sova_ai/
fake.rs

1//! In-memory fake [`LanguageModel`] for tests.
2
3use aisdk::core::capabilities::{
4    ReasoningSupport, StructuredOutputSupport, TextInputSupport, TextOutputSupport,
5    ToolCallSupport,
6};
7use aisdk::core::language_model::{
8    LanguageModel, LanguageModelOptions, LanguageModelResponse, LanguageModelStreamChunk,
9    LanguageModelStreamChunkType,
10};
11use aisdk::core::Message;
12use aisdk::Result as AisdkResult;
13use async_trait::async_trait;
14use futures_util::stream;
15use std::pin::Pin;
16use std::sync::{Arc, Mutex};
17
18/// Recorded call snapshot (system + user prompts extracted from options).
19#[derive(Debug, Clone)]
20pub struct FakeCall {
21    /// System prompt, if any.
22    pub system: Option<String>,
23    /// Flattened user/text prompts from the conversation.
24    pub prompts: Vec<String>,
25}
26
27#[derive(Debug, Default)]
28struct FakeInner {
29    stubs: Vec<String>,
30    default: String,
31    calls: Vec<FakeCall>,
32}
33
34/// Laravel-style fake model: stub replies, assert prompts.
35#[derive(Clone, Debug, Default)]
36pub struct FakeAi {
37    inner: Arc<Mutex<FakeInner>>,
38}
39
40impl FakeAi {
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    /// Queue a one-shot stub (FIFO). Falls back to [`Self::stub_text`] default.
46    pub fn push_text(self, text: impl Into<String>) -> Self {
47        self.inner.lock().unwrap().stubs.push(text.into());
48        self
49    }
50
51    /// Default reply when the stub queue is empty (also sets initial queue empty).
52    pub fn stub_text(self, text: impl Into<String>) -> Self {
53        self.inner.lock().unwrap().default = text.into();
54        self
55    }
56
57    fn next_text(&self) -> String {
58        let mut g = self.inner.lock().unwrap();
59        if !g.stubs.is_empty() {
60            g.stubs.remove(0)
61        } else {
62            g.default.clone()
63        }
64    }
65
66    fn record(&self, options: &LanguageModelOptions) {
67        let prompts = options
68            .messages()
69            .into_iter()
70            .filter_map(|m| match m {
71                Message::User(u) => Some(u.content),
72                Message::System(s) => Some(s.content),
73                _ => None,
74            })
75            .collect();
76        self.inner.lock().unwrap().calls.push(FakeCall {
77            system: options.system.clone(),
78            prompts,
79        });
80    }
81
82    /// All recorded calls (oldest first).
83    pub fn calls(&self) -> Vec<FakeCall> {
84        self.inner.lock().unwrap().calls.clone()
85    }
86
87    /// Flattened user/system prompt strings across calls.
88    pub fn prompts(&self) -> Vec<String> {
89        self.calls()
90            .into_iter()
91            .flat_map(|c| c.prompts)
92            .collect()
93    }
94
95    pub fn call_count(&self) -> usize {
96        self.inner.lock().unwrap().calls.len()
97    }
98
99    pub fn assert_called(&self) {
100        assert!(
101            self.call_count() > 0,
102            "FakeAi: expected at least one generate/stream call"
103        );
104    }
105
106    pub fn assert_called_times(&self, n: usize) {
107        assert_eq!(
108            self.call_count(),
109            n,
110            "FakeAi: expected {n} calls, got {}",
111            self.call_count()
112        );
113    }
114}
115
116impl TextInputSupport for FakeAi {}
117impl TextOutputSupport for FakeAi {}
118impl ToolCallSupport for FakeAi {}
119impl StructuredOutputSupport for FakeAi {}
120impl ReasoningSupport for FakeAi {}
121
122#[async_trait]
123impl LanguageModel for FakeAi {
124    fn name(&self) -> String {
125        "fake".into()
126    }
127
128    async fn generate_text(
129        &mut self,
130        options: LanguageModelOptions,
131    ) -> AisdkResult<LanguageModelResponse> {
132        self.record(&options);
133        Ok(LanguageModelResponse::new(self.next_text()))
134    }
135
136    async fn stream_text(
137        &mut self,
138        options: LanguageModelOptions,
139    ) -> AisdkResult<
140        Pin<Box<dyn futures_util::Stream<Item = AisdkResult<Vec<LanguageModelStreamChunk>>> + Send>>,
141    > {
142        self.record(&options);
143        let text = self.next_text();
144        let chunks = vec![
145            LanguageModelStreamChunk::Delta(LanguageModelStreamChunkType::Start),
146            LanguageModelStreamChunk::Delta(LanguageModelStreamChunkType::Text(text)),
147            LanguageModelStreamChunk::Delta(LanguageModelStreamChunkType::End(Default::default())),
148        ];
149        Ok(Box::pin(stream::once(async move { Ok(chunks) })))
150    }
151}
sova_ai/fake.rs

sova_ai/
fake.rs
