1use super::compactor::{LLMCompactorConfig, ToolCompaction};
4
5#[derive(Debug, Clone, PartialEq)]
7pub enum LLMProvider {
8 Anthropic,
9 OpenAI,
10 Google,
11 Cohere,
12 Bedrock,
13}
14
15#[derive(Debug, Clone)]
17pub struct CompactionConfig {
18 pub threshold: f64,
21 pub keep_recent_turns: usize,
23 pub tool_compaction: ToolCompaction,
25}
26
27impl Default for CompactionConfig {
28 fn default() -> Self {
29 Self {
30 threshold: 0.75,
31 keep_recent_turns: 5,
32 tool_compaction: ToolCompaction::Summarize,
33 }
34 }
35}
36
37#[derive(Debug, Clone)]
39pub enum CompactorType {
40 Threshold(CompactionConfig),
42 LLM(LLMCompactorConfig),
44}
45
46impl Default for CompactorType {
47 fn default() -> Self {
48 CompactorType::Threshold(CompactionConfig::default())
49 }
50}
51
52#[derive(Debug, Clone)]
54pub struct LLMSessionConfig {
55 pub provider: LLMProvider,
57 pub api_key: String,
59 pub model: String,
61 pub base_url: Option<String>,
64 pub max_tokens: Option<u32>,
66 pub system_prompt: Option<String>,
68 pub temperature: Option<f32>,
70 pub streaming: bool,
72 pub context_limit: i32,
74 pub compaction: Option<CompactorType>,
76 pub azure_resource: Option<String>,
79 pub azure_deployment: Option<String>,
81 pub azure_api_version: Option<String>,
83 pub bedrock_region: Option<String>,
85 pub bedrock_access_key_id: Option<String>,
87 pub bedrock_secret_access_key: Option<String>,
89 pub bedrock_session_token: Option<String>,
91}
92
93impl LLMSessionConfig {
94 pub fn anthropic(api_key: impl Into<String>, model: impl Into<String>) -> Self {
96 Self {
97 provider: LLMProvider::Anthropic,
98 api_key: api_key.into(),
99 model: model.into(),
100 base_url: None,
101 max_tokens: Some(4096),
102 system_prompt: None,
103 temperature: None,
104 streaming: true,
105 context_limit: 200_000, compaction: Some(CompactorType::default()),
107 azure_resource: None,
108 azure_deployment: None,
109 azure_api_version: None,
110 bedrock_region: None,
111 bedrock_access_key_id: None,
112 bedrock_secret_access_key: None,
113 bedrock_session_token: None,
114 }
115 }
116
117 pub fn openai(api_key: impl Into<String>, model: impl Into<String>) -> Self {
119 Self {
120 provider: LLMProvider::OpenAI,
121 api_key: api_key.into(),
122 model: model.into(),
123 base_url: None,
124 max_tokens: Some(4096),
125 system_prompt: None,
126 temperature: None,
127 streaming: true,
128 context_limit: 128_000, compaction: Some(CompactorType::default()),
130 azure_resource: None,
131 azure_deployment: None,
132 azure_api_version: None,
133 bedrock_region: None,
134 bedrock_access_key_id: None,
135 bedrock_secret_access_key: None,
136 bedrock_session_token: None,
137 }
138 }
139
140 pub fn openai_compatible(
145 api_key: impl Into<String>,
146 model: impl Into<String>,
147 base_url: impl Into<String>,
148 context_limit: i32,
149 ) -> Self {
150 Self {
151 provider: LLMProvider::OpenAI,
152 api_key: api_key.into(),
153 model: model.into(),
154 base_url: Some(base_url.into()),
155 max_tokens: Some(4096),
156 system_prompt: None,
157 temperature: None,
158 streaming: true,
159 context_limit,
160 compaction: Some(CompactorType::default()),
161 azure_resource: None,
162 azure_deployment: None,
163 azure_api_version: None,
164 bedrock_region: None,
165 bedrock_access_key_id: None,
166 bedrock_secret_access_key: None,
167 bedrock_session_token: None,
168 }
169 }
170
171 pub fn google(api_key: impl Into<String>, model: impl Into<String>) -> Self {
173 Self {
174 provider: LLMProvider::Google,
175 api_key: api_key.into(),
176 model: model.into(),
177 base_url: None,
178 max_tokens: Some(4096),
179 system_prompt: None,
180 temperature: None,
181 streaming: true,
182 context_limit: 1_000_000, compaction: Some(CompactorType::default()),
184 azure_resource: None,
185 azure_deployment: None,
186 azure_api_version: None,
187 bedrock_region: None,
188 bedrock_access_key_id: None,
189 bedrock_secret_access_key: None,
190 bedrock_session_token: None,
191 }
192 }
193
194 pub fn azure_openai(
204 api_key: impl Into<String>,
205 resource: impl Into<String>,
206 deployment: impl Into<String>,
207 ) -> Self {
208 Self {
209 provider: LLMProvider::OpenAI,
210 api_key: api_key.into(),
211 model: String::new(), base_url: None,
213 max_tokens: Some(4096),
214 system_prompt: None,
215 temperature: None,
216 streaming: true,
217 context_limit: 128_000, compaction: Some(CompactorType::default()),
219 azure_resource: Some(resource.into()),
220 azure_deployment: Some(deployment.into()),
221 azure_api_version: Some("2024-10-21".to_string()), bedrock_region: None,
223 bedrock_access_key_id: None,
224 bedrock_secret_access_key: None,
225 bedrock_session_token: None,
226 }
227 }
228
229 pub fn with_azure_api_version(mut self, version: impl Into<String>) -> Self {
231 self.azure_api_version = Some(version.into());
232 self
233 }
234
235 pub fn cohere(api_key: impl Into<String>, model: impl Into<String>) -> Self {
237 Self {
238 provider: LLMProvider::Cohere,
239 api_key: api_key.into(),
240 model: model.into(),
241 base_url: None,
242 max_tokens: Some(4096),
243 system_prompt: None,
244 temperature: None,
245 streaming: true,
246 context_limit: 128_000, compaction: Some(CompactorType::default()),
248 azure_resource: None,
249 azure_deployment: None,
250 azure_api_version: None,
251 bedrock_region: None,
252 bedrock_access_key_id: None,
253 bedrock_secret_access_key: None,
254 bedrock_session_token: None,
255 }
256 }
257
258 pub fn bedrock(
266 access_key_id: impl Into<String>,
267 secret_access_key: impl Into<String>,
268 region: impl Into<String>,
269 model: impl Into<String>,
270 ) -> Self {
271 Self {
272 provider: LLMProvider::Bedrock,
273 api_key: String::new(), model: model.into(),
275 base_url: None,
276 max_tokens: Some(4096),
277 system_prompt: None,
278 temperature: None,
279 streaming: true,
280 context_limit: 200_000, compaction: Some(CompactorType::default()),
282 azure_resource: None,
283 azure_deployment: None,
284 azure_api_version: None,
285 bedrock_region: Some(region.into()),
286 bedrock_access_key_id: Some(access_key_id.into()),
287 bedrock_secret_access_key: Some(secret_access_key.into()),
288 bedrock_session_token: None,
289 }
290 }
291
292 pub fn with_bedrock_session_token(mut self, token: impl Into<String>) -> Self {
294 self.bedrock_session_token = Some(token.into());
295 self
296 }
297
298 pub fn with_streaming(mut self, streaming: bool) -> Self {
300 self.streaming = streaming;
301 self
302 }
303
304 pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
306 self.max_tokens = Some(max_tokens);
307 self
308 }
309
310 pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
312 self.system_prompt = Some(prompt.into());
313 self
314 }
315
316 pub fn with_temperature(mut self, temperature: f32) -> Self {
318 self.temperature = Some(temperature);
319 self
320 }
321
322 pub fn with_context_limit(mut self, context_limit: i32) -> Self {
324 self.context_limit = context_limit;
325 self
326 }
327
328 pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
330 self.base_url = Some(base_url.into());
331 self
332 }
333
334 pub fn with_threshold_compaction(mut self, config: CompactionConfig) -> Self {
336 self.compaction = Some(CompactorType::Threshold(config));
337 self
338 }
339
340 pub fn with_llm_compaction(mut self, config: LLMCompactorConfig) -> Self {
342 self.compaction = Some(CompactorType::LLM(config));
343 self
344 }
345
346 pub fn with_compaction(mut self, compactor_type: CompactorType) -> Self {
348 self.compaction = Some(compactor_type);
349 self
350 }
351
352 pub fn without_compaction(mut self) -> Self {
354 self.compaction = None;
355 self
356 }
357}
358
359#[cfg(test)]
360mod tests {
361 use super::*;
362
363 #[test]
364 fn test_anthropic_config() {
365 let config = LLMSessionConfig::anthropic("test-key", "claude-3-sonnet")
366 .with_max_tokens(2048)
367 .with_system_prompt("You are helpful.");
368
369 assert_eq!(config.provider, LLMProvider::Anthropic);
370 assert_eq!(config.api_key, "test-key");
371 assert_eq!(config.model, "claude-3-sonnet");
372 assert_eq!(config.max_tokens, Some(2048));
373 assert_eq!(config.system_prompt, Some("You are helpful.".to_string()));
374 }
375
376 #[test]
377 fn test_openai_config() {
378 let config = LLMSessionConfig::openai("test-key", "gpt-4");
379
380 assert_eq!(config.provider, LLMProvider::OpenAI);
381 assert_eq!(config.model, "gpt-4");
382 }
383}