1use 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
13pub 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 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}