1use std::collections::HashMap;
2
3use tokio::sync::broadcast;
4use tokio_util::sync::CancellationToken;
5
6use crate::error::RuntimeError;
7use crate::event::{NodeEvent, Observable};
8use crate::message::{Message, MessageOrigin, MessagePart, MessageRole};
9use crate::provider::{
10 AssistantMessage, CallTiming, DEFAULT_STREAM_BUFFER, LlmRequest, Provider, StopReason,
11 TokenUsage, estimate_tokens,
12};
13use crate::tool::BoxFut;
14use crate::value::Value;
15
16pub struct MockProvider {
17 name: String,
18 by_model: HashMap<String, Value>,
19 by_prefix: Vec<(String, String, Value)>,
20 fallback: Option<Value>,
21 chunk_delay: Option<std::time::Duration>,
22}
23
24impl MockProvider {
25 pub fn new(name: impl Into<String>) -> Self {
26 Self {
27 name: name.into(),
28 by_model: HashMap::new(),
29 by_prefix: Vec::new(),
30 fallback: None,
31 chunk_delay: None,
32 }
33 }
34
35 pub fn with_chunk_delay(mut self, d: std::time::Duration) -> Self {
36 self.chunk_delay = Some(d);
37 self
38 }
39
40 pub fn with_model(mut self, model: impl Into<String>, value: Value) -> Self {
41 self.by_model.insert(model.into(), value);
42 self
43 }
44
45 pub fn with_prefix(
46 mut self,
47 model: impl Into<String>,
48 prompt_prefix: impl Into<String>,
49 value: Value,
50 ) -> Self {
51 self.by_prefix
52 .push((model.into(), prompt_prefix.into(), value));
53 self
54 }
55
56 pub fn with_fallback(mut self, value: Value) -> Self {
57 self.fallback = Some(value);
58 self
59 }
60}
61
62impl Provider for MockProvider {
63 fn name(&self) -> &str {
64 &self.name
65 }
66
67 fn call<'a>(&'a self, req: LlmRequest) -> BoxFut<'a, Result<AssistantMessage, RuntimeError>> {
68 let turn_id = turn_id_from_req(&req);
69 Box::pin(async move {
70 self.lookup(&req)
71 .map(|v| value_to_assistant_message(&v, turn_id))
72 })
73 }
74
75 fn call_streaming(&self, req: LlmRequest) -> Observable<AssistantMessage> {
76 let turn_id = turn_id_from_req(&req);
77 let (tx, events) = broadcast::channel(DEFAULT_STREAM_BUFFER);
78 let cancel = CancellationToken::new();
79 let cancel_for_task = cancel.clone();
80 let looked_up = self.lookup(&req);
81 let chunk_delay = self.chunk_delay;
82 let output: BoxFut<'static, Result<AssistantMessage, RuntimeError>> =
83 Box::pin(async move {
84 let value = match looked_up {
85 Ok(v) => v,
86 Err(e) => {
87 let _ = tx.send(NodeEvent::LlmDone { total_tokens: 0 });
88 return Err(e);
89 }
90 };
91 let text_form = value_to_stream_text(&value);
92 let chunks = split_for_stream(&text_form);
93 let mut running = 0u64;
94 for chunk in chunks {
95 if let Some(d) = chunk_delay {
96 tokio::select! {
97 biased;
98 _ = cancel_for_task.cancelled() => {
99 let _ = tx.send(NodeEvent::LlmDone { total_tokens: running });
100 return Err(RuntimeError::Cancelled("mock stream cancelled".into()));
101 }
102 _ = tokio::time::sleep(d) => {}
103 }
104 }
105 if cancel_for_task.is_cancelled() {
106 let _ = tx.send(NodeEvent::LlmDone {
107 total_tokens: running,
108 });
109 return Err(RuntimeError::Cancelled("mock stream cancelled".into()));
110 }
111 let inc = estimate_tokens(&chunk);
112 running += inc;
113 let _ = tx.send(NodeEvent::LlmChunk {
114 text: chunk,
115 cumulative_tokens: running,
116 });
117 }
118 let _ = tx.send(NodeEvent::LlmDone {
119 total_tokens: running,
120 });
121 Ok(value_to_assistant_message(&value, turn_id))
122 });
123 Observable {
124 output,
125 events,
126 cancel,
127 }
128 }
129}
130
131impl MockProvider {
132 fn lookup(&self, req: &LlmRequest) -> Result<Value, RuntimeError> {
133 let prompt_text = req
134 .messages
135 .last()
136 .map(|m| m.text_concat())
137 .unwrap_or_default();
138 for (model, prefix, value) in &self.by_prefix {
139 if req.model == *model && prompt_text.starts_with(prefix.as_str()) {
140 return Ok(value.clone());
141 }
142 }
143 if let Some(v) = self.by_model.get(&req.model) {
144 return Ok(v.clone());
145 }
146 if let Some(v) = &self.fallback {
147 return Ok(v.clone());
148 }
149 Err(RuntimeError::ToolFailed(format!(
150 "mock provider `{}` has no entry for model={} prompt.prefix={:?}",
151 self.name,
152 req.model,
153 prompt_text.chars().take(40).collect::<String>()
154 )))
155 }
156}
157
158fn value_to_assistant_message(v: &Value, turn_id: crate::event::TurnId) -> AssistantMessage {
159 let text = value_to_stream_text(v);
160 AssistantMessage {
161 message: Message {
162 role: MessageRole::Assistant,
163 parts: vec![MessagePart::Text { text: text.clone() }],
164 turn_id,
165 origin: MessageOrigin::User,
166 },
167 stop_reason: StopReason::End,
168 token_usage: TokenUsage {
169 output: estimate_tokens(&text),
170 ..Default::default()
171 },
172 timing: CallTiming::default(),
173 model: String::new(),
174 response_id: None,
175 }
176}
177
178fn turn_id_from_req(req: &LlmRequest) -> crate::event::TurnId {
179 req.messages
180 .first()
181 .map(|m| m.turn_id.clone())
182 .unwrap_or_else(crate::event::TurnId::now)
183}
184
185fn value_to_stream_text(v: &Value) -> String {
186 match v {
187 Value::Str(s) => s.clone(),
188 other => other.to_json().to_string(),
189 }
190}
191
192fn split_for_stream(s: &str) -> Vec<String> {
193 if s.len() > 8 {
194 s.as_bytes()
195 .chunks(s.len().div_ceil(3))
196 .map(|c| String::from_utf8_lossy(c).into_owned())
197 .collect()
198 } else {
199 vec![s.to_string()]
200 }
201}
202
203#[cfg(test)]
204mod tests {
205 use super::*;
206 use crate::provider::user_text_message;
207
208 fn req(model: &str, prompt: &str) -> LlmRequest {
209 LlmRequest {
210 model: model.into(),
211 messages: vec![user_text_message(prompt)],
212 system: None,
213 input: Value::Unit,
214 schema: None,
215 cache_prompt: false,
216 tools: Vec::new(),
217 thinking_enabled: false,
218 stall_timeout_secs: 0,
219 }
220 }
221
222 #[tokio::test]
223 async fn resolves_by_model_name() {
224 let p = MockProvider::new("mock").with_model("gpt-4o-mini", Value::Str("hi".into()));
225 let out = p.call(req("gpt-4o-mini", "anything")).await.unwrap();
226 assert_eq!(out.text_concat(), "hi");
227 }
228
229 #[tokio::test]
230 async fn prefix_wins_over_model() {
231 let p = MockProvider::new("mock")
232 .with_model("m", Value::Str("model-hit".into()))
233 .with_prefix("m", "review", Value::Str("prefix-hit".into()));
234 let out = p.call(req("m", "review please")).await.unwrap();
235 assert_eq!(out.text_concat(), "prefix-hit");
236 }
237
238 #[tokio::test]
239 async fn missing_entry_errors_with_hint() {
240 let p = MockProvider::new("mock");
241 let err = p.call(req("gpt", "hello")).await.unwrap_err();
242 assert!(matches!(err, RuntimeError::ToolFailed(msg) if msg.contains("gpt")));
243 }
244
245 #[tokio::test]
246 async fn fallback_captures_unmatched() {
247 let p = MockProvider::new("mock").with_fallback(Value::Str("fb".into()));
248 let out = p.call(req("anything", "")).await.unwrap();
249 assert_eq!(out.text_concat(), "fb");
250 }
251}