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