Skip to main content

ai_agents_runtime/spec/
llm.rs

1//! LLM configuration types
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6#[derive(Debug, Clone, Serialize, Deserialize, Default)]
7pub struct CliMetadata {
8    #[serde(default, skip_serializing_if = "Option::is_none")]
9    pub welcome: Option<String>,
10
11    #[serde(default, skip_serializing_if = "Vec::is_empty")]
12    pub hints: Vec<String>,
13
14    #[serde(default, skip_serializing_if = "Option::is_none")]
15    pub show_tools: Option<bool>,
16
17    #[serde(default, skip_serializing_if = "Option::is_none")]
18    pub show_state: Option<bool>,
19
20    #[serde(default, skip_serializing_if = "Option::is_none")]
21    pub show_timing: Option<bool>,
22
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub streaming: Option<bool>,
25
26    #[serde(default, skip_serializing_if = "Option::is_none")]
27    pub prompt_style: Option<CliPromptStyle>,
28
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub disable_builtin_commands: Option<bool>,
31
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub hitl: Option<CliHitlMetadata>,
34
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub theme: Option<String>,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
40#[serde(rename_all = "snake_case")]
41pub enum CliPromptStyle {
42    Simple,
43    WithState,
44}
45
46/// Controls how the CLI handles HITL approval requests at runtime.
47#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
48#[serde(rename_all = "snake_case")]
49pub enum CliHitlStyle {
50    /// Interactive y/N prompt in the terminal (default).
51    #[default]
52    Prompt,
53    /// Silently approve all requests.
54    AutoApprove,
55    /// Silently reject all requests.
56    AutoReject,
57}
58
59/// CLI-specific HITL display settings from `metadata.cli.hitl`.
60#[derive(Debug, Clone, Serialize, Deserialize, Default)]
61pub struct CliHitlMetadata {
62    #[serde(default, skip_serializing_if = "Option::is_none")]
63    pub style: Option<CliHitlStyle>,
64
65    #[serde(default, skip_serializing_if = "Option::is_none")]
66    pub show_context: Option<bool>,
67}
68
69/// Configuration for LLM provider
70#[derive(Debug, Clone, Serialize, Deserialize)]
71pub struct LLMConfig {
72    pub provider: String,
73
74    pub model: String,
75
76    #[serde(default = "default_temperature")]
77    pub temperature: f32,
78
79    #[serde(default = "default_max_tokens")]
80    pub max_tokens: u32,
81
82    #[serde(skip_serializing_if = "Option::is_none")]
83    pub top_p: Option<f32>,
84
85    /// Base URL for the LLM provider API.
86    /// Required for `openai-compatible`; optional override for other providers.
87    #[serde(default, skip_serializing_if = "Option::is_none")]
88    pub base_url: Option<String>,
89
90    /// Environment variable name containing the API key.
91    /// Overrides the provider's default env var (e.g. OPENAI_API_KEY).
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub api_key_env: Option<String>,
94
95    /// Request timeout in seconds.
96    #[serde(default, skip_serializing_if = "Option::is_none")]
97    pub timeout_seconds: Option<u64>,
98
99    /// Enable extended thinking / reasoning mode.
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub reasoning: Option<bool>,
102
103    /// Reasoning effort level: "low", "medium", or "high".
104    #[serde(default, skip_serializing_if = "Option::is_none")]
105    pub reasoning_effort: Option<String>,
106
107    /// Maximum token budget for reasoning.
108    #[serde(default, skip_serializing_if = "Option::is_none")]
109    pub reasoning_budget_tokens: Option<u32>,
110
111    /// Override whether the provider supports function/tool calling.
112    #[serde(default, skip_serializing_if = "Option::is_none")]
113    pub function_calling: Option<bool>,
114
115    /// Override whether the provider supports vision inputs.
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub vision: Option<bool>,
118
119    /// Override whether the provider supports JSON mode.
120    #[serde(default, skip_serializing_if = "Option::is_none")]
121    pub json_mode: Option<bool>,
122
123    /// Additional provider-specific configuration
124    #[serde(flatten)]
125    pub extra: HashMap<String, serde_json::Value>,
126}
127
128fn default_temperature() -> f32 {
129    0.7
130}
131
132fn default_max_tokens() -> u32 {
133    2000
134}
135
136impl Default for LLMConfig {
137    fn default() -> Self {
138        Self {
139            provider: "openai".to_string(),
140            model: "gpt-4".to_string(),
141            temperature: default_temperature(),
142            max_tokens: default_max_tokens(),
143            top_p: None,
144            base_url: None,
145            api_key_env: None,
146            timeout_seconds: None,
147            reasoning: None,
148            reasoning_effort: None,
149            reasoning_budget_tokens: None,
150            function_calling: None,
151            vision: None,
152            json_mode: None,
153            extra: HashMap::new(),
154        }
155    }
156}
157
158#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct LLMSelector {
160    #[serde(default = "default_alias")]
161    pub default: String,
162    #[serde(default)]
163    pub router: Option<String>,
164}
165
166fn default_alias() -> String {
167    "default".to_string()
168}
169
170impl Default for LLMSelector {
171    fn default() -> Self {
172        Self {
173            default: default_alias(),
174            router: None,
175        }
176    }
177}
178
179impl LLMSelector {
180    pub fn new(default: impl Into<String>) -> Self {
181        Self {
182            default: default.into(),
183            router: None,
184        }
185    }
186
187    pub fn with_router(mut self, router: impl Into<String>) -> Self {
188        self.router = Some(router.into());
189        self
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn test_cli_metadata_deserialize() {
199        let yaml = r#"
200welcome: "=== Demo ==="
201hints:
202  - "Try: hello"
203  - "Try: help"
204show_tools: true
205show_state: false
206show_timing: true
207streaming: true
208prompt_style: with_state
209disable_builtin_commands: false
210"#;
211        let metadata: CliMetadata = serde_yaml::from_str(yaml).unwrap();
212        assert_eq!(metadata.welcome.as_deref(), Some("=== Demo ==="));
213        assert_eq!(metadata.hints.len(), 2);
214        assert_eq!(metadata.show_tools, Some(true));
215        assert_eq!(metadata.show_state, Some(false));
216        assert_eq!(metadata.show_timing, Some(true));
217        assert_eq!(metadata.streaming, Some(true));
218        assert_eq!(metadata.prompt_style, Some(CliPromptStyle::WithState));
219        assert_eq!(metadata.disable_builtin_commands, Some(false));
220        assert!(metadata.hitl.is_none());
221    }
222
223    #[test]
224    fn test_llm_config_default() {
225        let config = LLMConfig::default();
226        assert_eq!(config.provider, "openai");
227        assert_eq!(config.model, "gpt-4");
228        assert_eq!(config.temperature, 0.7);
229        assert_eq!(config.max_tokens, 2000);
230        assert_eq!(config.base_url, None);
231        assert_eq!(config.api_key_env, None);
232    }
233
234    #[test]
235    fn test_llm_config_with_base_url() {
236        let yaml = r#"
237provider: openai-compatible
238model: llama3.2
239base_url: http://localhost:1234/v1
240"#;
241        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
242        assert_eq!(config.provider, "openai-compatible");
243        assert_eq!(
244            config.base_url,
245            Some("http://localhost:1234/v1".to_string())
246        );
247    }
248
249    #[test]
250    fn test_llm_config_with_api_key_env() {
251        let yaml = r#"
252provider: openai-compatible
253model: my-model
254base_url: http://my-server:8080/v1
255api_key_env: MY_SERVER_KEY
256"#;
257        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
258        assert_eq!(config.api_key_env, Some("MY_SERVER_KEY".to_string()));
259    }
260
261    #[test]
262    fn test_llm_config_base_url_does_not_leak_to_extra() {
263        let yaml = r#"
264provider: openai-compatible
265model: my-model
266base_url: http://localhost:1234/v1
267"#;
268        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
269        assert!(!config.extra.contains_key("base_url"));
270    }
271
272    #[test]
273    fn test_llm_config_deserialize() {
274        let yaml = r#"
275provider: openai
276model: gpt-3.5-turbo
277temperature: 0.5
278max_tokens: 1000
279"#;
280        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
281        assert_eq!(config.provider, "openai");
282        assert_eq!(config.model, "gpt-3.5-turbo");
283        assert_eq!(config.temperature, 0.5);
284        assert_eq!(config.max_tokens, 1000);
285    }
286
287    #[test]
288    fn test_llm_config_with_defaults() {
289        let yaml = r#"
290provider: openai
291model: gpt-4
292"#;
293        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
294        assert_eq!(config.temperature, 0.7); // default
295        assert_eq!(config.max_tokens, 2000); // default
296    }
297
298    #[test]
299    fn test_llm_config_extra_fields() {
300        let yaml = r#"
301provider: openai
302model: gpt-4
303custom_field: "value"
304another_field: 123
305"#;
306        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
307        assert!(config.extra.contains_key("custom_field"));
308        assert!(config.extra.contains_key("another_field"));
309    }
310
311    #[test]
312    fn test_llm_selector_default() {
313        let selector = LLMSelector::default();
314        assert_eq!(selector.default, "default");
315        assert!(selector.router.is_none());
316    }
317
318    #[test]
319    fn test_llm_selector_with_router() {
320        let selector = LLMSelector::new("main").with_router("cheap");
321        assert_eq!(selector.default, "main");
322        assert_eq!(selector.router, Some("cheap".to_string()));
323    }
324
325    #[test]
326    fn test_cli_hitl_metadata_deserialize() {
327        let yaml = r#"
328style: auto_approve
329show_context: false
330"#;
331        let meta: CliHitlMetadata = serde_yaml::from_str(yaml).unwrap();
332        assert_eq!(meta.style, Some(CliHitlStyle::AutoApprove));
333        assert_eq!(meta.show_context, Some(false));
334    }
335
336    #[test]
337    fn test_cli_hitl_style_default() {
338        assert_eq!(CliHitlStyle::default(), CliHitlStyle::Prompt);
339    }
340
341    #[test]
342    fn test_cli_metadata_with_hitl() {
343        let yaml = r#"
344welcome: "Hello"
345hints: []
346hitl:
347  style: prompt
348  show_context: true
349"#;
350        let meta: CliMetadata = serde_yaml::from_str(yaml).unwrap();
351        let hitl = meta.hitl.unwrap();
352        assert_eq!(hitl.style, Some(CliHitlStyle::Prompt));
353        assert_eq!(hitl.show_context, Some(true));
354    }
355
356    #[test]
357    fn test_llm_selector_deserialize() {
358        let yaml = r#"
359default: main
360router: router_llm
361"#;
362        let selector: LLMSelector = serde_yaml::from_str(yaml).unwrap();
363        assert_eq!(selector.default, "main");
364        assert_eq!(selector.router, Some("router_llm".to_string()));
365    }
366
367    #[test]
368    fn test_llm_config_reasoning_fields_deser() {
369        let yaml = r#"
370provider: openai
371model: o3
372timeout_seconds: 120
373reasoning: true
374reasoning_effort: high
375reasoning_budget_tokens: 16384
376"#;
377        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
378        assert_eq!(config.timeout_seconds, Some(120));
379        assert_eq!(config.reasoning, Some(true));
380        assert_eq!(config.reasoning_effort.as_deref(), Some("high"));
381        assert_eq!(config.reasoning_budget_tokens, Some(16384));
382        // Must NOT leak into extra
383        assert!(!config.extra.contains_key("timeout_seconds"));
384        assert!(!config.extra.contains_key("reasoning"));
385        assert!(!config.extra.contains_key("reasoning_effort"));
386        assert!(!config.extra.contains_key("reasoning_budget_tokens"));
387    }
388
389    #[test]
390    fn test_ollama_named_fields_land_in_extra() {
391        let yaml = r#"
392provider: ollama
393model: llama3.1
394num_ctx: 8192
395num_gpu: -1
396keep_alive: 5m
397"#;
398        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
399        assert_eq!(config.extra.get("num_ctx"), Some(&serde_json::json!(8192)));
400        assert_eq!(config.extra.get("num_gpu"), Some(&serde_json::json!(-1)));
401        assert_eq!(
402            config.extra.get("keep_alive"),
403            Some(&serde_json::json!("5m"))
404        );
405    }
406
407    #[test]
408    fn test_llm_config_feature_override_fields_deser() {
409        let yaml = r#"
410provider: openai-compatible
411model: qwen3:8b
412base_url: http://localhost:11434/v1
413function_calling: true
414vision: false
415json_mode: true
416"#;
417        let config: LLMConfig = serde_yaml::from_str(yaml).unwrap();
418        assert_eq!(config.function_calling, Some(true));
419        assert_eq!(config.vision, Some(false));
420        assert_eq!(config.json_mode, Some(true));
421        assert!(!config.extra.contains_key("function_calling"));
422        assert!(!config.extra.contains_key("vision"));
423        assert!(!config.extra.contains_key("json_mode"));
424    }
425}