Skip to content

Commit c567167

Browse files
timothycarambattimothycarambat
committed
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'
1 parent 8cdadd8 commit c567167

File tree

7 files changed

+1125
-82
lines changed

7 files changed

+1125
-82
lines changed

‎frontend/src/components/EmbeddingSelection/CohereOptions/index.jsx‎

Lines changed: 75 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,10 @@
1+
import { useState, useEffect } from "react";
2+
import System from "@/models/system";
3+
14
export default function CohereEmbeddingOptions({ settings }) {
5+
const [inputValue, setInputValue] = useState(settings?.CohereApiKey);
6+
const [cohereApiKey, setCohereApiKey] = useState(settings?.CohereApiKey);
7+
28
return (
39
<div className="w-full flex flex-col gap-y-4">
410
<div className="w-full flex items-center gap-[36px] mt-1.5">
@@ -15,41 +21,78 @@ export default function CohereEmbeddingOptions({ settings }) {
1521
required={true}
1622
autoComplete="off"
1723
spellCheck={false}
24+
onChange={(e) => setInputValue(e.target.value)}
25+
onBlur={() => setCohereApiKey(inputValue)}
1826
/>
1927
</div>
20-
<div className="flex flex-col w-60">
21-
<label className="text-white text-sm font-semibold block mb-3">
22-
Model Preference
23-
</label>
24-
<select
25-
name="EmbeddingModelPref"
26-
required={true}
27-
className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
28-
>
29-
<optgroup label="Available embedding models">
30-
{[
31-
"embed-english-v3.0",
32-
"embed-multilingual-v3.0",
33-
"embed-english-light-v3.0",
34-
"embed-multilingual-light-v3.0",
35-
"embed-english-v2.0",
36-
"embed-english-light-v2.0",
37-
"embed-multilingual-v2.0",
38-
].map((model) => {
39-
return (
40-
<option
41-
key={model}
42-
value={model}
43-
selected={settings?.EmbeddingModelPref === model}
44-
>
45-
{model}
46-
</option>
47-
);
48-
})}
49-
</optgroup>
50-
</select>
51-
</div>
28+
<CohereModelSelection settings={settings} apiKey={cohereApiKey} />
5229
</div>
5330
</div>
5431
);
5532
}
33+
34+
function CohereModelSelection({ apiKey, settings }) {
35+
const [models, setModels] = useState([]);
36+
const [loading, setLoading] = useState(true);
37+
38+
useEffect(() => {
39+
async function findCustomModels() {
40+
if (!apiKey) {
41+
setModels([]);
42+
setLoading(true);
43+
return;
44+
}
45+
46+
setLoading(true);
47+
const { models } = await System.customModels(
48+
"cohere-embedder",
49+
typeof apiKey === "boolean" ? null : apiKey
50+
);
51+
setModels(models || []);
52+
setLoading(false);
53+
}
54+
findCustomModels();
55+
}, [apiKey]);
56+
57+
if (loading) {
58+
return (
59+
<div className="flex flex-col w-60">
60+
<label className="text-white text-sm font-semibold block mb-3">
61+
Model Preference
62+
</label>
63+
<select
64+
name="EmbeddingModelPref"
65+
disabled={true}
66+
className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
67+
>
68+
<option disabled={true} selected={true}>
69+
-- loading available models --
70+
</option>
71+
</select>
72+
</div>
73+
);
74+
}
75+
76+
return (
77+
<div className="flex flex-col w-60">
78+
<label className="text-white text-sm font-semibold block mb-3">
79+
Model Preference
80+
</label>
81+
<select
82+
name="EmbeddingModelPref"
83+
required={true}
84+
className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
85+
>
86+
{models.map((model) => (
87+
<option
88+
key={model.id}
89+
value={model.id}
90+
selected={settings?.EmbeddingModelPref === model.id}
91+
>
92+
{model.name}
93+
</option>
94+
))}
95+
</select>
96+
</div>
97+
);
98+
}

‎frontend/src/components/LLMSelection/CohereAiOptions/index.jsx‎

Lines changed: 77 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,10 @@
1+
import { useState, useEffect } from "react";
2+
import System from "@/models/system";
3+
14
export default function CohereAiOptions({ settings }) {
5+
const [inputValue, setInputValue] = useState(settings?.CohereApiKey);
6+
const [cohereApiKey, setCohereApiKey] = useState(settings?.CohereApiKey);
7+
28
return (
39
<div className="w-full flex flex-col">
410
<div className="w-full flex items-center gap-[36px] mt-1.5">
@@ -15,35 +21,80 @@ export default function CohereAiOptions({ settings }) {
1521
required={true}
1622
autoComplete="off"
1723
spellCheck={false}
24+
onChange={(e) => setInputValue(e.target.value)}
25+
onBlur={() => setCohereApiKey(inputValue)}
1826
/>
1927
</div>
20-
<div className="flex flex-col w-60">
21-
<label className="text-white text-sm font-semibold block mb-3">
22-
Chat Model Selection
23-
</label>
24-
<select
25-
name="CohereModelPref"
26-
defaultValue={settings?.CohereModelPref || "command-r"}
27-
required={true}
28-
className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
29-
>
30-
{[
31-
"command-r",
32-
"command-r-plus",
33-
"command",
34-
"command-light",
35-
"command-nightly",
36-
"command-light-nightly",
37-
].map((model) => {
38-
return (
39-
<option key={model} value={model}>
40-
{model}
41-
</option>
42-
);
43-
})}
44-
</select>
45-
</div>
28+
{!settings?.credentialsOnly && (
29+
<CohereModelSelection settings={settings} apiKey={cohereApiKey} />
30+
)}
4631
</div>
4732
</div>
4833
);
4934
}
35+
36+
function CohereModelSelection({ apiKey, settings }) {
37+
const [models, setModels] = useState([]);
38+
const [loading, setLoading] = useState(true);
39+
40+
useEffect(() => {
41+
async function findCustomModels() {
42+
if (!apiKey) {
43+
setModels([]);
44+
setLoading(true);
45+
return;
46+
}
47+
48+
setLoading(true);
49+
const { models } = await System.customModels(
50+
"cohere",
51+
typeof apiKey === "boolean" ? null : apiKey
52+
);
53+
setModels(models || []);
54+
setLoading(false);
55+
}
56+
findCustomModels();
57+
}, [apiKey]);
58+
59+
if (loading) {
60+
return (
61+
<div className="flex flex-col w-60">
62+
<label className="text-white text-sm font-semibold block mb-3">
63+
Chat Model Selection
64+
</label>
65+
<select
66+
name="CohereModelPref"
67+
disabled={true}
68+
className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
69+
>
70+
<option disabled={true} selected={true}>
71+
-- loading available models --
72+
</option>
73+
</select>
74+
</div>
75+
);
76+
}
77+
78+
return (
79+
<div className="flex flex-col w-60">
80+
<label className="text-white text-sm font-semibold block mb-3">
81+
Chat Model Selection
82+
</label>
83+
<select
84+
name="CohereModelPref"
85+
required={true}
86+
className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
87+
>
88+
{models.map((model) => (
89+
<option
90+
key={model.id}
91+
value={model.id}
92+
selected={settings?.CohereModelPref === model.id}
93+
>
94+
{model.name}
95+
</option>
96+
))}
97+
</select>
98+
</div>
99+
);
100+
}

‎server/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@
4646
"check-disk-space": "^3.4.0",
4747
"cheerio": "^1.0.0",
4848
"chromadb": "^2.0.1",
49-
"cohere-ai": "^7.9.5",
49+
"cohere-ai": "^7.19.0",
5050
"cors": "^2.8.5",
5151
"dotenv": "^16.0.3",
5252
"elevenlabs": "^0.5.0",

‎server/utils/AiProviders/cohere/index.js‎

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ const {
88

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

2728
this.embedder = embedder ?? new NativeEmbedder();
29+
this.#log(
30+
`Initialized CohereLLM with model ${this.model}. ctx: ${this.promptWindowLimit()}`
31+
);
32+
}
33+
34+
#log(text, ...args) {
35+
console.log(`\x1b[32m[${this.className}]\x1b[0m ${text}`, ...args);
2836
}
2937

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

