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