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}