improve gemini stream to avoid word cutting

harmony-one · Mar 28, 2024 · 8d52ff7 · 8d52ff7
1 parent cbe63e8
commit 8d52ff7
Showing 1 changed file with 22 additions and 24 deletions.
diff --git a/src/modules/llms/api/vertex.ts b/src/modules/llms/api/vertex.ts
@@ -7,7 +7,7 @@ import { GrammyError } from 'grammy'
 import { pino } from 'pino'
 import { LlmsModelsEnum } from '../types'
 
-const API_ENDPOINT = config.llms.apiEndpoint // config.llms.apiEndpoint // config.llms.apiEndpoint // 'http://127.0.0.1:5000' // config.llms.apiEndpoint
+const API_ENDPOINT = config.llms.apiEndpoint // config.llms.apiEndpoint  // 'http://127.0.0.1:5000' // config.llms.apiEndpoint
 
 const logger = pino({
   name: 'Gemini - llmsBot',
@@ -76,31 +76,29 @@ export const vertexStreamCompletion = async (
   for await (const chunk of completionStream) {
     const msg = chunk.toString()
     if (msg) {
-      if (msg.startsWith('Text')) {
-        completion += msg.split('Text: ')[1]
-        if (msg.includes('Input Token:')) {
-          const tokenMsg = msg.split('Input Token: ')[1]
-          inputTokens = tokenMsg.split('Output Tokens: ')[0]
-          outputTokens = tokenMsg.split('Output Tokens: ')[1]
-          completion = completion.split('Input Token: ')[0]
-        }
-        completion = completion.replaceAll('...', '')
-        completion += '...'
-        if (ctx.chat?.id) {
-          await ctx.api
-            .editMessageText(ctx.chat?.id, msgId, completion)
-            .catch(async (e: any) => {
-              if (e instanceof GrammyError) {
-                if (e.error_code !== 400) {
-                  throw e
-                } else {
-                  logger.error(e)
-                }
-              } else {
+      completion += msg // .split('Text: ')[1]
+      if (msg.includes('Input Token:')) {
+        const tokenMsg = msg.split('Input Token: ')[1]
+        inputTokens = tokenMsg.split('Output Tokens: ')[0]
+        outputTokens = tokenMsg.split('Output Tokens: ')[1]
+        completion = completion.split('Input Token: ')[0]
+      }
+      completion = completion.replaceAll('...', '')
+      completion += '...'
+      if (ctx.chat?.id) {
+        await ctx.api
+          .editMessageText(ctx.chat?.id, msgId, completion)
+          .catch(async (e: any) => {
+            if (e instanceof GrammyError) {
+              if (e.error_code !== 400) {
                 throw e
+              } else {
+                logger.error(e)
               }
-            })
-        }
+            } else {
+              throw e
+            }
+          })
       }
     }
   }