Merge pull request #301 from TBXark/dev

perf: support cohere v2 api

package.json

@@ -36,7 +36,7 @@
         "start:debug": "wrangler dev --local",
         "prepare:vercel": "tsx ./scripts/plugins/vercel/setenv.ts",
         "wrangler": "wrangler",
-        "test": "tsx ./src/config/env.test.ts"
+        "test": "tsx ./src/agent/index.test.ts"
     },
     "dependencies": {
         "cloudflare-worker-adapter": "^1.3.3"

src/agent/cohere.ts

@@ -1,19 +1,12 @@
 import type { AgentUserConfig } from '../config/env';
 import type { SseChatCompatibleOptions } from './request';
-import type { SSEMessage, SSEParserResult } from './stream';
-import type { ChatAgent, ChatStreamTextHandler, HistoryItem, LLMChatParams } from './types';
+import type { ChatAgent, ChatStreamTextHandler, LLMChatParams } from './types';
 import { requestChatCompletions } from './request';
-import { Stream } from './stream';
 
 export class Cohere implements ChatAgent {
     readonly name = 'cohere';
     readonly modelKey = 'COHERE_CHAT_MODEL';
 
-    static COHERE_ROLE_MAP: Record<string, string> = {
-        assistant: 'CHATBOT',
-        user: 'USER',
-    };
-
     readonly enable = (context: AgentUserConfig): boolean => {
         return !!(context.COHERE_API_KEY);
     };
@@ -22,37 +15,6 @@ export class Cohere implements ChatAgent {
         return ctx.COHERE_CHAT_MODEL;
     };
 
-    private render = (item: HistoryItem): any => {
-        return {
-            role: Cohere.COHERE_ROLE_MAP[item.role] || 'USER',
-            content: item.content,
-        };
-    };
-
-    static parser(sse: SSEMessage): SSEParserResult {
-        // example:
-        //      event: text-generation
-        //      data: {"is_finished":false,"event_type":"text-generation","text":"?"}
-        //
-        //      event: stream-end
-        //      data: {"is_finished":true,...}
-        switch (sse.event) {
-            case 'text-generation':
-                try {
-                    return { data: JSON.parse(sse.data || '') };
-                } catch (e) {
-                    console.error(e, sse.data);
-                    return {};
-                }
-            case 'stream-start':
-                return {};
-            case 'stream-end':
-                return { finish: true };
-            default:
-                return {};
-        }
-    }
-
     readonly request = async (params: LLMChatParams, context: AgentUserConfig, onStream: ChatStreamTextHandler | null): Promise<string> => {
         const { message, prompt, history } = params;
         const url = `${context.COHERE_API_BASE}/chat`;
@@ -63,27 +25,22 @@ export class Cohere implements ChatAgent {
         };
 
         const messages = [...history || [], { role: 'user', content: message }];
+        if (prompt) {
+            messages.unshift({ role: 'assistant', content: prompt });
+        }
 
         const body = {
-            message,
+            messages,
             model: context.COHERE_CHAT_MODEL,
             stream: onStream != null,
-            preamble: prompt,
-            chat_history: messages.map(this.render),
         };
-        if (!body.preamble) {
-            delete body.preamble;
-        }
 
         const options: SseChatCompatibleOptions = {};
-        options.streamBuilder = function (r, c) {
-            return new Stream(r, c, Cohere.parser);
-        };
         options.contentExtractor = function (data: any) {
-            return data?.text;
+            return data?.delta?.message?.content?.text;
         };
         options.fullContentExtractor = function (data: any) {
-            return data?.text;
+            return data?.messages[0].content;
         };
         options.errorExtractor = function (data: any) {
             return data?.message;

src/agent/index.test.ts

@@ -6,7 +6,7 @@ import '../config/env.test';
 {
     const agent = loadChatLLM({
         ...ENV.USER_CONFIG,
-        AI_PROVIDER: 'anthropic',
+        AI_PROVIDER: 'cohere',
     });
     const params: LLMChatParams = {
         prompt: 'You are a useful assistant.',

src/agent/stream.ts

@@ -138,7 +138,7 @@ function defaultSSEJsonParser(sse: SSEMessage): SSEParserResult {
     if (sse.data?.startsWith('[DONE]')) {
         return { finish: true };
     }
-    if (sse.event === null && sse.data) {
+    if (sse.data) {
         try {
             return { data: JSON.parse(sse.data) };
         } catch (e) {

src/config/config.ts

@@ -169,7 +169,7 @@ export class CohereConfig {
     // cohere api key
     COHERE_API_KEY: string | null = null;
     // cohere api base
-    COHERE_API_BASE = 'https://api.cohere.com/v1';
+    COHERE_API_BASE = 'https://api.cohere.com/v2';
     // cohere api model
     COHERE_CHAT_MODEL = 'command-r-plus';
 }