Skip to main content

agent_core_runtime/controller/session/
config.rs

1// Session configuration types
2
3use super::compactor::{LLMCompactorConfig, ToolCompaction};
4
5/// LLM provider type
6#[derive(Debug, Clone, PartialEq)]
7pub enum LLMProvider {
8    Anthropic,
9    OpenAI,
10    Google,
11    Cohere,
12    Bedrock,
13}
14
15/// Configuration for conversation compaction
16#[derive(Debug, Clone)]
17pub struct CompactionConfig {
18    /// Context utilization threshold (0.0-1.0) that triggers compaction.
19    /// For example, 0.75 means compact when 75% of context is used.
20    pub threshold: f64,
21    /// Number of recent turns to preserve during compaction.
22    pub keep_recent_turns: usize,
23    /// Strategy for handling old tool results.
24    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/// Type of compaction strategy to use.
38#[derive(Debug, Clone)]
39pub enum CompactorType {
40    /// Simple threshold-based compaction that summarizes/redacts old tool results.
41    Threshold(CompactionConfig),
42    /// LLM-based conversation summarization that uses an LLM to create intelligent summaries.
43    LLM(LLMCompactorConfig),
44}
45
46impl Default for CompactorType {
47    fn default() -> Self {
48        CompactorType::Threshold(CompactionConfig::default())
49    }
50}
51
52/// Configuration for creating an LLM session
53#[derive(Debug, Clone)]
54pub struct LLMSessionConfig {
55    /// The LLM provider to use
56    pub provider: LLMProvider,
57    /// API key for the provider
58    pub api_key: String,
59    /// Model to use (e.g., "claude-3-sonnet-20240229", "gpt-4")
60    pub model: String,
61    /// Custom base URL for OpenAI-compatible providers.
62    /// Only used when provider is OpenAI. If None, uses default OpenAI endpoint.
63    pub base_url: Option<String>,
64    /// Default maximum tokens for responses
65    pub max_tokens: Option<u32>,
66    /// Default system prompt
67    pub system_prompt: Option<String>,
68    /// Default temperature
69    pub temperature: Option<f32>,
70    /// Enable streaming responses
71    pub streaming: bool,
72    /// Model's context window size (for compaction decisions)
73    pub context_limit: i32,
74    /// Compaction configuration (None to disable compaction)
75    pub compaction: Option<CompactorType>,
76    /// Azure OpenAI resource name (e.g., "my-resource").
77    /// When set, the provider uses Azure OpenAI instead of standard OpenAI.
78    pub azure_resource: Option<String>,
79    /// Azure OpenAI deployment name (e.g., "gpt-4-deployment").
80    pub azure_deployment: Option<String>,
81    /// Azure OpenAI API version (e.g., "2024-10-21").
82    pub azure_api_version: Option<String>,
83    /// AWS region for Bedrock (e.g., "us-east-1").
84    pub bedrock_region: Option<String>,
85    /// AWS access key ID for Bedrock.
86    pub bedrock_access_key_id: Option<String>,
87    /// AWS secret access key for Bedrock.
88    pub bedrock_secret_access_key: Option<String>,
89    /// AWS session token for Bedrock (optional, for temporary credentials).
90    pub bedrock_session_token: Option<String>,
91}
92
93impl LLMSessionConfig {
94    /// Creates a new Anthropic session config
95    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, // Claude default
106            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    /// Creates a new OpenAI session config
118    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, // GPT-4 default
129            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    /// Creates a new OpenAI-compatible session config with a custom base URL.
141    ///
142    /// Use this for providers like Groq, Together, Fireworks, etc. that have
143    /// OpenAI-compatible APIs.
144    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    /// Creates a new Google (Gemini) session config
172    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, // Gemini 2.5 default
183            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    /// Creates a new Azure OpenAI session config.
195    ///
196    /// Azure OpenAI uses a different URL format and authentication method.
197    /// The endpoint is: https://{resource}.openai.azure.com/openai/deployments/{deployment}/chat/completions?api-version={version}
198    ///
199    /// # Arguments
200    /// * `api_key` - Azure OpenAI API key
201    /// * `resource` - Azure resource name (e.g., "my-openai-resource")
202    /// * `deployment` - Deployment name (e.g., "gpt-4-deployment")
203    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(), // Not used for Azure - deployment determines model
212            base_url: None,
213            max_tokens: Some(4096),
214            system_prompt: None,
215            temperature: None,
216            streaming: true,
217            context_limit: 128_000, // Azure OpenAI default
218            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()), // Latest stable version
222            bedrock_region: None,
223            bedrock_access_key_id: None,
224            bedrock_secret_access_key: None,
225            bedrock_session_token: None,
226        }
227    }
228
229    /// Sets the Azure API version.
230    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    /// Creates a new Cohere session config
236    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, // Command-R context limit
247            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    /// Creates a new Amazon Bedrock session config.
259    ///
260    /// # Arguments
261    /// * `access_key_id` - AWS access key ID
262    /// * `secret_access_key` - AWS secret access key
263    /// * `region` - AWS region (e.g., "us-east-1")
264    /// * `model` - Bedrock model ID (e.g., "anthropic.claude-3-sonnet-20240229-v1:0")
265    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(), // Not used for Bedrock
274            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, // Claude on Bedrock default
281            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    /// Sets the Bedrock session token for temporary credentials.
293    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    /// Enable or disable streaming
299    pub fn with_streaming(mut self, streaming: bool) -> Self {
300        self.streaming = streaming;
301        self
302    }
303
304    /// Sets the default max tokens
305    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
306        self.max_tokens = Some(max_tokens);
307        self
308    }
309
310    /// Sets the default system prompt
311    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
312        self.system_prompt = Some(prompt.into());
313        self
314    }
315
316    /// Sets the default temperature
317    pub fn with_temperature(mut self, temperature: f32) -> Self {
318        self.temperature = Some(temperature);
319        self
320    }
321
322    /// Sets the model's context window size
323    pub fn with_context_limit(mut self, context_limit: i32) -> Self {
324        self.context_limit = context_limit;
325        self
326    }
327
328    /// Sets a custom base URL for OpenAI-compatible providers
329    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    /// Enables threshold compaction with custom configuration
335    pub fn with_threshold_compaction(mut self, config: CompactionConfig) -> Self {
336        self.compaction = Some(CompactorType::Threshold(config));
337        self
338    }
339
340    /// Enables LLM-based compaction with custom configuration
341    pub fn with_llm_compaction(mut self, config: LLMCompactorConfig) -> Self {
342        self.compaction = Some(CompactorType::LLM(config));
343        self
344    }
345
346    /// Enables compaction with the specified compactor type
347    pub fn with_compaction(mut self, compactor_type: CompactorType) -> Self {
348        self.compaction = Some(compactor_type);
349        self
350    }
351
352    /// Disables compaction
353    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}