Skip to main content

vv_agent/llm/
base.rs

1use std::collections::BTreeMap;
2use std::path::Path;
3use std::sync::Arc;
4
5use serde_json::Value;
6use thiserror::Error;
7
8use crate::memory::CompactionExhaustedError;
9use crate::model_settings::ModelSettings;
10use crate::prompt::PromptBundle;
11use crate::types::{LLMResponse, Message};
12
13pub type LlmStreamCallback = Arc<dyn Fn(&BTreeMap<String, Value>) + Send + Sync + 'static>;
14
15#[derive(Debug, Clone, PartialEq)]
16pub struct LlmRequest {
17    pub model: String,
18    pub messages: Vec<Message>,
19    pub tools: Vec<Value>,
20    pub metadata: Value,
21    pub model_settings: Option<ModelSettings>,
22    pub prompt_bundle: PromptBundle,
23}
24
25impl LlmRequest {
26    pub fn new(
27        model: impl Into<String>,
28        messages: Vec<Message>,
29        prompt_bundle: PromptBundle,
30    ) -> Self {
31        Self {
32            model: model.into(),
33            messages,
34            tools: Vec::new(),
35            metadata: Value::Null,
36            model_settings: None,
37            prompt_bundle,
38        }
39    }
40}
41
42#[derive(Debug, Error)]
43pub enum LlmError {
44    #[error("scripted response queue is empty")]
45    ScriptExhausted,
46    #[error("{0}")]
47    CompactionExhausted(CompactionExhaustedError),
48    #[error("llm request failed: {0}")]
49    Request(String),
50}
51
52pub trait LlmClient: Send + Sync {
53    fn complete(&self, request: LlmRequest) -> Result<LLMResponse, LlmError>;
54
55    fn set_debug_dump_dir(&self, _debug_dump_dir: &Path) {}
56
57    fn clone_with_debug_dump_dir(&self, _debug_dump_dir: &Path) -> Option<Arc<dyn LlmClient>> {
58        None
59    }
60
61    fn complete_with_stream(
62        &self,
63        request: LlmRequest,
64        stream_callback: Option<LlmStreamCallback>,
65    ) -> Result<LLMResponse, LlmError> {
66        let _ = stream_callback;
67        self.complete(request)
68    }
69}
70
71impl<T> LlmClient for Arc<T>
72where
73    T: LlmClient + ?Sized,
74{
75    fn complete(&self, request: LlmRequest) -> Result<LLMResponse, LlmError> {
76        (**self).complete(request)
77    }
78
79    fn set_debug_dump_dir(&self, debug_dump_dir: &Path) {
80        (**self).set_debug_dump_dir(debug_dump_dir);
81    }
82
83    fn clone_with_debug_dump_dir(&self, debug_dump_dir: &Path) -> Option<Arc<dyn LlmClient>> {
84        (**self).clone_with_debug_dump_dir(debug_dump_dir)
85    }
86
87    fn complete_with_stream(
88        &self,
89        request: LlmRequest,
90        stream_callback: Option<LlmStreamCallback>,
91    ) -> Result<LLMResponse, LlmError> {
92        (**self).complete_with_stream(request, stream_callback)
93    }
94}
95
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct EndpointTarget {
98    pub endpoint_id: String,
99    pub api_key: String,
100    pub api_base: String,
101    pub endpoint_type: String,
102    pub model_id: String,
103}
104
105impl EndpointTarget {
106    pub fn new(
107        endpoint_id: impl Into<String>,
108        api_key: impl Into<String>,
109        api_base: impl Into<String>,
110        endpoint_type: impl Into<String>,
111        model_id: impl Into<String>,
112    ) -> Self {
113        Self {
114            endpoint_id: endpoint_id.into(),
115            api_key: api_key.into(),
116            api_base: api_base.into(),
117            endpoint_type: endpoint_type.into(),
118            model_id: model_id.into(),
119        }
120    }
121}