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}