Merge pull request #63 from andyccliao/ollama-chat-reimplementation

Ollama chat reimplementation
semperai · Dec 23, 2023 · 7edaee1 · 7edaee1
2 parents b1437df + 269bd6c
commit 7edaee1
Showing 1 changed file with 12 additions and 21 deletions.
diff --git a/src/features/chat/ollamaChat.ts b/src/features/chat/ollamaChat.ts
@@ -6,13 +6,12 @@ export async function getOllamaChatResponseStream(messages: Message[]) {
   const headers: Record<string, string> = {
     "Content-Type": "application/json",
   };
-  const prompt = buildPrompt(messages);
-  const res = await fetch(`${config("ollama_url")}/api/generate`, {
+  const res = await fetch(`${config("ollama_url")}/api/chat`, {
     headers: headers,
     method: "POST",
     body: JSON.stringify({
       model: config("ollama_model"),
-      prompt,
+      messages,
     }),
   });
 
@@ -25,27 +24,20 @@ export async function getOllamaChatResponseStream(messages: Message[]) {
     async start(controller: ReadableStreamDefaultController) {
       const decoder = new TextDecoder("utf-8");
       try {
-        // sometimes the response is chunked, so we need to combine the chunks
-        let combined = "";
+        // Ollama sends chunks of multiple complete JSON objects separated by newlines
         while (true) {
           const { done, value } = await reader.read();
           if (done) break;
           const data = decoder.decode(value);
-          const chunks = data
-            .split("data:")
-            .filter((val) => !!val && val.trim() !== "[DONE]");
-
-          for (const chunk of chunks) {
-            // skip comments
-            if (chunk.length > 0 && chunk[0] === ":") {
-              continue;
-            }
-            combined += chunk;
+          const jsonResponses = data
+            .trim() // Ollama sends an empty line after the final JSON message...
+            .split("\n")
+            //.filter((val) => !!val) 
 
+          for (const jsonResponse of jsonResponses) {
             try {
-              const json = JSON.parse(combined);
-              const messagePiece = json.response;
-              combined = "";
+              const json = JSON.parse(jsonResponse);
+              const messagePiece = json.message.content;
               if (!!messagePiece) {
                 controller.enqueue(messagePiece);
               }
@@ -75,13 +67,12 @@ export async function getOllamaVisionChatResponse(messages: Message[], imageData
   const headers: Record<string, string> = {
     "Content-Type": "application/json",
   };
-  const prompt = buildPrompt(messages);
-  const res = await fetch(`${config("vision_ollama_url")}/api/generate`, {
+  const res = await fetch(`${config("vision_ollama_url")}/api/chat`, {
     headers: headers,
     method: "POST",
     body: JSON.stringify({
       model: config("vision_ollama_model"),
-      prompt,
+      messages,
       images: [imageData],
       stream: false,
     }),