rig_core/providers/cohere/
wire.rs

1//! Cohere's configuration, and its chat wire: the OpenAI Compatibility API
2//! by default, or the native chat API when the wire's route opts in.
3//!
4//! ```
5//! use rig_core::providers::cohere::{ChatRoute, CohereConfig};
6//! let config = CohereConfig::new("key").with_base_url("https://api.cohere.ai/");
7//! assert_eq!(config.base_url, "https://api.cohere.ai");
8//! let chat = config.completion("command-a-03-2025");
9//! assert_eq!(chat.route, ChatRoute::Compatibility);
10//! let chat = chat.with_route(ChatRoute::Auto);
11//! assert_eq!(chat.route, ChatRoute::Auto);
12//! ```
13
14use crate::client::env::{self, EnvError};
15use crate::completion::{CompletionRequest, ReplayTarget};
16use crate::error::{EncodeError, ProviderError};
17use crate::operation::Completion;
18use crate::providers::openai::wire::{COHERE, Chat, OpenAIConfig};
19use crate::wire::{
20    Decoder, Descriptor, Encoded, Flow, Mode, Out, Secret, Wire, WireEvent, WireFrame,
21};
22use serde::{Deserialize, Serialize};
23
24use super::NativeChat;
25
26/// Cohere's API root.
27const BASE_URL: &str = "https://api.cohere.ai";
28
29/// Where Cohere's OpenAI Compatibility API sits under the API root.
30const COMPATIBILITY_PATH: &str = "/compatibility/v1";
31
32/// The environment variable carrying the API key.
33const API_KEY_ENV: &str = "COHERE_API_KEY";
34
35/// The settings of a Cohere provider: serializable, and the credential is
36/// never serialized. [`connect`](Self::connect) puts it on a transport as a
37/// [`Cohere`](super::Cohere) client.
38#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
39pub struct CohereConfig {
40    /// The API key, sent as `Authorization: Bearer`.
41    pub api_key: Secret,
42    /// The API root, without a trailing slash.
43    pub base_url: String,
44}
45
46impl CohereConfig {
47    /// Cohere with default settings.
48    pub fn new(api_key: impl Into<Secret>) -> Self {
49        Self {
50            api_key: api_key.into(),
51            base_url: BASE_URL.to_owned(),
52        }
53    }
54
55    /// Cohere from `COHERE_API_KEY`.
56    pub fn from_env() -> Result<Self, EnvError> {
57        Ok(Self::new(env::required(API_KEY_ENV)?))
58    }
59
60    /// Point the wires at another API root.
61    pub fn with_base_url(mut self, base_url: impl AsRef<str>) -> Self {
62        self.base_url = base_url.as_ref().trim_end_matches('/').to_owned();
63        self
64    }
65
66    /// The chat wire for `model` under this API root, on the
67    /// [`ChatRoute::Compatibility`] route.
68    pub fn completion(&self, model: impl Into<String>) -> CohereChat {
69        let model = model.into();
70        CohereChat {
71            compatibility_api: OpenAIConfig::with_key(&COHERE, self.api_key.clone())
72                .with_base_url(format!("{}{COMPATIBILITY_PATH}", self.base_url))
73                .chat(model.clone()),
74            native_api: NativeChat::new(self.clone(), model),
75            route: ChatRoute::Compatibility,
76        }
77    }
78}
79
80/// Which Cohere API a chat request goes to. The default is
81/// [`Compatibility`](Self::Compatibility); the native API is opt-in.
82#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
83#[serde(rename_all = "snake_case")]
84pub enum ChatRoute {
85    /// The native API for a request that carries documents, so Cohere
86    /// grounds the answer in them and cites them, and the Compatibility
87    /// API for any other.
88    Auto,
89    /// The native API for every request.
90    Native,
91    /// The Compatibility API for every request, with documents sent as
92    /// text in the history.
93    #[default]
94    Compatibility,
95}
96
97/// Cohere chat: each request goes to the OpenAI Compatibility API or the
98/// native chat API, as [`ChatRoute`] says. Requests stay on the
99/// Compatibility API unless [`with_route`](Self::with_route) opts in to
100/// the native one.
101///
102/// The two are different APIs to history replay. A turn made on one route
103/// replays on the other from its canonical fields only, so its citations
104/// and tool plan stay behind. [`ChatRoute::Native`] keeps a whole
105/// conversation on the native API.
106#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
107pub struct CohereChat {
108    /// The Compatibility API's wire.
109    pub compatibility_api: Chat,
110    /// The native API's wire.
111    pub native_api: NativeChat,
112    /// Which API each request goes to.
113    pub route: ChatRoute,
114}
115
116impl CohereChat {
117    /// Send requests by `route`.
118    pub fn with_route(mut self, route: ChatRoute) -> Self {
119        self.route = route;
120        self
121    }
122
123    /// The same wire, both APIs encoding with `facts` and pricing their
124    /// replies by them.
125    pub fn with_facts(mut self, facts: crate::catalog::ModelFacts) -> Self {
126        self.compatibility_api = self.compatibility_api.with_facts(facts.clone());
127        self.native_api = self.native_api.with_facts(facts);
128        self
129    }
130
131    /// Ask both APIs to hold every tool call to its tool's schema.
132    pub fn with_strict_tools(mut self) -> Self {
133        self.compatibility_api = self.compatibility_api.with_strict_tools();
134        self.native_api = self.native_api.with_strict_tools();
135        self
136    }
137
138    /// Whether `request` goes to the native API.
139    fn native_for(&self, request: &CompletionRequest) -> bool {
140        match self.route {
141            ChatRoute::Auto => !request.documents.is_empty(),
142            ChatRoute::Native => true,
143            ChatRoute::Compatibility => false,
144        }
145    }
146}
147
148impl Wire for CohereChat {
149    type Op = Completion;
150    type Payload = Encoded;
151    type Frame = WireFrame;
152    type Event = CohereEvent;
153    type Decoder<'id> = CohereDecoder;
154    type Reassembler = document::Routed;
155
156    /// The Compatibility API's description, naming this wire as the replay
157    /// target that routes each request.
158    fn describe(&self) -> Descriptor<'_> {
159        self.compatibility_api.describe().replay(self)
160    }
161
162    fn encode(&self, request: CompletionRequest, mode: Mode) -> Result<Encoded, EncodeError> {
163        if self.native_for(&request) {
164            self.native_api.encode(request, mode)
165        } else {
166            self.compatibility_api.encode(request, mode)
167        }
168    }
169
170    fn decoder<'id>(&self) -> Self::Decoder<'id> {
171        CohereDecoder {
172            native_api: Wire::decoder(&self.native_api),
173            compatibility_api: Wire::decoder(&self.compatibility_api),
174            native_seen: false,
175        }
176    }
177
178    fn reassembler(&self) -> Self::Reassembler {
179        document::Routed::new(
180            Wire::reassembler(&self.native_api),
181            Wire::reassembler(&self.compatibility_api),
182        )
183    }
184}
185
186/// Every request names its route, so these are only what the replay of a
187/// history adapted to the wire as a whole reads: the Compatibility API's.
188impl ReplayTarget for CohereChat {
189    /// The answer of the API `request` goes to.
190    fn map_options(
191        &self,
192        request: &CompletionRequest,
193        fields: crate::completion::options::OptionFields<'_>,
194    ) -> crate::completion::options::OptionMap {
195        if self.native_for(request) {
196            self.native_api.map_options(request, fields)
197        } else {
198            self.compatibility_api.map_options(request, fields)
199        }
200    }
201
202    fn facts(&self) -> Option<&crate::catalog::ModelFacts> {
203        self.compatibility_api.facts()
204    }
205
206    fn api(&self) -> crate::message::Api {
207        self.compatibility_api.api()
208    }
209
210    fn provider(&self) -> &str {
211        ReplayTarget::provider(&self.compatibility_api)
212    }
213
214    fn model(&self) -> &str {
215        ReplayTarget::model(&self.compatibility_api)
216    }
217
218    fn accepts(&self, model: &str) -> crate::completion::Accepts {
219        self.compatibility_api.accepts(model)
220    }
221
222    fn route(&self, request: &CompletionRequest) -> Option<&dyn ReplayTarget> {
223        Some(if self.native_for(request) {
224            &self.native_api
225        } else {
226            &self.compatibility_api
227        })
228    }
229}
230
231/// One frame of a Cohere chat reply, from whichever API answered.
232pub enum CohereEvent {
233    /// A native chat event or reply.
234    Native(super::streaming::ChatEvent),
235    /// A Compatibility API chunk, reply or signal.
236    Compatibility(crate::providers::openai::wire::chat::ChatEvent),
237    /// Cohere's error envelope delivered with a success status, which fails
238    /// the turn.
239    Failure(ProviderError),
240}
241
242/// Decodes a Cohere chat reply from either API. The decoder is not told
243/// which API a request went to, so each frame is read by its shape: a
244/// native event carries a `type`, a native reply a `message` object, and
245/// Cohere's error envelope a `message` string.
246pub struct CohereDecoder {
247    native_api: super::streaming::ChatDecoder,
248    compatibility_api: crate::providers::openai::wire::chat::ChatDecoder,
249    /// Whether a native frame arrived, so the end of the frames is read as
250    /// the native API's.
251    native_seen: bool,
252}
253
254/// How a frame's shape names the API that sent it.
255#[derive(Debug, Clone, Copy)]
256enum Shape {
257    Native,
258    Compatibility,
259    /// `{"id": ..., "message": "..."}`, Cohere's error body.
260    Error,
261}
262
263/// Which API sent `data`, read off its top-level keys in one skim that
264/// builds no value: the API's own decoder parses the frame.
265fn shape(data: &str) -> Shape {
266    let starts = |value: &Option<&serde_json::value::RawValue>, first| {
267        value.is_some_and(|value| value.get().starts_with(first))
268    };
269    match crate::providers::internal::wire::top_level(data, &["type", "choices", "message"])
270        .as_deref()
271    {
272        Some([kind, _, _]) if starts(kind, '"') => Shape::Native,
273        Some([_, Some(_), _]) => Shape::Compatibility,
274        Some([_, _, message]) if starts(message, '{') => Shape::Native,
275        Some([_, _, message]) if starts(message, '"') => Shape::Error,
276        _ => Shape::Compatibility,
277    }
278}
279
280impl<'id> Decoder<'id, Completion> for CohereDecoder {
281    type Event = CohereEvent;
282
283    fn classify(&self, frame: WireFrame) -> WireEvent<CohereEvent> {
284        let data = frame.as_str();
285        match shape(&data) {
286            Shape::Error => WireEvent::Known(CohereEvent::Failure(
287                ProviderError::from_provider_body(data.into_owned()),
288            )),
289            Shape::Native => self.native_api.classify(frame).map(CohereEvent::Native),
290            Shape::Compatibility => self
291                .compatibility_api
292                .classify(frame)
293                .map(CohereEvent::Compatibility),
294        }
295    }
296
297    fn decode(
298        &mut self,
299        event: CohereEvent,
300        out: Out<'id, Completion>,
301    ) -> Result<Flow, ProviderError> {
302        match event {
303            CohereEvent::Native(event) => {
304                self.native_seen = true;
305                self.native_api.decode(event, out)
306            }
307            CohereEvent::Compatibility(event) => self.compatibility_api.decode(event, out),
308            CohereEvent::Failure(error) => Err(error),
309        }
310    }
311
312    fn eof(&mut self, out: Out<'id, Completion>) -> Result<Flow, ProviderError> {
313        if self.native_seen {
314            self.native_api.eof(out)
315        } else {
316            self.compatibility_api.eof(out)
317        }
318    }
319}
320
321mod document;
322
323#[cfg(test)]
324mod tests;
rig_core/providers/cohere/wire.rs

rig_core/providers/cohere/
wire.rs
