add in manual stop sequence detection

Author: charleslien

backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts

@@ -9,6 +9,7 @@ import { buildArray } from '@codebuff/common/util/array'
 import { convertCbToModelMessages } from '@codebuff/common/util/messages'
 import { errorToObject } from '@codebuff/common/util/object'
 import { withTimeout } from '@codebuff/common/util/promise'
+import { StopSequenceHandler } from '@codebuff/common/util/stop-sequence'
 import { generateCompactId } from '@codebuff/common/util/string'
 import { APICallError, generateObject, generateText, streamText } from 'ai'
 
@@ -99,10 +100,6 @@ export const promptAiSdkStream = async function* (
       },
       'Skipping stream due to canceled user input',
     )
-    yield {
-      type: 'text',
-      text: '',
-    }
     return
   }
   const startTime = Date.now()
@@ -117,8 +114,18 @@ export const promptAiSdkStream = async function* (
   })
 
   let content = ''
+  const stopSequenceHandler = new StopSequenceHandler(options.stopSequences)
 
   for await (const chunk of response.fullStream) {
+    if (chunk.type !== 'text-delta') {
+      const flushed = stopSequenceHandler.flush()
+      if (flushed) {
+        yield {
+          type: 'text',
+          text: flushed,
+        }
+      }
+    }
     if (chunk.type === 'error') {
       logger.error(
         {
@@ -161,13 +168,33 @@ export const promptAiSdkStream = async function* (
       }
     }
     if (chunk.type === 'text-delta') {
-      content += chunk.text
-      yield {
-        type: 'text',
-        text: chunk.text,
+      if (!options.stopSequences) {
+        content += chunk.text
+        if (chunk.text) {
+          yield {
+            type: 'text',
+            text: chunk.text,
+          }
+        }
+        continue
+      }
+
+      const stopSequenceResult = stopSequenceHandler.process(chunk.text)
+      if (stopSequenceResult.text) {
+        yield {
+          type: 'text',
+          text: stopSequenceResult.text,
+        }
       }
     }
   }
+  const flushed = stopSequenceHandler.flush()
+  if (flushed) {
+    yield {
+      type: 'text',
+      text: flushed,
+    }
+  }
 
   const messageId = (await response.response).id
   if (options.resolveMessageId) {

common/src/util/stop-sequence.ts

@@ -0,0 +1,61 @@
+import { suffixPrefixOverlap } from './string'
+
+export class StopSequenceHandler {
+  private buffer: string = ''
+  private finished: boolean = false
+  private stopSequences: string[]
+
+  constructor(stopSequences?: string[]) {
+    this.stopSequences = stopSequences ?? []
+  }
+
+  public process(
+    text: string,
+  ):
+    | { text: string; endOfStream: boolean }
+    | { text: null; endOfStream: true } {
+    if (this.finished) {
+      return {
+        text: null,
+        endOfStream: true,
+      }
+    }
+    this.buffer += text
+    let longestOverlap = ''
+
+    for (const stopSequence of this.stopSequences) {
+      const index = this.buffer.indexOf(stopSequence)
+      if (index !== -1) {
+        this.finished = true
+        return {
+          text: this.buffer.slice(0, index),
+          endOfStream: true,
+        }
+      }
+    }
+
+    for (const stopSequence of this.stopSequences) {
+      const overlap = suffixPrefixOverlap(this.buffer, stopSequence)
+      longestOverlap =
+        overlap.length > longestOverlap.length ? overlap : longestOverlap
+    }
+
+    const index = this.buffer.length - longestOverlap.length
+    const processed = this.buffer.slice(0, index)
+    this.buffer = this.buffer.slice(index)
+
+    return {
+      text: processed,
+      endOfStream: false,
+    }
+  }
+
+  public flush(): string {
+    if (this.finished) {
+      return ''
+    }
+    const b = this.buffer
+    this.buffer = ''
+    return b
+  }
+}