yapi_ai/
faux.rs

1//! A scripted provider for tests: replays canned responses as streams.
2
3use std::collections::VecDeque;
4use std::sync::{Arc, Mutex};
5
6use serde_json::{Map, Value};
7use yapi_types::message::{
8    AssistantMessage, ContentBlock, Message, StopReason, ThinkingContent, ToolCall,
9};
10
11use crate::stream::{EventStream, Provider, Request, StreamEvent, new_output, now_ms};
12
13/// A canned response.
14#[derive(Clone, Debug, Default)]
15pub struct Response {
16    /// Content blocks, streamed in order.
17    pub content: Vec<ContentBlock>,
18    /// Final stop reason. `None` means `toolUse` when the content has tool
19    /// calls, else `stop`.
20    pub stop_reason: Option<StopReason>,
21    /// Error message for `error` and `aborted`.
22    pub error_message: Option<String>,
23}
24
25impl Response {
26    /// A text response.
27    pub fn text(text: &str) -> Response {
28        Response {
29            content: vec![ContentBlock::text(text)],
30            ..Response::default()
31        }
32    }
33
34    /// A response with one tool call.
35    pub fn tool_call(id: &str, name: &str, arguments: Value) -> Response {
36        Response {
37            content: vec![ContentBlock::tool_call(
38                id,
39                name,
40                arguments.as_object().cloned().unwrap_or_else(Map::new),
41            )],
42            ..Response::default()
43        }
44    }
45
46    /// A failed response.
47    pub fn error(message: &str) -> Response {
48        Response {
49            stop_reason: Some(StopReason::Error),
50            error_message: Some(message.to_owned()),
51            ..Response::default()
52        }
53    }
54}
55
56/// Replays responses in order and records the requests it receives. When the
57/// script runs out it answers with an error.
58#[derive(Clone, Debug, Default)]
59pub struct Faux {
60    script: Arc<Mutex<VecDeque<Response>>>,
61    requests: Arc<Mutex<Vec<Vec<Message>>>>,
62}
63
64impl Faux {
65    /// A provider that will give `responses`, one per request.
66    pub fn new(responses: impl IntoIterator<Item = Response>) -> Faux {
67        Faux {
68            script: Arc::new(Mutex::new(responses.into_iter().collect())),
69            requests: Arc::default(),
70        }
71    }
72
73    /// The transcripts received so far.
74    pub fn requests(&self) -> Vec<Vec<Message>> {
75        self.requests.lock().map(|r| r.clone()).unwrap_or_default()
76    }
77}
78
79impl Provider for Faux {
80    fn api(&self) -> &str {
81        "faux"
82    }
83
84    fn stream(&self, request: Request) -> EventStream {
85        if let Ok(mut requests) = self.requests.lock() {
86            requests.push(request.messages.clone());
87        }
88        let response = self
89            .script
90            .lock()
91            .ok()
92            .and_then(|mut script| script.pop_front())
93            .unwrap_or_else(|| Response::error("faux: no scripted response left"));
94        let (sender, stream) = EventStream::channel();
95        let mut output: AssistantMessage = new_output(&request.model, now_ms());
96        output.response_id = Some("faux".into());
97        let input = crate::transcript::estimate_context_tokens(&request.messages);
98        output.usage.input = input;
99        output.usage.total_tokens = Some(input);
100        sender.send(StreamEvent::Start(output.clone()));
101
102        for block in response.content {
103            let (empty, delta) = match &block {
104                ContentBlock::Text(text) => (ContentBlock::text(""), text.text.clone()),
105                ContentBlock::Thinking(thinking) => (
106                    ContentBlock::Thinking(ThinkingContent {
107                        thinking: String::new(),
108                        ..thinking.clone()
109                    }),
110                    thinking.thinking.clone(),
111                ),
112                ContentBlock::ToolCall(call) => (
113                    ContentBlock::ToolCall(ToolCall {
114                        arguments: Map::new(),
115                        ..call.clone()
116                    }),
117                    yapi_types::json::stringify(&call.arguments),
118                ),
119                ContentBlock::Image(_) => continue,
120            };
121            let index = sender.start(&mut output, empty);
122            sender.delta(&mut output, index, &delta);
123            output.content[index] = block;
124            sender.end(&output, index);
125        }
126
127        let has_calls = output
128            .content
129            .iter()
130            .any(|block| matches!(block, ContentBlock::ToolCall(_)));
131        output.stop_reason = response.stop_reason.unwrap_or(if has_calls {
132            StopReason::ToolUse
133        } else {
134            StopReason::Stop
135        });
136        let produced = crate::transcript::estimate_message_tokens(&Message::Assistant(Box::new(
137            output.clone(),
138        )));
139        output.usage.output = produced;
140        output.usage.total_tokens = Some(input + produced);
141        match output.stop_reason {
142            StopReason::Error | StopReason::Aborted => {
143                output.error_message = response.error_message;
144                sender.send(StreamEvent::Error(output));
145            }
146            _ => sender.send(StreamEvent::Done(output)),
147        }
148        stream
149    }
150}
yapi_ai/faux.rs

yapi_ai/
faux.rs
