diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index fa22f92c1a..85151d141d 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -1823,6 +1823,49 @@ impl StreamingProcessor { .map_err(|_| "Client disconnected".to_string()) } + /// Stop the open content block, if any, so the next block gets the next + /// index: reasoning, text and tool calls can alternate. + async fn stop_open_block( + tx: &SseSender, + buffer: &mut Vec, + index: &mut u32, + open: [&mut bool; 3], + ) -> Result<(), String> { + if open + .into_iter() + .fold(false, |any, open| std::mem::take(open) | any) + { + let stop = MessageStreamEvent::ContentBlockStop { index: *index }; + Self::send_messages_event(tx, buffer, &stop).await?; + *index += 1; + } + Ok(()) + } + + /// Send tool call arguments to the open `tool_use` block. Reasoning can + /// stop that block in the middle of a call, and the arguments after it + /// have no block to go to, so they are dropped. + async fn send_tool_arguments( + tx: &SseSender, + buffer: &mut Vec, + index: u32, + tool_block_open: bool, + partial_json: String, + ) -> Result<(), String> { + if partial_json.is_empty() { + return Ok(()); + } + if !tool_block_open { + debug!("Dropping tool arguments without an open tool_use block"); + return Ok(()); + } + let delta = MessageStreamEvent::ContentBlockDelta { + index, + delta: ContentBlockDelta::InputJsonDelta { partial_json }, + }; + Self::send_messages_event(tx, buffer, &delta).await + } + /// Process reasoning content in Messages streaming mode (n=1 only). /// /// Returns `(normal_text, reasoning_text, in_reasoning)`. @@ -2275,6 +2318,17 @@ impl StreamingProcessor { // Emit thinking content block deltas if !reasoning_chunk_text.is_empty() { if !thinking_block_open { + Self::stop_open_block( + tx, + &mut sse_buffer, + &mut current_block_index, + [ + &mut thinking_block_open, + &mut text_block_open, + &mut tool_block_open, + ], + ) + .await?; Self::send_messages_event( tx, &mut sse_buffer, @@ -2326,19 +2380,18 @@ impl StreamingProcessor { // Specific function: entire output is arguments for one tool if !has_tool_calls { has_tool_calls = true; - // Close text block if open before starting tool block - if text_block_open { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { - index: current_block_index, - }, - ) - .await?; - text_block_open = false; - current_block_index += 1; - } + // Close the open block before starting tool block + Self::stop_open_block( + tx, + &mut sse_buffer, + &mut current_block_index, + [ + &mut thinking_block_open, + &mut text_block_open, + &mut tool_block_open, + ], + ) + .await?; // Emit content_block_start for the tool_use let tool_name = match &original_request.tool_choice { Some(messages::ToolChoice::Tool { name, .. }) => name.clone(), @@ -2365,30 +2418,57 @@ impl StreamingProcessor { .await?; tool_block_open = true; } - // Emit arguments delta - if !normal_text.is_empty() { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockDelta { - index: current_block_index, - delta: ContentBlockDelta::InputJsonDelta { - partial_json: normal_text, - }, - }, - ) - .await?; - } + // Emit arguments delta, unless reasoning stopped the block + Self::send_tool_arguments( + tx, + &mut sse_buffer, + current_block_index, + tool_block_open, + normal_text, + ) + .await?; } else if let Some(ref mut parser) = streaming_tool_parser { // Regular/required tool choice: use incremental parser match parser.parse_incremental(&normal_text, chat_tools).await { Ok(StreamingParseResult { normal_text: text, - calls, + mut calls, }) => { + // Arguments that finish the open call come before + // the text after it in the same chunk. + let finishing = if tool_block_open { + calls + .iter() + .position(|call| call.name.is_some()) + .unwrap_or(calls.len()) + } else { + 0 + }; + for tool_call_item in calls.drain(..finishing) { + Self::send_tool_arguments( + tx, + &mut sse_buffer, + current_block_index, + tool_block_open, + tool_call_item.parameters, + ) + .await?; + } + // Emit normal text from parser as text content blocks if !text.is_empty() { if !text_block_open { + Self::stop_open_block( + tx, + &mut sse_buffer, + &mut current_block_index, + [ + &mut thinking_block_open, + &mut text_block_open, + &mut tool_block_open, + ], + ) + .await?; Self::send_messages_event( tx, &mut sse_buffer, @@ -2420,29 +2500,17 @@ impl StreamingProcessor { if let Some(ref name) = tool_call_item.name { // New tool call: close previous blocks, emit start - if text_block_open { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { - index: current_block_index, - }, - ) - .await?; - text_block_open = false; - current_block_index += 1; - } - if tool_block_open { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { - index: current_block_index, - }, - ) - .await?; - current_block_index += 1; - } + Self::stop_open_block( + tx, + &mut sse_buffer, + &mut current_block_index, + [ + &mut thinking_block_open, + &mut text_block_open, + &mut tool_block_open, + ], + ) + .await?; let tool_call_id = utils::generate_tool_call_id( model, @@ -2469,19 +2537,14 @@ impl StreamingProcessor { } // Emit incremental arguments - if !tool_call_item.parameters.is_empty() { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockDelta { - index: current_block_index, - delta: ContentBlockDelta::InputJsonDelta { - partial_json: tool_call_item.parameters, - }, - }, - ) - .await?; - } + Self::send_tool_arguments( + tx, + &mut sse_buffer, + current_block_index, + tool_block_open, + tool_call_item.parameters, + ) + .await?; } } Err(e) => { @@ -2495,6 +2558,17 @@ impl StreamingProcessor { // Regular text emission (no tools active) if !normal_text.is_empty() { if !text_block_open { + Self::stop_open_block( + tx, + &mut sse_buffer, + &mut current_block_index, + [ + &mut thinking_block_open, + &mut text_block_open, + &mut tool_block_open, + ], + ) + .await?; Self::send_messages_event( tx, &mut sse_buffer, @@ -2529,6 +2603,17 @@ impl StreamingProcessor { let leftover_text = parser.take_unstreamed_normal_text(); if !leftover_text.is_empty() { if !text_block_open { + Self::stop_open_block( + tx, + &mut sse_buffer, + &mut current_block_index, + [ + &mut thinking_block_open, + &mut text_block_open, + &mut tool_block_open, + ], + ) + .await?; Self::send_messages_event( tx, &mut sse_buffer, @@ -2563,30 +2648,18 @@ impl StreamingProcessor { has_tool_calls = true; if let Some(ref name) = tool_call_item.name { - // Close text block if open before starting tool block - if text_block_open { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { - index: current_block_index, - }, - ) - .await?; - text_block_open = false; - current_block_index += 1; - } - if tool_block_open { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { - index: current_block_index, - }, - ) - .await?; - current_block_index += 1; - } + // Close the open block before starting tool block + Self::stop_open_block( + tx, + &mut sse_buffer, + &mut current_block_index, + [ + &mut thinking_block_open, + &mut text_block_open, + &mut tool_block_open, + ], + ) + .await?; let tool_call_id = utils::generate_tool_call_id( model, @@ -2610,19 +2683,14 @@ impl StreamingProcessor { tool_block_open = true; } - if !tool_call_item.parameters.is_empty() { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockDelta { - index: current_block_index, - delta: ContentBlockDelta::InputJsonDelta { - partial_json: tool_call_item.parameters, - }, - }, - ) - .await?; - } + Self::send_tool_arguments( + tx, + &mut sse_buffer, + current_block_index, + tool_block_open, + tool_call_item.parameters, + ) + .await?; } } } diff --git a/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs b/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs index 2787a00d01..5f31623d33 100644 --- a/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs +++ b/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs @@ -15,6 +15,7 @@ use smg_grpc_client::vllm_engine::{ }; use tokio::{net::TcpListener, task::JoinHandle}; use tonic::codec::Codec; +use tool_parser::types::ToolCallItem; use super::*; use crate::{routers::common::sse::SseReceiver, worker::WorkerRegistry}; @@ -578,3 +579,440 @@ async fn deepseek_does_not_emit_aggregate_usage_without_complete_frames() { ); } } + +/// Reasoning for every chunk but "text"; `is_in_reasoning` reports `.0`. +struct ReasoningButText(bool); + +impl ReasoningParser for ReasoningButText { + fn detect_and_parse_reasoning( + &mut self, + text: &str, + ) -> Result { + Ok(ParserResult::normal(text.to_string())) + } + + fn parse_reasoning_streaming_incremental( + &mut self, + text: &str, + ) -> Result { + Ok(match text { + "text" => ParserResult::normal(text.to_string()), + _ => ParserResult::reasoning(text.to_string()), + }) + } + + fn reset(&mut self) {} + + fn model_type(&self) -> &str { + "reasoning-but-text" + } + + fn is_in_reasoning(&self) -> bool { + self.0 + } + + fn mark_reasoning_started(&mut self) {} + + fn mark_think_start_stripped(&mut self) {} +} + +/// Returns the chunk as normal text and reports one whole `lookup` call on +/// its first parse. +#[derive(Default)] +struct CallFirst { + called: bool, +} + +#[async_trait::async_trait] +impl ToolParser for CallFirst { + async fn parse_complete( + &self, + output: &str, + ) -> tool_parser::errors::ParserResult<(String, Vec)> { + Ok((output.to_string(), Vec::new())) + } + + async fn parse_incremental( + &mut self, + chunk: &str, + _tools: &[Tool], + ) -> tool_parser::errors::ParserResult { + let calls = (!std::mem::replace(&mut self.called, true)).then(|| ToolCallItem { + tool_index: 0, + name: Some("lookup".to_string()), + parameters: "{}".to_string(), + }); + Ok(StreamingParseResult { + normal_text: chunk.to_string(), + calls: calls.into_iter().collect(), + }) + } + + fn has_tool_markers(&self, _text: &str) -> bool { + false + } +} + +/// A processor with the stub parsers above. +fn stub_processor(in_reasoning: bool) -> StreamingProcessor { + let reasoning = ReasoningParserFactory::new(); + reasoning + .registry() + .register_parser("reasoning-but-text", move || { + Box::new(ReasoningButText(in_reasoning)) + }); + let tools = ToolParserFactory::new(); + tools + .registry() + .register_parser("call-first", || Box::new(CallFirst::default())); + let resolver = utils::ParserResolver::new( + Arc::new(WorkerRegistry::new()), + Some("call-first".to_string()), + Some("reasoning-but-text".to_string()), + ); + StreamingProcessor::new(tools, reasoning, resolver, "vllm") +} + +/// The content block starts and stops from `messages_blocks_and_inputs`. +async fn messages_blocks( + processor: StreamingProcessor, + tool_choice: Option, + texts: &[&str], +) -> Vec { + messages_blocks_and_inputs(processor, tool_choice, texts) + .await + .0 +} + +/// The content block starts and stops of a Messages stream of `texts`, and +/// the input of each `tool_use` block joined from its `input_json_delta`s as +/// a client SDK builds it (the joined text if it is not complete JSON). A +/// block starts only when none is open, and every delta goes to the open +/// block and matches its type. +async fn messages_blocks_and_inputs( + processor: StreamingProcessor, + tool_choice: Option, + texts: &[&str], +) -> (Vec, Vec) { + let spec = MessagesResponseSpec { + thinking: Some(messages::ThinkingConfig::Enabled { + budget_tokens: 1024, + display: None, + }), + tool_choice, + has_tools: true, + history_tool_calls_count: 0, + chat_tools: chat_spec(true).tools.unwrap(), + stop_sequences: None, + }; + let mut frames: Vec<_> = texts.iter().map(|text| chunk(0, text)).collect(); + frames.push(complete(0, "stop")); + let (stream, server) = scripted_stream(frames, "0").await; + let (tx, rx) = sse_channel(); + let result = processor + .process_messages_streaming_chunks( + stream, + dispatch(), + Arc::new(CharacterTokenizer::default()), + (None, None, false, false, false), + spec, + &tx, + None, + ) + .await; + drop(tx); + let events = collect_events(rx).await; + server.abort(); + assert!(result.is_ok(), "{result:?}"); + let mut open = None; + let mut blocks = Vec::new(); + let mut inputs = Vec::new(); + for event in &events { + let index = &event["index"]; + match event["type"].as_str() { + Some("content_block_start") => { + let kind = event["content_block"]["type"].as_str().unwrap_or_default(); + blocks.push(format!("start {index} {kind:?}")); + assert!(open.replace((index, kind)).is_none(), "{blocks:?}"); + if kind == "tool_use" { + inputs.push(String::new()); + } + } + Some("content_block_delta") => { + let delta = &event["delta"]; + let kind = match delta["type"].as_str() { + Some("text_delta") => "text", + Some("input_json_delta") => "tool_use", + Some("thinking_delta" | "signature_delta") => "thinking", + other => panic!("unexpected delta {other:?}"), + }; + assert_eq!(open, Some((index, kind)), "delta {delta} after {blocks:?}"); + if let Some(json) = delta["partial_json"].as_str() { + inputs.last_mut().expect("open tool_use").push_str(json); + } + } + Some("content_block_stop") => { + blocks.push(format!("stop {index}")); + assert_eq!(open.take().map(|open| open.0), Some(index), "{blocks:?}"); + } + _ => {} + } + } + let inputs = inputs + .iter() + .map(|json| match json.as_str() { + "" => serde_json::json!({}), + json => serde_json::from_str(json).unwrap_or_else(|_| json.into()), + }) + .collect(); + (blocks, inputs) +} + +#[tokio::test] +async fn messages_blocks_do_not_overlap_when_reasoning_calls_and_text_alternate() { + // Reasoning, then a call with no text between; text; reasoning again. + assert_eq!( + messages_blocks(stub_processor(false), None, &["a", "text", "b"]).await, + [ + "start 0 \"thinking\"", + "stop 0", + "start 1 \"tool_use\"", + "stop 1", + "start 2 \"text\"", + "stop 2", + "start 3 \"thinking\"", + "stop 3", + ] + ); + // Text the reasoning parser returns while it stays in reasoning. Such text + // skips the tool parser and follows the chunk's reasoning, so a chunk like + // `answermore` still comes out in the wrong order; only the + // block boundaries are checked here. + assert_eq!( + messages_blocks(stub_processor(true), None, &["a", "text", "b"]).await, + [ + "start 0 \"thinking\"", + "stop 0", + "start 1 \"text\"", + "stop 1", + "start 2 \"thinking\"", + "stop 2", + ] + ); + // Reasoning right after a call. + assert_eq!( + messages_blocks(stub_processor(false), None, &["a", "b"]).await, + [ + "start 0 \"thinking\"", + "stop 0", + "start 1 \"tool_use\"", + "stop 1", + "start 2 \"thinking\"", + "stop 2", + ] + ); +} + +/// A processor with the registered parsers `tool` and `reasoning`. +fn named_processor(tool: &str, reasoning: &str) -> StreamingProcessor { + StreamingProcessor::new( + ToolParserFactory::new(), + ReasoningParserFactory::new(), + utils::ParserResolver::new( + Arc::new(WorkerRegistry::new()), + Some(tool.to_string()), + Some(reasoning.to_string()), + ), + "vllm", + ) +} + +#[tokio::test] +async fn messages_blocks_do_not_overlap_with_deepseek_parsers() { + // `` ends reasoning without text, and the arguments of a specific + // tool follow in the next chunk. + let tool = messages::ToolChoice::Tool { + name: "lookup".to_string(), + disable_parallel_tool_use: None, + }; + assert_eq!( + messages_blocks( + named_processor("deepseek", "deepseek_r1"), + Some(tool), + &["plan", "", "{}"], + ) + .await, + [ + "start 0 \"thinking\"", + "stop 0", + "start 1 \"tool_use\"", + "stop 1", + ] + ); + // A reasoning parser that enters reasoning again after text. + assert_eq!( + messages_blocks( + named_processor("deepseek_v41", "deepseek_v41"), + None, + &[ + "plan", + "", + "answer", + "", + "more", + "", + "done" + ], + ) + .await, + [ + "start 0 \"thinking\"", + "stop 0", + "start 1 \"text\"", + "stop 1", + "start 2 \"thinking\"", + "stop 2", + "start 3 \"text\"", + "stop 3", + ] + ); + // Reasoning again while the tool parser holds text that it releases at + // the end of the stream. + assert_eq!( + messages_blocks( + named_processor("json", "deepseek_v41"), + None, + &["plan", "", "{", "", "more"], + ) + .await, + [ + "start 0 \"thinking\"", + "stop 0", + "start 1 \"thinking\"", + "stop 1", + "start 2 \"text\"", + "stop 2", + ] + ); +} + +#[tokio::test] +async fn messages_tool_arguments_precede_text_in_the_same_chunk() { + // With multi-token chunks, qwen_xml returns the arguments that finish a + // call together with the text after the call. + let (blocks, inputs) = messages_blocks_and_inputs( + named_processor("qwen_xml", "qwen3"), + None, + &[ + "\n\n\n1", + "\n\n\n\n\n\n", + "\n2\n\n\n", + ], + ) + .await; + assert_eq!( + blocks, + [ + "start 0 \"tool_use\"", + "stop 0", + "start 1 \"text\"", + "stop 1", + "start 2 \"tool_use\"", + "stop 2", + ] + ); + assert_eq!( + inputs, + [serde_json::json!({"q": 1}), serde_json::json!({"q": 2})] + ); + for (texts, input) in [ + ( + [ + "\n\n\nPar", + "is\n\n\n\nDone.", + ], + serde_json::json!({"q": "Paris"}), + ), + ( + [ + "\n\n\nx\n\n", + "\ny\n\n\n\n", + ], + serde_json::json!({"a": "x", "b": "y"}), + ), + ] { + let (blocks, inputs) = + messages_blocks_and_inputs(named_processor("qwen_xml", "qwen3"), None, &texts).await; + assert_eq!( + blocks, + [ + "start 0 \"tool_use\"", + "stop 0", + "start 1 \"text\"", + "stop 1" + ] + ); + assert_eq!(inputs, [input]); + } +} + +#[tokio::test] +async fn messages_tool_arguments_need_an_open_block() { + // Reasoning that starts only after the specific tool's block stops that + // block; the arguments after it have no block to go to and are dropped. + let tool = messages::ToolChoice::Tool { + name: "lookup".to_string(), + disable_parallel_tool_use: None, + }; + let (blocks, inputs) = messages_blocks_and_inputs( + named_processor("qwen", "qwen3"), + Some(tool), + &["plan", "", "{}"], + ) + .await; + assert_eq!( + blocks, + [ + "start 0 \"tool_use\"", + "stop 0", + "start 1 \"thinking\"", + "stop 1", + ] + ); + assert_eq!(inputs, [serde_json::json!({})]); + // Until qwen3 strips a ``, it takes the first one anywhere as the + // start of reasoning, so one inside an argument stops the call's block. + // The rest of the arguments are dropped, in the middle of the stream and + // at its end, where qwen_xml releases the closing brace. + for (texts, input) in [ + ( + [ + "\n\n\nA ", + "x B\n\n", + "\n", + ], + serde_json::json!({}), + ), + ( + [ + "\n\n\nA\n\n", + "\nB ", + "x", + ], + serde_json::json!("{\"q\": \"A\""), + ), + ] { + let (blocks, inputs) = + messages_blocks_and_inputs(named_processor("qwen_xml", "qwen3"), None, &texts).await; + assert_eq!( + blocks, + [ + "start 0 \"tool_use\"", + "stop 0", + "start 1 \"thinking\"", + "stop 1", + ] + ); + assert_eq!(inputs, [input]); + } +}