vv_agent/llm/
scripted.rs

1use std::collections::VecDeque;
2use std::sync::{Arc, Mutex};
3
4use crate::types::LLMResponse;
5
6use super::{LlmClient, LlmError, LlmRequest};
7
8pub type ScriptStepCallback =
9    Arc<dyn Fn(&LlmRequest) -> Result<LLMResponse, LlmError> + Send + Sync + 'static>;
10
11#[derive(Clone)]
12// Boxing Response would break direct construction of this public test API.
13#[allow(clippy::large_enum_variant)]
14pub enum ScriptStep {
15    Response(LLMResponse),
16    Callback(ScriptStepCallback),
17}
18
19impl ScriptStep {
20    pub fn response(response: LLMResponse) -> Self {
21        Self::Response(response)
22    }
23
24    pub fn callback(
25        callback: impl Fn(&LlmRequest) -> Result<LLMResponse, LlmError> + Send + Sync + 'static,
26    ) -> Self {
27        Self::Callback(Arc::new(callback))
28    }
29}
30
31impl From<LLMResponse> for ScriptStep {
32    fn from(response: LLMResponse) -> Self {
33        Self::Response(response)
34    }
35}
36
37#[derive(Clone)]
38pub struct ScriptedLlmClient {
39    steps: Arc<Mutex<VecDeque<ScriptStep>>>,
40}
41
42impl ScriptedLlmClient {
43    pub fn new(responses: Vec<LLMResponse>) -> Self {
44        Self::from_steps(responses.into_iter().map(ScriptStep::from).collect())
45    }
46
47    pub fn from_steps(steps: Vec<ScriptStep>) -> Self {
48        Self {
49            steps: Arc::new(Mutex::new(VecDeque::from(steps))),
50        }
51    }
52
53    pub fn push_response(&self, response: LLMResponse) {
54        self.push_step(ScriptStep::Response(response));
55    }
56
57    pub fn push_step(&self, step: ScriptStep) {
58        if let Ok(mut queue) = self.steps.lock() {
59            queue.push_back(step);
60        }
61    }
62}
63
64impl std::fmt::Debug for ScriptedLlmClient {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        f.debug_struct("ScriptedLlmClient").finish_non_exhaustive()
67    }
68}
69
70impl LlmClient for ScriptedLlmClient {
71    fn complete(&self, request: LlmRequest) -> Result<LLMResponse, LlmError> {
72        let mut queue = self
73            .steps
74            .lock()
75            .map_err(|_| LlmError::Request("scripted response queue poisoned".to_string()))?;
76        let Some(step) = queue.pop_front() else {
77            return Err(LlmError::ScriptExhausted);
78        };
79        drop(queue);
80        match step {
81            ScriptStep::Response(response) => Ok(response),
82            ScriptStep::Callback(callback) => callback(&request),
83        }
84    }
85}
vv_agent/llm/scripted.rs

vv_agent/llm/
scripted.rs
