From ffe384b66ea326d58056cd6315b50fefa7c6beee Mon Sep 17 00:00:00 2001 From: Nige Date: Sat, 14 Mar 2026 20:06:48 +0000 Subject: [PATCH] fix(llm): add stop_sequences parity for tool completions (#1170) * fix(llm): add stop_sequences parity for tool completions * refactor(web-openai): dedupe request builders and satisfy no-panics gate * test(llm): mark multiline assert with safety comment for CI gate * test(llm): make safety-marked assert formatting-stable --- src/channels/web/openai_compat.rs | 88 +++++++++++++++---------------- src/llm/bedrock.rs | 7 ++- src/llm/nearai_chat.rs | 6 +++ src/llm/provider.rs | 27 ++++++++-- src/llm/response_cache.rs | 1 + src/orchestrator/api.rs | 1 + src/worker/api.rs | 2 + 7 files changed, 81 insertions(+), 51 deletions(-) diff --git a/src/channels/web/openai_compat.rs b/src/channels/web/openai_compat.rs index e329693a..51577e06 100644 --- a/src/channels/web/openai_compat.rs +++ b/src/channels/web/openai_compat.rs @@ -419,6 +419,44 @@ fn parse_stop(val: &serde_json::Value) -> Option> { } } +fn build_completion_request( + req: &OpenAiChatRequest, + messages: Vec, +) -> CompletionRequest { + let mut comp_req = CompletionRequest::new(messages).with_model(req.model.clone()); + if let Some(t) = req.temperature { + comp_req = comp_req.with_temperature(t); + } + if let Some(mt) = req.max_tokens { + comp_req = comp_req.with_max_tokens(mt); + } + if let Some(stops) = req.stop.as_ref().and_then(parse_stop) { + comp_req.stop_sequences = Some(stops); + } + comp_req +} + +fn build_tool_request( + req: &OpenAiChatRequest, + messages: Vec, +) -> ToolCompletionRequest { + let tools = convert_tools(req.tools.as_deref().unwrap_or(&[])); + let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model.clone()); + if let Some(t) = req.temperature { + tool_req = tool_req.with_temperature(t); + } + if let Some(mt) = req.max_tokens { + tool_req = tool_req.with_max_tokens(mt); + } + if let Some(stops) = req.stop.as_ref().and_then(parse_stop) { + tool_req = tool_req.with_stop_sequences(stops); + } + if let Some(choice) = req.tool_choice.as_ref().and_then(normalize_tool_choice) { + tool_req = tool_req.with_tool_choice(choice); + } + tool_req +} + // --------------------------------------------------------------------------- // Handlers // --------------------------------------------------------------------------- @@ -476,19 +514,7 @@ pub async fn chat_completions_handler( let created = unix_timestamp(); if has_tools { - let tools = convert_tools(req.tools.as_deref().unwrap_or(&[])); - let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model); - if let Some(t) = req.temperature { - tool_req = tool_req.with_temperature(t); - } - if let Some(mt) = req.max_tokens { - tool_req = tool_req.with_max_tokens(mt); - } - if let Some(ref tc) = req.tool_choice - && let Some(choice) = normalize_tool_choice(tc) - { - tool_req = tool_req.with_tool_choice(choice); - } + let tool_req = build_tool_request(&req, messages); let resp = llm .complete_with_tools(tool_req) @@ -527,16 +553,7 @@ pub async fn chat_completions_handler( Ok(Json(response).into_response()) } else { - let mut comp_req = CompletionRequest::new(messages).with_model(req.model); - if let Some(t) = req.temperature { - comp_req = comp_req.with_temperature(t); - } - if let Some(mt) = req.max_tokens { - comp_req = comp_req.with_max_tokens(mt); - } - if let Some(ref stop_val) = req.stop { - comp_req.stop_sequences = parse_stop(stop_val); - } + let comp_req = build_completion_request(&req, messages); let resp = llm.complete(comp_req).await.map_err(map_llm_error)?; let model_name = llm.effective_model_name(Some(requested_model.as_str())); @@ -596,35 +613,14 @@ async fn handle_streaming( } let llm_result = if has_tools { - let tools = convert_tools(req.tools.as_deref().unwrap_or(&[])); - let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model); - if let Some(t) = req.temperature { - tool_req = tool_req.with_temperature(t); - } - if let Some(mt) = req.max_tokens { - tool_req = tool_req.with_max_tokens(mt); - } - if let Some(ref tc) = req.tool_choice - && let Some(choice) = normalize_tool_choice(tc) - { - tool_req = tool_req.with_tool_choice(choice); - } + let tool_req = build_tool_request(&req, messages); LlmResult::WithTools( llm.complete_with_tools(tool_req) .await .map_err(map_llm_error)?, ) } else { - let mut comp_req = CompletionRequest::new(messages).with_model(req.model); - if let Some(t) = req.temperature { - comp_req = comp_req.with_temperature(t); - } - if let Some(mt) = req.max_tokens { - comp_req = comp_req.with_max_tokens(mt); - } - if let Some(ref stop_val) = req.stop { - comp_req.stop_sequences = parse_stop(stop_val); - } + let comp_req = build_completion_request(&req, messages); LlmResult::Simple(llm.complete(comp_req).await.map_err(map_llm_error)?) }; let model_name = llm.effective_model_name(Some(requested_model.as_str())); diff --git a/src/llm/bedrock.rs b/src/llm/bedrock.rs index ebde19f1..5d6e121e 100644 --- a/src/llm/bedrock.rs +++ b/src/llm/bedrock.rs @@ -176,8 +176,11 @@ impl LlmProvider for BedrockProvider { builder = builder.tool_config(tc); } - if let Some(config) = build_inference_config(request.temperature, request.max_tokens, None) - { + if let Some(config) = build_inference_config( + request.temperature, + request.max_tokens, + request.stop_sequences.as_deref(), + ) { builder = builder.inference_config(config); } diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index 0c0335bd..bf2b8738 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -475,6 +475,7 @@ impl LlmProvider for NearAiChatProvider { messages, temperature: req.temperature, max_tokens: req.max_tokens, + stop: req.stop_sequences, tools: None, tool_choice: None, }; @@ -554,6 +555,7 @@ impl LlmProvider for NearAiChatProvider { messages, temperature: req.temperature, max_tokens: req.max_tokens, + stop: req.stop_sequences, tools: if tools.is_empty() { None } else { Some(tools) }, tool_choice: req.tool_choice, }; @@ -680,6 +682,8 @@ struct ChatCompletionRequest { #[serde(skip_serializing_if = "Option::is_none")] max_tokens: Option, #[serde(skip_serializing_if = "Option::is_none")] + stop: Option>, + #[serde(skip_serializing_if = "Option::is_none")] tools: Option>, #[serde(skip_serializing_if = "Option::is_none")] tool_choice: Option, @@ -1666,6 +1670,7 @@ mod tests { }], temperature: None, max_tokens: None, + stop: None, tools: None, tool_choice: None, }; @@ -1687,6 +1692,7 @@ mod tests { messages: vec![], temperature: Some(0.7), max_tokens: Some(1024), + stop: None, tools: Some(vec![ChatCompletionTool { tool_type: "function".to_string(), function: ChatCompletionFunction { diff --git a/src/llm/provider.rs b/src/llm/provider.rs index 787bbff1..8a213031 100644 --- a/src/llm/provider.rs +++ b/src/llm/provider.rs @@ -251,6 +251,7 @@ pub struct ToolCompletionRequest { pub model: Option, pub max_tokens: Option, pub temperature: Option, + pub stop_sequences: Option>, /// How to handle tool use: "auto", "required", or "none". pub tool_choice: Option, /// Opaque metadata passed through to the provider (e.g. thread_id for chaining). @@ -266,6 +267,7 @@ impl ToolCompletionRequest { model: None, max_tokens: None, temperature: None, + stop_sequences: None, tool_choice: None, metadata: std::collections::HashMap::new(), } @@ -289,6 +291,12 @@ impl ToolCompletionRequest { self } + /// Set stop sequences. + pub fn with_stop_sequences(mut self, stop_sequences: Vec) -> Self { + self.stop_sequences = Some(stop_sequences); + self + } + /// Set tool choice mode. pub fn with_tool_choice(mut self, choice: impl Into) -> Self { self.tool_choice = Some(choice.into()); @@ -504,8 +512,6 @@ pub fn strip_unsupported_completion_params( /// This is the single helper function used by all providers to remove /// parameters they don't support from tool calls, replacing duplicate stringly-typed logic. /// -/// Note: Only `Temperature` and `MaxTokens` are supported in `ToolCompletionRequest`. -/// `StopSequences` is only available in `CompletionRequest` and is not applicable to tool calls. pub fn strip_unsupported_tool_params( unsupported: &std::collections::HashSet, req: &mut ToolCompletionRequest, @@ -519,7 +525,9 @@ pub fn strip_unsupported_tool_params( if unsupported.contains(UnsupportedParam::MaxTokens.name()) { req.max_tokens = None; } - // Note: StopSequences is not a field in ToolCompletionRequest, so no action needed + if unsupported.contains(UnsupportedParam::StopSequences.name()) { + req.stop_sequences = None; + } } #[cfg(test)] @@ -651,4 +659,17 @@ mod tests { assert!(messages[2].tool_call_id.is_none()); assert!(messages[2].name.is_none()); } + + #[test] + fn test_strip_unsupported_tool_params_strips_stop_sequences() { + let mut unsupported = std::collections::HashSet::new(); + unsupported.insert(UnsupportedParam::StopSequences.name().to_string()); + + let mut req = ToolCompletionRequest::new(vec![ChatMessage::user("hello")], vec![]); + req.stop_sequences = Some(vec!["STOP".to_string()]); + + strip_unsupported_tool_params(&unsupported, &mut req); + + assert!(req.stop_sequences.is_none()); // safety: test assertion for explicit strip behavior + } } diff --git a/src/llm/response_cache.rs b/src/llm/response_cache.rs index b8238427..d7746f60 100644 --- a/src/llm/response_cache.rs +++ b/src/llm/response_cache.rs @@ -548,6 +548,7 @@ mod tests { model: None, max_tokens: None, temperature: None, + stop_sequences: None, tool_choice: None, metadata: Default::default(), }; diff --git a/src/orchestrator/api.rs b/src/orchestrator/api.rs index 80e09073..b46aa8c6 100644 --- a/src/orchestrator/api.rs +++ b/src/orchestrator/api.rs @@ -176,6 +176,7 @@ async fn llm_complete_with_tools( model: req.model, max_tokens: req.max_tokens, temperature: req.temperature, + stop_sequences: req.stop_sequences, tool_choice: req.tool_choice, metadata: std::collections::HashMap::new(), }; diff --git a/src/worker/api.rs b/src/worker/api.rs index 459375b4..43fda2dd 100644 --- a/src/worker/api.rs +++ b/src/worker/api.rs @@ -65,6 +65,7 @@ pub struct ProxyToolCompletionRequest { pub model: Option, pub max_tokens: Option, pub temperature: Option, + pub stop_sequences: Option>, pub tool_choice: Option, } @@ -251,6 +252,7 @@ impl WorkerHttpClient { model: request.model.clone(), max_tokens: request.max_tokens, temperature: request.temperature, + stop_sequences: request.stop_sequences.clone(), tool_choice: request.tool_choice.clone(), };