rig_core/providers/cohere/
chat.rs

1//! Cohere's native chat API, `POST /v2/chat`: the request built as JSON,
2//! with the request's documents sent as Cohere `documents`. Its replies
3//! are read by [`ChatDecoder`].
4//!
5//! ```
6//! use rig_core::providers::cohere::{COMMAND_A_03_2025, CohereConfig, NativeChat};
7//!
8//! let wire = NativeChat::new(CohereConfig::new("key"), COMMAND_A_03_2025);
9//! assert_eq!(wire.model, COMMAND_A_03_2025);
10//! ```
11
12use std::collections::BTreeMap;
13use std::sync::Arc;
14
15use serde::{Deserialize, Serialize};
16use serde_json::{Map, Value, json};
17
18use super::CohereConfig;
19use super::streaming::ChatDecoder;
20use crate::catalog::ModelFacts;
21use crate::completion::options::{BaseInput, FinalBody, RawAt, request_params};
22use crate::completion::{CompletionRequest, Document, ProviderCapabilities, Replay};
23use crate::error::EncodeError;
24use crate::json_utils::Lenient;
25use crate::message::{
26    AssistantContent, AssistantMessage, DocumentData, DocumentSourceKind as Source, Message,
27    MimeType, ToolChoice, ToolResult, ToolResultContent, UserContent,
28};
29use crate::operation::Completion;
30use crate::providers::internal::object;
31use crate::providers::internal::wire_ids::WireIds;
32use crate::wire::{Capabilities, Descriptor, Encoded, Framing, Mode, Wire};
33
34/// Where the native chat endpoint sits under the API root.
35const CHAT_PATH: &str = "/v2/chat";
36
37/// The replay API of turns made on the native chat endpoint.
38pub(crate) const API: &str = "cohere.chat";
39
40/// The native chat wire: a provider configuration, a model, and the
41/// per-turn options the endpoint takes.
42#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
43pub struct NativeChat {
44    /// Which provider, and how to reach it.
45    pub provider: CohereConfig,
46    /// The model this wire addresses.
47    pub model: String,
48    /// Whether requests ask Cohere to hold tool calls to their schemas
49    /// (`strict_tools`).
50    pub strict_tools: bool,
51    /// The model facts the encoder reads and replies are priced by.
52    #[serde(skip)]
53    pub facts: ModelFacts,
54}
55
56impl NativeChat {
57    /// The wire for `model` on `provider`, with every option off.
58    pub fn new(provider: CohereConfig, model: impl Into<String>) -> Self {
59        Self {
60            provider,
61            model: model.into(),
62            strict_tools: false,
63            facts: ModelFacts::default(),
64        }
65    }
66
67    /// The same wire, encoding with `facts` and pricing its replies by
68    /// them.
69    pub fn with_facts(mut self, facts: ModelFacts) -> Self {
70        self.facts = facts;
71        self
72    }
73
74    /// Ask Cohere to hold every tool call to its tool's schema.
75    pub fn with_strict_tools(mut self) -> Self {
76        self.strict_tools = true;
77        self
78    }
79
80    /// The request body: the wire's encoding of `request`, then the mapped
81    /// options, then `additional_params`, merged key by key.
82    fn body(&self, request: &CompletionRequest, mode: Mode) -> Result<FinalBody, EncodeError> {
83        request_params(
84            self,
85            request,
86            |input| self.base(request, mode, input),
87            RawAt::Top,
88            &[],
89        )
90    }
91
92    /// The wire's own encoding of `request`.
93    fn base(
94        &self,
95        request: &CompletionRequest,
96        mode: Mode,
97        input: &mut BaseInput<'_>,
98    ) -> Result<Map<String, Value>, EncodeError> {
99        let model = request.model.clone().unwrap_or_else(|| self.model.clone());
100        let messages = self.messages(&request.chat_history, &model)?;
101        let mut tools: Vec<Value> = request
102            .tools
103            .iter()
104            .filter(|tool| match &request.tool_choice {
105                // Cohere cannot name the tool it must call: only the named
106                // tools are offered, and one of them is required.
107                Some(ToolChoice::Specific { function_names }) => {
108                    function_names.contains(&tool.name)
109                }
110                _ => true,
111            })
112            .map(|tool| {
113                json!({"type": "function", "function": {
114                    "name": tool.name,
115                    "description": tool.description,
116                    "parameters": tool.parameters,
117                }})
118            })
119            .collect();
120        tools.extend(input.raw_tools()?);
121        let tool_choice = match &request.tool_choice {
122            None | Some(ToolChoice::Auto) => None,
123            Some(ToolChoice::None) => Some("NONE"),
124            Some(ToolChoice::Required | ToolChoice::Specific { .. }) => Some("REQUIRED"),
125        }
126        .filter(|_| !tools.is_empty());
127        let documents: Vec<Value> = request
128            .documents
129            .iter()
130            .enumerate()
131            .map(|(position, document)| document_value(position, document))
132            .collect();
133        let response_format = request
134            .output_schema
135            .clone()
136            .map(|schema| json!({"type": "json_object", "schema": schema.to_value()}));
137        Ok(crate::providers::internal::set_fields([
138            ("model", Some(Value::String(model))),
139            ("messages", Some(Value::Array(messages))),
140            (
141                "documents",
142                (!documents.is_empty()).then_some(Value::Array(documents)),
143            ),
144            ("tools", (!tools.is_empty()).then_some(Value::Array(tools))),
145            ("tool_choice", tool_choice.map(Value::from)),
146            (
147                "strict_tools",
148                self.strict_tools.then_some(Value::Bool(true)),
149            ),
150            ("response_format", response_format),
151            ("temperature", request.temperature.map(Value::from)),
152            ("max_tokens", request.max_tokens.map(Value::from)),
153            (
154                "stream",
155                (mode == Mode::Streaming).then_some(Value::Bool(true)),
156            ),
157        ]))
158    }
159
160    /// The history as Cohere messages, each call and result spelled by one
161    /// [`WireIds`].
162    fn messages(&self, history: &[Arc<Message>], model: &str) -> Result<Vec<Value>, EncodeError> {
163        let ids = WireIds::for_target(history, self, model);
164        let mut messages = Vec::new();
165        for message in history {
166            match &**message {
167                Message::System { content } => {
168                    messages.push(json!({"role": "system", "content": content}));
169                }
170                Message::User { content } => {
171                    let mut parts = Vec::new();
172                    for part in content {
173                        if let UserContent::ToolResult(result) = part {
174                            user_message(&mut messages, &mut parts);
175                            messages.push(tool_message(result, &ids));
176                        } else {
177                            parts.push(user_part(part)?);
178                        }
179                    }
180                    user_message(&mut messages, &mut parts);
181                }
182                Message::Assistant(turn) => messages.extend(self.assistant(turn, &ids)),
183            }
184        }
185        if messages.is_empty() {
186            return Err(EncodeError::request(
187                "Cohere chat request has no messages after conversion",
188            ));
189        }
190        Ok(messages)
191    }
192
193    /// One assistant turn: text, thinking and unknown parts as content
194    /// parts, a tool plan as `tool_plan`, each call under `tool_calls`, and
195    /// the citations each block's item holds, pointed at the part they
196    /// cite. A current item goes back as it came, less its citations. A
197    /// turn with nothing to send is `None`.
198    fn assistant(&self, turn: &AssistantMessage, ids: &WireIds) -> Option<Value> {
199        let (mut content, mut plan, mut calls, mut citations) =
200            (Vec::new(), String::new(), Vec::new(), Vec::new());
201        for block in &turn.content {
202            let replay = block.replay(self, ids);
203            let kind = replay_kind(&replay);
204            let mut item = match replay {
205                Replay::Item(item) => match item.into_owned() {
206                    Value::Object(item) => item,
207                    _ => Map::new(),
208                },
209                Replay::Identity(_) | Replay::Rebuild => Map::new(),
210            };
211            let cited = match item.shift_remove("citations") {
212                Some(Value::Array(cited)) => cited,
213                _ => Vec::new(),
214            };
215            let part = match block {
216                AssistantContent::Text(text) if !text.text.is_empty() => {
217                    json!({"type": "text", "text": text.text})
218                }
219                AssistantContent::Reasoning(reasoning) if kind.as_deref() == Some(PLAN) => {
220                    citations.extend(pointed(cited, None));
221                    plan.push_str(&reasoning.text);
222                    continue;
223                }
224                AssistantContent::Reasoning(reasoning) if !reasoning.text.is_empty() => {
225                    json!({"type": "thinking", "thinking": reasoning.text})
226                }
227                AssistantContent::Opaque(opaque)
228                    if opaque.replay && opaque.item.get("type").is_some() =>
229                {
230                    content.push(opaque.item.clone());
231                    continue;
232                }
233                AssistantContent::ToolCall(call) => {
234                    item.entry("type")
235                        .or_insert_with(|| Value::from("function"));
236                    item.insert("id".to_owned(), Value::String(ids.spell(&call.id)));
237                    let function = item
238                        .entry("function")
239                        .or_insert_with(|| Value::Object(Map::new()));
240                    if !function.is_object() {
241                        *function = Value::Object(Map::new());
242                    }
243                    if let Value::Object(function) = function {
244                        function
245                            .insert("name".to_owned(), Value::from(call.function.name.as_str()));
246                        function.insert(
247                            "arguments".to_owned(),
248                            Value::String(
249                                serde_json::to_string(&call.function.arguments).unwrap_or_default(),
250                            ),
251                        );
252                    }
253                    calls.push(Value::Object(item));
254                    continue;
255                }
256                AssistantContent::Text(_)
257                | AssistantContent::Reasoning(_)
258                | AssistantContent::Image(_)
259                | AssistantContent::Opaque(_) => continue,
260            };
261            citations.extend(pointed(cited, Some(content.len())));
262            content.push(if item.is_empty() {
263                part
264            } else {
265                Value::Object(item)
266            });
267        }
268        if content.is_empty() && calls.is_empty() {
269            return None;
270        }
271        Some(Value::Object(crate::providers::internal::set_fields([
272            ("role", Some(Value::from("assistant"))),
273            (
274                "content",
275                (!content.is_empty()).then_some(Value::Array(content)),
276            ),
277            (
278                "tool_plan",
279                (!plan.is_empty()).then_some(Value::String(plan)),
280            ),
281            (
282                "tool_calls",
283                (!calls.is_empty()).then_some(Value::Array(calls)),
284            ),
285            (
286                "citations",
287                (!citations.is_empty()).then_some(Value::Array(citations)),
288            ),
289        ])))
290    }
291}
292
293/// The item kind a tool plan's reasoning block holds.
294pub(crate) const PLAN: &str = "tool_plan";
295
296/// The kind of item replay hands the encoder for a block: the current
297/// item's `type`, or the one an edited block kept.
298fn replay_kind(replay: &Replay<'_>) -> Option<String> {
299    match replay {
300        Replay::Item(item) => item.str("type").map(str::to_owned),
301        Replay::Identity(identity) => identity
302            .get("type")
303            .and_then(Value::as_str)
304            .map(str::to_owned),
305        Replay::Rebuild => None,
306    }
307}
308
309/// The kind of item `block` replays as on `target`.
310fn kind(block: &AssistantContent, target: &NativeChat) -> Option<String> {
311    replay_kind(&block.replay(target, &WireIds::default()))
312}
313
314/// `citations` as a request sends them back: each at `content_index`, the
315/// position of the part it cites, or with none for a tool plan's.
316fn pointed(citations: Vec<Value>, content_index: Option<usize>) -> Vec<Value> {
317    citations
318        .into_iter()
319        .map(|mut citation| {
320            if let (Some(citation), Some(index)) = (citation.as_object_mut(), content_index) {
321                citation.insert("content_index".to_owned(), Value::from(index));
322            }
323            citation
324        })
325        .collect()
326}
327
328/// `document` as a Cohere document: its metadata and `text` under `data`,
329/// keys sorted so the same document always sends the same bytes, under its
330/// id, or `doc_<position>` when it has none.
331fn document_value(position: usize, document: &Document) -> Value {
332    let mut data: BTreeMap<&str, &str> = document
333        .additional_props
334        .iter()
335        .map(|(key, value)| (key.as_str(), value.as_str()))
336        .collect();
337    data.insert("text", &document.text);
338    let id = if document.id.is_empty() {
339        format!("doc_{position}")
340    } else {
341        document.id.clone()
342    };
343    json!({"id": id, "data": data})
344}
345
346/// One user content part as Cohere carries it: text, a text document's
347/// text, or an image by URL. Replay leaves no other part.
348fn user_part(part: &UserContent) -> Result<Value, EncodeError> {
349    Ok(match part {
350        UserContent::Text(text) => json!({"type": "text", "text": text.text}),
351        UserContent::Image(image) => {
352            let mime = image.media_type.as_ref().map(MimeType::to_mime_type);
353            let url = match (&image.data, mime) {
354                (Source::Url(url), _) => url.clone(),
355                (Source::Base64(data), Some(mime)) => format!("data:{mime};base64,{data}"),
356                _ => return Err(unsendable("an image")),
357            };
358            object! {"type": "image_url", "image_url": object! {"url": url}}
359        }
360        UserContent::Document(document) => match &document.data {
361            DocumentData::Text(text) => json!({"type": "text", "text": text}),
362            DocumentData::File(_) => return Err(unsendable("a document")),
363        },
364        UserContent::Audio(_) => return Err(unsendable("audio")),
365        UserContent::Video(_) => return Err(unsendable("a video")),
366        UserContent::ToolResult(_) => return Err(unsendable("a tool result as a content part")),
367    })
368}
369
370/// The error for content Cohere chat cannot carry, which replay replaces
371/// before a request is encoded.
372fn unsendable(what: &str) -> EncodeError {
373    EncodeError::request(format!("Cohere chat cannot carry {what} in this form"))
374}
375
376/// Push the user parts gathered so far as one message.
377fn user_message(messages: &mut Vec<Value>, parts: &mut Vec<Value>) {
378    if parts.is_empty() {
379        return;
380    }
381    messages.push(object! {"role": "user", "content": std::mem::take(parts)});
382}
383
384/// A result as the `tool` message that answers its call, its parts as text.
385fn tool_message(result: &ToolResult, ids: &WireIds) -> Value {
386    let content: Vec<Value> = result
387        .content
388        .iter()
389        .filter_map(|part| match part {
390            ToolResultContent::Text(text) => Some(json!({"type": "text", "text": text.text})),
391            ToolResultContent::Json { value } => {
392                Some(json!({"type": "text", "text": value.to_string()}))
393            }
394            ToolResultContent::Image(_) => None,
395        })
396        .collect();
397    object! {"role": "tool", "tool_call_id": ids.spell(&result.call), "content": content}
398}
399
400impl Wire for NativeChat {
401    type Op = Completion;
402    type Payload = Encoded;
403    type Frame = crate::wire::WireFrame;
404    type Event = super::streaming::ChatEvent;
405    type Decoder<'id> = ChatDecoder;
406    type Reassembler = super::streaming::document::ChatResponse;
407
408    fn describe(&self) -> Descriptor<'_> {
409        Descriptor::new(super::PROVIDER_NAME)
410            .model(self.model.as_str())
411            .facts(&self.facts)
412            .capabilities(Capabilities::completion(
413                ProviderCapabilities::default().with_native_output_tool_composition(true),
414            ))
415            .replay(self)
416    }
417
418    fn encode(&self, request: CompletionRequest, mode: Mode) -> Result<Encoded, EncodeError> {
419        let body = self.body(&request, mode)?;
420        crate::providers::internal::trace_json(
421            crate::providers::internal::LogTarget::Completions,
422            "Cohere chat request",
423            &body,
424        );
425        let request = self.provider.post(CHAT_PATH).body(body.into_body())?;
426        let framing = match mode {
427            Mode::Streaming => Framing::Sse,
428            Mode::Unary => Framing::Whole,
429        };
430        Ok(Encoded::new(request, framing)
431            .with_request_id_header(Some(REQUEST_ID_HEADER))
432            .with_projection(ChatDecoder::project)
433            .with_route(Some(CHAT_PATH)))
434    }
435
436    fn decoder<'id>(&self) -> Self::Decoder<'id> {
437        ChatDecoder::default()
438    }
439}
440
441/// The reply header carrying Cohere's request id.
442const REQUEST_ID_HEADER: &str = "x-debug-trace-id";
443
444impl crate::completion::ReplayTarget for NativeChat {
445    /// Section 6.5 of the typed-options design, for the native API.
446    fn map_options(
447        &self,
448        request: &CompletionRequest,
449        fields: crate::completion::options::OptionFields<'_>,
450    ) -> crate::completion::options::OptionMap {
451        use crate::completion::options::{Mapping, OptionMap};
452        use crate::completion::{CacheRetention, Effort, Reasoning};
453        let model = request.model.as_deref().unwrap_or(&self.model);
454        // A model that thinks does so by default. An id the catalog does
455        // not list thinks when its name says `reasoning`.
456        let thinks = super::thinks(&self.facts, model);
457        let reasons = thinks.unwrap_or_else(|| model.contains("reasoning"));
458        const NO_FIELD: &str = "Cohere's chat API has no such field";
459        OptionMap {
460            reasoning: Mapping::of(fields.reasoning, |reasoning| match reasoning {
461                Reasoning::Off if reasons => {
462                    Mapping::Send(json!({"thinking": {"type": "disabled"}}))
463                }
464                Reasoning::Off => Mapping::Omit("the model does not think"),
465                Reasoning::Effort(_) | Reasoning::Budget { .. } if thinks == Some(false) => {
466                    Mapping::unsupported("the model does not think")
467                }
468                Reasoning::Effort(Effort::High) => {
469                    Mapping::Send(json!({"thinking": {"type": "enabled"}}))
470                }
471                Reasoning::Effort(effort) => Mapping::unsupported(format!(
472                    "Cohere takes thinking on or a token budget, not `{}`",
473                    effort.as_str()
474                )),
475                Reasoning::Budget { tokens } => Mapping::Send(json!({
476                    "thinking": {"type": "enabled", "token_budget": tokens},
477                })),
478            }),
479            cache: Mapping::of(fields.cache, |cache| match cache {
480                CacheRetention::None => Mapping::Omit("Cohere does not cache prompts"),
481                CacheRetention::Short | CacheRetention::Long => {
482                    Mapping::unsupported("Cohere has no prompt cache")
483                }
484            }),
485            service_tier: Mapping::of(fields.service_tier, |_| {
486                Mapping::unsupported("Cohere's `priority` is a queue position, not a tier")
487            }),
488            verbosity: Mapping::of(fields.verbosity, |_| Mapping::unsupported(NO_FIELD)),
489            parallel_tool_calls: Mapping::of(fields.parallel_tool_calls, |_| {
490                Mapping::unsupported(NO_FIELD)
491            }),
492            top_p: Mapping::of(fields.top_p, |top_p| {
493                if (0.01..=0.99).contains(&top_p) {
494                    Mapping::Send(json!({ "p": top_p }))
495                } else {
496                    Mapping::unsupported("Cohere takes `p` from 0.01 to 0.99")
497                }
498            }),
499            seed: Mapping::of(fields.seed, |seed| Mapping::Send(json!({ "seed": seed }))),
500            stop: Mapping::of_stop(fields.stop, |stop| match stop.len() {
501                0..=5 => Mapping::Send(json!({ "stop_sequences": stop })),
502                _ => Mapping::unsupported("Cohere takes at most 5 stop sequences"),
503            }),
504            cache_key: Mapping::unrouted(fields.cache_key),
505        }
506    }
507
508    fn api(&self) -> crate::message::Api {
509        crate::message::Api::from_static(API)
510    }
511
512    fn facts(&self) -> Option<&ModelFacts> {
513        Some(&self.facts)
514    }
515
516    fn provider(&self) -> &str {
517        super::PROVIDER_NAME
518    }
519
520    fn model(&self) -> &str {
521        &self.model
522    }
523
524    /// Cohere's vision models read user images; no model reads images in
525    /// assistant turns or tool results.
526    fn accepts(&self, model: &str) -> crate::completion::Accepts {
527        crate::completion::Accepts {
528            user_images: self.facts.reads_images_or(
529                super::PROVIDER_NAME,
530                model,
531                super::reads_images,
532            ),
533            assistant_images: false,
534            tool_result_images: false,
535            tools: true,
536            foreign_reasoning: false,
537        }
538    }
539
540    /// The encoder carries a user image by URL or typed data, and a text
541    /// document as its text. It carries no other media.
542    fn encodes(&self, _model: &str, media: crate::completion::Media<'_>) -> bool {
543        use crate::completion::{Media, Place};
544        match media {
545            Media::Image(image, Place::User) => match &image.data {
546                Source::Url(_) => true,
547                Source::Base64(_) => image.media_type.is_some(),
548                _ => false,
549            },
550            Media::Document(document) => matches!(document.data, DocumentData::Text(_)),
551            Media::Image(..) | Media::Audio(_) | Media::Video(_) => false,
552        }
553    }
554
555    /// An edited block keeps its kind: a tool plan stays a plan.
556    fn identity(&self, item: &Value) -> Map<String, Value> {
557        item.str("type")
558            .filter(|kind| *kind == PLAN)
559            .map(|kind| Map::from_iter([("type".to_owned(), Value::from(kind))]))
560            .unwrap_or_default()
561    }
562
563    fn call_id_slot(&self) -> Option<&'static str> {
564        Some("/id")
565    }
566
567    fn takes_documents(&self) -> bool {
568        true
569    }
570
571    /// Text and thinking with text are content parts, an opaque item with
572    /// a `type` is a part, and a call is a call. A tool plan rides with its
573    /// calls, and images are never sent.
574    fn sends_alone(&self, block: &AssistantContent) -> bool {
575        match block {
576            AssistantContent::Text(text) => !text.text.is_empty(),
577            AssistantContent::Reasoning(reasoning) => {
578                !reasoning.text.is_empty() && kind(block, self).as_deref() != Some(PLAN)
579            }
580            AssistantContent::Opaque(opaque) => opaque.item.get("type").is_some(),
581            AssistantContent::ToolCall(_) => true,
582            AssistantContent::Image(_) => false,
583        }
584    }
585}
586
587#[cfg(test)]
588mod tests;
rig_core/providers/cohere/chat.rs

rig_core/providers/cohere/
chat.rs