73-
async isValidChatCompletionModel(model = "") {
74-
const validModels = [
75-
"command-r",
76-
"command-r-plus",
77-
"command",
78-
"command-light",
79-
"command-nightly",
80-
"command-light-nightly",
81-
];
82-
return validModels.includes(model);
81+
async isValidChatCompletionModel() {
82+
return true;
8383
}
8484

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

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

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

136131
async streamGetChatCompletion(messages = null, { temperature = 0.7 }) {
137-
if (!(await this.isValidChatCompletionModel(this.model)))
138-
throw new Error(
139-
`Cohere chat: ${this.model} is not valid for chat completion!`
140-
);
141-
142132
const message = messages[messages.length - 1].content; // Get the last message
143133
const cohereHistory = this.#convertChatHistoryCohere(messages.slice(0, -1)); // Remove the last message and convert to Cohere
144134
const measuredStreamRequest = await LLMPerformanceMonitor.measureStream(

‎server/utils/AiProviders/modelMap/legacy.js‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,18 @@ const LEGACY_MODEL_MAP = {
2222
"command-light": 4096,
2323
"command-nightly": 8192,
2424
"command-light-nightly": 8192,
25+
"command-r-plus-08-2024": 132096,
26+
"command-a-03-2025": 288000,
27+
"c4ai-aya-vision-32b": 16384,
28+
"command-a-reasoning-08-2025": 288768,
29+
"command-r-08-2024": 132096,
30+
"c4ai-aya-vision-8b": 16384,
31+
"command-r7b-12-2024": 132000,
32+
"command-r7b-arabic-02-2025": 128000,
33+
"command-a-vision-07-2025": 128000,
34+
"c4ai-aya-expanse-8b": 8192,
35+
"c4ai-aya-expanse-32b": 128000,
36+
"command-a-translate-08-2025": 8992,
2537
},
2638
gemini: {
2739
"gemini-1.5-pro-001": 2000000,

‎server/utils/helpers/customModels.js‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,10 @@ const SUPPORT_CUSTOM_MODELS = [
3737
"dpais",
3838
"moonshotai",
3939
"foundry",
40+
"cohere",
4041
// Embedding Engines
4142
"native-embedder",
43+
"cohere-embedder",
4244
];
4345

4446
async function getCustomModels(provider = "", apiKey = null, basePath = null) {
@@ -96,8 +98,12 @@ async function getCustomModels(provider = "", apiKey = null, basePath = null) {
9698
return await getMoonshotAiModels(apiKey);
9799
case "foundry":
98100
return await getFoundryModels(basePath);
101+
case "cohere":
102+
return await getCohereModels(apiKey, "chat");
99103
case "native-embedder":
100104
return await getNativeEmbedderModels();
105+
case "cohere-embedder":
106+
return await getCohereModels(apiKey, "embed");
101107
default:
102108
return { models: [], error: "Invalid provider for custom models" };
103109
}
@@ -759,6 +765,39 @@ async function getFoundryModels(basePath = null) {
759765
}
760766
}
761767

768+
/**
769+
* Get Cohere models
770+
* @param {string} _apiKey - The API key to use
771+
* @param {'chat' | 'embed'} type - The type of model to get
772+
* @returns {Promise<{models: Array<{id: string, organization: string, name: string}>, error: string | null}>}
773+
*/
774+
async function getCohereModels(_apiKey = null, type = "chat") {
775+
const apiKey =
776+
_apiKey === true
777+
? process.env.COHERE_API_KEY
778+
: _apiKey || process.env.COHERE_API_KEY || null;
779+
780+
const { CohereClient } = require("cohere-ai");
781+
const cohere = new CohereClient({
782+
token: apiKey,
783+
});
784+
const models = await cohere.models
785+
.list({ pageSize: 1000, endpoint: type })
786+
.then((results) => results.models)
787+
.then((models) =>
788+
models.map((model) => ({
789+
id: model.id,
790+
name: model.name,
791+
}))
792+
)
793+
.catch((e) => {
794+
console.error(`Cohere:listModels`, e.message);
795+
return [];
796+
});
797+
798+
return { models, error: null };
799+
}
800+
762801
module.exports = {
763802
getCustomModels,
764803
SUPPORT_CUSTOM_MODELS,

0 commit comments

Comments
 (0)