|
| 1 | +const { CohereClient } = require("cohere-ai"); |
| 2 | +const Provider = require("./ai-provider"); |
| 3 | +const InheritMultiple = require("./helpers/classes"); |
| 4 | +const UnTooled = require("./helpers/untooled"); |
| 5 | +const { v4 } = require("uuid"); |
| 6 | +const { safeJsonParse } = require("../../../http"); |
| 7 | + |
| 8 | +class CohereProvider extends InheritMultiple([Provider, UnTooled]) { |
| 9 | + model; |
| 10 | + |
| 11 | + constructor(config = {}) { |
| 12 | + const { model = process.env.COHERE_MODEL_PREF || "command-r-08-2024" } = |
| 13 | + config; |
| 14 | + super(); |
| 15 | + const client = new CohereClient({ |
| 16 | + token: process.env.COHERE_API_KEY, |
| 17 | + }); |
| 18 | + this._client = client; |
| 19 | + this.model = model; |
| 20 | + this.verbose = true; |
| 21 | + } |
| 22 | + |
| 23 | + get client() { |
| 24 | + return this._client; |
| 25 | + } |
| 26 | + |
| 27 | + get supportsAgentStreaming() { |
| 28 | + return true; |
| 29 | + } |
| 30 | + |
| 31 | + #convertChatHistoryCohere(chatHistory = []) { |
| 32 | + let cohereHistory = []; |
| 33 | + chatHistory.forEach((message) => { |
| 34 | + switch (message.role) { |
| 35 | + case "SYSTEM": |
| 36 | + case "system": |
| 37 | + cohereHistory.push({ role: "SYSTEM", message: message.content }); |
| 38 | + break; |
| 39 | + case "USER": |
| 40 | + case "user": |
| 41 | + cohereHistory.push({ role: "USER", message: message.content }); |
| 42 | + break; |
| 43 | + case "CHATBOT": |
| 44 | + case "assistant": |
| 45 | + cohereHistory.push({ role: "CHATBOT", message: message.content }); |
| 46 | + break; |
| 47 | + } |
| 48 | + }); |
| 49 | + |
| 50 | + return cohereHistory; |
| 51 | + } |
| 52 | + |
| 53 | + async #handleFunctionCallStream({ messages = [] }) { |
| 54 | + const userPrompt = messages[messages.length - 1]?.content || ""; |
| 55 | + const history = messages.slice(0, -1); |
| 56 | + return await this.client.chatStream({ |
| 57 | + model: this.model, |
| 58 | + chatHistory: this.#convertChatHistoryCohere(history), |
| 59 | + message: userPrompt, |
| 60 | + }); |
| 61 | + } |
| 62 | + |
| 63 | + async stream(messages, functions = [], eventHandler = null) { |
| 64 | + return await UnTooled.prototype.stream.call( |
| 65 | + this, |
| 66 | + messages, |
| 67 | + functions, |
| 68 | + this.#handleFunctionCallStream.bind(this), |
| 69 | + eventHandler |
| 70 | + ); |
| 71 | + } |
| 72 | + |
| 73 | + async streamingFunctionCall( |
| 74 | + messages, |
| 75 | + functions, |
| 76 | + chatCb = null, |
| 77 | + eventHandler = null |
| 78 | + ) { |
| 79 | + const history = [...messages].filter((msg) => |
| 80 | + ["user", "assistant"].includes(msg.role) |
| 81 | + ); |
| 82 | + if (history[history.length - 1]?.role !== "user") return null; |
| 83 | + |
| 84 | + const msgUUID = v4(); |
| 85 | + let textResponse = ""; |
| 86 | + const historyMessages = this.buildToolCallMessages(history, functions); |
| 87 | + const stream = await chatCb({ messages: historyMessages }); |
| 88 | + |
| 89 | + eventHandler?.("reportStreamEvent", { |
| 90 | + type: "statusResponse", |
| 91 | + uuid: v4(), |
| 92 | + content: "Agent is thinking...", |
| 93 | + }); |
| 94 | + |
| 95 | + for await (const event of stream) { |
| 96 | + if (event.eventType !== "text-generation") continue; |
| 97 | + textResponse += event.text; |
| 98 | + eventHandler?.("reportStreamEvent", { |
| 99 | + type: "statusResponse", |
| 100 | + uuid: msgUUID, |
| 101 | + content: event.text, |
| 102 | + }); |
| 103 | + } |
| 104 | + |
| 105 | + const call = safeJsonParse(textResponse, null); |
| 106 | + if (call === null) |
| 107 | + return { toolCall: null, text: textResponse, uuid: msgUUID }; |
| 108 | + |
| 109 | + const { valid, reason } = this.validFuncCall(call, functions); |
| 110 | + if (!valid) { |
| 111 | + this.providerLog(`Invalid function tool call: ${reason}.`); |
| 112 | + eventHandler?.("reportStreamEvent", { |
| 113 | + type: "removeStatusResponse", |
| 114 | + uuid: msgUUID, |
| 115 | + content: |
| 116 | + "The model attempted to make an invalid function call - it was ignored.", |
| 117 | + }); |
| 118 | + return { toolCall: null, text: null, uuid: msgUUID }; |
| 119 | + } |
| 120 | + |
| 121 | + const { isDuplicate, reason: duplicateReason } = |
| 122 | + this.deduplicator.isDuplicate(call.name, call.arguments); |
| 123 | + if (isDuplicate) { |
| 124 | + this.providerLog( |
| 125 | + `Cannot call ${call.name} again because ${duplicateReason}.` |
| 126 | + ); |
| 127 | + eventHandler?.("reportStreamEvent", { |
| 128 | + type: "removeStatusResponse", |
| 129 | + uuid: msgUUID, |
| 130 | + content: |
| 131 | + "The model tried to call a function with the same arguments as a previous call - it was ignored.", |
| 132 | + }); |
| 133 | + return { toolCall: null, text: null, uuid: msgUUID }; |
| 134 | + } |
| 135 | + |
| 136 | + eventHandler?.("reportStreamEvent", { |
| 137 | + uuid: `${msgUUID}:tool_call_invocation`, |
| 138 | + type: "toolCallInvocation", |
| 139 | + content: `Parsed Tool Call: ${call.name}(${JSON.stringify(call.arguments)})`, |
| 140 | + }); |
| 141 | + return { toolCall: call, text: null, uuid: msgUUID }; |
| 142 | + } |
| 143 | + |
| 144 | + /** |
| 145 | + * Stream a chat completion from the LLM with tool calling |
| 146 | + * Override the inherited `stream` method since Cohere uses a different API format. |
| 147 | + * |
| 148 | + * @param {any[]} messages - The messages to send to the LLM. |
| 149 | + * @param {any[]} functions - The functions to use in the LLM. |
| 150 | + * @param {function} eventHandler - The event handler to use to report stream events. |
| 151 | + * @returns {Promise<{ functionCall: any, textResponse: string }>} - The result of the chat completion. |
| 152 | + */ |
| 153 | + async stream(messages, functions = [], eventHandler = null) { |
| 154 | + this.providerLog( |
| 155 | + "CohereProvider.stream - will process this chat completion." |
| 156 | + ); |
| 157 | + try { |
| 158 | + let completion = { content: "" }; |
| 159 | + if (functions.length > 0) { |
| 160 | + const { |
| 161 | + toolCall, |
| 162 | + text, |
| 163 | + uuid: msgUUID, |
| 164 | + } = await this.streamingFunctionCall( |
| 165 | + messages, |
| 166 | + functions, |
| 167 | + this.#handleFunctionCallStream.bind(this), |
| 168 | + eventHandler |
| 169 | + ); |
| 170 | + |
| 171 | + if (toolCall !== null) { |
| 172 | + this.providerLog(`Valid tool call found - running ${toolCall.name}.`); |
| 173 | + this.deduplicator.trackRun(toolCall.name, toolCall.arguments, { |
| 174 | + cooldown: this.isMCPTool(toolCall, functions), |
| 175 | + }); |
| 176 | + return { |
| 177 | + result: null, |
| 178 | + functionCall: { |
| 179 | + name: toolCall.name, |
| 180 | + arguments: toolCall.arguments, |
| 181 | + }, |
| 182 | + cost: 0, |
| 183 | + }; |
| 184 | + } |
| 185 | + |
| 186 | + if (text) { |
| 187 | + this.providerLog( |
| 188 | + `No tool call found in the response - will send as a full text response.` |
| 189 | + ); |
| 190 | + completion.content = text; |
| 191 | + eventHandler?.("reportStreamEvent", { |
| 192 | + type: "removeStatusResponse", |
| 193 | + uuid: msgUUID, |
| 194 | + content: "No tool call found in the response", |
| 195 | + }); |
| 196 | + eventHandler?.("reportStreamEvent", { |
| 197 | + type: "statusResponse", |
| 198 | + uuid: v4(), |
| 199 | + content: "Done thinking.", |
| 200 | + }); |
| 201 | + eventHandler?.("reportStreamEvent", { |
| 202 | + type: "fullTextResponse", |
| 203 | + uuid: v4(), |
| 204 | + content: text, |
| 205 | + }); |
| 206 | + } |
| 207 | + } |
| 208 | + |
| 209 | + if (!completion?.content) { |
| 210 | + eventHandler?.("reportStreamEvent", { |
| 211 | + type: "statusResponse", |
| 212 | + uuid: v4(), |
| 213 | + content: "Done thinking.", |
| 214 | + }); |
| 215 | + |
| 216 | + this.providerLog( |
| 217 | + "Will assume chat completion without tool call inputs." |
| 218 | + ); |
| 219 | + const msgUUID = v4(); |
| 220 | + completion = { content: "" }; |
| 221 | + const stream = await this.#handleFunctionCallStream({ |
| 222 | + messages: this.cleanMsgs(messages), |
| 223 | + }); |
| 224 | + |
| 225 | + for await (const chunk of stream) { |
| 226 | + if (chunk.eventType !== "text-generation") continue; |
| 227 | + completion.content += chunk.text; |
| 228 | + eventHandler?.("reportStreamEvent", { |
| 229 | + type: "textResponseChunk", |
| 230 | + uuid: msgUUID, |
| 231 | + content: chunk.text, |
| 232 | + }); |
| 233 | + } |
| 234 | + } |
| 235 | + |
| 236 | + this.deduplicator.reset("runs"); |
| 237 | + return { |
| 238 | + textResponse: completion.content, |
| 239 | + cost: 0, |
| 240 | + }; |
| 241 | + } catch (error) { |
| 242 | + throw error; |
| 243 | + } |
| 244 | + } |
| 245 | + |
| 246 | + getCost(_usage) { |
| 247 | + return 0; |
| 248 | + } |
| 249 | +} |
| 250 | + |
| 251 | +module.exports = CohereProvider; |
0 commit comments