Merge pull request #6054 from maxbrunet/feat/cohere-refresh-2025

✨ refresh Cohere support

Author: sestinj

core/llm/autodetect.ts

@@ -52,6 +52,7 @@ const PROVIDER_HANDLES_TEMPLATING: string[] = [
   "msty",
   "anthropic",
   "bedrock",
+  "cohere",
   "sagemaker",
   "continue-proxy",
   "mistral",
@@ -65,6 +66,7 @@ const PROVIDER_HANDLES_TEMPLATING: string[] = [
 const PROVIDER_SUPPORTS_IMAGES: string[] = [
   "openai",
   "ollama",
+  "cohere",
   "gemini",
   "msty",
   "anthropic",
@@ -89,6 +91,8 @@ const MODEL_SUPPORTS_IMAGES: string[] = [
   "gpt-4o-mini",
   "gpt-4-vision",
   "claude-3",
+  "c4ai-aya-vision-8b",
+  "c4ai-aya-vision-32b",
   "gemini-ultra",
   "gemini-1.5-pro",
   "gemini-1.5-flash",
@@ -140,6 +144,7 @@ function modelSupportsImages(
 const PARALLEL_PROVIDERS: string[] = [
   "anthropic",
   "bedrock",
+  "cohere",
   "sagemaker",
   "deepinfra",
   "gemini",
@@ -176,6 +181,7 @@ function autodetectTemplateType(model: string): TemplateType | undefined {
   if (
     lower.includes("gpt") ||
     lower.includes("command") ||
+    lower.includes("aya") ||
     lower.includes("chat-bison") ||
     lower.includes("pplx") ||
     lower.includes("gemini") ||

core/llm/llms/Cohere.ts

@@ -1,31 +1,100 @@
-import { streamJSON } from "@continuedev/fetch";
+import { streamSse } from "@continuedev/fetch";
 import {
   ChatMessage,
   Chunk,
   CompletionOptions,
   LLMOptions,
+  MessageContent,
 } from "../../index.js";
 import { renderChatMessage, stripImages } from "../../util/messageContent.js";
 import { BaseLLM } from "../index.js";
 
 class Cohere extends BaseLLM {
   static providerName = "cohere";
   static defaultOptions: Partial<LLMOptions> = {
-    apiBase: "https://api.cohere.ai/v1",
+    apiBase: "https://api.cohere.ai/v2",
     maxEmbeddingBatchSize: 96,
   };
   static maxStopSequences = 5;
 
   private _convertMessages(msgs: ChatMessage[]): any[] {
     const messages = [];
+    let lastToolPlan: MessageContent | undefined;
     for (const m of msgs) {
-      if (m.role === "system" || !m.content) {
+      if (!m.content) {
         continue;
       }
-      messages.push({
-        role: m.role === "assistant" ? "chatbot" : m.role,
-        message: m.content,
-      });
+      switch (m.role) {
+        case "user":
+          if (typeof m.content === "string") {
+            messages.push({
+              role: m.role,
+              content: m.content,
+            });
+            break;
+          }
+
+          messages.push({
+            role: m.role,
+            content: m.content.map((part) => {
+              if (part.type === "imageUrl") {
+                return {
+                  type: "image_url",
+                  image_url: { url: part.imageUrl.url },
+                };
+              }
+              return part;
+            }),
+          });
+          break;
+        case "thinking":
+          lastToolPlan = m.content;
+          break;
+        case "assistant":
+          if (m.toolCalls) {
+            if (!lastToolPlan) {
+              throw new Error("No tool plan found");
+            }
+            messages.push({
+              role: m.role,
+              tool_calls: m.toolCalls.map((toolCall) => ({
+                id: toolCall.id,
+                type: "function",
+                function: {
+                  name: toolCall.function?.name,
+                  arguments: toolCall.function?.arguments,
+                },
+              })),
+              // Ideally the tool plan would be in this message, but it is
+              // split in another, usually the previous, this one's content is
+              // a space.
+              // tool_plan: m.content,
+              tool_plan: lastToolPlan,
+            });
+            lastToolPlan = undefined;
+            break;
+          }
+          messages.push({
+            role: m.role,
+            content: m.content,
+          });
+          break;
+        case "system":
+          messages.push({
+            role: m.role,
+            content: stripImages(m.content),
+          });
+          break;
+        case "tool":
+          messages.push({
+            role: m.role,
+            content: m.content,
+            tool_call_id: m.toolCallId,
+          });
+          break;
+        default:
+          break;
+      }
     }
     return messages;
   }
@@ -41,7 +110,14 @@ class Cohere extends BaseLLM {
       stop_sequences: options.stop?.slice(0, Cohere.maxStopSequences),
       frequency_penalty: options.frequencyPenalty,
       presence_penalty: options.presencePenalty,
-      raw_prompting: options.raw,
+      tools: options.tools?.map((tool) => ({
+        type: "function",
+        function: {
+          name: tool.function.name,
+          parameters: tool.function.parameters,
+          description: tool.function.description,
+        },
+      })),
     };
   }
 
@@ -67,19 +143,12 @@ class Cohere extends BaseLLM {
       ...this.requestOptions?.headers,
     };
 
-    let preamble: string | undefined = undefined;
-    const systemMessage = messages.find((m) => m.role === "system")?.content;
-    if (systemMessage) {
-      preamble = stripImages(systemMessage);
-    }
     const resp = await this.fetch(new URL("chat", this.apiBase), {
       method: "POST",
       headers,
       body: JSON.stringify({
         ...this._convertArgs(options),
-        message: messages.pop()?.content,
-        chat_history: this._convertMessages(messages),
-        preamble,
+        messages: this._convertMessages(messages),
       }),
       signal,
     });
@@ -90,13 +159,97 @@ class Cohere extends BaseLLM {
 
     if (options.stream === false) {
       const data = await resp.json();
-      yield { role: "assistant", content: data.text };
+      if (data.message.tool_calls) {
+        yield {
+          // Use the "thinking" role for `tool_plan`, since there is no such
+          // role in the Cohere API at the moment and it is a "a
+          // chain-of-thought style reflection".
+          role: "thinking",
+          content: data.message.tool_plan,
+        };
+        yield {
+          role: "assistant",
+          content: "",
+          toolCalls: data.message.tool_calls.map((toolCall: any) => ({
+            id: toolCall.id,
+            type: "function",
+            function: {
+              name: toolCall.function?.name,
+              arguments: toolCall.function?.arguments,
+            },
+          })),
+        };
+        return;
+      }
+      yield { role: "assistant", content: data.message.content[0].text };
       return;
     }
 
-    for await (const value of streamJSON(resp)) {
-      if (value.event_type === "text-generation") {
-        yield { role: "assistant", content: value.text };
+    let lastToolUseId: string | undefined;
+    let lastToolUseName: string | undefined;
+    for await (const value of streamSse(resp)) {
+      // https://docs.cohere.com/v2/docs/streaming#stream-events
+      switch (value.type) {
+        // https://docs.cohere.com/v2/docs/streaming#content-delta
+        case "content-delta":
+          yield {
+            role: "assistant",
+            content: value.delta.message.content.text,
+          };
+          break;
+        // https://docs.cohere.com/reference/chat-stream#request.body.messages.assistant.tool_plan
+        case "tool-plan-delta":
+          // Use the "thinking" role for `tool_plan`, since there is no such
+          // role in the Cohere API at the moment and it is a "a
+          // chain-of-thought style reflection".
+          yield {
+            role: "thinking",
+            content: value.delta.message.tool_plan,
+          };
+          break;
+        case "tool-call-start":
+          lastToolUseId = value.delta.message.tool_calls.id;
+          lastToolUseName = value.delta.message.tool_calls.function.name;
+          yield {
+            role: "assistant",
+            content: "",
+            toolCalls: [
+              {
+                id: lastToolUseId,
+                type: "function",
+                function: {
+                  name: lastToolUseName,
+                  arguments: value.delta.message.tool_calls.function.arguments,
+                },
+              },
+            ],
+          };
+          break;
+        case "tool-call-delta":
+          if (!lastToolUseId || !lastToolUseName) {
+            throw new Error("No tool use found");
+          }
+          yield {
+            role: "assistant",
+            content: "",
+            toolCalls: [
+              {
+                id: lastToolUseId,
+                type: "function",
+                function: {
+                  name: lastToolUseName,
+                  arguments: value.delta.message.tool_calls.function.arguments,
+                },
+              },
+            ],
+          };
+          break;
+        case "tool-call-end":
+          lastToolUseId = undefined;
+          lastToolUseName = undefined;
+          break;
+        default:
+          break;
       }
     }
   }

core/llm/toolSupport.test.ts

@@ -109,6 +109,19 @@ describe("PROVIDER_TOOL_SUPPORT", () => {
     });
   });
 
+  describe("cohere", () => {
+    const supportsFn = PROVIDER_TOOL_SUPPORT["cohere"];
+
+    it("should return true for Command models", () => {
+      expect(supportsFn("command-r")).toBe(true);
+      expect(supportsFn("command-a")).toBe(true);
+    });
+
+    it("should return false for other models", () => {
+      expect(supportsFn("c4ai-aya-expanse-32b")).toBe(false);
+    });
+  });
+
   describe("gemini", () => {
     const supportsFn = PROVIDER_TOOL_SUPPORT["gemini"];
 
@@ -221,6 +234,7 @@ describe("PROVIDER_TOOL_SUPPORT", () => {
       expect(supportsFn("qwen2")).toBe(true);
       expect(supportsFn("mixtral-8x7b")).toBe(true);
       expect(supportsFn("command-r")).toBe(true);
+      expect(supportsFn("command-a")).toBe(true);
       expect(supportsFn("smollm2")).toBe(true);
       expect(supportsFn("hermes3")).toBe(true);
       expect(supportsFn("athene-v2")).toBe(true);

core/llm/toolSupport.ts

@@ -71,6 +71,9 @@ export const PROVIDER_TOOL_SUPPORT: Record<string, (model: string) => boolean> =
 
       return false;
     },
+    cohere: (model) => {
+      return model.toLowerCase().startsWith("command");
+    },
     gemini: (model) => {
       // All gemini models support function calling
       return model.toLowerCase().includes("gemini");
@@ -144,6 +147,7 @@ export const PROVIDER_TOOL_SUPPORT: Record<string, (model: string) => boolean> =
           "qwen3",
           "mixtral",
           "command-r",
+          "command-a",
           "smollm2",
           "hermes3",
           "athene-v2",
@@ -226,6 +230,7 @@ export const PROVIDER_TOOL_SUPPORT: Record<string, (model: string) => boolean> =
         "qwen/qwen3",
         "qwen/qwen-",
         "cohere/command-r",
+        "cohere/command-a",
         "ai21/jamba-1.6",
         "mistralai/mistral",
         "mistralai/ministral",