fix: streaming anthropic models (janhq#3079)

Signed-off-by: James <[email protected]>
hansongwang · Jun 21, 2024 · a91fe5c · a91fe5c
1 parent af92c1d
commit a91fe5c
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 7 deletions.
diff --git a/core/src/types/api/index.ts b/core/src/types/api/index.ts
@@ -169,5 +169,5 @@ export const APIEvents = [
 export type PayloadType = {
   messages: ChatCompletionMessage[]
   model: string
-  stream: Boolean
+  stream: boolean
 }
diff --git a/extensions/inference-anthropic-extension/package.json b/extensions/inference-anthropic-extension/package.json
@@ -1,7 +1,7 @@
 {
   "name": "@janhq/inference-anthropic-extension",
   "productName": "Anthropic Inference Engine",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "description": "This extension enables Anthropic chat completion API calls",
   "main": "dist/index.js",
   "module": "dist/module.js",

diff --git a/extensions/inference-anthropic-extension/src/index.ts b/extensions/inference-anthropic-extension/src/index.ts
@@ -19,6 +19,7 @@ enum Settings {
 }
 
 type AnthropicPayloadType = {
+  stream: boolean
   model?: string
   max_tokens?: number
   messages?: Array<{ role: string; content: string }>
@@ -86,16 +87,22 @@ export default class JanInferenceAnthropicExtension extends RemoteOAIEngine {
   // Override the transformPayload method to convert the payload to the required format
   transformPayload = (payload: PayloadType): AnthropicPayloadType => {
     if (!payload.messages || payload.messages.length === 0) {
-      return { max_tokens: this.maxTokens, messages: [], model: payload.model }
+      return {
+        max_tokens: this.maxTokens,
+        messages: [],
+        model: payload.model,
+        stream: payload.stream,
+      }
     }
 
     const convertedData: AnthropicPayloadType = {
       max_tokens: this.maxTokens,
       messages: [],
       model: payload.model,
+      stream: payload.stream,
     }
 
-    payload.messages.forEach((item, index) => {
+    payload.messages.forEach((item) => {
       if (item.role === ChatCompletionRole.User) {
         convertedData.messages.push({
           role: 'user',
@@ -112,13 +119,30 @@ export default class JanInferenceAnthropicExtension extends RemoteOAIEngine {
     return convertedData
   }
 
+  // Sample returned stream data from anthropic
+  // {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}         }
+  // {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}              }
+  // {"type":"content_block_stop","index":0        }
+  // {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":12}   }
+
   // Override the transformResponse method to convert the response to the required format
   transformResponse = (data: any): string => {
+    // handling stream response
+    if (typeof data === 'string' && data.trim().length === 0) return ''
+    if (typeof data === 'string' && data.startsWith('event: ')) return ''
+    if (typeof data === 'string' && data.startsWith('data: ')) {
+      data = data.replace('data: ', '')
+      const parsedData = JSON.parse(data)
+      if (parsedData.type !== 'content_block_delta') return ''
+      return parsedData.delta?.text ?? ''
+    }
+
+    // non stream response
     if (data.content && data.content.length > 0 && data.content[0].text) {
       return data.content[0].text
-    } else {
-      console.error('Invalid response format:', data)
-      return ''
     }
+
+    console.error('Invalid response format:', data)
+    return ''
   }
 }