From 7a7dffeb7a8bf5abd8db20e0ad9a0dbdb6895850 Mon Sep 17 00:00:00 2001 From: ajianaz Date: Tue, 7 Jul 2026 08:34:22 +0700 Subject: [PATCH] fix(llm): add conditional max_tokens parameter naming per provider LLM API providers use different parameter names for max output tokens: - OpenAI (legacy): max_tokens - OpenAI (new): max_completion_tokens - Anthropic: max_tokens - Google Gemini/Vertex: max_output_tokens cora-cli previously hardcoded 'max_tokens' everywhere, causing HTTP 400 errors ('Unsupported parameter: max_output_tokens') when providers don't accept it. Changes: - Add max_tokens_param to LlmSection config schema (supports 'auto' for provider-based detection, or explicit override) - Add resolve_max_tokens_param() in loader.rs for auto-detection logic - Replace hardcoded 'max_tokens' JSON key with dynamic key in both streaming and non-streaming LLM request paths - Add max_tokens_param to LLMConfig struct for runtime use - Display max_tokens_param in 'cora config' output Auto-detection maps: - gemini/google/vertex -> max_output_tokens - All other providers -> max_tokens Co-authored-by: Hermes Agent --- src/commands/config_cmd.rs | 3 ++ src/config/loader.rs | 74 +++++++++++++++++++++++++++++++++ src/config/schema.rs | 67 ++++++++++++++++++++++++++++++ src/engine/llm.rs | 83 ++++++++++++++++++++++++++++---------- src/engine/types.rs | 4 ++ 5 files changed, 210 insertions(+), 21 deletions(-) diff --git a/src/commands/config_cmd.rs b/src/commands/config_cmd.rs index 612bfea..4ccca69 100644 --- a/src/commands/config_cmd.rs +++ b/src/commands/config_cmd.rs @@ -326,6 +326,9 @@ fn print_cora_file(cora: &CoraFile) { if let Some(v) = llm.max_tokens { parts.push(format!("max_tokens={}", v)); } + if let Some(ref v) = llm.max_tokens_param { + parts.push(format!("max_tokens_param={}", v)); + } if let Some(v) = llm.timeout { parts.push(format!("timeout={}", v)); } diff --git a/src/config/loader.rs b/src/config/loader.rs index d402128..53543b0 100644 --- a/src/config/loader.rs +++ b/src/config/loader.rs @@ -220,6 +220,22 @@ fn migrate_old_config() { } } +/// Resolve the max tokens parameter name based on provider and config setting. +/// +/// - If `config_value` is not `"auto"`, use it directly (explicit override). +/// - Otherwise, detect from provider name: +/// - `gemini`, `google`, `vertex` → `"max_output_tokens"` +/// - Everything else → `"max_tokens"` +pub fn resolve_max_tokens_param(provider: &str, config_value: &str) -> String { + match config_value { + v if v != "auto" => v.to_string(), + _ => match provider { + "gemini" | "google" | "vertex" => "max_output_tokens".to_string(), + _ => "max_tokens".to_string(), + }, + } +} + /// Load the full resolved config: defaults ← global config ← .cora.yaml ← CLI overrides. /// /// `cli_provider`, `cli_model`, `cli_api_key`, and `cli_format` are `None` @@ -374,6 +390,8 @@ pub fn build_llm_config( }) .unwrap_or_else(|| config.provider.base_url.clone()); + let max_tokens_param = resolve_max_tokens_param(&provider, &config.max_tokens_param); + Ok(LLMConfig { api_key, base_url, @@ -381,6 +399,7 @@ pub fn build_llm_config( provider, temperature: config.temperature, max_tokens: config.max_tokens, + max_tokens_param, timeout: config.timeout, }) } @@ -738,3 +757,58 @@ pub fn remove_provider_info() -> std::result::Result<(), CoraError> { Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn resolve_max_tokens_param_auto_gemini() { + assert_eq!( + resolve_max_tokens_param("gemini", "auto"), + "max_output_tokens" + ); + } + + #[test] + fn resolve_max_tokens_param_auto_google() { + assert_eq!( + resolve_max_tokens_param("google", "auto"), + "max_output_tokens" + ); + } + + #[test] + fn resolve_max_tokens_param_auto_vertex() { + assert_eq!( + resolve_max_tokens_param("vertex", "auto"), + "max_output_tokens" + ); + } + + #[test] + fn resolve_max_tokens_param_auto_openai() { + assert_eq!(resolve_max_tokens_param("openai", "auto"), "max_tokens"); + } + + #[test] + fn resolve_max_tokens_param_auto_anthropic() { + assert_eq!(resolve_max_tokens_param("anthropic", "auto"), "max_tokens"); + } + + #[test] + fn resolve_max_tokens_param_explicit_override() { + assert_eq!( + resolve_max_tokens_param("gemini", "max_tokens"), + "max_tokens" + ); + assert_eq!( + resolve_max_tokens_param("openai", "max_output_tokens"), + "max_output_tokens" + ); + assert_eq!( + resolve_max_tokens_param("anthropic", "max_completion_tokens"), + "max_completion_tokens" + ); + } +} diff --git a/src/config/schema.rs b/src/config/schema.rs index 5fdb9db..fba49e7 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -35,6 +35,8 @@ pub struct Config { pub temperature: f32, /// Max tokens for LLM responses. pub max_tokens: u32, + /// JSON parameter name for max tokens (e.g. "max_tokens" or "max_output_tokens"). + pub max_tokens_param: String, /// HTTP timeout in seconds for LLM requests. pub timeout: u64, /// Cache TTL in minutes for review caching. @@ -126,6 +128,7 @@ impl Default for Config { scan_system_prompt_file: None, temperature: 0.0, max_tokens: 4096, + max_tokens_param: "auto".to_string(), timeout: 600, cache_ttl: 1440, // 24h in minutes static_analysis: StaticAnalysisConfig::default(), @@ -297,6 +300,11 @@ pub struct LlmSection { pub temperature: Option, #[serde(skip_serializing_if = "Option::is_none")] pub max_tokens: Option, + /// Name of the max tokens parameter to use in API requests. + /// Supported values: `"auto"` (detect from provider), `"max_tokens"`, + /// `"max_output_tokens"`, `"max_completion_tokens"`. + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens_param: Option, #[serde(skip_serializing_if = "Option::is_none")] pub timeout: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -449,6 +457,9 @@ impl CoraFile { if let Some(v) = llm.cache_ttl { config.cache_ttl = v; } + if let Some(ref v) = llm.max_tokens_param { + config.max_tokens_param = v.clone(); + } } if let Some(re) = &self.rules_engine { config.rules_config.enabled = re.enabled; @@ -1038,6 +1049,7 @@ llm: llm: Some(LlmSection { temperature: Some(0.7), max_tokens: None, + max_tokens_param: None, timeout: None, cache_ttl: None, }), @@ -1058,6 +1070,7 @@ llm: llm: Some(LlmSection { temperature: None, max_tokens: Some(2048), + max_tokens_param: None, timeout: None, cache_ttl: None, }), @@ -1074,6 +1087,7 @@ llm: llm: Some(LlmSection { temperature: None, max_tokens: None, + max_tokens_param: None, timeout: Some(300), cache_ttl: None, }), @@ -1090,6 +1104,7 @@ llm: llm: Some(LlmSection { temperature: Some(1.0), max_tokens: Some(16384), + max_tokens_param: None, timeout: Some(240), cache_ttl: Some(2880), }), @@ -1134,6 +1149,7 @@ provider: llm: Some(LlmSection { temperature: Some(0.5), max_tokens: Some(8192), + max_tokens_param: None, timeout: Some(60), cache_ttl: None, }), @@ -1278,4 +1294,55 @@ bundling: Some(crate::engine::bundling::GroupingStrategy::Smart) ); } + + // ─── max_tokens_param ─── + + #[test] + fn config_default_max_tokens_param() { + let cfg = Config::default(); + assert_eq!(cfg.max_tokens_param, "auto"); + } + + #[test] + fn merge_llm_max_tokens_param_explicit() { + let mut cfg = Config::default(); + let cora = CoraFile { + llm: Some(LlmSection { + max_tokens_param: Some("max_output_tokens".to_string()), + ..Default::default() + }), + ..Default::default() + }; + cora.merge_into(&mut cfg).unwrap(); + assert_eq!(cfg.max_tokens_param, "max_output_tokens"); + } + + #[test] + fn merge_llm_max_tokens_param_absent_leaves_default() { + let mut cfg = Config::default(); + let cora = CoraFile { + llm: Some(LlmSection { + temperature: Some(0.5), + ..Default::default() + }), + ..Default::default() + }; + cora.merge_into(&mut cfg).unwrap(); + assert_eq!(cfg.max_tokens_param, "auto"); + } + + #[test] + fn llm_section_yaml_roundtrip_with_max_tokens_param() { + let section = LlmSection { + temperature: Some(0.7), + max_tokens: Some(8192), + max_tokens_param: Some("max_output_tokens".to_string()), + timeout: Some(300), + cache_ttl: Some(60), + }; + let yaml = serde_yaml_ng::to_string(§ion).unwrap(); + let back: LlmSection = serde_yaml_ng::from_str(&yaml).unwrap(); + assert_eq!(back.max_tokens_param, Some("max_output_tokens".to_string())); + assert_eq!(back.max_tokens, Some(8192)); + } } diff --git a/src/engine/llm.rs b/src/engine/llm.rs index 8325054..cf2d5cc 100644 --- a/src/engine/llm.rs +++ b/src/engine/llm.rs @@ -54,7 +54,8 @@ struct ChatMessage { content: String, } -/// Request body for /chat/completions. +/// Request body for /chat/completions (kept for reference; unused after migration to dynamic json!). +#[allow(dead_code)] #[derive(Debug, Clone, Serialize)] struct ChatRequest { model: String, @@ -285,26 +286,19 @@ async fn chat_completion( )); } - let request = ChatRequest { - model: config.model.clone(), - messages: vec![ - ChatMessage { - role: "system".into(), - content: system_prompt.into(), - }, - ChatMessage { - role: "user".into(), - content: user_message.into(), - }, + let mut request = serde_json::json!({ + "model": config.model, + "messages": [ + { "role": "system", "content": system_prompt }, + { "role": "user", "content": user_message } ], - temperature: config.temperature, - max_tokens: config.max_tokens, - response_format: if response_format == "json_object" { - Some(serde_json::json!({"type": "json_object"})) - } else { - None - }, - }; + "temperature": config.temperature, + }); + request[config.max_tokens_param.clone()] = serde_json::json!(config.max_tokens); + + if response_format == "json_object" { + request["response_format"] = serde_json::json!({"type": "json_object"}); + } debug!(model = %config.model, url = %url, "sending LLM request"); @@ -541,12 +535,12 @@ async fn chat_completion_stream( { "role": "user", "content": user_message } ], "temperature": config.temperature, - "max_tokens": config.max_tokens, "stream": true, // Ask OpenAI-compatible providers to include token usage in the final // SSE chunk. Providers that don't recognise this field simply ignore it. "stream_options": { "include_usage": true } }); + request_body[config.max_tokens_param.clone()] = serde_json::json!(config.max_tokens); if response_format == "json_object" { request_body["response_format"] = serde_json::json!({"type": "json_object"}); @@ -1885,4 +1879,51 @@ mod tests { let messy = "hello\n\t world\n\n"; assert_eq!(preview_raw(messy), "hello world"); } + + // ─── max_tokens_param JSON key naming ─── + + #[test] + fn chat_request_uses_max_output_tokens() { + // Verify that when max_tokens_param is "max_output_tokens", the JSON + // body contains "max_output_tokens" (not "max_tokens") as the key. + let mut body = serde_json::json!({ + "model": "gemini-pro", + "messages": [ + { "role": "system", "content": "test" }, + { "role": "user", "content": "hello" } + ], + "temperature": 0.0, + }); + let param_name = "max_output_tokens"; + body[param_name] = serde_json::json!(4096); + + let serialized = serde_json::to_string(&body).unwrap(); + assert!( + serialized.contains(r#""max_output_tokens":4096"#), + "Expected max_output_tokens key in JSON, got: {serialized}" + ); + assert!( + !serialized.contains(r#""max_tokens":"#), + "Should NOT contain hardcoded max_tokens key, got: {serialized}" + ); + } + + #[test] + fn chat_request_uses_max_tokens_default() { + let mut body = serde_json::json!({ + "model": "gpt-4o-mini", + "messages": [ + { "role": "system", "content": "test" }, + { "role": "user", "content": "hello" } + ], + "temperature": 0.0, + }); + body["max_tokens"] = serde_json::json!(8192); + + let serialized = serde_json::to_string(&body).unwrap(); + assert!( + serialized.contains(r#""max_tokens":8192"#), + "Expected max_tokens key in JSON, got: {serialized}" + ); + } } diff --git a/src/engine/types.rs b/src/engine/types.rs index 326b71a..8a60ce1 100644 --- a/src/engine/types.rs +++ b/src/engine/types.rs @@ -306,6 +306,7 @@ mod tests { assert_eq!(cfg.provider, "openai"); assert_eq!(cfg.temperature, 0.0); assert_eq!(cfg.max_tokens, 4096); + assert_eq!(cfg.max_tokens_param, "max_tokens"); assert_eq!(cfg.timeout, 600); } @@ -423,6 +424,8 @@ pub struct LLMConfig { pub provider: String, pub temperature: f32, pub max_tokens: u32, + /// JSON parameter name for max tokens (resolved from config). + pub max_tokens_param: String, pub timeout: u64, } @@ -435,6 +438,7 @@ impl Default for LLMConfig { provider: "openai".to_string(), temperature: 0.0, max_tokens: 4096, + max_tokens_param: "max_tokens".to_string(), timeout: 600, } }