llm/testing/
fake_llm.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::{Arc, Mutex};
4
5use tokio::spawn;
6use tokio::sync::{Notify, mpsc};
7use tokio_stream::wrappers::UnboundedReceiverStream;
8
9use crate::{Context, LlmError, LlmModel, LlmResponse, LlmResponseStream, StreamingModelProvider};
10
11pub struct FakeLlmProvider {
12    responses: Vec<Vec<Result<LlmResponse, LlmError>>>,
13    pauses: HashMap<usize, HashMap<usize, Arc<Notify>>>,
14    call_count: AtomicUsize,
15    /// Captured contexts from each call to `stream_response`
16    captured_contexts: Arc<Mutex<Vec<Context>>>,
17    display_name: String,
18    model: Option<LlmModel>,
19    context_window: Option<u32>,
20}
21
22impl FakeLlmProvider {
23    pub fn new(responses: Vec<Vec<LlmResponse>>) -> Self {
24        let wrapped = responses.into_iter().map(|turn| turn.into_iter().map(Ok).collect()).collect();
25        Self::from_results(wrapped)
26    }
27
28    pub fn with_single_response(chunks: Vec<LlmResponse>) -> Self {
29        Self::new(vec![chunks])
30    }
31
32    pub fn from_results(responses: Vec<Vec<Result<LlmResponse, LlmError>>>) -> Self {
33        Self {
34            responses,
35            pauses: HashMap::new(),
36            call_count: AtomicUsize::new(0),
37            captured_contexts: Arc::new(Mutex::new(Vec::new())),
38            display_name: "Fake LLM".to_string(),
39            model: None,
40            context_window: None,
41        }
42    }
43
44    pub fn with_display_name(mut self, name: &str) -> Self {
45        self.display_name = name.to_string();
46        self
47    }
48
49    pub fn with_model(mut self, model: LlmModel) -> Self {
50        self.model = Some(model);
51        self
52    }
53
54    pub fn with_context_window(mut self, window: Option<u32>) -> Self {
55        self.context_window = window;
56        self
57    }
58
59    /// Pause the stream for `turn_index` after emitting the chunk at `chunk_index`
60    /// until the returned `Notify` is notified. Enables deterministic mid-stream
61    /// tests (e.g. send a user message after the provider has emitted some text).
62    pub fn pause_turn_after(mut self, turn_index: usize, chunk_index: usize, notify: Arc<Notify>) -> Self {
63        self.pauses.entry(turn_index).or_default().insert(chunk_index, notify);
64        self
65    }
66
67    /// Returns a handle to the captured contexts that can be used to verify
68    /// what contexts were passed to the LLM.
69    pub fn captured_contexts(&self) -> Arc<Mutex<Vec<Context>>> {
70        Arc::clone(&self.captured_contexts)
71    }
72}
73
74impl StreamingModelProvider for FakeLlmProvider {
75    fn stream_response(&self, context: &Context) -> LlmResponseStream {
76        if let Ok(mut contexts) = self.captured_contexts.lock() {
77            contexts.push(context.clone());
78        }
79
80        let current_call = self.call_count.fetch_add(1, Ordering::SeqCst);
81
82        let response = if current_call < self.responses.len() {
83            self.responses[current_call].clone()
84        } else {
85            vec![Ok(LlmResponse::done())]
86        };
87
88        let pauses = self.pauses.get(&current_call).cloned().unwrap_or_default();
89        if pauses.is_empty() {
90            return Box::pin(tokio_stream::iter(response));
91        }
92
93        let (tx, rx) = mpsc::unbounded_channel();
94        spawn(async move {
95            for (index, chunk) in response.into_iter().enumerate() {
96                if tx.send(chunk).is_err() {
97                    return;
98                }
99                if let Some(notify) = pauses.get(&index) {
100                    notify.notified().await;
101                }
102            }
103        });
104        Box::pin(UnboundedReceiverStream::new(rx))
105    }
106
107    fn display_name(&self) -> String {
108        self.display_name.clone()
109    }
110
111    fn context_window(&self) -> Option<u32> {
112        self.context_window
113    }
114
115    fn model(&self) -> Option<LlmModel> {
116        self.model.clone()
117    }
118}
llm/testing/fake_llm.rs

llm/testing/
fake_llm.rs
