Skip to main content

cairn_runtime/
testing.rs

1//! Deterministic models for tests and offline examples.
2
3use std::collections::VecDeque;
4use std::sync::Mutex;
5
6use async_trait::async_trait;
7use serde_json::Value;
8
9use crate::message::Content;
10use crate::model::{Model, ModelError, ModelRequest, ModelResponse, StopReason, Usage};
11use crate::pricing::Pricing;
12
13/// Returns pre-scripted responses in order and records every request.
14pub struct ScriptedModel {
15    responses: Mutex<VecDeque<ModelResponse>>,
16    requests: Mutex<Vec<ModelRequest>>,
17}
18
19impl ScriptedModel {
20    pub fn new(responses: Vec<ModelResponse>) -> Self {
21        Self { responses: Mutex::new(responses.into()), requests: Mutex::new(Vec::new()) }
22    }
23
24    pub fn requests(&self) -> Vec<ModelRequest> {
25        self.requests.lock().unwrap().clone()
26    }
27}
28
29#[async_trait]
30impl Model for ScriptedModel {
31    fn id(&self) -> String {
32        "scripted:test".into()
33    }
34
35    async fn complete(&self, request: &ModelRequest) -> Result<ModelResponse, ModelError> {
36        self.requests.lock().unwrap().push(request.clone());
37        self.responses
38            .lock()
39            .unwrap()
40            .pop_front()
41            .ok_or_else(|| ModelError::Other("scripted model has no responses left".into()))
42    }
43
44    /// $1 per million tokens of every kind, so tests can exercise USD budgets.
45    fn pricing(&self) -> Option<Pricing> {
46        Some(Pricing { input: 1.0, output: 1.0, cache_read: 1.0, cache_write: 1.0 })
47    }
48}
49
50pub fn text_response(text: &str) -> ModelResponse {
51    ModelResponse {
52        model: "test".into(),
53        content: vec![Content::Text { text: text.into() }],
54        stop_reason: StopReason::EndTurn,
55        usage: Usage { input_tokens: 100, output_tokens: 20, ..Default::default() },
56        ..Default::default()
57    }
58}
59
60pub fn tool_call_response(calls: &[(&str, &str, Value)]) -> ModelResponse {
61    ModelResponse {
62        model: "test".into(),
63        content: calls
64            .iter()
65            .map(|(id, name, input)| Content::ToolUse { id: (*id).into(), name: (*name).into(), input: input.clone() })
66            .collect(),
67        stop_reason: StopReason::ToolUse,
68        usage: Usage { input_tokens: 100, output_tokens: 20, ..Default::default() },
69        ..Default::default()
70    }
71}