Update utils/server/index.ts

utils/server/index.ts

@@ -74,31 +74,54 @@ export const LLMStream = async (
   const stream = new ReadableStream({
     async start(controller) {
       const onParse = (event: ParsedEvent | ReconnectInterval) => {
-        console.log(event);
-        if (event.type === 'event') {
-          const data = event.data;
-          if (data === '<|end|>') {
-            controller.close();
-            return; 
-          }
-
-          try {
+      console.log(event);
+      if (event.type === 'event') {
+        const data = event.data;
+        if (data === '') {
+          controller.close();
+          return; 
+        }
+    
+        try {
+          const json = JSON.parse(data);
+    
+          // New condition to handle Python API format
+          if (json.choices && json.choices[0].message) {
+            const text = json.choices[0].message.content;
+            const finishReason = json.choices[0].finish_reason;
+    
+            if (finishReason === 'stop') {
+              controller.close();
+              return;
+            }
             
-            const queue = encoder.encode(data);
+            const queue = encoder.encode(text);
+            controller.enqueue(queue);
+          }
+          // Old condition to handle existing format
+          else if (json.choices && json.choices[0].delta) {
+            if (json.choices[0].finish_reason != null) {
+              controller.close();
+              return;
+            }
             
+            const text = json.choices[0].delta.content;
+            const queue = encoder.encode(text);
             controller.enqueue(queue);
-          } catch (e) {
-            controller.error(e);
           }
+    
+        } catch (e) {
+          controller.error(e);
         }
-      };
+      }
+    };
 
-      const parser = createParser(onParse);
+    const parser = createParser(onParse);
 
-      for await (const chunk of res.body as any) {
-        parser.feed(decoder.decode(chunk));
-      }
-    },
+    for await (const chunk of res.body as any) {
+      parser.feed(decoder.decode(chunk));
+    }
+  }
   });
 
   return stream;