Reimplement Cohere models for basic chat (#4489)

* Reimplement Cohere models - Redo LLM implementation to grab models from endpoint and pre-filter - Migrate embedding models to also grab from remote - Add records for easy context window lookup' * fix comment
2025-10-03 18:28:20 -07:00 · 2025-10-03 18:28:20 -07:00 · c2e7ccc00f
commit c2e7ccc00f
parent 988a14e67e
7 changed files with 1125 additions and 82 deletions
--- a/frontend/src/components/EmbeddingSelection/CohereOptions/index.jsx
+++ b/frontend/src/components/EmbeddingSelection/CohereOptions/index.jsx
@ -1,4 +1,10 @@
+import { useState, useEffect } from "react";
+import System from "@/models/system";
+
 export default function CohereEmbeddingOptions({ settings }) {
+  const [inputValue, setInputValue] = useState(settings?.CohereApiKey);
+  const [cohereApiKey, setCohereApiKey] = useState(settings?.CohereApiKey);
+
  return (
    <div className="w-full flex flex-col gap-y-4">
      <div className="w-full flex items-center gap-[36px] mt-1.5">
@ -15,41 +21,78 @@ export default function CohereEmbeddingOptions({ settings }) {
            required={true}
            autoComplete="off"
            spellCheck={false}
+            onChange={(e) => setInputValue(e.target.value)}
+            onBlur={() => setCohereApiKey(inputValue)}
          />
        </div>
-        <div className="flex flex-col w-60">
-          <label className="text-white text-sm font-semibold block mb-3">
-            Model Preference
-          </label>
-          <select
-            name="EmbeddingModelPref"
-            required={true}
-            className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
-          >
-            <optgroup label="Available embedding models">
-              {[
-                "embed-english-v3.0",
-                "embed-multilingual-v3.0",
-                "embed-english-light-v3.0",
-                "embed-multilingual-light-v3.0",
-                "embed-english-v2.0",
-                "embed-english-light-v2.0",
-                "embed-multilingual-v2.0",
-              ].map((model) => {
-                return (
-                  <option
-                    key={model}
-                    value={model}
-                    selected={settings?.EmbeddingModelPref === model}
-                  >
-                    {model}
-                  </option>
-                );
-              })}
-            </optgroup>
-          </select>
-        </div>
+        <CohereModelSelection settings={settings} apiKey={cohereApiKey} />
      </div>
    </div>
  );
 }
+
+function CohereModelSelection({ apiKey, settings }) {
+  const [models, setModels] = useState([]);
+  const [loading, setLoading] = useState(true);
+
+  useEffect(() => {
+    async function findCustomModels() {
+      if (!apiKey) {
+        setModels([]);
+        setLoading(true);
+        return;
+      }
+
+      setLoading(true);
+      const { models } = await System.customModels(
+        "cohere-embedder",
+        typeof apiKey === "boolean" ? null : apiKey
+      );
+      setModels(models || []);
+      setLoading(false);
+    }
+    findCustomModels();
+  }, [apiKey]);
+
+  if (loading) {
+    return (
+      <div className="flex flex-col w-60">
+        <label className="text-white text-sm font-semibold block mb-3">
+          Model Preference
+        </label>
+        <select
+          name="EmbeddingModelPref"
+          disabled={true}
+          className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
+        >
+          <option disabled={true} selected={true}>
+            -- loading available models --
+          </option>
+        </select>
+      </div>
+    );
+  }
+
+  return (
+    <div className="flex flex-col w-60">
+      <label className="text-white text-sm font-semibold block mb-3">
+        Model Preference
+      </label>
+      <select
+        name="EmbeddingModelPref"
+        required={true}
+        className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
+      >
+        {models.map((model) => (
+          <option
+            key={model.id}
+            value={model.id}
+            selected={settings?.EmbeddingModelPref === model.id}
+          >
+            {model.name}
+          </option>
+        ))}
+      </select>
+    </div>
+  );
+}
--- a/frontend/src/components/LLMSelection/CohereAiOptions/index.jsx
+++ b/frontend/src/components/LLMSelection/CohereAiOptions/index.jsx
@ -1,4 +1,10 @@
+import { useState, useEffect } from "react";
+import System from "@/models/system";
+
 export default function CohereAiOptions({ settings }) {
+  const [inputValue, setInputValue] = useState(settings?.CohereApiKey);
+  const [cohereApiKey, setCohereApiKey] = useState(settings?.CohereApiKey);
+
  return (
    <div className="w-full flex flex-col">
      <div className="w-full flex items-center gap-[36px] mt-1.5">
@ -15,35 +21,80 @@ export default function CohereAiOptions({ settings }) {
            required={true}
            autoComplete="off"
            spellCheck={false}
+            onChange={(e) => setInputValue(e.target.value)}
+            onBlur={() => setCohereApiKey(inputValue)}
          />
        </div>
-        <div className="flex flex-col w-60">
-          <label className="text-white text-sm font-semibold block mb-3">
-            Chat Model Selection
-          </label>
-          <select
-            name="CohereModelPref"
-            defaultValue={settings?.CohereModelPref || "command-r"}
-            required={true}
-            className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
-          >
-            {[
-              "command-r",
-              "command-r-plus",
-              "command",
-              "command-light",
-              "command-nightly",
-              "command-light-nightly",
-            ].map((model) => {
-              return (
-                <option key={model} value={model}>
-                  {model}
-                </option>
-              );
-            })}
-          </select>
-        </div>
+        {!settings?.credentialsOnly && (
+          <CohereModelSelection settings={settings} apiKey={cohereApiKey} />
+        )}
      </div>
    </div>
  );
 }
+
+function CohereModelSelection({ apiKey, settings }) {
+  const [models, setModels] = useState([]);
+  const [loading, setLoading] = useState(true);
+
+  useEffect(() => {
+    async function findCustomModels() {
+      if (!apiKey) {
+        setModels([]);
+        setLoading(true);
+        return;
+      }
+
+      setLoading(true);
+      const { models } = await System.customModels(
+        "cohere",
+        typeof apiKey === "boolean" ? null : apiKey
+      );
+      setModels(models || []);
+      setLoading(false);
+    }
+    findCustomModels();
+  }, [apiKey]);
+
+  if (loading) {
+    return (
+      <div className="flex flex-col w-60">
+        <label className="text-white text-sm font-semibold block mb-3">
+          Chat Model Selection
+        </label>
+        <select
+          name="CohereModelPref"
+          disabled={true}
+          className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
+        >
+          <option disabled={true} selected={true}>
+            -- loading available models --
+          </option>
+        </select>
+      </div>
+    );
+  }
+
+  return (
+    <div className="flex flex-col w-60">
+      <label className="text-white text-sm font-semibold block mb-3">
+        Chat Model Selection
+      </label>
+      <select
+        name="CohereModelPref"
+        required={true}
+        className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
+      >
+        {models.map((model) => (
+          <option
+            key={model.id}
+            value={model.id}
+            selected={settings?.CohereModelPref === model.id}
+          >
+            {model.name}
+          </option>
+        ))}
+      </select>
+    </div>
+  );
+}
--- a/server/package.json
+++ b/server/package.json
@ -46,7 +46,7 @@
    "check-disk-space": "^3.4.0",
    "cheerio": "^1.0.0",
    "chromadb": "^2.0.1",
-    "cohere-ai": "^7.9.5",
+    "cohere-ai": "^7.19.0",
    "cors": "^2.8.5",
    "dotenv": "^16.0.3",
    "elevenlabs": "^0.5.0",
--- a/server/utils/AiProviders/cohere/index.js
+++ b/server/utils/AiProviders/cohere/index.js
@ -8,6 +8,7 @@ const {

 class CohereLLM {
  constructor(embedder = null) {
+    this.className = "CohereLLM";
    const { CohereClient } = require("cohere-ai");
    if (!process.env.COHERE_API_KEY)
      throw new Error("No Cohere API key was set.");
@ -25,6 +26,13 @@ class CohereLLM {
    };

    this.embedder = embedder ?? new NativeEmbedder();
+    this.#log(
+      `Initialized with model ${this.model}. ctx: ${this.promptWindowLimit()}`
+    );
+  }
+
+  #log(text, ...args) {
+    console.log(`\x1b[32m[${this.className}]\x1b[0m ${text}`, ...args);
  }

  #appendContext(contextTexts = []) {
@ -70,16 +78,8 @@ class CohereLLM {
    return MODEL_MAP.get("cohere", this.model) ?? 4_096;
  }

-  async isValidChatCompletionModel(model = "") {
-    const validModels = [
-      "command-r",
-      "command-r-plus",
-      "command",
-      "command-light",
-      "command-nightly",
-      "command-light-nightly",
-    ];
-    return validModels.includes(model);
+  async isValidChatCompletionModel() {
+    return true;
  }

  constructPrompt({
@ -96,11 +96,6 @@ class CohereLLM {
  }

  async getChatCompletion(messages = null, { temperature = 0.7 }) {
-    if (!(await this.isValidChatCompletionModel(this.model)))
-      throw new Error(
-        `Cohere chat: ${this.model} is not valid for chat completion!`
-      );
-
    const message = messages[messages.length - 1].content; // Get the last message
    const cohereHistory = this.#convertChatHistoryCohere(messages.slice(0, -1)); // Remove the last message and convert to Cohere

@ -134,11 +129,6 @@ class CohereLLM {
  }

  async streamGetChatCompletion(messages = null, { temperature = 0.7 }) {
-    if (!(await this.isValidChatCompletionModel(this.model)))
-      throw new Error(
-        `Cohere chat: ${this.model} is not valid for chat completion!`
-      );
-
    const message = messages[messages.length - 1].content; // Get the last message
    const cohereHistory = this.#convertChatHistoryCohere(messages.slice(0, -1)); // Remove the last message and convert to Cohere
    const measuredStreamRequest = await LLMPerformanceMonitor.measureStream(
--- a/server/utils/AiProviders/modelMap/legacy.js
+++ b/server/utils/AiProviders/modelMap/legacy.js
@ -22,6 +22,18 @@ const LEGACY_MODEL_MAP = {
    "command-light": 4096,
    "command-nightly": 8192,
    "command-light-nightly": 8192,
+    "command-r-plus-08-2024": 132096,
+    "command-a-03-2025": 288000,
+    "c4ai-aya-vision-32b": 16384,
+    "command-a-reasoning-08-2025": 288768,
+    "command-r-08-2024": 132096,
+    "c4ai-aya-vision-8b": 16384,
+    "command-r7b-12-2024": 132000,
+    "command-r7b-arabic-02-2025": 128000,
+    "command-a-vision-07-2025": 128000,
+    "c4ai-aya-expanse-8b": 8192,
+    "c4ai-aya-expanse-32b": 128000,
+    "command-a-translate-08-2025": 8992,
  },
  gemini: {
    "gemini-1.5-pro-001": 2000000,
--- a/server/utils/helpers/customModels.js
+++ b/server/utils/helpers/customModels.js
@ -37,8 +37,10 @@ const SUPPORT_CUSTOM_MODELS = [
  "dpais",
  "moonshotai",
  "foundry",
+  "cohere",
  // Embedding Engines
  "native-embedder",
+  "cohere-embedder",
 ];

 async function getCustomModels(provider = "", apiKey = null, basePath = null) {
@ -96,8 +98,12 @@ async function getCustomModels(provider = "", apiKey = null, basePath = null) {
      return await getMoonshotAiModels(apiKey);
    case "foundry":
      return await getFoundryModels(basePath);
+    case "cohere":
+      return await getCohereModels(apiKey, "chat");
    case "native-embedder":
      return await getNativeEmbedderModels();
+    case "cohere-embedder":
+      return await getCohereModels(apiKey, "embed");
    default:
      return { models: [], error: "Invalid provider for custom models" };
  }
@ -759,6 +765,39 @@ async function getFoundryModels(basePath = null) {
  }
 }

+/**
+ * Get Cohere models
+ * @param {string} _apiKey - The API key to use
+ * @param {'chat' | 'embed'} type - The type of model to get
+ * @returns {Promise<{models: Array<{id: string, organization: string, name: string}>, error: string | null}>}
+ */
+async function getCohereModels(_apiKey = null, type = "chat") {
+  const apiKey =
+    _apiKey === true
+      ? process.env.COHERE_API_KEY
+      : _apiKey || process.env.COHERE_API_KEY || null;
+
+  const { CohereClient } = require("cohere-ai");
+  const cohere = new CohereClient({
+    token: apiKey,
+  });
+  const models = await cohere.models
+    .list({ pageSize: 1000, endpoint: type })
+    .then((results) => results.models)
+    .then((models) =>
+      models.map((model) => ({
+        id: model.id,
+        name: model.name,
+      }))
+    )
+    .catch((e) => {
+      console.error(`Cohere:listModels`, e.message);
+      return [];
+    });
+
+  return { models, error: null };
+}
+
 module.exports = {
  getCustomModels,
  SUPPORT_CUSTOM_MODELS,
--- a/server/yarn.lock
+++ b/server/yarn.lock