llm_universal/
cohere.rs

1use async_trait::async_trait;
2use serde_json::{json, Value};
3
4use crate::http::{secret_header, Http};
5use crate::openai::{
6    float_vec, parse_arguments, parse_tool_calls, usage, validate_embed_args, OpenAiProvider,
7};
8use crate::stream::{self, parse_data, Accumulator, ChatStream, StreamEvent};
9use crate::{
10    ChatRequest, ChatResponse, EmbeddingProvider, Embeddings, LlmError, LlmProvider, Message,
11    ModelInfo, Part, Result, Role, StopReason,
12};
13
14/// Client for the Cohere Chat API v2 (`POST {base_url}/chat`).
15///
16/// Requests use the same message, image and tool shapes as OpenAI; responses differ.
17#[derive(Clone)]
18pub struct CohereProvider {
19    http: Http,
20    api_key: String,
21    base_url: String,
22}
23
24impl std::fmt::Debug for CohereProvider {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        f.debug_struct("CohereProvider")
27            .field("api_key", &"<redacted>")
28            .field("base_url", &self.base_url)
29            .field("max_retries", &self.http.max_retries)
30            .finish()
31    }
32}
33
34impl CohereProvider {
35    pub fn new(api_key: impl Into<String>) -> Self {
36        Self {
37            http: Http::default(),
38            api_key: api_key.into(),
39            base_url: "https://api.cohere.com/v2".to_string(),
40        }
41    }
42
43    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
44        self.base_url = base_url.into().trim_end_matches('/').to_string();
45        self
46    }
47
48    /// Uses your own `reqwest::Client`, for example to share a connection pool or set a proxy.
49    pub fn with_client(mut self, client: reqwest::Client) -> Self {
50        self.http.client = client;
51        self
52    }
53
54    /// How many times a failed request is retried (default [`DEFAULT_MAX_RETRIES`](crate::DEFAULT_MAX_RETRIES); 0 disables).
55    pub fn with_max_retries(mut self, max_retries: u32) -> Self {
56        self.http.max_retries = max_retries;
57        self
58    }
59
60    fn build_body(req: &ChatRequest) -> Value {
61        let mut body = OpenAiProvider::build_body(req, true);
62        if let Some(obj) = body.as_object_mut() {
63            // Cohere names the limit `max_tokens`, takes `json_object` + schema and its own thinking switch.
64            if let Some(max_tokens) = obj.remove("max_completion_tokens") {
65                obj.insert("max_tokens".into(), max_tokens);
66            }
67            obj.remove("response_format");
68            obj.remove("reasoning_effort");
69        }
70        if let Some(format) = &req.response_format {
71            body["response_format"] =
72                json!({ "type": "json_object", "json_schema": format.schema });
73        }
74        if req.reasoning.is_some() {
75            body["thinking"] = json!({ "type": "enabled" });
76        }
77        body
78    }
79
80    fn post(&self, path: &str) -> Result<reqwest::RequestBuilder> {
81        Ok(self
82            .http
83            .client
84            .post(format!("{}{path}", self.base_url))
85            .header(
86                reqwest::header::AUTHORIZATION,
87                secret_header("CohereProvider", &format!("Bearer {}", self.api_key))?,
88            ))
89    }
90
91    /// Model listing lives on the v1 API.
92    fn models_url(&self) -> String {
93        match self.base_url.strip_suffix("/v2") {
94            Some(root) => format!("{root}/v1/models?page_size=1000"),
95            None => format!("{}/models?page_size=1000", self.base_url),
96        }
97    }
98}
99
100fn parse_response(v: &Value, requested_model: &str) -> Result<ChatResponse> {
101    let message = &v["message"];
102    if !message.is_object() {
103        return Err(LlmError::InvalidResponse("missing field `message`".into()));
104    }
105
106    let mut parts = Vec::new();
107    for block in message["content"].as_array().into_iter().flatten() {
108        match block["type"].as_str() {
109            Some("text") => {
110                if let Some(text) = block["text"].as_str().filter(|t| !t.is_empty()) {
111                    parts.push(Part::text(text));
112                }
113            }
114            Some("thinking") => parts.push(Part::Reasoning {
115                text: block["thinking"].as_str().unwrap_or_default().to_string(),
116                raw: Value::Null,
117            }),
118            _ => {}
119        }
120    }
121    parts.extend(parse_tool_calls(&message["tool_calls"]));
122
123    let stop_reason = match v["finish_reason"].as_str() {
124        Some("COMPLETE" | "STOP_SEQUENCE") => StopReason::EndTurn,
125        Some("TOOL_CALL") => StopReason::ToolUse,
126        Some("MAX_TOKENS") => StopReason::MaxTokens,
127        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
128    };
129
130    Ok(ChatResponse {
131        // The response does not echo the model.
132        model: requested_model.to_string(),
133        message: Message::new(Role::Assistant, parts),
134        stop_reason,
135        usage: usage(&v["usage"]["tokens"], "input_tokens", "output_tokens"),
136    })
137}
138
139/// Rebuilds the response from `content-delta`, `tool-call-*` and `message-end` events.
140struct CohereStream {
141    model: String,
142    text: String,
143    thinking: String,
144    /// (id, name, arguments)
145    calls: Vec<(String, String, String)>,
146    end: Value,
147    done: bool,
148}
149
150impl Accumulator for CohereStream {
151    fn on_event(&mut self, _event: &str, data: &str) -> Result<Vec<StreamEvent>> {
152        let v = parse_data(data)?;
153        let message = &v["delta"]["message"];
154        let mut out = Vec::new();
155        match v["type"].as_str() {
156            Some("content-delta") => {
157                if let Some(text) = message["content"]["text"]
158                    .as_str()
159                    .filter(|t| !t.is_empty())
160                {
161                    self.text.push_str(text);
162                    out.push(StreamEvent::TextDelta(text.to_string()));
163                }
164                if let Some(text) = message["content"]["thinking"]
165                    .as_str()
166                    .filter(|t| !t.is_empty())
167                {
168                    self.thinking.push_str(text);
169                    out.push(StreamEvent::ReasoningDelta(text.to_string()));
170                }
171            }
172            Some("tool-call-start") => {
173                let call = &message["tool_calls"];
174                let id = call["id"].as_str().unwrap_or_default().to_string();
175                let name = call["function"]["name"]
176                    .as_str()
177                    .unwrap_or_default()
178                    .to_string();
179                let args = call["function"]["arguments"]
180                    .as_str()
181                    .unwrap_or_default()
182                    .to_string();
183                out.push(StreamEvent::ToolCallStart {
184                    id: id.clone(),
185                    name: name.clone(),
186                });
187                self.calls.push((id, name, args));
188            }
189            Some("tool-call-delta") => {
190                let args = message["tool_calls"]["function"]["arguments"]
191                    .as_str()
192                    .unwrap_or_default();
193                if let Some(call) = self.calls.last_mut() {
194                    call.2.push_str(args);
195                    out.push(StreamEvent::ToolCallDelta {
196                        id: call.0.clone(),
197                        arguments: args.to_string(),
198                    });
199                }
200            }
201            Some("message-end") => {
202                self.end = v["delta"].clone();
203                self.done = true;
204            }
205            _ => {}
206        }
207        Ok(out)
208    }
209
210    fn finish(&mut self) -> Result<ChatResponse> {
211        if !self.done {
212            return Err(crate::stream::ended_early());
213        }
214        let mut parts = Vec::new();
215        if !self.thinking.is_empty() {
216            parts.push(Part::Reasoning {
217                text: std::mem::take(&mut self.thinking),
218                raw: Value::Null,
219            });
220        }
221        if !self.text.is_empty() {
222            parts.push(Part::text(std::mem::take(&mut self.text)));
223        }
224        for (id, name, args) in self.calls.drain(..) {
225            parts.push(Part::ToolCall {
226                id,
227                name,
228                arguments: parse_arguments(&args),
229                signature: None,
230            });
231        }
232        let response = parse_response(
233            &json!({ "message": {}, "finish_reason": self.end["finish_reason"], "usage": self.end["usage"] }),
234            &self.model,
235        )?;
236        Ok(ChatResponse {
237            message: Message::new(Role::Assistant, parts),
238            ..response
239        })
240    }
241}
242
243#[async_trait]
244impl LlmProvider for CohereProvider {
245    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
246        crate::validate(
247            "CohereProvider",
248            &self.api_key,
249            true,
250            &self.base_url,
251            64,
252            req,
253        )?;
254        let builder = self.post("/chat")?.json(&Self::build_body(req));
255        parse_response(&self.http.json(builder).await?, &req.model)
256    }
257
258    async fn chat_stream(&self, req: &ChatRequest) -> Result<ChatStream> {
259        crate::validate(
260            "CohereProvider",
261            &self.api_key,
262            true,
263            &self.base_url,
264            64,
265            req,
266        )?;
267        let mut body = Self::build_body(req);
268        body["stream"] = json!(true);
269        let response = self.http.send(self.post("/chat")?.json(&body)).await?;
270        Ok(stream::drive(
271            response,
272            Box::new(CohereStream {
273                model: req.model.clone(),
274                text: String::new(),
275                thinking: String::new(),
276                calls: Vec::new(),
277                end: Value::Null,
278                done: false,
279            }),
280        ))
281    }
282
283    async fn list_models(&self) -> Result<Vec<ModelInfo>> {
284        crate::validate_endpoint("CohereProvider", &self.api_key, true, &self.base_url)?;
285        let builder = self.http.client.get(self.models_url()).header(
286            reqwest::header::AUTHORIZATION,
287            secret_header("CohereProvider", &format!("Bearer {}", self.api_key))?,
288        );
289        let v = self.http.json(builder).await?;
290        Ok(v["models"]
291            .as_array()
292            .into_iter()
293            .flatten()
294            .filter_map(|m| m["name"].as_str())
295            .map(|name| ModelInfo {
296                id: name.to_string(),
297                display_name: None,
298            })
299            .collect())
300    }
301}
302
303#[async_trait]
304impl EmbeddingProvider for CohereProvider {
305    async fn embed(&self, model: &str, inputs: &[&str]) -> Result<Embeddings> {
306        crate::validate_endpoint("CohereProvider", &self.api_key, true, &self.base_url)?;
307        validate_embed_args(model, inputs)?;
308        let builder = self.post("/embed")?.json(&json!({
309            "model": model,
310            "texts": inputs,
311            "input_type": "search_document",
312            "embedding_types": ["float"],
313        }));
314        let v = self.http.json(builder).await?;
315        Ok(Embeddings {
316            vectors: v["embeddings"]["float"]
317                .as_array()
318                .into_iter()
319                .flatten()
320                .map(float_vec)
321                .collect(),
322            usage: usage(&v["meta"]["billed_units"], "input_tokens", "output_tokens"),
323        })
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330    use crate::ReasoningEffort;
331
332    #[test]
333    fn builds_cohere_specific_fields() {
334        let req = ChatRequest::new("command-a")
335            .max_tokens(50)
336            .json_schema("answer", json!({"type": "object"}))
337            .reasoning(ReasoningEffort::High)
338            .message(Message::user("hi"));
339        let body = CohereProvider::build_body(&req);
340        assert_eq!(body["max_tokens"], 50);
341        assert!(body.get("max_completion_tokens").is_none());
342        assert!(body.get("reasoning_effort").is_none());
343        assert_eq!(body["response_format"]["type"], "json_object");
344        assert_eq!(body["thinking"]["type"], "enabled");
345        assert_eq!(body["messages"][0]["content"], "hi");
346    }
347
348    #[test]
349    fn parses_tool_call_response() {
350        let v = json!({
351            "finish_reason": "TOOL_CALL",
352            "message": {
353                "role": "assistant",
354                "tool_plan": "I will look up the weather.",
355                "tool_calls": [{
356                    "id": "get_weather_abc",
357                    "type": "function",
358                    "function": { "name": "get_weather", "arguments": "{\"city\":\"Berlin\"}" }
359                }]
360            },
361            "usage": { "tokens": { "input_tokens": 12, "output_tokens": 8 } }
362        });
363        let resp = parse_response(&v, "command-a").unwrap();
364        assert_eq!(resp.stop_reason, StopReason::ToolUse);
365        assert_eq!(resp.tool_calls()[0].arguments["city"], "Berlin");
366        assert_eq!(resp.usage.input_tokens, 12);
367        assert_eq!(resp.model, "command-a");
368    }
369
370    #[test]
371    fn parses_text_response() {
372        let v = json!({
373            "finish_reason": "COMPLETE",
374            "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] }
375        });
376        let resp = parse_response(&v, "command-a").unwrap();
377        assert_eq!(resp.text(), "Hello");
378        assert_eq!(resp.stop_reason, StopReason::EndTurn);
379    }
380
381    #[test]
382    fn stream_rebuilds_text_and_tool_calls() {
383        let mut acc = CohereStream {
384            model: "command-a".into(),
385            text: String::new(),
386            thinking: String::new(),
387            calls: Vec::new(),
388            end: Value::Null,
389            done: false,
390        };
391        let events = [
392            r#"{"type":"message-start","id":"x"}"#,
393            r#"{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":"Hi"}}}}"#,
394            r#"{"type":"tool-call-start","index":0,"delta":{"message":{"tool_calls":{"id":"c1","type":"function","function":{"name":"f","arguments":""}}}}}"#,
395            r#"{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":"{\"x\":1}"}}}}}"#,
396            r#"{"type":"message-end","delta":{"finish_reason":"TOOL_CALL","usage":{"tokens":{"input_tokens":3,"output_tokens":5}}}}"#,
397        ];
398        for e in events {
399            acc.on_event("", e).unwrap();
400        }
401        let resp = acc.finish().unwrap();
402        assert_eq!(resp.text(), "Hi");
403        assert_eq!(resp.tool_calls()[0].arguments["x"], 1);
404        assert_eq!(resp.stop_reason, StopReason::ToolUse);
405        assert_eq!(resp.usage.output_tokens, 5);
406    }
407
408    #[test]
409    fn models_use_the_v1_api() {
410        assert_eq!(
411            CohereProvider::new("k").models_url(),
412            "https://api.cohere.com/v1/models?page_size=1000"
413        );
414    }
415}
llm_universal/cohere.rs

llm_universal/
cohere.rs
