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}