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

vv_agent/llm/
scripted.rs
