telos_agent/model/
mock.rs1use async_trait::async_trait;
10use std::collections::VecDeque;
11use tokio::sync::Mutex;
12
13use crate::error::AgentError;
14use crate::model::provider::{CompletionRequest, CompletionResponse, ModelProvider};
15
16pub struct MockProvider {
18 responses: Mutex<VecDeque<CompletionResponse>>,
19 pub requests: Mutex<Vec<CompletionRequest>>,
21}
22
23impl MockProvider {
24 pub fn new(responses: Vec<CompletionResponse>) -> Self {
26 Self { responses: Mutex::new(responses.into()), requests: Mutex::new(Vec::new()) }
27 }
28}
29
30#[async_trait]
31impl ModelProvider for MockProvider {
32 async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, AgentError> {
33 self.requests.lock().await.push(request);
34 self.responses.lock().await.pop_front().ok_or_else(|| {
35 AgentError::Provider(crate::error::ProviderError::Other(
36 "mock provider has no more responses".into(),
37 ))
38 })
39 }
40}