gemini_adk_rs/llm/
mock.rs

1//! [`MockLlm`]: a scripted, inspectable model for tests.
2
3use std::collections::VecDeque;
4use std::fmt;
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use parking_lot::Mutex;
9
10use super::{BaseLlm, LlmError, LlmRequest, LlmResponse, LlmStream};
11
12type ReplyFn = dyn Fn(&LlmRequest) -> Result<LlmResponse, LlmError> + Send + Sync;
13
14enum Replies {
15    Repeat(LlmResponse),
16    Script(VecDeque<Result<LlmResponse, LlmError>>),
17    Function(Box<ReplyFn>),
18}
19
20struct Inner {
21    replies: Mutex<Replies>,
22    requests: Mutex<Vec<LlmRequest>>,
23}
24
25/// A [`BaseLlm`] that replies from a script or a closure and records every
26/// request, for tests that must not reach a provider.
27///
28/// It replies three ways, one per kind of test:
29///
30/// | Constructor | Replies with | Use it to test |
31/// |---|---|---|
32/// | [`MockLlm::text`] | the same text, every call | wiring, prompts, state flow |
33/// | [`MockLlm::script`] | each response in turn, then an error | tool rounds, retries, multi-turn |
34/// | [`MockLlm::from_fn`] | whatever the closure returns | replies that depend on the request |
35///
36/// Every call's [`LlmRequest`] is recorded, so a test can assert on what the
37/// agent actually sent — instructions, history, tool declarations, sampling —
38/// not only on what came back.
39///
40/// `MockLlm` is a cheap handle: clones share the script and the recording.
41/// Give one clone to the agent and keep one to inspect.
42///
43/// ```
44/// use gemini_adk_rs::llm::{BaseLlm, LlmRequest, LlmResponse, MockLlm};
45///
46/// # tokio_test::block_on(async {
47/// let llm = MockLlm::script([
48///     LlmResponse::tool_call("get_weather", serde_json::json!({ "city": "Paris" })),
49///     LlmResponse::from_text("It is sunny in Paris."),
50/// ]);
51///
52/// let first = llm.generate(LlmRequest::from_text("Weather in Paris?")).await.unwrap();
53/// assert_eq!(first.function_calls()[0].name, "get_weather");
54///
55/// let second = llm.generate(LlmRequest::from_text("…")).await.unwrap();
56/// assert_eq!(second.text(), "It is sunny in Paris.");
57///
58/// // The script is spent: a third call is a test failure, not a silent reply.
59/// assert!(llm.generate(LlmRequest::from_text("again")).await.is_err());
60/// assert_eq!(llm.requests().len(), 3);
61/// # });
62/// ```
63#[derive(Clone)]
64pub struct MockLlm {
65    model_id: String,
66    inner: Arc<Inner>,
67}
68
69impl MockLlm {
70    fn with_replies(replies: Replies) -> Self {
71        Self {
72            model_id: "mock".into(),
73            inner: Arc::new(Inner {
74                replies: Mutex::new(replies),
75                requests: Mutex::new(Vec::new()),
76            }),
77        }
78    }
79
80    /// Reply with `text` to every request.
81    pub fn text(text: impl Into<String>) -> Self {
82        Self::with_replies(Replies::Repeat(LlmResponse::from_text(text)))
83    }
84
85    /// Reply with each response in order.
86    ///
87    /// A call after the script is spent returns an error naming the call, so
88    /// an agent that makes one model call more than the test expected fails
89    /// instead of receiving a plausible reply.
90    pub fn script(responses: impl IntoIterator<Item = LlmResponse>) -> Self {
91        Self::with_replies(Replies::Script(responses.into_iter().map(Ok).collect()))
92    }
93
94    /// Reply by calling `reply` with the request.
95    ///
96    /// This is the programmable mock: branch on the conversation so far, echo
97    /// the prompt, or return an error for a particular input.
98    pub fn from_fn<F>(reply: F) -> Self
99    where
100        F: Fn(&LlmRequest) -> Result<LlmResponse, LlmError> + Send + Sync + 'static,
101    {
102        Self::with_replies(Replies::Function(Box::new(reply)))
103    }
104
105    /// Append a failure to a script, to test how an agent handles a provider
106    /// error at that point in the conversation.
107    ///
108    /// # Panics
109    ///
110    /// On a mock built with [`text`](Self::text) or [`from_fn`](Self::from_fn):
111    /// those have no sequence to append to, and a closure can return the error
112    /// itself.
113    pub fn then_fail(self, error: LlmError) -> Self {
114        match &mut *self.inner.replies.lock() {
115            Replies::Script(queue) => queue.push_back(Err(error)),
116            _ => panic!("MockLlm::then_fail applies to MockLlm::script only"),
117        }
118        self
119    }
120
121    /// Report `model_id` from [`BaseLlm::model_id`] instead of `"mock"`.
122    pub fn with_model_id(mut self, model_id: impl Into<String>) -> Self {
123        self.model_id = model_id.into();
124        self
125    }
126
127    /// Every request received so far, oldest first.
128    pub fn requests(&self) -> Vec<LlmRequest> {
129        self.inner.requests.lock().clone()
130    }
131
132    /// The most recent request, if any.
133    pub fn last_request(&self) -> Option<LlmRequest> {
134        self.inner.requests.lock().last().cloned()
135    }
136
137    /// How many times the model has been called.
138    pub fn call_count(&self) -> usize {
139        self.inner.requests.lock().len()
140    }
141}
142
143impl fmt::Debug for MockLlm {
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        let replies = match &*self.inner.replies.lock() {
146            Replies::Repeat(_) => "repeat".to_string(),
147            Replies::Script(queue) => format!("script({} left)", queue.len()),
148            Replies::Function(_) => "fn".to_string(),
149        };
150        f.debug_struct("MockLlm")
151            .field("model_id", &self.model_id)
152            .field("replies", &replies)
153            .field("calls", &self.call_count())
154            .finish()
155    }
156}
157
158#[async_trait]
159impl BaseLlm for MockLlm {
160    fn model_id(&self) -> &str {
161        &self.model_id
162    }
163
164    /// Streams a text reply one word at a time (usage and the finish reason
165    /// ride on the last chunk), so tests exercise the multi-chunk path; a
166    /// reply with tool calls arrives as one chunk.
167    async fn generate_stream(&self, request: LlmRequest) -> Result<LlmStream, LlmError> {
168        use futures_util::StreamExt;
169
170        let reply = self.generate(request).await?;
171        let words: Vec<String> = reply
172            .text()
173            .split_inclusive(' ')
174            .map(str::to_owned)
175            .collect();
176        if !reply.function_calls().is_empty() || words.len() < 2 {
177            return Ok(futures_util::stream::once(async move { Ok(reply) }).boxed());
178        }
179        let last = words.len() - 1;
180        let chunks: Vec<Result<LlmResponse, LlmError>> = words
181            .into_iter()
182            .enumerate()
183            .map(|(i, word)| {
184                let mut chunk = LlmResponse::from_text(word);
185                chunk.finish_reason = None;
186                if i == last {
187                    chunk.finish_reason = reply.finish_reason.clone();
188                    chunk.usage = reply.usage;
189                }
190                Ok(chunk)
191            })
192            .collect();
193        Ok(futures_util::stream::iter(chunks).boxed())
194    }
195
196    async fn generate(&self, request: LlmRequest) -> Result<LlmResponse, LlmError> {
197        let call = {
198            let mut requests = self.inner.requests.lock();
199            requests.push(request.clone());
200            requests.len()
201        };
202        match &mut *self.inner.replies.lock() {
203            Replies::Repeat(response) => Ok(response.clone()),
204            Replies::Script(queue) => queue.pop_front().unwrap_or_else(|| {
205                Err(LlmError::Other(format!(
206                    "MockLlm script exhausted: call {call} has no scripted reply \
207                     (the agent called the model more times than the test expected)"
208                )))
209            }),
210            Replies::Function(reply) => reply(&request),
211        }
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    #[tokio::test]
220    async fn text_replies_the_same_to_every_call() {
221        let llm = MockLlm::text("hi");
222        for _ in 0..3 {
223            let reply = llm.generate(LlmRequest::from_text("x")).await.unwrap();
224            assert_eq!(reply.text(), "hi");
225        }
226        assert_eq!(llm.call_count(), 3);
227    }
228
229    #[tokio::test]
230    async fn script_replies_in_order_then_fails_loudly() {
231        let llm = MockLlm::script([LlmResponse::from_text("one"), LlmResponse::from_text("two")]);
232        assert_eq!(
233            llm.generate(LlmRequest::default()).await.unwrap().text(),
234            "one"
235        );
236        assert_eq!(
237            llm.generate(LlmRequest::default()).await.unwrap().text(),
238            "two"
239        );
240        let spent = llm.generate(LlmRequest::default()).await.unwrap_err();
241        assert!(spent.to_string().contains("call 3"), "{spent}");
242    }
243
244    #[tokio::test]
245    async fn then_fail_injects_a_provider_error_at_that_step() {
246        let llm = MockLlm::script([LlmResponse::from_text("ok")]).then_fail(LlmError::RateLimited);
247        assert!(llm.generate(LlmRequest::default()).await.is_ok());
248        assert!(matches!(
249            llm.generate(LlmRequest::default()).await,
250            Err(LlmError::RateLimited)
251        ));
252    }
253
254    #[tokio::test]
255    async fn from_fn_sees_the_request() {
256        let llm = MockLlm::from_fn(|req| {
257            Ok(LlmResponse::from_text(format!(
258                "{} turns",
259                req.contents.len()
260            )))
261        });
262        let reply = llm.generate(LlmRequest::from_text("x")).await.unwrap();
263        assert_eq!(reply.text(), "1 turns");
264    }
265
266    #[tokio::test]
267    async fn clones_share_the_script_and_the_recording() {
268        let llm = MockLlm::script([LlmResponse::from_text("a"), LlmResponse::from_text("b")]);
269        let handed_to_agent = llm.clone();
270        handed_to_agent
271            .generate(LlmRequest::from_text("first"))
272            .await
273            .unwrap();
274        assert_eq!(
275            llm.generate(LlmRequest::default()).await.unwrap().text(),
276            "b"
277        );
278        assert_eq!(llm.call_count(), 2);
279        assert_eq!(llm.requests()[0].contents.len(), 1);
280    }
281
282    #[tokio::test]
283    async fn with_model_id_keeps_the_replies() {
284        let llm = MockLlm::text("hi").with_model_id("gemini-test");
285        assert_eq!(llm.model_id(), "gemini-test");
286        assert_eq!(
287            llm.generate(LlmRequest::default()).await.unwrap().text(),
288            "hi"
289        );
290    }
291
292    #[test]
293    #[should_panic(expected = "script only")]
294    fn then_fail_on_a_repeating_mock_is_a_test_bug() {
295        let _ = MockLlm::text("hi").then_fail(LlmError::RateLimited);
296    }
297}
gemini_adk_rs/llm/mock.rs

gemini_adk_rs/llm/
mock.rs
