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

rig_core/providers/cohere/
chat.rs
