cairn_runtime/
model.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use futures::stream::{self, BoxStream};
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8
9use crate::message::{Content, Message};
10use crate::pricing::Pricing;
11use crate::tool::ToolSpec;
12
13#[derive(Debug, Clone, Default)]
14pub struct ModelRequest {
15    pub system: Option<String>,
16    pub messages: Vec<Message>,
17    pub tools: Vec<ToolSpec>,
18    pub max_tokens: u32,
19    pub temperature: Option<f32>,
20    /// Provider-specific options, keyed by provider id (e.g. `"anthropic"`,
21    /// `"openai"`). Adapters merge their own entry into the request body and
22    /// ignore the rest, so one agent definition can carry tuned options for
23    /// every provider it may fall back to.
24    pub provider_options: BTreeMap<String, Value>,
25}
26
27#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(rename_all = "snake_case")]
29pub enum StopReason {
30    #[default]
31    EndTurn,
32    ToolUse,
33    MaxTokens,
34    Other(String),
35}
36
37#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
38pub struct Usage {
39    pub input_tokens: u64,
40    pub output_tokens: u64,
41    #[serde(default)]
42    pub cache_read_tokens: u64,
43    #[serde(default)]
44    pub cache_write_tokens: u64,
45}
46
47impl Usage {
48    pub fn total(&self) -> u64 {
49        self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_write_tokens
50    }
51
52    pub fn add(&mut self, other: &Usage) {
53        self.input_tokens += other.input_tokens;
54        self.output_tokens += other.output_tokens;
55        self.cache_read_tokens += other.cache_read_tokens;
56        self.cache_write_tokens += other.cache_write_tokens;
57    }
58}
59
60#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
61pub struct ModelResponse {
62    /// Model name as reported by the provider.
63    pub model: String,
64    pub content: Vec<Content>,
65    pub stop_reason: StopReason,
66    pub usage: Usage,
67    /// [`Model::id`] of the model that served this response, when a
68    /// composite model (e.g. [`FallbackModel`]) chose among several.
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub served_by: Option<String>,
71    /// Earlier answers in this call that were discarded as unusable. Their
72    /// tokens were still paid for and count toward budgets.
73    #[serde(default, skip_serializing_if = "Vec::is_empty")]
74    pub rejected: Vec<RejectedAttempt>,
75}
76
77/// A model answer that was received, paid for, and discarded.
78#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
79pub struct RejectedAttempt {
80    pub model: String,
81    pub reason: String,
82    pub usage: Usage,
83    pub cost_usd: Option<f64>,
84}
85
86impl ModelResponse {
87    /// Tokens for this response plus any rejected attempts.
88    pub fn total_usage(&self) -> Usage {
89        let mut u = self.usage;
90        for r in &self.rejected {
91            u.add(&r.usage);
92        }
93        u
94    }
95}
96
97#[derive(Debug, Clone, thiserror::Error)]
98pub enum ModelError {
99    #[error("transport error: {0}")]
100    Transport(String),
101    #[error("provider returned {status}: {body}")]
102    Api { status: u16, body: String },
103    #[error("could not decode provider response: {0}")]
104    Decode(String),
105    /// Every model answered, but no answer passed the usability check.
106    #[error("no usable response: {reason}")]
107    Unusable { reason: String, attempts: Vec<RejectedAttempt> },
108    #[error("{0}")]
109    Other(String),
110}
111
112impl ModelError {
113    /// Whether another attempt (or another provider) may succeed.
114    pub fn retryable(&self) -> bool {
115        match self {
116            ModelError::Transport(_) => true,
117            ModelError::Api { status, .. } => *status == 408 || *status == 429 || *status >= 500,
118            _ => false,
119        }
120    }
121
122    /// Tokens and USD already spent before this error (e.g. on answers that
123    /// were rejected as unusable). `cost` is `None` if any attempt was unpriced.
124    pub fn spent(&self) -> (Usage, Option<f64>) {
125        match self {
126            ModelError::Unusable { attempts, .. } => {
127                let mut usage = Usage::default();
128                let mut cost = Some(0.0);
129                for a in attempts {
130                    usage.add(&a.usage);
131                    cost = cost.zip(a.cost_usd).map(|(x, y)| x + y);
132                }
133                (usage, cost)
134            }
135            _ => (Usage::default(), Some(0.0)),
136        }
137    }
138}
139
140#[async_trait]
141pub trait Model: Send + Sync {
142    /// Identifier such as `anthropic:claude-opus-5-5`.
143    fn id(&self) -> String;
144
145    async fn complete(&self, request: &ModelRequest) -> Result<ModelResponse, ModelError>;
146
147    /// Stream the response as it is generated. The stream always ends with
148    /// [`StreamEvent::Completed`] carrying the full response. The default
149    /// implementation calls [`Model::complete`] and replays it as deltas.
150    async fn stream(&self, request: &ModelRequest) -> Result<ModelStream, ModelError> {
151        let response = self.complete(request).await?;
152        Ok(Box::pin(stream::iter(synthesize_stream(response).into_iter().map(Ok))))
153    }
154
155    /// Prices for this model, if configured. Required for USD budgets.
156    fn pricing(&self) -> Option<Pricing> {
157        None
158    }
159
160    /// Cost of a response in USD, if pricing is configured.
161    fn cost_usd(&self, response: &ModelResponse) -> Option<f64> {
162        self.pricing().map(|p| p.cost(&response.usage))
163    }
164
165    /// Upper bound in USD for one call with this many input and output
166    /// tokens, counting every attempt a composite model might make. Input is
167    /// priced at the higher of the input and cache-write rates.
168    fn max_call_cost(&self, input_tokens: u64, output_tokens: u64) -> Option<f64> {
169        self.pricing()
170            .map(|p| (input_tokens as f64 * p.input.max(p.cache_write) + output_tokens as f64 * p.output) / 1_000_000.0)
171    }
172
173    /// Ids of models (including any inside a composite) that have no
174    /// pricing. A USD budget cannot be enforced while this is non-empty.
175    fn unpriced_models(&self) -> Vec<String> {
176        if self.pricing().is_none() { vec![self.id()] } else { Vec::new() }
177    }
178}
179
180/// Incremental output from a streaming model call.
181#[derive(Debug, Clone, PartialEq)]
182pub enum StreamEvent {
183    TextDelta(String),
184    ThinkingDelta(String),
185    /// Signature for the current thinking block (Anthropic).
186    ThinkingSignature(String),
187    /// A new tool call; `index` counts tool calls within this response.
188    ToolCallStart {
189        index: usize,
190        id: String,
191        name: String,
192    },
193    ToolInputDelta {
194        index: usize,
195        partial_json: String,
196    },
197    /// Final event of every successful stream.
198    Completed(ModelResponse),
199}
200
201pub type ModelStream = BoxStream<'static, Result<StreamEvent, ModelError>>;
202
203/// Express a complete response as the deltas a streaming call would yield.
204pub fn synthesize_stream(response: ModelResponse) -> Vec<StreamEvent> {
205    let mut events = Vec::new();
206    let mut tool_index = 0;
207    for c in &response.content {
208        match c {
209            Content::Text { text } => events.push(StreamEvent::TextDelta(text.clone())),
210            Content::Thinking { text, provider_data } => {
211                events.push(StreamEvent::ThinkingDelta(text.clone()));
212                if let Some(sig) = provider_data.as_ref().and_then(|d| d["signature"].as_str()) {
213                    events.push(StreamEvent::ThinkingSignature(sig.to_owned()));
214                }
215            }
216            Content::ToolUse { id, name, input } => {
217                events.push(StreamEvent::ToolCallStart { index: tool_index, id: id.clone(), name: name.clone() });
218                events.push(StreamEvent::ToolInputDelta { index: tool_index, partial_json: input.to_string() });
219                tool_index += 1;
220            }
221            Content::ToolResult { .. } => {}
222        }
223    }
224    events.push(StreamEvent::Completed(response));
225    events
226}
227
228type UsableFn = Arc<dyn Fn(&ModelResponse) -> Result<(), String> + Send + Sync>;
229
230/// Tries each model in order. Moves on when a request fails with a
231/// retryable error **or** when the answer is unusable:
232///
233/// * cut off by the token limit (unless [`accept_truncated`](Self::accept_truncated)),
234/// * empty — no text and no tool calls (unless [`accept_empty`](Self::accept_empty)),
235/// * a final answer that is not JSON with the required fields
236///   ([`require_json_fields`](Self::require_json_fields)),
237/// * rejected by a custom check ([`require`](Self::require)).
238///
239/// Rejected answers are recorded in [`ModelResponse::rejected`] and their
240/// cost still counts toward budgets. With usability checks enabled,
241/// streaming is buffered per attempt so a rejected answer never reaches
242/// the caller.
243pub struct FallbackModel {
244    models: Vec<Arc<dyn Model>>,
245    allow_truncated: bool,
246    allow_empty: bool,
247    json_fields: Option<Vec<String>>,
248    custom: Option<UsableFn>,
249    keep_last: bool,
250}
251
252impl FallbackModel {
253    pub fn new(models: Vec<Arc<dyn Model>>) -> Self {
254        assert!(!models.is_empty(), "FallbackModel needs at least one model");
255        Self { models, allow_truncated: false, allow_empty: false, json_fields: None, custom: None, keep_last: false }
256    }
257
258    /// When every model's answer is unusable, return the last one anyway
259    /// instead of failing the call.
260    ///
261    /// For agents with a [`finish_schema`](crate::AgentBuilder::finish_schema):
262    /// a final answer that is not quite valid JSON then reaches the agent's
263    /// corrective retry, which usually fixes it on the same cheap model,
264    /// rather than ending the run. Transport and API errors still fail.
265    pub fn keep_last_answer(mut self) -> Self {
266        self.keep_last = true;
267        self
268    }
269
270    /// Keep answers that hit the token limit.
271    pub fn accept_truncated(mut self) -> Self {
272        self.allow_truncated = true;
273        self
274    }
275
276    /// Keep answers with no text and no tool calls.
277    pub fn accept_empty(mut self) -> Self {
278        self.allow_empty = true;
279        self
280    }
281
282    /// Final answers (no tool calls) must contain a JSON object with these
283    /// top-level fields, e.g. `["summary", "work", "risks", "confidence"]`.
284    pub fn require_json_fields(mut self, fields: &[&str]) -> Self {
285        self.json_fields = Some(fields.iter().map(|f| (*f).to_owned()).collect());
286        self
287    }
288
289    /// Custom usability check; return `Err(reason)` to try the next model.
290    pub fn require(mut self, check: impl Fn(&ModelResponse) -> Result<(), String> + Send + Sync + 'static) -> Self {
291        self.custom = Some(Arc::new(check));
292        self
293    }
294
295    fn has_checks(&self) -> bool {
296        !self.allow_truncated || !self.allow_empty || self.json_fields.is_some() || self.custom.is_some()
297    }
298
299    /// Why a response is unusable, if it is.
300    pub fn check(&self, r: &ModelResponse) -> Result<(), String> {
301        if !self.allow_truncated && r.stop_reason == StopReason::MaxTokens {
302            return Err("reply was cut off at the token limit".into());
303        }
304        let has_tools = r.content.iter().any(|c| matches!(c, Content::ToolUse { .. }));
305        let text = crate::message::content_text(&r.content);
306        if !self.allow_empty && !has_tools && text.trim().is_empty() {
307            return Err("reply is empty".into());
308        }
309        if let (Some(fields), false) = (&self.json_fields, has_tools) {
310            let Some(serde_json::Value::Object(obj)) = crate::output::extract_json(&text) else {
311                return Err("reply is not a JSON object".into());
312            };
313            let missing: Vec<&str> = fields.iter().filter(|f| !obj.contains_key(*f)).map(String::as_str).collect();
314            if !missing.is_empty() {
315                return Err(format!("JSON is missing required field(s): {}", missing.join(", ")));
316            }
317        }
318        if let Some(check) = &self.custom {
319            check(r)?;
320        }
321        Ok(())
322    }
323}
324
325#[async_trait]
326impl Model for FallbackModel {
327    fn id(&self) -> String {
328        let ids: Vec<_> = self.models.iter().map(|m| m.id()).collect();
329        format!("fallback({})", ids.join(","))
330    }
331
332    async fn complete(&self, request: &ModelRequest) -> Result<ModelResponse, ModelError> {
333        let mut last_err = None;
334        let mut rejected: Vec<RejectedAttempt> = Vec::new();
335        // The most recent unusable answer, kept only for `keep_last_answer`.
336        let mut last_unusable: Option<ModelResponse> = None;
337        for model in &self.models {
338            match model.complete(request).await {
339                Ok(mut resp) => {
340                    resp.served_by.get_or_insert_with(|| model.id());
341                    match self.check(&resp) {
342                        Ok(()) => {
343                            rejected.append(&mut resp.rejected);
344                            resp.rejected = rejected;
345                            return Ok(resp);
346                        }
347                        Err(reason) => {
348                            tracing::warn!(model = %model.id(), %reason, "unusable answer, falling back");
349                            rejected.append(&mut resp.rejected);
350                            if self.keep_last {
351                                last_unusable = Some(resp.clone());
352                                last_err = None;
353                            }
354                            rejected.push(RejectedAttempt {
355                                model: model.id(),
356                                reason,
357                                usage: resp.usage,
358                                cost_usd: model.cost_usd(&resp),
359                            });
360                        }
361                    }
362                }
363                Err(e) if e.retryable() => {
364                    tracing::warn!(model = %model.id(), error = %e, "model failed, falling back");
365                    last_err = Some(e);
366                }
367                Err(ModelError::Unusable { reason, attempts }) => {
368                    rejected.extend(attempts);
369                    last_err = Some(ModelError::Other(reason));
370                }
371                Err(e) if rejected.is_empty() => return Err(e),
372                Err(e) => {
373                    last_err = Some(e);
374                    break;
375                }
376            }
377        }
378        if rejected.is_empty() {
379            return Err(last_err.expect("at least one model"));
380        }
381        // Every model answered and none usably: hand back the last answer, with
382        // the earlier attempts attached so budgets still count them. Its own
383        // tokens are counted once, as the response, not again as a rejection.
384        if let (Some(mut resp), None) = (last_unusable, &last_err) {
385            rejected.pop();
386            resp.rejected = rejected;
387            tracing::warn!(model = resp.served_by.as_deref().unwrap_or(""), "no usable answer; keeping the last one");
388            return Ok(resp);
389        }
390        // Keep the paid-for attempts on the error so budgets still see them.
391        let mut reason = rejected.iter().map(|r| format!("{}: {}", r.model, r.reason)).collect::<Vec<_>>().join("; ");
392        if let Some(e) = last_err {
393            reason = format!("{reason}; last error: {e}");
394        }
395        Err(ModelError::Unusable { reason, attempts: rejected })
396    }
397
398    /// Fails over only while opening the stream; once tokens flow, errors
399    /// are surfaced to the caller.
400    async fn stream(&self, request: &ModelRequest) -> Result<ModelStream, ModelError> {
401        if self.has_checks() {
402            // An answer must be validated before any of it is released.
403            let response = self.complete(request).await?;
404            return Ok(Box::pin(stream::iter(synthesize_stream(response).into_iter().map(Ok))));
405        }
406        let mut last_err = None;
407        for model in &self.models {
408            match model.stream(request).await {
409                Ok(s) => {
410                    let id = model.id();
411                    return Ok(Box::pin(futures::StreamExt::map(s, move |ev| match ev {
412                        Ok(StreamEvent::Completed(mut resp)) => {
413                            resp.served_by.get_or_insert_with(|| id.clone());
414                            Ok(StreamEvent::Completed(resp))
415                        }
416                        other => other,
417                    })));
418                }
419                Err(e) if e.retryable() => {
420                    tracing::warn!(model = %model.id(), error = %e, "model stream failed, falling back");
421                    last_err = Some(e);
422                }
423                Err(e) => return Err(e),
424            }
425        }
426        Err(last_err.expect("at least one model"))
427    }
428
429    /// Conservative (field-wise maximum) pricing across all models, or
430    /// `None` if any model is unpriced.
431    fn pricing(&self) -> Option<Pricing> {
432        self.models.iter().map(|m| m.pricing()).try_fold(Pricing::default(), |acc, p| p.map(|p| acc.max(p)))
433    }
434
435    /// Prices the response with the model that actually served it; falls
436    /// back to the conservative maximum rather than ever under-counting.
437    /// Cost of the answer (priced on the model that served it, else the
438    /// conservative maximum) plus every rejected attempt.
439    fn cost_usd(&self, response: &ModelResponse) -> Option<f64> {
440        let served = response.served_by.as_deref();
441        let answer = match self.models.iter().find(|m| Some(m.id().as_str()) == served) {
442            Some(m) => m.cost_usd(&ModelResponse { rejected: Vec::new(), ..response.clone() }),
443            None => self.pricing().map(|p| p.cost(&response.usage)),
444        }?;
445        let conservative = self.pricing();
446        response
447            .rejected
448            .iter()
449            .try_fold(answer, |acc, r| r.cost_usd.or_else(|| conservative.map(|p| p.cost(&r.usage))).map(|c| acc + c))
450    }
451
452    fn unpriced_models(&self) -> Vec<String> {
453        self.models.iter().flat_map(|m| m.unpriced_models()).collect()
454    }
455
456    /// Every model may be tried (and paid for) in one call.
457    fn max_call_cost(&self, input_tokens: u64, output_tokens: u64) -> Option<f64> {
458        self.models.iter().map(|m| m.max_call_cost(input_tokens, output_tokens)).sum()
459    }
460}
461
462#[async_trait]
463impl<M: Model + ?Sized> Model for Arc<M> {
464    fn id(&self) -> String {
465        (**self).id()
466    }
467
468    async fn complete(&self, request: &ModelRequest) -> Result<ModelResponse, ModelError> {
469        (**self).complete(request).await
470    }
471
472    async fn stream(&self, request: &ModelRequest) -> Result<ModelStream, ModelError> {
473        (**self).stream(request).await
474    }
475
476    fn pricing(&self) -> Option<Pricing> {
477        (**self).pricing()
478    }
479
480    fn cost_usd(&self, response: &ModelResponse) -> Option<f64> {
481        (**self).cost_usd(response)
482    }
483
484    fn unpriced_models(&self) -> Vec<String> {
485        (**self).unpriced_models()
486    }
487
488    fn max_call_cost(&self, input_tokens: u64, output_tokens: u64) -> Option<f64> {
489        (**self).max_call_cost(input_tokens, output_tokens)
490    }
491}
cairn_runtime/model.rs

cairn_runtime/
model.rs
