fix word cut in streaming

harmony-one · Mar 27, 2024 · cbe63e8 · cbe63e8
1 parent 4588b41
commit cbe63e8
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/src/modules/llms/api/athropic.ts b/src/modules/llms/api/athropic.ts
@@ -88,9 +88,16 @@ export const anthropicStreamCompletion = async (
     if (msg) {
       if (msg.startsWith('Input Token')) {
         inputTokens = msg.split('Input Token: ')[1]
-      } else if (msg.startsWith('Text')) {
+      } else if (msg.startsWith('Output Tokens')) {
+        outputTokens = msg.split('Output Tokens: ')[1]
+      } else {
         wordCount++
-        completion += msg.split('Text: ')[1]
+        completion += msg // .split('Text: ')[1]
+        if (msg.includes('Output Tokens:')) {
+          const tokenMsg = msg.split('Output Tokens: ')[1]
+          outputTokens = tokenMsg.split('Output Tokens: ')[1]
+          completion = completion.split('Output Tokens: ')[0]
+        }
         if (wordCount > wordCountMinimum) { // if (chunck === '.' && wordCount > wordCountMinimum) {
           if (wordCountMinimum < 64) {
             wordCountMinimum *= 2
@@ -114,8 +121,6 @@ export const anthropicStreamCompletion = async (
               })
           }
         }
-      } else if (msg.startsWith('Output Tokens')) {
-        outputTokens = msg.split('Output Tokens: ')[1]
       }
     }
   }