From e8cf6ee9e7260a9c545d572d5d8fac0c81ff28ad Mon Sep 17 00:00:00 2001
From: w1gs <root@w1gs.dev>
Date: Sat, 21 Sep 2024 00:36:11 -0400
Subject: [PATCH] Replaced fetch with Ollama client

---
 extensions/void/package-lock.json            | 14 +++++
 extensions/void/package.json                 |  1 +
 extensions/void/src/common/sendLLMMessage.ts | 66 ++++++--------------
 3 files changed, 33 insertions(+), 48 deletions(-)

diff --git a/extensions/void/package-lock.json b/extensions/void/package-lock.json
index 1c49e24a..354cc7e7 100644
--- a/extensions/void/package-lock.json
+++ b/extensions/void/package-lock.json
@@ -9,6 +9,7 @@
       "version": "0.0.1",
       "dependencies": {
         "@anthropic-ai/sdk": "^0.27.1",
+        "ollama": "^0.5.9",
         "openai": "^4.57.0"
       },
       "devDependencies": {
@@ -5183,6 +5184,14 @@
         "url": "https://github.com/sponsors/ljharb"
       }
     },
+    "node_modules/ollama": {
+      "version": "0.5.9",
+      "resolved": "https://registry.npmjs.org/ollama/-/ollama-0.5.9.tgz",
+      "integrity": "sha512-F/KZuDRC+ZsVCuMvcOYuQ6zj42/idzCkkuknGyyGVmNStMZ/sU3jQpvhnl4SyC0+zBzLiKNZJnJeuPFuieWZvQ==",
+      "dependencies": {
+        "whatwg-fetch": "^3.6.20"
+      }
+    },
     "node_modules/once": {
       "version": "1.4.0",
       "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz",
@@ -7231,6 +7240,11 @@
       "resolved": "https://registry.npmjs.org/webidl-conversions/-/webidl-conversions-3.0.1.tgz",
       "integrity": "sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ=="
     },
+    "node_modules/whatwg-fetch": {
+      "version": "3.6.20",
+      "resolved": "https://registry.npmjs.org/whatwg-fetch/-/whatwg-fetch-3.6.20.tgz",
+      "integrity": "sha512-EqhiFU6daOA8kpjOWTL0olhVOF3i7OrFzSYiGsEMB8GcXS+RrzauAERX65xMeNWVqxA6HXH2m69Z9LaKKdisfg=="
+    },
     "node_modules/whatwg-url": {
       "version": "5.0.0",
       "resolved": "https://registry.npmjs.org/whatwg-url/-/whatwg-url-5.0.0.tgz",
diff --git a/extensions/void/package.json b/extensions/void/package.json
index a0267896..4e66383d 100644
--- a/extensions/void/package.json
+++ b/extensions/void/package.json
@@ -142,6 +142,7 @@
   },
   "dependencies": {
     "@anthropic-ai/sdk": "^0.27.1",
+    "ollama": "^0.5.9",
     "openai": "^4.57.0"
   }
 }
diff --git a/extensions/void/src/common/sendLLMMessage.ts b/extensions/void/src/common/sendLLMMessage.ts
index bf87f40b..2fb3dd3d 100644
--- a/extensions/void/src/common/sendLLMMessage.ts
+++ b/extensions/void/src/common/sendLLMMessage.ts
@@ -1,5 +1,6 @@
 import Anthropic from '@anthropic-ai/sdk';
 import OpenAI from 'openai';
+import { Ollama } from 'ollama/browser'
 
 export type ApiConfig = {
 	anthropic: {
@@ -239,6 +240,8 @@ export const sendLLMMessage: SendLLMMessageFnTypeExternal = ({ messages, onText,
 
 // Ollama
 export const sendOllamaMsg: SendLLMMessageFnTypeInternal = ({ messages, onText, onFinalMessage, apiConfig }) => {
+	const ollamaClient = new Ollama({ host: apiConfig.ollama.endpoint })
+
 	let didAbort = false;
 	let fullText = "";
 
@@ -247,57 +250,24 @@ export const sendOllamaMsg: SendLLMMessageFnTypeInternal = ({ messages, onText,
 		didAbort = true;
 	};
 
-	const handleError = (error: any) => {
-		console.error('Error:', error);
-		onFinalMessage(fullText);
-	};
-
-	if (apiConfig.ollama.endpoint.endsWith('/')) {
-		apiConfig.ollama.endpoint = apiConfig.ollama.endpoint.slice(0, -1);
-	}
-
-	fetch(`${apiConfig.ollama.endpoint}/api/chat`, {
-		method: "POST",
-		headers: {
-			"Content-Type": "application/json",
-		},
-		body: JSON.stringify({
-			model: apiConfig.ollama.model,
-			messages: messages,
-			stream: true,
-		}),
+	ollamaClient.chat({
+		model: apiConfig.ollama.model,
+		messages: messages,
+		stream: true,
 	})
-		.then(response => {
-			if (didAbort) return;
-			const reader = response.body?.getReader();
-			if (!reader) {
-				onFinalMessage(fullText);
-				return;
+		.then(async (stream) => {
+			for await (const chunk of stream) {
+				if (didAbort) return;
+				const newText = chunk.message.content;
+				fullText += newText;
+				onText(newText, fullText);
 			}
-			return reader;
+			onFinalMessage(fullText);
 		})
-		.then(reader => {
-			if (!reader) return;
-
-			const readStream = async () => {
-				try {
-					let done, value;
-					while ({ done, value } = await reader.read(), !done) {
-						if (didAbort) return;
-						const stringedResponse = new TextDecoder().decode(value);
-						const newText = JSON.parse(stringedResponse).message.content;
-						fullText += newText;
-						onText(newText, fullText);
-					}
-					onFinalMessage(fullText);
-				} catch (error) {
-					handleError(error);
-				}
-			};
-
-			readStream();
-		})
-		.catch(handleError);
+		.catch((error) => {
+			console.error('Error:', error);
+			onFinalMessage(fullText);
+		});
 
 	return { abort };
 };