use std::sync::Mutex;
use super::*;
use crate::llm::Embedding;
#[derive(Default)]
pub struct MockProvider {
pub chat_calls: Mutex<Vec<Vec<ChatMessage>>>,
pub chat_replies: Mutex<Vec<Result<ChatResponse, LlmError>>>,
pub embed_calls: Mutex<Vec<Vec<String>>>,
pub embed_replies: Mutex<Vec<Result<EmbedResult, LlmError>>>,
}
impl MockProvider {
pub fn new() -> Self {
Self {
chat_calls: Mutex::new(Vec::new()),
chat_replies: Mutex::new(Vec::new()),
embed_calls: Mutex::new(Vec::new()),
embed_replies: Mutex::new(Vec::new()),
}
}
pub fn push_chat(&self, reply: Result<ChatResponse, LlmError>) {
self.chat_replies.lock().unwrap().push(reply);
}
pub fn push_embed(&self, reply: Result<EmbedResult, LlmError>) {
self.embed_replies.lock().unwrap().push(reply);
}
pub fn chat_call_count(&self) -> usize {
self.chat_calls.lock().unwrap().len()
}
pub fn embed_call_count(&self) -> usize {
self.embed_calls.lock().unwrap().len()
}
pub fn last_chat_call(&self) -> Option<Vec<ChatMessage>> {
self.chat_calls.lock().unwrap().last().cloned()
}
pub fn canned_embed(dim: usize, count: usize, model: &str) -> EmbedResult {
EmbedResult {
model: model.to_string(),
embeddings: (0..count).map(|_| Embedding(vec![0.1; dim])).collect(),
}
}
}
impl Provider for MockProvider {
fn name(&self) -> &str {
"mock"
}
fn chat<'a>(
&'a self,
messages: &'a [ChatMessage],
) -> Pin<Box<dyn Future<Output = Result<ChatResponse, LlmError>> + Send + 'a>> {
self.chat_calls.lock().unwrap().push(messages.to_vec());
let reply = self.chat_replies.lock().unwrap().pop().unwrap_or_else(|| {
Ok(ChatResponse {
model: "mock-default".into(),
content: "[mock-default reply]".into(),
})
});
Box::pin(async move { reply })
}
}
impl EmbeddingProvider for MockProvider {
fn embed<'a>(
&'a self,
texts: &'a [String],
) -> Pin<Box<dyn Future<Output = Result<EmbedResult, LlmError>> + Send + 'a>> {
self.embed_calls.lock().unwrap().push(texts.to_vec());
let reply = self.embed_replies.lock().unwrap().pop().unwrap_or_else(|| {
Ok(EmbedResult {
model: "mock-default".into(),
embeddings: texts.iter().map(|_| Embedding(vec![0.0; 4])).collect(),
})
});
Box::pin(async move { reply })
}
}