llm_universal/
cohere.rs

1use async_trait::async_trait;
2use serde_json::Value;
3
4use crate::openai::{parse_tool_calls, OpenAiProvider};
5use crate::{
6    send_json, ChatRequest, ChatResponse, LlmError, LlmProvider, Message, Part, Result, Role,
7    StopReason, Usage,
8};
9
10/// Client for the Cohere Chat API v2 (`POST {base_url}/chat`).
11///
12/// Requests use the same message and tool shapes as OpenAI; responses differ.
13#[derive(Clone)]
14pub struct CohereProvider {
15    client: reqwest::Client,
16    api_key: String,
17    base_url: String,
18}
19
20impl std::fmt::Debug for CohereProvider {
21    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22        f.debug_struct("CohereProvider")
23            .field("api_key", &"<redacted>")
24            .field("base_url", &self.base_url)
25            .finish()
26    }
27}
28
29impl CohereProvider {
30    pub fn new(api_key: impl Into<String>) -> Self {
31        Self {
32            client: crate::http_client(),
33            api_key: api_key.into(),
34            base_url: "https://api.cohere.com/v2".to_string(),
35        }
36    }
37
38    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
39        self.base_url = base_url.into().trim_end_matches('/').to_string();
40        self
41    }
42
43    /// Uses your own `reqwest::Client`, for example to share a connection pool or set a proxy.
44    pub fn with_client(mut self, client: reqwest::Client) -> Self {
45        self.client = client;
46        self
47    }
48
49    fn build_body(req: &ChatRequest) -> Value {
50        let mut body = OpenAiProvider::build_body(req);
51        // Cohere names the limit `max_tokens`.
52        if let Some(max_tokens) = body
53            .as_object_mut()
54            .and_then(|o| o.remove("max_completion_tokens"))
55        {
56            body["max_tokens"] = max_tokens;
57        }
58        body
59    }
60}
61
62fn parse_response(v: &Value, requested_model: &str) -> Result<ChatResponse> {
63    let message = &v["message"];
64    if !message.is_object() {
65        return Err(LlmError::InvalidResponse("missing field `message`".into()));
66    }
67
68    let mut parts: Vec<Part> = message["content"]
69        .as_array()
70        .into_iter()
71        .flatten()
72        .filter(|block| block["type"] == "text")
73        .filter_map(|block| block["text"].as_str())
74        .filter(|text| !text.is_empty())
75        .map(|text| Part::Text {
76            text: text.to_string(),
77        })
78        .collect();
79    parts.extend(parse_tool_calls(&message["tool_calls"]));
80
81    let stop_reason = match v["finish_reason"].as_str() {
82        Some("COMPLETE" | "STOP_SEQUENCE") => StopReason::EndTurn,
83        Some("TOOL_CALL") => StopReason::ToolUse,
84        Some("MAX_TOKENS") => StopReason::MaxTokens,
85        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
86    };
87
88    let tokens = |field: &str| {
89        u32::try_from(v["usage"]["tokens"][field].as_u64().unwrap_or(0)).unwrap_or(u32::MAX)
90    };
91    let usage = Usage {
92        input_tokens: tokens("input_tokens"),
93        output_tokens: tokens("output_tokens"),
94    };
95
96    Ok(ChatResponse {
97        // The response does not echo the model.
98        model: requested_model.to_string(),
99        message: Message::new(Role::Assistant, parts),
100        stop_reason,
101        usage,
102    })
103}
104
105#[async_trait]
106impl LlmProvider for CohereProvider {
107    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
108        crate::validate(
109            "CohereProvider",
110            &self.api_key,
111            true,
112            &self.base_url,
113            64,
114            req,
115        )?;
116        let url = format!("{}/chat", self.base_url);
117        let builder = self
118            .client
119            .post(url)
120            .header(
121                reqwest::header::AUTHORIZATION,
122                crate::secret_header("CohereProvider", &format!("Bearer {}", self.api_key))?,
123            )
124            .json(&Self::build_body(req));
125        let body = send_json(builder).await?;
126        parse_response(&body, &req.model)
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use serde_json::json;
134
135    #[test]
136    fn renames_max_tokens() {
137        let req = ChatRequest::new("command-a")
138            .max_tokens(50)
139            .message(Message::user("hi"));
140        let body = CohereProvider::build_body(&req);
141        assert_eq!(body["max_tokens"], 50);
142        assert!(body.get("max_completion_tokens").is_none());
143        assert_eq!(body["messages"][0]["content"], "hi");
144    }
145
146    #[test]
147    fn parses_tool_call_response() {
148        let v = json!({
149            "finish_reason": "TOOL_CALL",
150            "message": {
151                "role": "assistant",
152                "tool_plan": "I will look up the weather.",
153                "tool_calls": [{
154                    "id": "get_weather_abc",
155                    "type": "function",
156                    "function": { "name": "get_weather", "arguments": "{\"city\":\"Berlin\"}" }
157                }]
158            },
159            "usage": { "tokens": { "input_tokens": 12, "output_tokens": 8 } }
160        });
161        let resp = parse_response(&v, "command-a").unwrap();
162        assert_eq!(resp.stop_reason, StopReason::ToolUse);
163        assert_eq!(resp.tool_calls()[0].2["city"], "Berlin");
164        assert_eq!(resp.usage.input_tokens, 12);
165        assert_eq!(resp.model, "command-a");
166    }
167
168    #[test]
169    fn parses_text_response() {
170        let v = json!({
171            "finish_reason": "COMPLETE",
172            "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] }
173        });
174        let resp = parse_response(&v, "command-a").unwrap();
175        assert_eq!(resp.text(), "Hello");
176        assert_eq!(resp.stop_reason, StopReason::EndTurn);
177    }
178}
llm_universal/cohere.rs

llm_universal/
cohere.rs
