diff --git a/sgl-router/Cargo.toml b/sgl-router/Cargo.toml index 28058109c..c980f9913 100644 --- a/sgl-router/Cargo.toml +++ b/sgl-router/Cargo.toml @@ -83,6 +83,7 @@ oracle = { version = "0.6.3", features = ["chrono"] } subtle = "2.6" rustpython-parser = "0.4.0" num-traits = "0.2" +openai-harmony = { git = "https://github.com/openai/harmony", tag = "v0.0.4" } # gRPC and Protobuf dependencies tonic = { version = "0.14.2", features = ["gzip", "transport"] } diff --git a/sgl-router/Makefile b/sgl-router/Makefile index c220b26f1..60cc72e9c 100644 --- a/sgl-router/Makefile +++ b/sgl-router/Makefile @@ -1,4 +1,4 @@ -# SGLang Router Makefile +# Model Gateway Makefile # Provides convenient shortcuts for common development tasks # Check if sccache is available and set RUSTC_WRAPPER accordingly @@ -13,14 +13,14 @@ endif .PHONY: help bench bench-quick bench-baseline bench-compare test build clean help: ## Show this help message - @echo "SGLang Router Development Commands" + @echo "Model Gateway Development Commands" @echo "==================================" @echo "" @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-20s\033[0m %s\n", $$1, $$2}' @echo "" build: ## Build the project in release mode - @echo "Building SGLang Router..." + @echo "Building SGLang Model Gateway..." @cargo build --release test: ## Run all tests @@ -59,11 +59,11 @@ check: ## Run cargo check and clippy @echo "Running cargo check..." @cargo check @echo "Running clippy..." - @cargo clippy + @cargo clippy --all-targets --all-features -- -D warnings fmt: ## Format code with rustfmt @echo "Formatting code..." - @cargo fmt + @rustup run nightly cargo fmt # Development workflow shortcuts dev-setup: build test ## Set up development environment diff --git a/sgl-router/src/grpc_client/sglang_scheduler.rs b/sgl-router/src/grpc_client/sglang_scheduler.rs index 53377b8e3..244678461 100644 --- a/sgl-router/src/grpc_client/sglang_scheduler.rs +++ b/sgl-router/src/grpc_client/sglang_scheduler.rs @@ -16,6 +16,7 @@ use crate::protocols::{ chat::ChatCompletionRequest, common::{ResponseFormat, StringOrArray, ToolChoice, ToolChoiceValue}, generate::GenerateRequest, + responses::ResponsesRequest, sampling_params::SamplingParams as GenerateSamplingParams, }; @@ -301,6 +302,42 @@ impl SglangSchedulerClient { Ok(grpc_request) } + /// Build a GenerateRequest from ResponsesRequest (OpenAI Responses API) + pub fn build_generate_request_from_responses( + &self, + request_id: String, + body: &ResponsesRequest, + processed_text: String, + token_ids: Vec, + harmony_stop_ids: Option>, + ) -> Result { + // Build sampling params from ResponsesRequest + let mut sampling_params = self.build_grpc_sampling_params_from_responses(body)?; + + // Inject Harmony stop token IDs if provided + if let Some(stop_ids) = harmony_stop_ids { + sampling_params.stop_token_ids = stop_ids; + } + + let grpc_request = proto::GenerateRequest { + request_id, + tokenized: Some(proto::TokenizedInput { + original_text: processed_text, + input_ids: token_ids, + }), + mm_inputs: None, // Responses API doesn't support multimodal yet + sampling_params: Some(sampling_params), + return_logprob: false, // Responses API uses top_logprobs field instead + logprob_start_len: -1, + top_logprobs_num: body.top_logprobs.unwrap_or(0) as i32, + return_hidden_states: false, + stream: body.stream.unwrap_or(false), + ..Default::default() + }; + + Ok(grpc_request) + } + /// Build gRPC SamplingParams from OpenAI request fn build_grpc_sampling_params( &self, @@ -400,6 +437,37 @@ impl SglangSchedulerClient { } } + /// Build gRPC SamplingParams from ResponsesRequest + fn build_grpc_sampling_params_from_responses( + &self, + request: &ResponsesRequest, + ) -> Result { + // ResponsesRequest doesn't have stop sequences in the same way + // Tools are handled externally by MCP loop, not via constraints + + let max_new_tokens = request.max_output_tokens.map(|v| v as i32); + + Ok(proto::SamplingParams { + temperature: request.temperature.unwrap_or(1.0), + top_p: request.top_p.unwrap_or(1.0), + top_k: -1, // ResponsesRequest doesn't expose top_k + min_p: 0.0, // ResponsesRequest doesn't expose min_p + frequency_penalty: 0.0, // ResponsesRequest doesn't expose frequency_penalty + presence_penalty: 0.0, // ResponsesRequest doesn't expose presence_penalty + repetition_penalty: 1.0, // ResponsesRequest doesn't expose repetition_penalty + max_new_tokens, + stop: vec![], // No stop sequences in Responses API + stop_token_ids: vec![], // Handled by Harmony stop tokens + skip_special_tokens: false, // Keep special tokens for Harmony + spaces_between_special_tokens: true, + ignore_eos: false, + no_stop_trim: false, + n: 1, // Responses API doesn't support n>1 + constraint: None, // No constraints - tools handled by MCP + ..Default::default() + }) + } + fn build_single_constraint_from_plain( params: &GenerateSamplingParams, ) -> Result, String> { diff --git a/sgl-router/src/protocols/chat.rs b/sgl-router/src/protocols/chat.rs index aea9bdd0f..ee73c44d5 100644 --- a/sgl-router/src/protocols/chat.rs +++ b/sgl-router/src/protocols/chat.rs @@ -674,7 +674,6 @@ pub struct ChatMessageDelta { pub struct ChatStreamChoice { pub index: u32, pub delta: ChatMessageDelta, - #[serde(skip_serializing_if = "Option::is_none")] pub logprobs: Option, pub finish_reason: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/sgl-router/src/protocols/responses.rs b/sgl-router/src/protocols/responses.rs index f48e6f9ef..138277d79 100644 --- a/sgl-router/src/protocols/responses.rs +++ b/sgl-router/src/protocols/responses.rs @@ -21,6 +21,9 @@ use super::common::{ pub struct ResponseTool { #[serde(rename = "type")] pub r#type: ResponseToolType, + // Function tool fields (used when type == "function") + #[serde(skip_serializing_if = "Option::is_none")] + pub function: Option, // MCP-specific fields (used when type == "mcp") #[serde(skip_serializing_if = "Option::is_none")] pub server_url: Option, @@ -40,6 +43,7 @@ impl Default for ResponseTool { fn default() -> Self { Self { r#type: ResponseToolType::WebSearchPreview, + function: None, server_url: None, authorization: None, server_label: None, @@ -53,6 +57,7 @@ impl Default for ResponseTool { #[derive(Debug, Clone, Deserialize, Serialize)] #[serde(rename_all = "snake_case")] pub enum ResponseToolType { + Function, WebSearchPreview, CodeInterpreter, Mcp, @@ -134,6 +139,13 @@ pub enum ResponseInputOutputItem { #[serde(skip_serializing_if = "Option::is_none")] status: Option, }, + #[serde(rename = "function_call_output")] + FunctionCallOutput { + call_id: String, + output: String, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + }, #[serde(untagged)] SimpleInputMessage { content: StringOrContentParts, @@ -499,7 +511,7 @@ pub struct ResponsesRequest { pub store: Option, /// Whether to stream the response - #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] pub stream: Option, /// Temperature for sampling @@ -678,6 +690,9 @@ impl GenerationRequest for ResponsesRequest { ResponseInputOutputItem::FunctionToolCall { arguments, .. } => { Some(arguments.clone()) } + ResponseInputOutputItem::FunctionCallOutput { output, .. } => { + Some(output.clone()) + } }) .collect::>() .join(" "), diff --git a/sgl-router/src/routers/grpc/context.rs b/sgl-router/src/routers/grpc/context.rs index 1d251dd64..e6f9c1725 100644 --- a/sgl-router/src/routers/grpc/context.rs +++ b/sgl-router/src/routers/grpc/context.rs @@ -15,6 +15,7 @@ use crate::{ protocols::{ chat::{ChatCompletionRequest, ChatCompletionResponse}, generate::{GenerateRequest, GenerateResponse}, + responses::ResponsesRequest, }, reasoning_parser::ParserFactory as ReasoningParserFactory, tokenizer::{stop::StopSequenceDecoder, traits::Tokenizer}, @@ -53,6 +54,7 @@ pub struct RequestInput { pub enum RequestType { Chat(Arc), Generate(Arc), + Responses(Arc), } /// Shared components (injected once at creation) @@ -104,6 +106,19 @@ pub struct PreparationOutput { /// Filtered request (if tools were filtered) pub filtered_request: Option, + + // Harmony-specific fields + /// Whether this is a Harmony request (default: false) + pub harmony_mode: bool, + + /// Selection text for worker routing (Harmony only) + pub selection_text: Option, + + /// Harmony messages for history tracking (Harmony only) + pub harmony_messages: Option>, + + /// Stop token IDs for Harmony models + pub harmony_stop_ids: Option>, } /// Worker selection (Step 2) @@ -155,6 +170,16 @@ pub struct ResponseState { /// Final processed response pub final_response: Option, + + /// Responses API iteration result (Harmony only, for tool loop orchestration) + pub responses_iteration_result: Option, + + // Harmony-specific parser state + /// Harmony parser for non-streaming (single parser for all indices) + pub harmony_parser: Option, + + /// Harmony parsers for streaming (one per index for n>1 support) + pub harmony_parser_per_index: Option>, } /// Streaming state (per-choice tracking) @@ -217,6 +242,24 @@ impl RequestContext { } } + /// Create context for Responses API request + pub fn for_responses( + request: Arc, + headers: Option, + model_id: Option, + components: Arc, + ) -> Self { + Self { + input: RequestInput { + request_type: RequestType::Responses(request), + headers, + model_id, + }, + components, + state: ProcessingState::default(), + } + } + /// Get reference to original request (type-safe) pub fn request(&self) -> &RequestType { &self.input.request_type @@ -254,11 +297,28 @@ impl RequestContext { } } + /// Get responses request (panics if not responses) + pub fn responses_request(&self) -> &ResponsesRequest { + match &self.input.request_type { + RequestType::Responses(req) => req.as_ref(), + _ => panic!("Expected responses request"), + } + } + + /// Get Arc clone of responses request (panics if not responses) + pub fn responses_request_arc(&self) -> Arc { + match &self.input.request_type { + RequestType::Responses(req) => Arc::clone(req), + _ => panic!("Expected responses request"), + } + } + /// Check if request is streaming pub fn is_streaming(&self) -> bool { match &self.input.request_type { RequestType::Chat(req) => req.stream, RequestType::Generate(req) => req.stream, + RequestType::Responses(req) => req.stream.unwrap_or(false), } } } diff --git a/sgl-router/src/routers/grpc/harmony/builder.rs b/sgl-router/src/routers/grpc/harmony/builder.rs new file mode 100644 index 000000000..03753eead --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/builder.rs @@ -0,0 +1,880 @@ +//! Harmony request builder +//! +//! Handles encoding of Chat/Responses requests into Harmony format using openai-harmony library. + +use std::sync::OnceLock; + +use chrono::Local; +use openai_harmony::{ + chat::{ + Author, ChannelConfig, Content, Conversation, DeveloperContent, Message as HarmonyMessage, + ReasoningEffort, Role, SystemContent, TextContent, ToolDescription, + }, + HarmonyEncoding, HarmonyEncodingName, +}; +use tracing::debug; + +use super::types::HarmonyBuildOutput; +use crate::protocols::{ + chat::{ChatCompletionRequest, ChatMessage, UserMessageContent}, + common::{ContentPart, Tool}, + responses::{ + ReasoningEffort as ResponsesReasoningEffort, ResponseContentPart, ResponseInput, + ResponseInputOutputItem, ResponseReasoningContent, ResponseTool, ResponseToolType, + ResponsesRequest, StringOrContentParts, + }, +}; + +/// Global Harmony encoding (lazy-initialized) +static HARMONY_ENCODING: OnceLock = OnceLock::new(); + +/// Get or initialize the Harmony encoding +/// +/// Uses HarmonyGptOss encoding which supports the gpt-oss model family. +pub(super) fn get_harmony_encoding() -> &'static HarmonyEncoding { + HARMONY_ENCODING.get_or_init(|| { + openai_harmony::load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss) + .expect("Failed to load Harmony encoding") + }) +} + +/// Built-in tools that are added to the system message +const BUILTIN_TOOLS: &[&str] = &["web_search_preview", "code_interpreter", "container"]; + +/// Trait for tool-like objects that can be converted to Harmony ToolDescription +trait ToolLike { + /// Check if this is a built-in tool (should be skipped in developer message) + #[allow(dead_code)] + fn is_builtin(&self) -> bool; + + /// Check if this is a custom tool (function or MCP) + fn is_custom(&self) -> bool; + + /// Convert to ToolDescription + fn to_tool_description(&self) -> Option; +} + +/// Implement ToolLike for Chat Completion Tool +impl ToolLike for Tool { + fn is_builtin(&self) -> bool { + matches!( + self.tool_type.as_str(), + "web_search_preview" | "code_interpreter" | "container" + ) + } + + fn is_custom(&self) -> bool { + matches!(self.tool_type.as_str(), "mcp" | "function") + } + + fn to_tool_description(&self) -> Option { + Some(ToolDescription::new( + self.function.name.clone(), + self.function.description.clone().unwrap_or_default(), + Some(self.function.parameters.clone()), + )) + } +} + +/// Implement ToolLike for Responses API Tool +impl ToolLike for ResponseTool { + fn is_builtin(&self) -> bool { + matches!( + self.r#type, + ResponseToolType::WebSearchPreview | ResponseToolType::CodeInterpreter + ) + } + + fn is_custom(&self) -> bool { + matches!( + self.r#type, + ResponseToolType::Mcp | ResponseToolType::Function + ) + } + + fn to_tool_description(&self) -> Option { + self.function.as_ref().map(|func| { + ToolDescription::new( + func.name.clone(), + func.description.clone().unwrap_or_default(), + Some(func.parameters.clone()), + ) + }) + } +} + +fn has_custom_tools(tool_types: &[&str]) -> bool { + !tool_types.iter().all(|t| BUILTIN_TOOLS.contains(t)) +} + +/// Harmony request builder +/// +/// Converts OpenAI-format requests into Harmony-encoded format with input_ids, +/// stop tokens, and selection text for worker routing. +pub struct HarmonyBuilder { + encoding: &'static HarmonyEncoding, +} + +impl HarmonyBuilder { + /// Create a new Harmony builder + pub fn new() -> Self { + Self { + encoding: get_harmony_encoding(), + } + } + + /// Build Harmony request from Chat Completion request + /// + /// # Arguments + /// + /// * `request` - The ChatCompletionRequest to encode + /// + /// # Returns + /// + /// HarmonyBuildOutput containing input_ids, stop_token_ids, selection_text, and messages + pub fn build_from_chat( + &self, + request: &ChatCompletionRequest, + ) -> Result { + let mut all_messages = Vec::new(); + + let sys_msg = self.build_system_message_from_chat(request); + all_messages.push(sys_msg); + + let dev_msg = self.build_developer_message_from_chat(request.tools.as_ref()); + all_messages.push(dev_msg); + + let mut user_messages = self.convert_chat_messages(&request.messages)?; + all_messages.append(&mut user_messages); + + let conversation = Conversation::from_messages(all_messages.clone()); + let token_ids = self + .encoding + .render_conversation_for_completion(&conversation, Role::Assistant, None) + .map_err(|e| format!("Failed to encode Harmony conversation: {}", e))?; + + let selection_text = self.extract_selection_text(&all_messages); + + // Get stop tokens for Harmony assistant actions (<|return|> and <|call|>) + let stop_token_ids: Vec = self + .encoding + .stop_tokens_for_assistant_actions() + .into_iter() + .flat_map(|set| set.into_iter()) + .collect(); + + Ok(HarmonyBuildOutput { + input_ids: token_ids, + stop_token_ids, + selection_text, + harmony_messages: all_messages + .into_iter() + .map(super::types::HarmonyMessage::from_openai_harmony) + .collect(), + }) + } + + /// Build Harmony request from Responses request + /// + /// # Arguments + /// + /// * `request` - The ResponsesRequest to encode + /// + /// # Returns + /// + /// HarmonyBuildOutput containing input_ids, stop_token_ids, selection_text, and messages + pub fn build_from_responses( + &self, + request: &ResponsesRequest, + ) -> Result { + let all_messages = self.construct_input_messages_with_harmony(request)?; + + let conversation = Conversation::from_messages(all_messages.clone()); + let token_ids = self + .encoding + .render_conversation_for_completion(&conversation, Role::Assistant, None) + .map_err(|e| format!("Failed to encode Harmony conversation: {}", e))?; + + let selection_text = self.extract_selection_text(&all_messages); + + // Get stop tokens for Harmony assistant actions (<|return|> and <|call|>) + let stop_token_ids: Vec = self + .encoding + .stop_tokens_for_assistant_actions() + .into_iter() + .flat_map(|set| set.into_iter()) + .collect(); + + // Decode tokens to see what the model actually receives + let decoded_text = self + .encoding + .tokenizer() + .decode_utf8(&token_ids) + .unwrap_or_else(|_| "".to_string()); + debug!( + token_count = token_ids.len(), + token_preview = ?&token_ids[..token_ids.len().min(20)], + decoded_length = decoded_text.len(), + "Encoded conversation to tokens - decoded text follows:" + ); + debug!("DECODED_TEXT_START\n{}\nDECODED_TEXT_END", decoded_text); + + Ok(HarmonyBuildOutput { + input_ids: token_ids, + stop_token_ids, + selection_text, + harmony_messages: all_messages + .into_iter() + .map(super::types::HarmonyMessage::from_openai_harmony) + .collect(), + }) + } + + /// Build system message from ChatCompletionRequest + /// Build system message with common logic + /// + /// # Arguments + /// * `reasoning_effort` - Optional reasoning effort level + /// * `has_tools` - Whether custom tools are present + fn build_system_message( + &self, + reasoning_effort: Option, + has_tools: bool, + ) -> HarmonyMessage { + let mut sys_content = SystemContent::new(); + + // Add reasoning_effort if provided + if let Some(effort) = reasoning_effort { + sys_content = sys_content.with_reasoning_effort(effort); + } + + // Set conversation start date (always current date) + sys_content = + sys_content.with_conversation_start_date(Local::now().format("%Y-%m-%d").to_string()); + + // If no tools, remove "commentary" from valid channels + if !has_tools { + if let Some(channel_config) = &sys_content.channel_config { + let valid_channels: Vec = channel_config + .valid_channels + .iter() + .filter(|c| c.as_str() != "commentary") + .cloned() + .collect(); + sys_content = sys_content + .with_channel_config(ChannelConfig::require_channels(valid_channels)); + } + } + + HarmonyMessage::from_role_and_content(Role::System, sys_content) + } + + fn build_system_message_from_chat(&self, request: &ChatCompletionRequest) -> HarmonyMessage { + let reasoning_effort = request + .reasoning_effort + .as_deref() + .map(|effort| match effort { + "high" => ReasoningEffort::High, + "medium" => ReasoningEffort::Medium, + "low" => ReasoningEffort::Low, + _ => ReasoningEffort::Medium, + }); + + let has_tools = request.tools.is_some(); + self.build_system_message(reasoning_effort, has_tools) + } + + /// Build system message from ResponsesRequest + /// + /// # Arguments + /// * `request` - The ResponsesRequest + /// * `with_custom_tools` - Whether custom tools (beyond built-ins) are present + fn build_system_message_from_responses( + &self, + request: &ResponsesRequest, + with_custom_tools: bool, + ) -> HarmonyMessage { + let reasoning_effort = request + .reasoning + .as_ref() + .and_then(|r| r.effort.as_ref()) + .map(|effort| match effort { + ResponsesReasoningEffort::High => ReasoningEffort::High, + ResponsesReasoningEffort::Medium => ReasoningEffort::Medium, + ResponsesReasoningEffort::Low => ReasoningEffort::Low, + }); + + self.build_system_message(reasoning_effort, with_custom_tools) + } + + /// Build developer message with common logic + /// + /// Filters out built-in tools and converts custom tools to ToolDescription + /// + /// # Arguments + /// * `tools` - Optional list of tools + /// * `instructions` - Optional instructions (Responses API only) + fn build_developer_message( + &self, + tools: Option<&Vec>, + instructions: Option<&str>, + ) -> HarmonyMessage { + let mut dev_content = DeveloperContent::new(); + + // Add instructions if provided (Responses API only) + if let Some(instructions) = instructions { + dev_content = dev_content.with_instructions(instructions.to_string()); + } + + // Early return if no tools + let Some(tools) = tools else { + return HarmonyMessage::from_role_and_content(Role::Developer, dev_content); + }; + + // Filter to custom tools and convert to ToolDescription + let tool_descriptions: Vec = tools + .iter() + .filter(|t| t.is_custom()) + .filter_map(|t| t.to_tool_description()) + .collect(); + + // Add function tools to developer content + if !tool_descriptions.is_empty() { + dev_content = dev_content.with_function_tools(tool_descriptions); + } + + HarmonyMessage::from_role_and_content(Role::Developer, dev_content) + } + + fn build_developer_message_from_chat(&self, tools: Option<&Vec>) -> HarmonyMessage { + self.build_developer_message(tools, None) + } + + /// Build developer message from Responses request + /// + /// # Arguments + /// * `instructions` - Optional instructions (Responses API specific) + /// * `tools` - Optional list of tools + fn build_developer_message_from_responses( + &self, + instructions: Option<&str>, + tools: Option<&Vec>, + ) -> HarmonyMessage { + self.build_developer_message(tools, instructions) + } + + /// Construct input messages for Responses API with Harmony + /// + /// Handles both new conversations and continuations of previous responses. + /// + /// This handles: + /// - New conversation: system message, developer message, and user input + /// - Continuing conversation: loads previous messages, cleans up chain-of-thoughts + /// - MCP tool allowlisting for special tool types + /// - Complex response input parsing with function call tracking + /// + /// # Arguments + /// * `request` - The ResponsesRequest + /// * `prev_response` - Optional previous response to continue from + fn construct_input_messages_with_harmony( + &self, + request: &ResponsesRequest, + ) -> Result, String> { + let mut all_messages = Vec::new(); + + // Handle new vs continuing conversation + if request.previous_response_id.is_none() { + // New conversation + + let tool_types: Vec<&str> = request + .tools + .as_ref() + .map(|tools| { + tools + .iter() + .map(|tool| match tool.r#type { + ResponseToolType::Function => "function", + ResponseToolType::WebSearchPreview => "web_search_preview", + ResponseToolType::CodeInterpreter => "code_interpreter", + ResponseToolType::Mcp => "mcp", + }) + .collect() + }) + .unwrap_or_default(); + + let with_custom_tools = has_custom_tools(&tool_types); + + // Add system message + let sys_msg = self.build_system_message_from_responses(request, with_custom_tools); + all_messages.push(sys_msg); + + // Add developer message only if we have custom tools + if with_custom_tools { + let dev_msg = self.build_developer_message_from_responses( + request.instructions.as_deref(), + request.tools.as_ref(), + ); + all_messages.push(dev_msg); + } + } else { + // Continue the previous conversation + // NOTE: Previous messages are loaded by serve_harmony_responses() before calling this method. + // The request.input will already contain the conversation history when previous_response_id was set. + // We just proceed with parsing the input items as normal. + debug!("Continuing conversation (history already loaded in request.input)"); + } + + // Append the new input + // Responses API supports simple text inputs without chat format + match &request.input { + ResponseInput::Text(text) => { + let user_msg = HarmonyMessage { + author: Author { + role: Role::User, + name: None, + }, + recipient: None, + content: vec![Content::Text(TextContent { text: text.clone() })], + channel: None, + content_type: None, + }; + all_messages.push(user_msg); + } + ResponseInput::Items(items) => { + // Track function calls for looking up call_id → name mapping + let mut prev_outputs: Vec<&ResponseInputOutputItem> = Vec::new(); + + for item in items { + let msg = self.parse_response_item_to_harmony_message(item, &prev_outputs)?; + all_messages.push(msg); + + // Track function tool calls so that function_call_output can find the name + if matches!(item, ResponseInputOutputItem::FunctionToolCall { .. }) { + prev_outputs.push(item); + } + } + } + } + + debug!( + message_count = all_messages.len(), + "Constructed Harmony messages for Responses API" + ); + Ok(all_messages) + } + + /// Parse a ResponseInputOutputItem into a HarmonyMessage + /// + /// Handles conversion of various response item types (messages, function calls, reasoning, etc.) + /// to Harmony message format. + /// + /// # Arguments + /// * `item` - The ResponseInputOutputItem to parse + /// * `prev_outputs` - Previous items for looking up function call names (for function_call_output) + fn parse_response_item_to_harmony_message( + &self, + item: &ResponseInputOutputItem, + prev_outputs: &[&ResponseInputOutputItem], + ) -> Result { + match item { + // Regular message (user or assistant) + ResponseInputOutputItem::Message { role, content, .. } => { + let harmony_role = match role.as_str() { + "user" => Role::User, + "assistant" => Role::Assistant, + "system" => Role::System, + _ => Role::User, // Default to user for unknown roles + }; + + // Extract text from content parts + let text_parts: Vec = content + .iter() + .filter_map(|part| match part { + ResponseContentPart::OutputText { text, .. } => Some(text.clone()), + ResponseContentPart::InputText { text } => Some(text.clone()), + ResponseContentPart::Unknown => None, + }) + .collect(); + + let text = text_parts.join("\n"); + + Ok(HarmonyMessage { + author: Author { + role: harmony_role, + name: None, + }, + recipient: None, + content: vec![Content::Text(TextContent { text })], + channel: None, + content_type: None, + }) + } + + // Reasoning content (chain-of-thought) + ResponseInputOutputItem::Reasoning { content, .. } => { + // Extract reasoning text + let reasoning_texts: Vec = content + .iter() + .map(|rc| match rc { + ResponseReasoningContent::ReasoningText { text } => text.clone(), + }) + .collect(); + + let text = reasoning_texts.join("\n"); + + // Reasoning goes in the "analysis" channel for Harmony + Ok(HarmonyMessage { + author: Author { + role: Role::Assistant, + name: None, + }, + recipient: None, + content: vec![Content::Text(TextContent { text })], + channel: Some("analysis".to_string()), + content_type: None, + }) + } + + // Function tool call (with optional output) + ResponseInputOutputItem::FunctionToolCall { + name, + arguments, + output, + .. + } => { + // If there's an output, this represents the tool result + // Otherwise, it's the tool call itself + if let Some(output_str) = output { + // Tool result - use Tool role with "functions.{name}" as author name + // IMPORTANT: Must include recipient="assistant" for parser to recognize it. + // We keep channel=None to minimize what the model might copy. + let author_name = format!("functions.{}", name); + debug!( + tool_name = %name, + author_name = %author_name, + output_preview = %output_str.chars().take(100).collect::(), + "Building tool result message with Tool role (recipient=assistant, no channel)" + ); + Ok(HarmonyMessage { + author: Author { + role: Role::Tool, + name: Some(author_name), + }, + recipient: Some("assistant".to_string()), + content: vec![Content::Text(TextContent { + text: output_str.clone(), + })], + channel: None, + content_type: None, + }) + } else { + // Tool call - assistant message in commentary channel with recipient + // msg.with_channel("commentary").with_recipient(f"functions.{name}") + let recipient = format!("functions.{}", name); + debug!( + tool_name = %name, + recipient = %recipient, + "Building tool call message with recipient" + ); + Ok(HarmonyMessage { + author: Author { + role: Role::Assistant, + name: None, + }, + recipient: Some(recipient), + content: vec![Content::Text(TextContent { + text: arguments.clone(), + })], + channel: Some("commentary".to_string()), + content_type: Some("json".to_string()), + }) + } + } + + // Function call output (separate from call) - requires looking up the original call + ResponseInputOutputItem::FunctionCallOutput { + call_id, output, .. + } => { + // Search prev_outputs in reverse order to find the matching function call + let call = prev_outputs + .iter() + .rev() + .find_map(|item| match item { + ResponseInputOutputItem::FunctionToolCall { id, name, .. } + if id == call_id => + { + Some(name.clone()) + } + _ => None, + }) + .ok_or_else(|| format!("No function call found for call_id: {}", call_id))?; + + // Create Tool message with "functions.{name}" prefix + // IMPORTANT: Must include recipient="assistant" for parser to recognize it. + // We keep channel=None to minimize what the model might copy. + Ok(HarmonyMessage { + author: Author { + role: Role::Tool, + name: Some(format!("functions.{}", call)), + }, + recipient: Some("assistant".to_string()), + content: vec![Content::Text(TextContent { + text: output.clone(), + })], + channel: None, + content_type: None, + }) + } + + // Simple input message (usually user message) + ResponseInputOutputItem::SimpleInputMessage { content, role, .. } => { + let harmony_role = match role.as_str() { + "user" => Role::User, + "assistant" => Role::Assistant, + "system" => Role::System, + _ => Role::User, + }; + + let text = match content { + StringOrContentParts::String(s) => s.clone(), + StringOrContentParts::Array(parts) => { + // Extract text from content parts + parts + .iter() + .filter_map(|part| match part { + ResponseContentPart::OutputText { text, .. } => Some(text.clone()), + ResponseContentPart::InputText { text } => Some(text.clone()), + ResponseContentPart::Unknown => None, + }) + .collect::>() + .join("\n") + } + }; + + Ok(HarmonyMessage { + author: Author { + role: harmony_role, + name: None, + }, + recipient: None, + content: vec![Content::Text(TextContent { text })], + channel: None, + content_type: None, + }) + } + } + } + + /// Convert OpenAI ChatMessage format to Harmony messages + /// + /// - Assistant messages with tool_calls create multiple messages (one per tool call) + /// - Tool role messages use Role::Tool with proper author + /// - Tool-related messages use channel="commentary" + fn convert_chat_messages( + &self, + messages: &[ChatMessage], + ) -> Result, String> { + let mut harmony_messages = Vec::new(); + + // Build a map of tool_call_id -> function_name for tool responses + let mut tool_call_map = std::collections::HashMap::new(); + for msg in messages { + if let ChatMessage::Assistant { + tool_calls: Some(calls), + .. + } = msg + { + for call in calls { + tool_call_map.insert(call.id.clone(), call.function.name.clone()); + } + } + } + + for msg in messages { + match msg { + ChatMessage::System { content, name } => { + // System messages stay as-is + let harmony_msg = HarmonyMessage { + author: Author { + role: Role::System, + name: name.clone(), + }, + recipient: None, + content: vec![Content::Text(TextContent { + text: content.clone(), + })], + channel: None, + content_type: None, + }; + harmony_messages.push(harmony_msg); + } + + ChatMessage::User { content, name } => { + // Extract text from user content + let text = match content { + UserMessageContent::Text(text) => text.clone(), + UserMessageContent::Parts(parts) => { + // For multimodal content, extract text parts + parts + .iter() + .filter_map(|part| { + if let ContentPart::Text { text } = part { + Some(text.as_str()) + } else { + None + } + }) + .collect::>() + .join("\n") + } + }; + + let harmony_msg = HarmonyMessage { + author: Author { + role: Role::User, + name: name.clone(), + }, + recipient: None, + content: vec![Content::Text(TextContent { text })], + channel: None, + content_type: None, + }; + harmony_messages.push(harmony_msg); + } + + ChatMessage::Assistant { + content, + name, + tool_calls, + reasoning_content, + } => { + if let Some(calls) = tool_calls { + // Create one message per tool call with channel="commentary" + for call in calls { + let function_name = &call.function.name; + let arguments = call.function.arguments.clone().unwrap_or_default(); + + let tool_call_msg = HarmonyMessage { + author: Author { + role: Role::Assistant, + name: name.clone(), + }, + recipient: Some(format!("functions.{}", function_name)), + content: vec![Content::Text(TextContent { text: arguments })], + channel: Some("commentary".to_string()), + content_type: Some("json".to_string()), + }; + harmony_messages.push(tool_call_msg); + } + } else { + // Regular assistant message with content + // Combine content with reasoning if present + let mut text = content.clone().unwrap_or_default(); + if let Some(reasoning) = reasoning_content { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(reasoning); + } + + let harmony_msg = HarmonyMessage { + author: Author { + role: Role::Assistant, + name: name.clone(), + }, + recipient: None, + content: vec![Content::Text(TextContent { text })], + channel: Some("final".to_string()), + content_type: None, + }; + harmony_messages.push(harmony_msg); + } + } + + ChatMessage::Tool { + content, + tool_call_id, + } => { + // Look up the function name from the tool_call_id + let function_name = tool_call_map + .get(tool_call_id) + .cloned() + .unwrap_or_else(|| tool_call_id.clone()); + + // Tool result - Must include recipient="assistant" for parser to recognize it. + // We keep channel=None to minimize what the model might copy. + let harmony_msg = HarmonyMessage { + author: Author { + role: Role::Tool, + name: Some(format!("functions.{}", function_name)), + }, + recipient: Some("assistant".to_string()), + content: vec![Content::Text(TextContent { + text: content.clone(), + })], + channel: None, + content_type: None, + }; + harmony_messages.push(harmony_msg); + } + + ChatMessage::Function { content, name } => { + // Function messages also use Role::Tool + // Tool result - Must include recipient="assistant" for parser to recognize it. + // We keep channel=None to minimize what the model might copy. + let harmony_msg = HarmonyMessage { + author: Author { + role: Role::Tool, + name: Some(format!("functions.{}", name)), + }, + recipient: Some("assistant".to_string()), + content: vec![Content::Text(TextContent { + text: content.clone(), + })], + channel: None, + content_type: None, + }; + harmony_messages.push(harmony_msg); + } + } + } + + Ok(harmony_messages) + } + + /// Extract selection text for worker routing + /// + /// Uses the last user message for load balancing + fn extract_selection_text(&self, messages: &[HarmonyMessage]) -> String { + // Find the last user message + if let Some(last_user_msg) = messages.iter().rev().find(|m| m.author.role == Role::User) { + // Extract full text from content + return last_user_msg + .content + .iter() + .filter_map(|c| match c { + Content::Text(tc) => Some(tc.text.as_str()), + _ => None, + }) + .collect::>() + .join(""); + } + + // Fallback: concatenate all text + messages + .iter() + .flat_map(|m| &m.content) + .filter_map(|c| match c { + Content::Text(tc) => Some(tc.text.as_str()), + _ => None, + }) + .collect::>() + .join(" ") + } +} + +impl Default for HarmonyBuilder { + fn default() -> Self { + Self::new() + } +} diff --git a/sgl-router/src/routers/grpc/harmony/detector.rs b/sgl-router/src/routers/grpc/harmony/detector.rs new file mode 100644 index 000000000..88ef6a6a5 --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/detector.rs @@ -0,0 +1,17 @@ +//! Harmony model detection + +/// Harmony model detector +/// +/// Detects if a model name indicates support for Harmony encoding/parsing. +pub struct HarmonyDetector; + +impl HarmonyDetector { + pub fn is_harmony_model(model_name: &str) -> bool { + // Case-insensitive substring search without heap allocation + // More efficient than to_lowercase() which allocates a new String + model_name + .as_bytes() + .windows(7) // "gpt-oss".len() + .any(|window| window.eq_ignore_ascii_case(b"gpt-oss")) + } +} diff --git a/sgl-router/src/routers/grpc/harmony/mod.rs b/sgl-router/src/routers/grpc/harmony/mod.rs new file mode 100644 index 000000000..9178d804c --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/mod.rs @@ -0,0 +1,54 @@ +//! Harmony pipeline implementation +//! +//! This module provides support for GPT-OSS models that use Harmony encoding/parsing. +//! The Harmony protocol uses a channel-based approach with three channels: +//! - **analysis**: Reasoning/thinking content (optional) +//! - **commentary**: Tool calls (optional) +//! - **final**: Final response text (required) +//! +//! ## Architecture +//! +//! The Harmony implementation is structured as follows: +//! +//! - **detector**: Model detection (is this a Harmony-capable model?) +//! - **builder**: Request encoding (convert Chat/Responses → input_ids) +//! - **parser**: Response parsing (output_ids → channels) +//! - **types**: Shared type definitions +//! +//! ## Usage +//! +//! ```ignore +//! use sglang_router_rs::routers::grpc::harmony::{HarmonyDetector, HarmonyBuilder}; +//! +//! // Detect if model supports Harmony +//! if HarmonyDetector::is_harmony_model("gpt-4o") { +//! // Build Harmony request +//! let builder = HarmonyBuilder::new(); +//! let output = builder.build_from_chat(&request)?; +//! // ... use output.input_ids for gRPC request +//! } +//! ``` + +pub mod builder; +pub mod detector; +pub mod parser; +pub mod processor; +pub mod responses; +pub mod stages; +pub mod streaming; +pub mod types; + +// Re-export main types for convenience +pub use builder::HarmonyBuilder; +pub use detector::HarmonyDetector; +pub use parser::HarmonyParserAdapter; +pub use processor::{HarmonyResponseProcessor, ResponsesIterationResult}; +pub use responses::{serve_harmony_responses, HarmonyResponsesContext}; +pub use stages::{ + HarmonyPreparationStage, HarmonyRequestBuildingStage, HarmonyResponseProcessingStage, +}; +pub use streaming::HarmonyStreamingProcessor; +pub use types::{ + FunctionDelta, HarmonyBuildOutput, HarmonyChannelDelta, HarmonyChannelOutput, HarmonyMessage, + ToolCallDelta, +}; diff --git a/sgl-router/src/routers/grpc/harmony/parser.rs b/sgl-router/src/routers/grpc/harmony/parser.rs new file mode 100644 index 000000000..dc45c4c38 --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/parser.rs @@ -0,0 +1,448 @@ +//! Harmony response parser +//! +//! Adapter for openai_harmony::StreamableParser that handles channel-based parsing. + +use openai_harmony::{chat::Role, HarmonyEncoding, StreamableParser}; +use uuid::Uuid; + +use super::types::{HarmonyChannelDelta, HarmonyChannelOutput}; +use crate::protocols::common::{FunctionCallResponse, ToolCall}; + +/// Get the global Harmony encoding +/// +/// References the same encoding used by the builder for consistency +fn get_harmony_encoding() -> &'static HarmonyEncoding { + use super::builder::get_harmony_encoding; + get_harmony_encoding() +} + +/// Harmony parser adapter +/// +/// Wraps openai_harmony::StreamableParser and provides methods for parsing +/// complete responses and streaming chunks. +pub struct HarmonyParserAdapter { + parser: StreamableParser, + prev_recipient: Option, +} + +impl HarmonyParserAdapter { + /// Create a new Harmony parser + pub fn new() -> Result { + let encoding = get_harmony_encoding(); + let parser = StreamableParser::new(encoding.clone(), Some(Role::Assistant)) + .map_err(|e| format!("Failed to create StreamableParser: {}", e))?; + + Ok(Self { + parser, + prev_recipient: None, + }) + } + + /// Extract text from message content (private helper) + /// + /// Filters text content from a message's content array and joins them into a single string. + /// + /// # Arguments + /// + /// * `content` - The content array from a Harmony message + /// + /// # Returns + /// + /// Joined text string from all text content items + fn extract_text_from_content(content: &[openai_harmony::chat::Content]) -> String { + content + .iter() + .filter_map(|c| match c { + openai_harmony::chat::Content::Text(tc) => Some(tc.text.as_str()), + _ => None, + }) + .collect::>() + .join("") + } + + /// Handle incomplete content from parser state (private helper) + /// + /// Checks for any remaining incomplete content in the parser and appends it + /// to the appropriate channel (analysis or final_text). + /// + /// # Arguments + /// + /// * `parser` - Reference to the StreamableParser + /// * `analysis` - Mutable reference to analysis content + /// * `final_text` - Mutable reference to final text content + fn handle_incomplete_content( + parser: &StreamableParser, + analysis: &mut Option, + final_text: &mut String, + ) { + if let Ok(current_content) = parser.current_content() { + if !current_content.is_empty() { + let current_channel = parser.current_channel(); + match current_channel.as_deref() { + Some("analysis") => { + *analysis = Some(current_content); + } + Some("final") | None => { + final_text.push_str(¤t_content); + } + _ => {} + } + } + } + } + + /// Parse messages into channel outputs (private helper) + /// + /// Extracts analysis, commentary (tool calls), and final text from Harmony messages. + /// This is the core parsing logic shared by both parse_complete and finalize. + /// + /// # Arguments + /// + /// * `messages` - The messages to parse from the Harmony parser + /// + /// # Returns + /// + /// Tuple of (analysis, commentary, final_text) + fn parse_messages( + messages: &[openai_harmony::chat::Message], + ) -> (Option, Option>, String) { + let mut analysis = None; + let mut commentary: Option> = None; + let mut final_text = String::new(); + + for msg in messages { + // Filter: Only process assistant messages + if msg.author.role != Role::Assistant { + continue; + } + + let channel = msg.channel.as_deref().unwrap_or(""); + let recipient = msg.recipient.as_deref(); + + match channel { + "analysis" => { + // Process each content item + // For Chat API, we join them into a single reasoning_content + let text = Self::extract_text_from_content(&msg.content); + + if !text.is_empty() { + analysis = Some(text); + } + } + "commentary" => { + // Handle different recipient types + if let Some(recipient_str) = recipient { + if recipient_str.starts_with("functions.") { + let function_name = recipient_str.strip_prefix("functions.").unwrap(); + + // Process each content item separately + for content in &msg.content { + if let openai_harmony::chat::Content::Text(tc) = content { + let call_id = format!("call_{}", Uuid::new_v4()); + let tool_call = ToolCall { + id: call_id, + tool_type: "function".to_string(), + function: FunctionCallResponse { + name: function_name.to_string(), + arguments: Some(tc.text.clone()), + }, + }; + + match commentary.as_mut() { + Some(calls) => calls.push(tool_call), + None => commentary = Some(vec![tool_call]), + } + } + } + } else if recipient_str.starts_with("python") + || recipient_str.starts_with("browser") + || recipient_str.starts_with("container") + { + // Built-in tools → treat as reasoning + // For Chat API, we add to analysis content + let text = Self::extract_text_from_content(&msg.content); + + if !text.is_empty() { + // Append to analysis (built-in tools are reasoning) + match analysis.as_mut() { + Some(existing) => { + existing.push('\n'); + existing.push_str(&text); + } + None => analysis = Some(text), + } + } + } + // Unknown recipient would raise ValueError + // For now, we silently ignore (can add logging later) + } + } + "final" => { + // Process final channel content + let text = Self::extract_text_from_content(&msg.content); + final_text.push_str(&text); + } + _ => { + // Unknown channel, append to final text as fallback + let text = Self::extract_text_from_content(&msg.content); + final_text.push_str(&text); + } + } + } + + (analysis, commentary, final_text) + } + + /// Parse complete response + /// + /// Parses all output token IDs and returns the complete channel output + /// containing analysis, commentary (tool calls), and final text. + /// + /// # Arguments + /// + /// * `output_ids` - The complete output token IDs from the model + /// * `finish_reason` - The finish reason from GenerateComplete ("stop", "length", etc.) + /// * `matched_stop` - Optional matched stop token information from GenerateComplete + /// + /// # Returns + /// + /// Complete HarmonyChannelOutput with all three channels parsed + pub fn parse_complete( + &mut self, + output_ids: &[u32], + finish_reason: String, + matched_stop: Option, + ) -> Result { + // Feed all tokens to the parser + for &token_id in output_ids { + self.parser.process(token_id).map_err(|e| { + // Log the full output_ids context on error + tracing::error!( + token_id = token_id, + output_ids = ?output_ids, + error = %e, + "Harmony parser failed to process token" + ); + format!("Failed to process token {}: {}", token_id, e) + })?; + } + + // Extract all completed messages from the parser + let messages = self.parser.messages(); + + // Parse messages into channel outputs using shared helper + let (mut analysis, commentary, mut final_text) = Self::parse_messages(messages); + + // Check for incomplete content in parser state + Self::handle_incomplete_content(&self.parser, &mut analysis, &mut final_text); + + // Determine finish reason: override to "tool_calls" if commentary has tool calls + let final_finish_reason = if commentary.is_some() { + "tool_calls".to_string() + } else { + finish_reason + }; + + Ok(HarmonyChannelOutput { + analysis, + commentary, + final_text, + finish_reason: final_finish_reason, + matched_stop, + }) + } + + /// Get all messages from the parser + /// + /// Returns the raw messages extracted by the Harmony parser. + /// Used for validation checks. + pub fn get_messages(&self) -> Vec { + self.parser.messages().to_vec() + } + + /// Parse streaming chunk + /// + /// Parses incremental token IDs and returns a delta with any new content + /// from the analysis, commentary, or final channels. + /// + /// # Arguments + /// + /// * `chunk_ids` - New token IDs from the current chunk + /// + /// # Returns + /// + /// Optional HarmonyChannelDelta if there's new content to emit + pub fn parse_chunk( + &mut self, + chunk_ids: &[u32], + ) -> Result, String> { + let mut has_delta = false; + let mut analysis_delta = None; + let mut commentary_delta = None; + let mut final_delta = None; + + // Track message count before processing + let prev_message_count = self.parser.messages().len(); + + // Accumulate delta text for commentary channel + let mut accumulated_delta = String::new(); + + // Process each token + for &token_id in chunk_ids { + self.parser + .process(token_id) + .map_err(|e| format!("Failed to process token {}: {}", token_id, e))?; + + // Check for content delta + if let Ok(Some(delta_text)) = self.parser.last_content_delta() { + has_delta = true; + + // Determine which channel this delta belongs to + let channel = self.parser.current_channel(); + match channel.as_deref() { + Some("analysis") => { + analysis_delta = Some(delta_text); + } + Some("final") | None => { + final_delta = Some(delta_text); + } + Some("commentary") => { + // Accumulate delta for commentary + accumulated_delta.push_str(&delta_text); + } + _ => {} + } + } + } + + // Handle commentary channel tool call deltas + if self.parser.current_channel().as_deref() == Some("commentary") { + if let Some(cur_recipient) = self.parser.current_recipient() { + if cur_recipient.starts_with("functions.") { + has_delta = true; + + // Count completed tool calls for index + let base_index = self + .parser + .messages() + .iter() + .filter(|msg| { + msg.channel.as_deref() == Some("commentary") + && msg + .recipient + .as_deref() + .is_some_and(|r| r.starts_with("functions.")) + }) + .count(); + + // Check if recipient changed (new tool call) + let recipient_changed = self.prev_recipient.as_deref() != Some(&cur_recipient); + + if recipient_changed { + // NEW tool call: emit name + id + let tool_name = cur_recipient.strip_prefix("functions.").unwrap(); + let call_id = format!("call_{}", Uuid::new_v4()); + + commentary_delta = Some(super::types::ToolCallDelta { + index: base_index, + id: Some(call_id), + function: Some(super::types::FunctionDelta { + name: Some(tool_name.to_string()), + arguments: Some(String::new()), + }), + }); + + // Update prev_recipient + self.prev_recipient = Some(cur_recipient); + } else if !accumulated_delta.is_empty() { + // CONTINUING tool call: emit arguments delta + commentary_delta = Some(super::types::ToolCallDelta { + index: base_index, + id: None, + function: Some(super::types::FunctionDelta { + name: None, + arguments: Some(accumulated_delta), + }), + }); + } + } + } + } + + // Check if new messages were completed + let current_message_count = self.parser.messages().len(); + let is_final = current_message_count > prev_message_count; + + if has_delta { + Ok(Some(HarmonyChannelDelta { + analysis_delta, + commentary_delta, + final_delta, + is_final, + })) + } else { + Ok(None) + } + } + + /// Finalize parsing + /// + /// Called at the end of streaming to get the final state and any + /// remaining content. + /// + /// # Arguments + /// + /// * `finish_reason` - The finish reason from GenerateComplete ("stop", "length", etc.) + /// * `matched_stop` - Optional matched stop token information from GenerateComplete + /// + /// # Returns + /// + /// Final HarmonyChannelOutput with complete parsed content + pub fn finalize( + &mut self, + finish_reason: String, + matched_stop: Option, + ) -> Result { + // Extract all completed messages + let messages = self.parser.messages(); + + // Parse messages into channel outputs using shared helper + let (mut analysis, commentary, mut final_text) = Self::parse_messages(messages); + + // Check for remaining incomplete content + Self::handle_incomplete_content(&self.parser, &mut analysis, &mut final_text); + + // Determine finish reason: override to "tool_calls" if commentary has tool calls + let final_finish_reason = if commentary.is_some() { + "tool_calls".to_string() + } else { + finish_reason + }; + + Ok(HarmonyChannelOutput { + analysis, + commentary, + final_text, + finish_reason: final_finish_reason, + matched_stop, + }) + } + + /// Reset parser state + /// + /// Resets the parser to initial state for reuse + pub fn reset(&mut self) -> Result<(), String> { + // Create a new parser instance (StreamableParser doesn't have a reset method) + let encoding = get_harmony_encoding(); + self.parser = StreamableParser::new(encoding.clone(), Some(Role::Assistant)) + .map_err(|e| format!("Failed to reset parser: {}", e))?; + self.prev_recipient = None; + Ok(()) + } +} + +impl Default for HarmonyParserAdapter { + fn default() -> Self { + Self::new().expect("Failed to create default parser") + } +} diff --git a/sgl-router/src/routers/grpc/harmony/processor.rs b/sgl-router/src/routers/grpc/harmony/processor.rs new file mode 100644 index 000000000..bb5be6128 --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/processor.rs @@ -0,0 +1,336 @@ +//! Harmony response processor for non-streaming responses + +use std::sync::Arc; + +use axum::response::Response; +use proto::generate_complete::MatchedStop::{MatchedStopStr, MatchedTokenId}; + +use super::HarmonyParserAdapter; +use crate::{ + grpc_client::proto, + protocols::{ + chat::{ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse}, + common::{ToolCall, Usage}, + responses::{ + ResponseContentPart, ResponseOutputItem, ResponseReasoningContent, ResponseStatus, + ResponseUsage, ResponsesRequest, ResponsesResponse, ResponsesUsage, + }, + }, + routers::grpc::{ + context::{DispatchMetadata, ExecutionResult}, + utils, + }, +}; + +/// Processor for non-streaming Harmony responses +/// +/// Collects all output tokens from execution and parses them using +/// HarmonyParserAdapter to extract the complete response. +pub struct HarmonyResponseProcessor; + +impl HarmonyResponseProcessor { + /// Create a new Harmony response processor + pub fn new() -> Self { + Self + } + + /// Collect responses from ExecutionResult (similar to regular processor) + async fn collect_responses( + execution_result: ExecutionResult, + ) -> Result, Response> { + match execution_result { + ExecutionResult::Single { mut stream } => { + let responses = utils::collect_stream_responses(&mut stream, "Single").await?; + stream.mark_completed(); + Ok(responses) + } + ExecutionResult::Dual { prefill, decode } => { + // For Harmony we currently rely only on decode stream for outputs + let mut decode_stream = *decode; + let responses = + utils::collect_stream_responses(&mut decode_stream, "Decode").await?; + prefill.mark_completed(); + decode_stream.mark_completed(); + Ok(responses) + } + } + } + + /// Process a non-streaming Harmony chat response + pub async fn process_non_streaming_chat_response( + &self, + execution_result: ExecutionResult, + chat_request: Arc, + dispatch: DispatchMetadata, + ) -> Result { + // Collect all completed responses (one per choice) + let all_responses = Self::collect_responses(execution_result).await?; + if all_responses.is_empty() { + return Err(utils::internal_error_static("No responses from server")); + } + + // Build choices by parsing output with HarmonyParserAdapter + let mut choices: Vec = Vec::new(); + for (index, complete) in all_responses.iter().enumerate() { + // Convert matched_stop from proto to JSON + let matched_stop = complete.matched_stop.as_ref().map(|m| match m { + MatchedTokenId(id) => { + serde_json::json!(id) + } + MatchedStopStr(s) => { + serde_json::json!(s) + } + }); + + // Parse Harmony channels with HarmonyParserAdapter + let mut parser = HarmonyParserAdapter::new().map_err(|e| { + utils::internal_error_message(format!("Failed to create Harmony parser: {}", e)) + })?; + + // Parse Harmony channels with finish_reason and matched_stop + let parsed = parser + .parse_complete( + &complete.output_ids, + complete.finish_reason.clone(), + matched_stop.clone(), + ) + .map_err(|e| { + utils::internal_error_message(format!("Harmony parsing failed: {}", e)) + })?; + + // Build response message (assistant) + let message = ChatCompletionMessage { + role: "assistant".to_string(), + content: (!parsed.final_text.is_empty()).then_some(parsed.final_text), + tool_calls: parsed.commentary, + reasoning_content: parsed.analysis, + }; + + let finish_reason = parsed.finish_reason; + + choices.push(ChatChoice { + index: index as u32, + message, + logprobs: None, + finish_reason: Some(finish_reason), + matched_stop, + hidden_states: None, + }); + } + + // Build usage from proto fields + let prompt_tokens: u32 = all_responses.iter().map(|r| r.prompt_tokens as u32).sum(); + let completion_tokens: u32 = all_responses + .iter() + .map(|r| r.completion_tokens as u32) + .sum(); + let usage = Usage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + completion_tokens_details: None, + }; + + // Final ChatCompletionResponse + let response = ChatCompletionResponse { + id: dispatch.request_id.clone(), + object: "chat.completion".to_string(), + created: dispatch.created, + model: chat_request.model.clone(), + choices, + usage: Some(usage), + system_fingerprint: dispatch.weight_version.clone(), + }; + + Ok(response) + } +} + +impl Default for HarmonyResponseProcessor { + fn default() -> Self { + Self::new() + } +} + +/// Result of processing a single Responses API iteration +/// +/// Used by the MCP tool loop to determine whether to continue +/// executing tools or return the final response. +pub enum ResponsesIterationResult { + /// Tool calls found in commentary channel - continue MCP loop + ToolCallsFound { + tool_calls: Vec, + analysis: Option, // For streaming emission + partial_text: String, // For streaming emission + }, + /// No tool calls - return final ResponsesResponse + Completed { + response: Box, + usage: Usage, + }, +} + +impl HarmonyResponseProcessor { + /// Process a single Responses API iteration + /// + /// Parses Harmony channels and determines if tool calls are present. + /// If tool calls found, returns ToolCallsFound for MCP loop to execute. + /// If no tool calls, builds final ResponsesResponse. + /// + /// # Arguments + /// + /// * `execution_result` - The execution result from the model + /// * `responses_request` - The original Responses API request + /// * `dispatch` - Dispatch metadata for request tracking + /// + /// # Returns + /// + /// ResponsesIterationResult indicating whether to continue loop or return + pub async fn process_responses_iteration( + &self, + execution_result: ExecutionResult, + responses_request: Arc, + dispatch: DispatchMetadata, + ) -> Result { + // Collect all completed responses + let all_responses = Self::collect_responses(execution_result).await?; + if all_responses.is_empty() { + return Err(utils::internal_error_static("No responses from server")); + } + + // For Responses API, we only process the first response (n=1) + let complete = all_responses + .first() + .ok_or_else(|| utils::internal_error_static("No complete response"))?; + + // Parse Harmony channels + let mut parser = HarmonyParserAdapter::new().map_err(|e| { + utils::internal_error_message(format!("Failed to create Harmony parser: {}", e)) + })?; + + // Convert matched_stop from proto to JSON + let matched_stop = complete.matched_stop.as_ref().map(|m| match m { + MatchedTokenId(id) => { + serde_json::json!(id) + } + MatchedStopStr(s) => { + serde_json::json!(s) + } + }); + + let parsed = parser + .parse_complete( + &complete.output_ids, + complete.finish_reason.clone(), + matched_stop, + ) + .map_err(|e| utils::internal_error_message(format!("Harmony parsing failed: {}", e)))?; + + // VALIDATION: Check if model incorrectly generated Tool role messages + // This happens when the model copies the format of tool result messages + // instead of continuing as assistant. This is a model hallucination bug. + let messages = parser.get_messages(); + let tool_messages_generated = messages.iter().any(|msg| { + msg.author.role == openai_harmony::chat::Role::Tool + && msg.recipient.as_deref() == Some("assistant") + }); + + if tool_messages_generated { + tracing::warn!( + "Model generated Tool->Assistant message instead of Assistant message. \ + This is a model hallucination bug where it copies tool result format." + ); + } + + // Check for tool calls in commentary channel + if let Some(tool_calls) = parsed.commentary { + // Tool calls found - return for MCP loop execution + return Ok(ResponsesIterationResult::ToolCallsFound { + tool_calls, + analysis: parsed.analysis, + partial_text: parsed.final_text, + }); + } + + // No tool calls - build final ResponsesResponse + let mut output: Vec = Vec::new(); + + // Map analysis channel → ResponseOutputItem::Reasoning + if let Some(analysis) = parsed.analysis { + let reasoning_item = ResponseOutputItem::Reasoning { + id: format!("reasoning_{}", dispatch.request_id), + summary: vec![], + content: vec![ResponseReasoningContent::ReasoningText { text: analysis }], + status: Some("completed".to_string()), + }; + output.push(reasoning_item); + } + + // Map final channel → ResponseOutputItem::Message + if !parsed.final_text.is_empty() { + let message_item = ResponseOutputItem::Message { + id: format!("msg_{}", dispatch.request_id), + role: "assistant".to_string(), + content: vec![ResponseContentPart::OutputText { + text: parsed.final_text, + annotations: vec![], + logprobs: None, + }], + status: "completed".to_string(), + }; + output.push(message_item); + } + + // Build usage + let prompt_tokens = complete.prompt_tokens as u32; + let completion_tokens = complete.completion_tokens as u32; + let usage = Usage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + completion_tokens_details: None, + }; + + // Build ResponsesResponse with all required fields + let response = ResponsesResponse { + id: dispatch.request_id.clone(), + object: "response".to_string(), + created_at: dispatch.created as i64, + status: ResponseStatus::Completed, + error: None, + incomplete_details: None, + instructions: responses_request.instructions.clone(), + max_output_tokens: responses_request.max_output_tokens, + model: responses_request.model.clone(), + output, + parallel_tool_calls: responses_request.parallel_tool_calls.unwrap_or(true), + previous_response_id: responses_request.previous_response_id.clone(), + reasoning: None, // Set by caller if needed + store: responses_request.store.unwrap_or(true), + temperature: responses_request.temperature, + text: None, + tool_choice: responses_request + .tool_choice + .as_ref() + .map(|tc| serde_json::to_string(tc).unwrap_or_else(|_| "auto".to_string())) + .unwrap_or_else(|| "auto".to_string()), + tools: responses_request.tools.clone().unwrap_or_default(), + top_p: responses_request.top_p, + truncation: None, + usage: Some(ResponsesUsage::Modern(ResponseUsage { + input_tokens: prompt_tokens, + output_tokens: completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + input_tokens_details: None, + output_tokens_details: None, + })), + user: None, + metadata: responses_request.metadata.clone().unwrap_or_default(), + }; + + Ok(ResponsesIterationResult::Completed { + response: Box::new(response), + usage, + }) + } +} diff --git a/sgl-router/src/routers/grpc/harmony/responses.rs b/sgl-router/src/routers/grpc/harmony/responses.rs new file mode 100644 index 000000000..fac7180f2 --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/responses.rs @@ -0,0 +1,879 @@ +//! Harmony Responses API implementation with multi-turn MCP tool support +//! +//! This module implements the Harmony Responses API orchestration logic, +//! coordinating full pipeline execution with MCP tool support for multi-turn conversations. +//! +//! ## Architecture +//! +//! Multi-turn pipeline orchestration (NOT just a tool loop): +//! - Serves Harmony Responses API requests end-to-end +//! - Each iteration executes FULL pipeline (worker selection + client acquisition + execution + parsing) +//! - Handles MCP tool execution and history building between iterations +//! - Clean separation: serving orchestration (this file) vs. pipeline stages (stages/) +//! +//! ## Flow +//! +//! ```text +//! loop { +//! // Execute through FULL pipeline +//! let result = pipeline.execute_harmony_responses(&request, &ctx).await?; +//! +//! match result { +//! ToolCallsFound { tool_calls, .. } => { +//! // Execute MCP tools +//! // Build next request with tool results +//! // Continue loop +//! } +//! Completed { response, .. } => { +//! return Ok(response); +//! } +//! } +//! } +//! ``` +//! +//! ## Design Reference +//! +//! See `/Users/simolin/workspace/sglang/.claude/docs/harmony_pipeline/tool_loop_design.md` +//! for complete architecture, rationale, and implementation details. + +use std::sync::Arc; + +use axum::response::Response; +use serde_json::Value as JsonValue; + +use crate::{ + data_connector::{ResponseId, ResponseStorage}, + mcp::McpManager, + protocols::{ + common::{Function, ToolCall}, + responses::{ + ResponseInput, ResponseInputOutputItem, ResponseTool, ResponsesRequest, + ResponsesResponse, StringOrContentParts, + }, + }, + routers::grpc::{ + context::SharedComponents, harmony::processor::ResponsesIterationResult, + pipeline::RequestPipeline, utils, + }, +}; + +/// Maximum number of tool execution iterations to prevent infinite loops +const MAX_TOOL_ITERATIONS: usize = 10; + +/// Record of a single MCP tool call execution +/// +/// Stores metadata needed to build mcp_call output items for Responses API format +#[derive(Debug, Clone)] +struct McpCallRecord { + /// Tool call ID (stored for potential future use, currently generate new IDs) + #[allow(dead_code)] + call_id: String, + /// Tool name + tool_name: String, + /// JSON-encoded arguments + arguments: String, + /// JSON-encoded output/result + output: String, + /// Whether execution succeeded + success: bool, + /// Error message if execution failed + error: Option, +} + +/// Tracking structure for MCP tool calls across iterations +/// +/// Accumulates all MCP tool call metadata during multi-turn conversation +/// so we can build proper mcp_list_tools and mcp_call output items. +#[derive(Debug, Clone)] +struct McpCallTracking { + /// MCP server label (e.g., "sglang-mcp") + server_label: String, + /// All tool call records across all iterations + tool_calls: Vec, +} + +impl McpCallTracking { + fn new(server_label: String) -> Self { + Self { + server_label, + tool_calls: Vec::new(), + } + } + + fn record_call( + &mut self, + call_id: String, + tool_name: String, + arguments: String, + output: String, + success: bool, + error: Option, + ) { + self.tool_calls.push(McpCallRecord { + call_id, + tool_name, + arguments, + output, + success, + error, + }); + } + + fn total_calls(&self) -> usize { + self.tool_calls.len() + } +} + +/// Context for Harmony Responses execution with MCP tool support +/// +/// Contains all dependencies needed for multi-turn Responses API execution. +/// Cheap to clone (all Arc references). +#[derive(Clone)] +pub struct HarmonyResponsesContext { + /// Pipeline for executing Harmony requests + pub pipeline: Arc, + + /// Shared components (tokenizer, parsers) + pub components: Arc, + + /// MCP manager for tool execution + pub mcp_manager: Arc, + + /// Response storage for loading conversation history + pub response_storage: Arc, + + /// Optional streaming sender (for future streaming support) + pub stream_tx: Option>>, +} + +impl HarmonyResponsesContext { + /// Create a new Harmony Responses context + pub fn new( + pipeline: Arc, + components: Arc, + mcp_manager: Arc, + response_storage: Arc, + ) -> Self { + Self { + pipeline, + components, + mcp_manager, + response_storage, + stream_tx: None, + } + } + + /// Create with streaming support + pub fn with_streaming( + pipeline: Arc, + components: Arc, + mcp_manager: Arc, + response_storage: Arc, + stream_tx: tokio::sync::mpsc::UnboundedSender>, + ) -> Self { + Self { + pipeline, + components, + mcp_manager, + response_storage, + stream_tx: Some(stream_tx), + } + } +} + +/// Execute Harmony Responses API request with multi-turn MCP tool support +/// +/// This function orchestrates the multi-turn conversation flow: +/// 1. Execute request through full pipeline +/// 2. Check for tool calls in commentary channel +/// 3. If tool calls found: +/// - Execute MCP tools +/// - Build next request with tool results +/// - Repeat from step 1 (full pipeline re-execution) +/// 4. If no tool calls, return final response +/// +/// # Architecture +/// +/// Uses **external loop pattern**: wraps full pipeline execution rather than +/// implementing loop inside pipeline. Each iteration goes through: +/// - Worker selection (fresh selection based on current context) +/// - Client acquisition (new gRPC client if worker changed) +/// - Request building (Harmony prefill with complete history) +/// - Execution (model generation) +/// - Response processing (parse channels, detect tool calls) +/// +/// # Arguments +/// +/// * `ctx` - Harmony Responses context with pipeline, components, MCP manager +/// * `request` - Initial Responses API request +/// +/// # Returns +/// +/// Final ResponsesResponse after all tool iterations complete +/// +/// # Errors +/// +/// Returns error if: +/// - Max iterations exceeded (10 iterations) +/// - Pipeline execution fails +/// - MCP tool execution fails +/// - Response building fails +pub async fn serve_harmony_responses( + ctx: &HarmonyResponsesContext, + request: ResponsesRequest, +) -> Result { + // Load previous conversation history if previous_response_id is set + let mut current_request = load_previous_messages(ctx, request).await?; + let mut iteration_count = 0; + + // Check if request has MCP tools - if so, ensure dynamic client is registered + // and add static MCP tools to the request + use crate::{ + protocols::responses::ResponseToolType, routers::openai::mcp::ensure_request_mcp_client, + }; + + let has_mcp_tools = current_request + .tools + .as_ref() + .map(|tools| { + tools + .iter() + .any(|t| matches!(t.r#type, ResponseToolType::Mcp)) + }) + .unwrap_or(false); + + // Initialize MCP call tracking (will be passed to processor for final response) + let mut mcp_tracking = if has_mcp_tools { + Some(McpCallTracking::new("sglang-mcp".to_string())) + } else { + None + }; + + if has_mcp_tools { + // Ensure dynamic MCP client is registered for request-scoped tools + if let Some(tools) = ¤t_request.tools { + ensure_request_mcp_client(&ctx.mcp_manager, tools).await; + } + + // Add static MCP tools from inventory to the request + // (similar to non-Harmony pipeline pattern) + let mcp_tools = ctx.mcp_manager.list_tools(); + if !mcp_tools.is_empty() { + let mcp_response_tools = convert_mcp_tools_to_response_tools(&mcp_tools); + + let mut all_tools = current_request.tools.clone().unwrap_or_default(); + all_tools.extend(mcp_response_tools); + current_request.tools = Some(all_tools); + + tracing::debug!( + mcp_tool_count = mcp_tools.len(), + total_tool_count = current_request.tools.as_ref().map(|t| t.len()).unwrap_or(0), + "Request has MCP tools - added static MCP tools to Harmony Responses request" + ); + } + } + + loop { + iteration_count += 1; + + // Safety check: prevent infinite loops + if iteration_count > MAX_TOOL_ITERATIONS { + return Err(utils::internal_error_message(format!( + "Maximum tool iterations ({}) exceeded", + MAX_TOOL_ITERATIONS + ))); + } + + tracing::debug!( + iteration = iteration_count, + "Harmony Responses serving iteration" + ); + + // Execute through full pipeline + // This includes: + // - HarmonyPreparationStage (builder.rs: construct_input_messages_with_harmony) + // - WorkerSelectionStage (FRESH selection based on current context) + // - ClientAcquisitionStage (NEW gRPC client if needed) + // - HarmonyRequestBuildingStage (encode to token_ids) + // - RequestExecutionStage (model generation) + // - HarmonyResponseProcessingStage (processor.rs: process_responses_iteration) + let iteration_result = ctx + .pipeline + .execute_harmony_responses(¤t_request, ctx) + .await?; + + match iteration_result { + ResponsesIterationResult::ToolCallsFound { + tool_calls, + analysis, + partial_text, + } => { + tracing::debug!( + tool_call_count = tool_calls.len(), + has_analysis = analysis.is_some(), + partial_text_len = partial_text.len(), + "Tool calls found in commentary channel" + ); + + // TODO: Streaming support - emit intermediate chunks + // if let Some(tx) = &ctx.stream_tx { + // emit_intermediate_chunks(tx, &analysis, &partial_text, iteration_count).await?; + // } + + // Execute MCP tools via MCP manager + // If tools don't exist, call_tool() will return error naturally + let tool_results = if let Some(ref mut tracking) = mcp_tracking { + execute_mcp_tools(&ctx.mcp_manager, &tool_calls, tracking).await? + } else { + // Should never happen (we only get tool_calls when has_mcp_tools=true) + return Err(utils::internal_error_static( + "Tool calls found but MCP tracking not initialized", + )); + }; + + // Build next request with appended history + current_request = build_next_request_with_tools( + current_request, + tool_calls, + tool_results, + analysis, + partial_text, + ) + .map_err(|e| *e)?; + + // Continue loop - next iteration will select workers and execute + } + ResponsesIterationResult::Completed { + mut response, + usage, + } => { + tracing::debug!( + output_items = response.output.len(), + input_tokens = usage.prompt_tokens, + output_tokens = usage.completion_tokens, + has_mcp_tracking = mcp_tracking.is_some(), + "Harmony Responses serving completed - no more tool calls" + ); + + // Inject MCP output items if MCP tools were available + // (even if no tools were called, we still list available tools) + if let Some(tracking) = mcp_tracking { + inject_mcp_metadata(&mut response, &tracking, &ctx.mcp_manager); + + tracing::debug!( + mcp_calls = tracking.total_calls(), + output_items_after = response.output.len(), + "Injected MCP metadata into final response" + ); + } + + // No tool calls - this is the final response + // TODO: Accumulate usage across all iterations if needed + return Ok(*response); + } + } + } +} + +/// Execute MCP tools and collect results +/// +/// Executes each tool call sequentially via the MCP manager. +/// Tool execution errors are returned as error results to the model +/// (allows model to handle gracefully). +/// +/// # Arguments +/// +/// * `mcp_manager` - MCP manager for tool execution +/// * `tool_calls` - Tool calls from commentary channel +/// +/// # Returns +/// +/// Vector of tool results (one per tool call) +async fn execute_mcp_tools( + mcp_manager: &Arc, + tool_calls: &[ToolCall], + tracking: &mut McpCallTracking, +) -> Result, Response> { + let mut results = Vec::new(); + + for tool_call in tool_calls { + tracing::debug!( + tool_name = %tool_call.function.name, + call_id = %tool_call.id, + "Executing MCP tool" + ); + + // Parse tool arguments from JSON string + let args_str = tool_call.function.arguments.as_deref().unwrap_or("{}"); + let args: JsonValue = serde_json::from_str(args_str).map_err(|e| { + utils::internal_error_message(format!( + "Invalid tool arguments JSON for tool '{}': {}", + tool_call.function.name, e + )) + })?; + + // Execute tool via MCP manager + // Convert JsonValue to ToolArgs via Option (MCP manager expects this) + let args_map = if let JsonValue::Object(map) = args { + Some(map) + } else { + None + }; + + match mcp_manager + .call_tool(&tool_call.function.name, args_map) + .await + { + Ok(mcp_result) => { + tracing::debug!( + tool_name = %tool_call.function.name, + call_id = %tool_call.id, + "Tool execution succeeded" + ); + + // Extract content from MCP result + let output = if let Some(content) = mcp_result.content.first() { + // TODO: Handle different content types (text, image, resource) + // For now, serialize the entire content item + serde_json::to_value(content).unwrap_or_else( + |_| serde_json::json!({"error": "Failed to serialize tool result"}), + ) + } else { + serde_json::json!({"result": "success"}) + }; + + let is_error = mcp_result.is_error.unwrap_or(false); + let output_str = serde_json::to_string(&output) + .unwrap_or_else(|_| r#"{"error": "Failed to serialize output"}"#.to_string()); + + // Record this call in tracking + tracking.record_call( + tool_call.id.clone(), + tool_call.function.name.clone(), + args_str.to_string(), + output_str.clone(), + !is_error, + if is_error { + Some(output_str.clone()) + } else { + None + }, + ); + + results.push(ToolResult { + call_id: tool_call.id.clone(), + tool_name: tool_call.function.name.clone(), + output, + is_error, + }); + } + Err(e) => { + tracing::warn!( + tool_name = %tool_call.function.name, + call_id = %tool_call.id, + error = %e, + "Tool execution failed" + ); + + let error_msg = format!("Tool execution failed: {}", e); + let error_output = serde_json::json!({ + "error": error_msg.clone() + }); + let error_output_str = serde_json::to_string(&error_output) + .unwrap_or_else(|_| format!(r#"{{"error": "{}"}}"#, error_msg)); + + // Record failed call in tracking + tracking.record_call( + tool_call.id.clone(), + tool_call.function.name.clone(), + args_str.to_string(), + error_output_str.clone(), + false, + Some(error_msg), + ); + + // Return error result to model (let it handle gracefully) + results.push(ToolResult { + call_id: tool_call.id.clone(), + tool_name: tool_call.function.name.clone(), + output: error_output, + is_error: true, + }); + } + } + } + + Ok(results) +} + +/// Build next request with tool results appended to history +/// +/// Constructs a new ResponsesRequest with: +/// 1. Original input items (preserved) +/// 2. Assistant message with analysis (reasoning) + partial_text + tool_calls +/// 3. Tool result messages for each tool execution +/// +/// # Arguments +/// +/// * `request` - Current request (contains original input) +/// * `tool_calls` - Tool calls from commentary channel +/// * `tool_results` - Results from MCP tool execution +/// * `analysis` - Analysis channel content (becomes reasoning content) +/// * `partial_text` - Final channel content (becomes message content) +/// +/// # Returns +/// +/// New ResponsesRequest with appended history +fn build_next_request_with_tools( + mut request: ResponsesRequest, + tool_calls: Vec, + tool_results: Vec, + analysis: Option, + partial_text: String, +) -> Result> { + use uuid::Uuid; + + use crate::protocols::responses::{ + ResponseContentPart, ResponseInputOutputItem, ResponseReasoningContent, + }; + + // Get current input items (or empty vec if Text variant) + let mut items = match request.input { + ResponseInput::Items(items) => items, + ResponseInput::Text(text) => { + // Convert text to items format + vec![ResponseInputOutputItem::SimpleInputMessage { + content: StringOrContentParts::String(text), + role: "user".to_string(), + r#type: None, + }] + } + }; + + // Build assistant response item with reasoning + content + tool calls + // This represents what the model generated in this iteration + let assistant_id = format!("msg_{}", Uuid::new_v4()); + + // Add reasoning if present (from analysis channel) + if let Some(analysis_text) = analysis { + items.push(ResponseInputOutputItem::Reasoning { + id: format!("reasoning_{}", assistant_id), + summary: vec![], + content: vec![ResponseReasoningContent::ReasoningText { + text: analysis_text, + }], + status: Some("completed".to_string()), + }); + } + + // Add message content if present (from final channel) + if !partial_text.is_empty() { + items.push(ResponseInputOutputItem::Message { + id: assistant_id.clone(), + role: "assistant".to_string(), + content: vec![ResponseContentPart::OutputText { + text: partial_text, + annotations: vec![], + logprobs: None, + }], + status: Some("completed".to_string()), + }); + } + + // Add function tool calls (from commentary channel) + for tool_call in tool_calls { + items.push(ResponseInputOutputItem::FunctionToolCall { + id: tool_call.id.clone(), + name: tool_call.function.name.clone(), + arguments: tool_call + .function + .arguments + .unwrap_or_else(|| "{}".to_string()), + output: None, // Output will be added next + status: Some("in_progress".to_string()), + }); + } + + // Add tool results + for tool_result in tool_results { + // Serialize tool output to string + let output_str = serde_json::to_string(&tool_result.output).unwrap_or_else(|e| { + format!("{{\"error\": \"Failed to serialize tool output: {}\"}}", e) + }); + + // Update the corresponding tool call with output and completed status + // Find and update the matching FunctionToolCall + if let Some(ResponseInputOutputItem::FunctionToolCall { + output, + status, + .. + }) = items + .iter_mut() + .find(|item| matches!(item, ResponseInputOutputItem::FunctionToolCall { id, .. } if id == &tool_result.call_id)) + { + *output = Some(output_str); + *status = if tool_result.is_error { + Some("failed".to_string()) + } else { + Some("completed".to_string()) + }; + } + } + + // Update request with new items + request.input = ResponseInput::Items(items); + + Ok(request) +} + +/// Tool execution result +/// +/// Contains the result of executing a single MCP tool. +struct ToolResult { + /// Tool call ID (for matching with request) + call_id: String, + + /// Tool name + #[allow(dead_code)] // Kept for documentation and future use + tool_name: String, + + /// Tool output (JSON value) + output: JsonValue, + + /// Whether this is an error result + is_error: bool, +} + +/// Convert MCP tools to Responses API tool format +/// +/// Converts MCP Tool entries (from rmcp SDK) to ResponseTool format so the model +/// knows about available MCP tools when making tool calls. +/// +/// # Arguments +/// +/// * `mcp_tools` - MCP tools from the MCP manager inventory (rmcp::model::Tool) +/// +/// # Returns +/// +/// Vector of ResponseTool entries in MCP format +fn convert_mcp_tools_to_response_tools(mcp_tools: &[crate::mcp::Tool]) -> Vec { + use serde_json::Value; + + use crate::protocols::responses::ResponseToolType; + + mcp_tools + .iter() + .map(|tool_info| ResponseTool { + r#type: ResponseToolType::Mcp, + function: Some(Function { + name: tool_info.name.to_string(), + description: tool_info.description.as_ref().map(|d| d.to_string()), + parameters: Value::Object((*tool_info.input_schema).clone()), + strict: None, + }), + server_url: None, // MCP tools from inventory don't have individual server URLs + authorization: None, + server_label: None, + server_description: tool_info.description.as_ref().map(|d| d.to_string()), + require_approval: None, + allowed_tools: None, + }) + .collect() +} + +/// Inject MCP metadata into final response +/// +/// Adds mcp_list_tools and mcp_call output items to the response output array. +/// Following non-Harmony pipeline pattern: +/// 1. Prepend mcp_list_tools at the beginning +/// 2. Append all mcp_call items at the end +/// +/// # Arguments +/// +/// * `response` - Final response to modify +/// * `tracking` - MCP call tracking data +/// * `mcp_manager` - MCP manager for listing tools +fn inject_mcp_metadata( + response: &mut ResponsesResponse, + tracking: &McpCallTracking, + mcp_manager: &Arc, +) { + use serde_json::{json, Value}; + use uuid::Uuid; + + use crate::protocols::responses::{McpToolInfo, ResponseOutputItem}; + + // Build mcp_list_tools item + let tools = mcp_manager.list_tools(); + let tools_info: Vec = tools + .iter() + .map(|t| McpToolInfo { + name: t.name.to_string(), + description: t.description.as_ref().map(|d| d.to_string()), + input_schema: Value::Object((*t.input_schema).clone()), + annotations: Some(json!({ + "read_only": false + })), + }) + .collect(); + + let mcp_list_tools = ResponseOutputItem::McpListTools { + id: format!("mcpl_{}", Uuid::new_v4()), + server_label: tracking.server_label.clone(), + tools: tools_info, + }; + + // Build mcp_call items for each tracked call + let mcp_call_items: Vec = tracking + .tool_calls + .iter() + .map(|record| ResponseOutputItem::McpCall { + id: format!("mcp_{}", Uuid::new_v4()), + status: if record.success { + "completed" + } else { + "failed" + } + .to_string(), + approval_request_id: None, + arguments: record.arguments.clone(), + error: record.error.clone(), + name: record.tool_name.clone(), + output: record.output.clone(), + server_label: tracking.server_label.clone(), + }) + .collect(); + + // Inject into response output: + // 1. Prepend mcp_list_tools at the beginning + response.output.insert(0, mcp_list_tools); + + // 2. Append all mcp_call items at the end + response.output.extend(mcp_call_items); +} + +/// Load previous conversation messages from storage +/// +/// If the request has `previous_response_id`, loads the response chain from storage +/// and prepends the conversation history to the request input items. +/// +/// # Arguments +/// +/// * `ctx` - Harmony Responses context with response_storage +/// * `request` - Current request (may have previous_response_id set) +/// +/// # Returns +/// +/// Modified request with conversation history prepended to input items +async fn load_previous_messages( + ctx: &HarmonyResponsesContext, + request: ResponsesRequest, +) -> Result { + let Some(ref prev_id_str) = request.previous_response_id else { + // No previous_response_id, return request as-is + return Ok(request); + }; + + let prev_id = ResponseId::from(prev_id_str.as_str()); + + // Load response chain from storage + let chain = ctx + .response_storage + .get_response_chain(&prev_id, None) + .await + .map_err(|e| { + utils::internal_error_message(format!( + "Failed to load previous response chain for {}: {}", + prev_id_str, e + )) + })?; + + // Build conversation history from stored responses + let mut history_items = Vec::new(); + + // Helper to deserialize and collect items from a JSON array + let deserialize_items = + |arr: &serde_json::Value, item_type: &str| -> Vec { + arr.as_array() + .into_iter() + .flat_map(|items| items.iter()) + .filter_map(|item| { + serde_json::from_value::(item.clone()) + .map_err(|e| { + tracing::warn!( + "Failed to deserialize stored {} item: {}. Item: {}", + item_type, + e, + item + ); + }) + .ok() + }) + .collect() + }; + + for stored in chain.responses.iter() { + history_items.extend(deserialize_items(&stored.input, "input")); + history_items.extend(deserialize_items(&stored.output, "output")); + } + + tracing::debug!( + previous_response_id = %prev_id_str, + history_items_count = history_items.len(), + "Loaded conversation history from previous response" + ); + + // Build modified request with history prepended + let mut modified_request = request; + + // Convert current input to items format + let all_items = match modified_request.input { + ResponseInput::Items(items) => { + // Prepend history to existing items + let mut combined = history_items; + combined.extend(items); + combined + } + ResponseInput::Text(text) => { + // Convert text to item and prepend history + history_items.push(ResponseInputOutputItem::SimpleInputMessage { + content: StringOrContentParts::String(text), + role: "user".to_string(), + r#type: None, + }); + history_items + } + }; + + // Update request with combined items and clear previous_response_id + modified_request.input = ResponseInput::Items(all_items); + modified_request.previous_response_id = None; + + Ok(modified_request) +} + +// TODO: Implement streaming support +// /// Emit intermediate streaming chunks for analysis and partial text +// /// +// /// Emits SSE chunks for Responses API streaming: +// /// - Reasoning chunks for analysis channel +// /// - Message chunks for partial text from final channel +// /// +// /// # Arguments +// /// +// /// * `tx` - Streaming sender +// /// * `analysis` - Analysis channel content +// /// * `partial_text` - Final channel content +// /// * `iteration` - Current iteration number +// async fn emit_intermediate_chunks( +// tx: &tokio::sync::mpsc::UnboundedSender>, +// analysis: &Option, +// partial_text: &str, +// iteration: usize, +// ) -> Result<(), Response> { +// // TODO: Implement streaming emission +// // - Emit reasoning chunks for analysis +// // - Emit message chunks for partial_text +// // - Follow OpenAI Responses streaming format (14 SSE event types) +// Ok(()) +// } diff --git a/sgl-router/src/routers/grpc/harmony/stages/mod.rs b/sgl-router/src/routers/grpc/harmony/stages/mod.rs new file mode 100644 index 000000000..8e130fffb --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/stages/mod.rs @@ -0,0 +1,14 @@ +//! Harmony-specific pipeline stages +//! +//! These stages replace their regular counterparts in the Harmony pipeline: +//! - HarmonyPreparationStage: Harmony encoding instead of chat template + tokenization +//! - HarmonyRequestBuildingStage: Token-based request building +//! - HarmonyResponseProcessingStage: Harmony channel parsing + +pub mod preparation; +pub mod request_building; +pub mod response_processing; + +pub use preparation::HarmonyPreparationStage; +pub use request_building::HarmonyRequestBuildingStage; +pub use response_processing::HarmonyResponseProcessingStage; diff --git a/sgl-router/src/routers/grpc/harmony/stages/preparation.rs b/sgl-router/src/routers/grpc/harmony/stages/preparation.rs new file mode 100644 index 000000000..7ab655b0c --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/stages/preparation.rs @@ -0,0 +1,245 @@ +//! Harmony Preparation Stage: Harmony encoding for chat and generate requests + +use async_trait::async_trait; +use axum::response::Response; +use serde_json::json; + +use super::super::HarmonyBuilder; +use crate::{ + protocols::{ + chat::ChatCompletionRequest, + common::{Tool, ToolChoice, ToolChoiceValue}, + responses::ResponsesRequest, + }, + routers::grpc::{ + context::{PreparationOutput, RequestContext, RequestType}, + stages::PipelineStage, + utils, + }, +}; + +/// Harmony Preparation stage: Encode requests using Harmony protocol +/// +/// Replaces the regular PreparationStage for Harmony models. +/// Converts chat/generate requests to Harmony-encoded token_ids and extraction_text. +pub struct HarmonyPreparationStage { + builder: HarmonyBuilder, +} + +impl HarmonyPreparationStage { + /// Create a new Harmony preparation stage + pub fn new() -> Self { + Self { + builder: HarmonyBuilder::new(), + } + } +} + +impl Default for HarmonyPreparationStage { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl PipelineStage for HarmonyPreparationStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + // Clone Arc before match to avoid borrow checker issues + // Arc clone is cheap (8 bytes) - avoids full request clone (15KB-200KB) + let is_chat = matches!(&ctx.input.request_type, RequestType::Chat(_)); + let is_responses = matches!(&ctx.input.request_type, RequestType::Responses(_)); + + if is_chat { + let request_arc = ctx.chat_request_arc(); + self.prepare_chat(ctx, &request_arc).await?; + } else if is_responses { + let request_arc = ctx.responses_request_arc(); + self.prepare_responses(ctx, &request_arc).await?; + } else { + return Err(utils::bad_request_error( + "Only Chat and Responses requests supported in Harmony pipeline".to_string(), + )); + } + + Ok(None) + } + + fn name(&self) -> &'static str { + "HarmonyPreparation" + } +} + +impl HarmonyPreparationStage { + /// Prepare a chat completion request using Harmony encoding + async fn prepare_chat( + &self, + ctx: &mut RequestContext, + request: &ChatCompletionRequest, + ) -> Result, Response> { + // Validate - reject logprobs + if request.logprobs { + return Err(utils::bad_request_error( + "logprobs are not supported for Harmony models".to_string(), + )); + } + + // Step 1: Filter tools if needed + let body_ref = utils::filter_tools_for_request(request); + + // Step 2: Build tool constraints + let tool_constraints = if let Some(tools) = body_ref.tools.as_ref() { + Self::generate_harmony_structural_tag(tools, &body_ref.tool_choice).map_err(|e| *e)? + } else { + None + }; + + // Step 3: Build via Harmony + let build_output = self + .builder + .build_from_chat(&body_ref) + .map_err(|e| utils::bad_request_error(format!("Harmony build failed: {}", e)))?; + + // Step 4: Store results + ctx.state.preparation = Some(PreparationOutput { + original_text: None, + token_ids: build_output.input_ids, + processed_messages: None, + tool_constraints, + filtered_request: if matches!(body_ref, std::borrow::Cow::Owned(_)) { + Some(body_ref.into_owned()) + } else { + None + }, + harmony_mode: true, + selection_text: Some(build_output.selection_text), + harmony_messages: Some(build_output.harmony_messages), + harmony_stop_ids: Some(build_output.stop_token_ids), + }); + + Ok(None) + } + + /// Prepare a responses API request using Harmony encoding + /// + /// For responses API, we build from conversation history using the same Harmony + /// encoding that the builder provides. This handles the MCP loop integration. + pub async fn prepare_responses( + &self, + ctx: &mut RequestContext, + request: &ResponsesRequest, + ) -> Result, Response> { + // Build via Harmony from responses API request + let build_output = self + .builder + .build_from_responses(request) + .map_err(|e| utils::bad_request_error(format!("Harmony build failed: {}", e)))?; + + // Store results in preparation output + ctx.state.preparation = Some(PreparationOutput { + original_text: None, + token_ids: build_output.input_ids, + processed_messages: None, + tool_constraints: None, + filtered_request: None, + harmony_mode: true, + selection_text: Some(build_output.selection_text), + harmony_messages: Some(build_output.harmony_messages), + harmony_stop_ids: Some(build_output.stop_token_ids), + }); + + Ok(None) + } + + /// Generate Harmony structural tag for tool constraints + /// + /// Uses structural tags with `triggered_tags` format to force Harmony format output. + /// This ensures the model outputs in Harmony format (with channels) even when constrained. + fn generate_harmony_structural_tag( + tools: &[Tool], + tool_choice: &Option, + ) -> Result, Box> { + let Some(choice) = tool_choice.as_ref() else { + return Ok(None); + }; + + match choice { + ToolChoice::Function { function, .. } => { + let tag = Self::build_harmony_structural_tag(tools, Some(&function.name))?; + Ok(Some(("structural_tag".to_string(), tag))) + } + ToolChoice::Value(ToolChoiceValue::Required) => { + let tag = Self::build_harmony_structural_tag(tools, None)?; + Ok(Some(("structural_tag".to_string(), tag))) + } + ToolChoice::AllowedTools { mode, .. } => { + if mode == "required" { + let tag = Self::build_harmony_structural_tag(tools, None)?; + Ok(Some(("structural_tag".to_string(), tag))) + } else { + Ok(None) + } + } + _ => Ok(None), + } + } + + /// Build Harmony structural tag for tool calling constraints + fn build_harmony_structural_tag( + tools: &[Tool], + specific_function: Option<&str>, + ) -> Result> { + let mut tags = Vec::new(); + + // Filter tools if specific function requested + let tools_to_use: Vec<&Tool> = if let Some(func_name) = specific_function { + tools + .iter() + .filter(|t| t.function.name == func_name) + .collect() + } else { + tools.iter().collect() + }; + + // Validate specific function exists + if specific_function.is_some() && tools_to_use.is_empty() { + return Err(Box::new(utils::bad_request_error(format!( + "Tool '{}' not found in tools list", + specific_function.unwrap() + )))); + } + + // Build tags for each tool + for tool in tools_to_use { + let tool_name = &tool.function.name; + let params_schema = &tool.function.parameters; + + tags.push(json!({ + "begin": format!("<|channel|>commentary to=functions.{}<|constrain|>json<|message|>", tool_name), + "content": { + "type": "json_schema", + "json_schema": params_schema + }, + "end": "" // `end` is empty because <|call|> comes naturally from Harmony stop tokens + })); + } + + let stop_after_first = specific_function.is_some(); + + let structural_tag = json!({ + "format": { + "type": "triggered_tags", + "triggers": ["<|channel|>commentary"], + "tags": tags, + "at_least_one": true, + "stop_after_first": stop_after_first + } + }); + + serde_json::to_string(&structural_tag).map_err(|e| { + Box::new(utils::internal_error_message(format!( + "Failed to serialize structural tag: {}", + e + ))) + }) + } +} diff --git a/sgl-router/src/routers/grpc/harmony/stages/request_building.rs b/sgl-router/src/routers/grpc/harmony/stages/request_building.rs new file mode 100644 index 000000000..b5526fec9 --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/stages/request_building.rs @@ -0,0 +1,159 @@ +//! Harmony Request Building Stage: Build gRPC request from Harmony-encoded tokens + +use std::sync::Arc; + +use async_trait::async_trait; +use axum::response::Response; +use rand::Rng; +use tracing::debug; +use uuid::Uuid; + +use crate::{ + core::Worker, + grpc_client::proto::{DisaggregatedParams, GenerateRequest}, + routers::grpc::{ + context::{ClientSelection, RequestContext, RequestType, WorkerSelection}, + stages::PipelineStage, + utils, + }, +}; + +/// Harmony Request Building stage: Convert Harmony tokens to gRPC request +/// +/// Takes the Harmony-encoded input_ids from preparation and builds a proto::GenerateRequest. +/// Unlike regular request building, this uses token_ids directly (Harmony encoding handles messages). +pub struct HarmonyRequestBuildingStage { + inject_pd_metadata: bool, +} + +impl HarmonyRequestBuildingStage { + /// Create a new Harmony request building stage + pub fn new(inject_pd_metadata: bool) -> Self { + Self { inject_pd_metadata } + } + + /// Inject PD (prefill-decode) bootstrap metadata + fn inject_bootstrap_metadata( + &self, + request: &mut GenerateRequest, + prefill_worker: &Arc, + ) { + let hostname = prefill_worker.bootstrap_host(); + let bootstrap_port = prefill_worker.bootstrap_port().unwrap_or(8998); + + // Generate room ID for bootstrap + let room_id = rand::rng().random_range(0..i32::MAX); + + // Create DisaggregatedParams + let disagg_params = DisaggregatedParams { + bootstrap_host: hostname.to_string(), + bootstrap_port: bootstrap_port as i32, + bootstrap_room: room_id, + }; + + // Inject metadata directly into request + request.disaggregated_params = Some(disagg_params); + + debug!( + "Injected Harmony bootstrap metadata: host={}, port={}, room={}", + hostname, bootstrap_port, room_id + ); + } +} + +#[async_trait] +impl PipelineStage for HarmonyRequestBuildingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + // Get preparation output + let prep = ctx + .state + .preparation + .as_ref() + .ok_or_else(|| utils::internal_error_static("Preparation not completed"))?; + + // Get clients + let clients = ctx + .state + .clients + .as_ref() + .ok_or_else(|| utils::internal_error_static("Client acquisition not completed"))?; + let builder_client = match clients { + ClientSelection::Single { client } => client, + ClientSelection::Dual { prefill, .. } => prefill, + }; + + // Generate request_id based on request type + let request_id = match &ctx.input.request_type { + RequestType::Chat(_) => format!("chatcmpl-{}", Uuid::new_v4()), + RequestType::Responses(_) => format!("responses-{}", Uuid::new_v4()), + RequestType::Generate(_) => { + return Err(utils::bad_request_error( + "Generate requests are not supported with Harmony models".to_string(), + )); + } + }; + + // Build gRPC request using token_ids directly (Harmony encoding already handled message rendering) + // Use a placeholder for original_text; Harmony uses input_ids for tokenization + let placeholder_processed_text = "[harmony]".to_string(); + + let mut proto_request = match &ctx.input.request_type { + RequestType::Chat(request) => { + // Use filtered request if present from preparation; otherwise original + let body = prep.filtered_request.as_ref().unwrap_or(request.as_ref()); + + builder_client + .build_generate_request( + request_id, + body, + placeholder_processed_text, + prep.token_ids.clone(), + None, + prep.tool_constraints.clone(), + ) + .map_err(|e| { + utils::bad_request_error(format!("Invalid request parameters: {}", e)) + })? + } + RequestType::Responses(request) => builder_client + .build_generate_request_from_responses( + request_id, + request.as_ref(), + placeholder_processed_text, + prep.token_ids.clone(), + prep.harmony_stop_ids.clone(), + ) + .map_err(|e| { + utils::bad_request_error(format!("Invalid request parameters: {}", e)) + })?, + _ => unreachable!(), + }; + + // Inject Harmony stop token IDs into sampling params for ALL Harmony requests + // These stop tokens (<|return|> and <|call|>) prevent the model from generating + // malformed Harmony sequences + if let Some(harmony_stops) = &prep.harmony_stop_ids { + if let Some(params) = proto_request.sampling_params.as_mut() { + params.stop_token_ids.extend_from_slice(harmony_stops); + debug!( + stop_token_count = harmony_stops.len(), + "Injected Harmony stop tokens into sampling params" + ); + } + } + + // Inject PD metadata if needed + if self.inject_pd_metadata { + if let Some(WorkerSelection::Dual { prefill, .. }) = ctx.state.workers.as_ref() { + self.inject_bootstrap_metadata(&mut proto_request, prefill); + } + } + + ctx.state.proto_request = Some(proto_request); + Ok(None) + } + + fn name(&self) -> &'static str { + "HarmonyRequestBuilding" + } +} diff --git a/sgl-router/src/routers/grpc/harmony/stages/response_processing.rs b/sgl-router/src/routers/grpc/harmony/stages/response_processing.rs new file mode 100644 index 000000000..cb96efe13 --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/stages/response_processing.rs @@ -0,0 +1,123 @@ +//! Harmony Response Processing Stage: Parse Harmony channels to ChatCompletionResponse + +use std::sync::Arc; + +use async_trait::async_trait; +use axum::response::Response; + +use super::super::{HarmonyResponseProcessor, HarmonyStreamingProcessor}; +use crate::routers::grpc::{ + context::{FinalResponse, RequestContext, RequestType}, + stages::PipelineStage, + utils, +}; + +/// Harmony Response Processing stage: Parse and format Harmony responses +/// +/// Takes output tokens from execution and parses them using HarmonyParserAdapter +/// to extract analysis, tool calls, and final response text from Harmony channels. +pub struct HarmonyResponseProcessingStage { + processor: HarmonyResponseProcessor, + streaming_processor: Arc, +} + +impl HarmonyResponseProcessingStage { + /// Create a new Harmony response processing stage + pub fn new() -> Self { + Self { + processor: HarmonyResponseProcessor::new(), + streaming_processor: Arc::new(HarmonyStreamingProcessor::new()), + } + } +} + +impl Default for HarmonyResponseProcessingStage { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl PipelineStage for HarmonyResponseProcessingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + // Get execution result (output tokens from model) + let execution_result = ctx + .state + .response + .execution_result + .take() + .ok_or_else(|| utils::internal_error_static("No execution result"))?; + + let is_streaming = ctx.is_streaming(); + let dispatch = ctx + .state + .dispatch + .as_ref() + .cloned() + .ok_or_else(|| utils::internal_error_static("Dispatch metadata not set"))?; + + // Check request type to determine which processor method to call + match &ctx.input.request_type { + RequestType::Chat(_) => { + // For streaming, delegate to streaming processor and return SSE response + if is_streaming { + return Ok(Some( + self.streaming_processor + .clone() + .process_streaming_chat_response( + execution_result, + ctx.chat_request_arc(), + dispatch, + ), + )); + } + + // For non-streaming, delegate to Harmony response processor to build ChatCompletionResponse + let chat_request = ctx.chat_request_arc(); + let response = self + .processor + .process_non_streaming_chat_response(execution_result, chat_request, dispatch) + .await?; + + ctx.state.response.final_response = Some(FinalResponse::Chat(response)); + Ok(None) + } + RequestType::Responses(_) => { + // For Responses API, process iteration and store result + // Streaming not yet supported for Responses API + if is_streaming { + return Err(utils::internal_error_static( + "Streaming not yet supported for Responses API", + )); + } + + let responses_request = ctx.responses_request_arc(); + let iteration_result = self + .processor + .process_responses_iteration(execution_result, responses_request, dispatch) + .await?; + + ctx.state.response.responses_iteration_result = Some(iteration_result); + Ok(None) + } + RequestType::Generate(_) => Err(utils::internal_error_static( + "Generate requests not supported in Harmony pipeline", + )), + } + } + + fn name(&self) -> &'static str { + "HarmonyResponseProcessing" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_response_processing_stage_creation() { + let stage = HarmonyResponseProcessingStage::new(); + assert_eq!(stage.name(), "HarmonyResponseProcessing"); + } +} diff --git a/sgl-router/src/routers/grpc/harmony/streaming.rs b/sgl-router/src/routers/grpc/harmony/streaming.rs new file mode 100644 index 000000000..c9b32459c --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/streaming.rs @@ -0,0 +1,551 @@ +//! Harmony streaming response processor + +use std::{ + collections::{hash_map::Entry::Vacant, HashMap}, + io, + sync::Arc, +}; + +use axum::{body::Body, http::StatusCode, response::Response}; +use bytes::Bytes; +use http::header::{HeaderValue, CONTENT_TYPE}; +use proto::{ + generate_complete::MatchedStop::{MatchedStopStr, MatchedTokenId}, + generate_response::Response::{Chunk, Complete}, +}; +use serde_json::json; +use tokio::sync::mpsc; +use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt}; +use tracing::error; + +use super::{types::HarmonyChannelDelta, HarmonyParserAdapter}; +use crate::{ + grpc_client::{proto, sglang_scheduler::AbortOnDropStream}, + protocols::{ + chat::{ + ChatCompletionRequest, ChatCompletionStreamResponse, ChatMessageDelta, ChatStreamChoice, + }, + common::{FunctionCallDelta, ToolCallDelta, Usage}, + }, + routers::grpc::context, +}; + +/// Processor for streaming Harmony responses +/// +/// Returns an SSE stream that parses Harmony tokens incrementally and +/// emits ChatCompletionChunk events for streaming responses. +pub struct HarmonyStreamingProcessor; + +impl HarmonyStreamingProcessor { + /// Create a new Harmony streaming processor + pub fn new() -> Self { + Self + } + + /// Process a streaming Harmony Chat Completion response + /// + /// Returns an SSE response with streaming token updates. + pub fn process_streaming_chat_response( + self: Arc, + execution_result: context::ExecutionResult, + chat_request: Arc, + dispatch: context::DispatchMetadata, + ) -> Response { + // Create SSE channel + let (tx, rx) = mpsc::unbounded_channel::>(); + + // Spawn background task based on execution mode + match execution_result { + context::ExecutionResult::Single { stream } => { + tokio::spawn(async move { + let result = + Self::process_single_stream(stream, dispatch, chat_request, &tx).await; + + if let Err(e) = result { + error!("Harmony streaming error: {}", e); + let error_chunk = format!( + "data: {}\n\n", + json!({ + "error": { + "message": e, + "type": "internal_error" + } + }) + ); + let _ = tx.send(Ok(Bytes::from(error_chunk))); + } + + let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + }); + } + context::ExecutionResult::Dual { prefill, decode } => { + tokio::spawn(async move { + let result = + Self::process_dual_stream(prefill, *decode, dispatch, chat_request, &tx) + .await; + + if let Err(e) = result { + error!("Harmony dual streaming error: {}", e); + let error_chunk = format!( + "data: {}\n\n", + json!({ + "error": { + "message": e, + "type": "internal_error" + } + }) + ); + let _ = tx.send(Ok(Bytes::from(error_chunk))); + } + + let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + }); + } + } + + // Return SSE response + Self::build_sse_response(rx) + } + + /// Process streaming chunks from a single stream + async fn process_single_stream( + mut grpc_stream: AbortOnDropStream, + dispatch: context::DispatchMetadata, + original_request: Arc, + tx: &mpsc::UnboundedSender>, + ) -> Result<(), String> { + // Per-index state management (for n>1 support) + let mut parsers: HashMap = HashMap::new(); + let mut is_firsts: HashMap = HashMap::new(); + let mut finish_reasons: HashMap> = HashMap::new(); + let mut matched_stops: HashMap> = HashMap::new(); + let mut prompt_tokens: HashMap = HashMap::new(); + let mut completion_tokens: HashMap = HashMap::new(); + + let stream_options = &original_request.stream_options; + + // Process stream + while let Some(result) = grpc_stream.next().await { + let response = result.map_err(|e| format!("Stream error: {}", e))?; + + match response.response { + Some(Chunk(chunk)) => { + let index = chunk.index; + + // Initialize parser for this index if needed + if let Vacant(e) = parsers.entry(index) { + e.insert( + HarmonyParserAdapter::new() + .map_err(|e| format!("Failed to create parser: {}", e))?, + ); + is_firsts.insert(index, true); + } + + // Track token counts + *completion_tokens.entry(index).or_insert(0) += 1; + + // Parse chunk via Harmony parser + let parser = parsers + .get_mut(&index) + .ok_or("Parser not found for index")?; + + let delta_result = parser + .parse_chunk(&chunk.token_ids) + .map_err(|e| format!("Parse error: {}", e))?; + + // Emit SSE event if there's a delta + if let Some(delta) = delta_result { + let is_first = is_firsts.get(&index).copied().unwrap_or(false); + Self::emit_chunk_delta( + &delta, + index, + is_first, + &dispatch, + &original_request, + tx, + )?; + + if is_first { + is_firsts.insert(index, false); + } + } + } + Some(Complete(complete)) => { + let index = complete.index; + + // Store final metadata + finish_reasons.insert(index, Some(complete.finish_reason.clone())); + matched_stops.insert( + index, + complete.matched_stop.as_ref().map(|m| match m { + MatchedTokenId(id) => { + serde_json::json!(id) + } + MatchedStopStr(s) => { + serde_json::json!(s) + } + }), + ); + prompt_tokens.insert(index, complete.prompt_tokens as u32); + *completion_tokens.entry(index).or_insert(0) = + complete.completion_tokens as u32; + + // Finalize parser and emit final chunk + if let Some(parser) = parsers.get_mut(&index) { + let matched_stop = matched_stops.get(&index).and_then(|m| m.clone()); + let final_output = parser + .finalize(complete.finish_reason.clone(), matched_stop.clone()) + .map_err(|e| format!("Finalize error: {}", e))?; + + Self::emit_final_chunk( + index, + &final_output.finish_reason, + final_output.matched_stop.as_ref(), + &dispatch, + &original_request, + tx, + )?; + } + } + Some(proto::generate_response::Response::Error(err)) => { + return Err(format!("Server error: {}", err.message)); + } + None => {} + } + } + + // Emit final usage if requested + if let Some(true) = stream_options.as_ref().and_then(|so| so.include_usage) { + let total_prompt: u32 = prompt_tokens.values().sum(); + let total_completion: u32 = completion_tokens.values().sum(); + + Self::emit_usage_chunk( + total_prompt, + total_completion, + &dispatch, + &original_request, + tx, + )?; + } + + // Mark stream as completed successfully to prevent abort on drop + grpc_stream.mark_completed(); + + Ok(()) + } + + /// Process streaming chunks from dual streams (prefill + decode) + async fn process_dual_stream( + mut prefill_stream: AbortOnDropStream, + mut decode_stream: AbortOnDropStream, + dispatch: context::DispatchMetadata, + original_request: Arc, + tx: &mpsc::UnboundedSender>, + ) -> Result<(), String> { + // Phase 1: Process prefill stream (collect metadata) + let mut prompt_tokens: HashMap = HashMap::new(); + + while let Some(result) = prefill_stream.next().await { + let response = result.map_err(|e| format!("Prefill stream error: {}", e))?; + + if let Some(Complete(complete)) = response.response { + prompt_tokens.insert(complete.index, complete.prompt_tokens as u32); + } + } + + // Phase 2: Process decode stream (same as single stream) + let mut parsers: HashMap = HashMap::new(); + let mut is_firsts: HashMap = HashMap::new(); + let mut finish_reasons: HashMap> = HashMap::new(); + let mut matched_stops: HashMap> = HashMap::new(); + let mut completion_tokens: HashMap = HashMap::new(); + + let stream_options = &original_request.stream_options; + + while let Some(result) = decode_stream.next().await { + let response = result.map_err(|e| format!("Decode stream error: {}", e))?; + + match response.response { + Some(Chunk(chunk)) => { + let index = chunk.index; + + // Initialize parser for this index if needed + if let Vacant(e) = parsers.entry(index) { + e.insert( + HarmonyParserAdapter::new() + .map_err(|e| format!("Failed to create parser: {}", e))?, + ); + is_firsts.insert(index, true); + } + + *completion_tokens.entry(index).or_insert(0) += 1; + + let parser = parsers + .get_mut(&index) + .ok_or("Parser not found for index")?; + + let delta_result = parser + .parse_chunk(&chunk.token_ids) + .map_err(|e| format!("Parse error: {}", e))?; + + if let Some(delta) = delta_result { + let is_first = is_firsts.get(&index).copied().unwrap_or(false); + Self::emit_chunk_delta( + &delta, + index, + is_first, + &dispatch, + &original_request, + tx, + )?; + + if is_first { + is_firsts.insert(index, false); + } + } + } + Some(Complete(complete)) => { + let index = complete.index; + + finish_reasons.insert(index, Some(complete.finish_reason.clone())); + matched_stops.insert( + index, + complete.matched_stop.as_ref().map(|m| match m { + MatchedTokenId(id) => { + json!(id) + } + MatchedStopStr(s) => { + json!(s) + } + }), + ); + *completion_tokens.entry(index).or_insert(0) = + complete.completion_tokens as u32; + + if let Some(parser) = parsers.get_mut(&index) { + let matched_stop = matched_stops.get(&index).and_then(|m| m.clone()); + let final_output = parser + .finalize(complete.finish_reason.clone(), matched_stop.clone()) + .map_err(|e| format!("Finalize error: {}", e))?; + + Self::emit_final_chunk( + index, + &final_output.finish_reason, + final_output.matched_stop.as_ref(), + &dispatch, + &original_request, + tx, + )?; + } + } + Some(proto::generate_response::Response::Error(err)) => { + return Err(format!("Server error: {}", err.message)); + } + None => {} + } + } + + decode_stream.mark_completed(); + + // Mark prefill stream as completed AFTER decode completes successfully + // This ensures that if client disconnects during decode, BOTH streams send abort + prefill_stream.mark_completed(); + + // Emit final usage if requested + if let Some(true) = stream_options.as_ref().and_then(|so| so.include_usage) { + let total_prompt: u32 = prompt_tokens.values().sum(); + let total_completion: u32 = completion_tokens.values().sum(); + + Self::emit_usage_chunk( + total_prompt, + total_completion, + &dispatch, + &original_request, + tx, + )?; + } + + Ok(()) + } + + /// Emit a chunk delta from Harmony channels + fn emit_chunk_delta( + delta: &HarmonyChannelDelta, + index: u32, + is_first: bool, + dispatch: &context::DispatchMetadata, + original_request: &ChatCompletionRequest, + tx: &mpsc::UnboundedSender>, + ) -> Result<(), String> { + // On first chunk, emit role announcement separately + if is_first { + let role_chunk = ChatCompletionStreamResponse { + id: dispatch.request_id.clone(), + object: "chat.completion.chunk".to_string(), + created: dispatch.created, + model: original_request.model.clone(), + system_fingerprint: dispatch.weight_version.clone(), + choices: vec![ChatStreamChoice { + index, + delta: ChatMessageDelta { + role: Some("assistant".to_string()), + content: Some(String::new()), + tool_calls: None, + reasoning_content: None, + }, + logprobs: None, + finish_reason: None, + matched_stop: None, + }], + usage: None, + }; + + let chunk_json = serde_json::to_string(&role_chunk) + .map_err(|e| format!("JSON serialization error: {}", e))?; + let sse_data = format!("data: {}\n\n", chunk_json); + + tx.send(Ok(Bytes::from(sse_data))) + .map_err(|_| "Failed to send role chunk".to_string())?; + } + + // Emit content delta (role is always None for content chunks) + let chat_delta = ChatMessageDelta { + role: None, + content: delta.final_delta.clone(), + tool_calls: delta.commentary_delta.as_ref().map(|tc_delta| { + vec![ToolCallDelta { + index: tc_delta.index as u32, + id: tc_delta.id.clone(), + tool_type: tc_delta.id.as_ref().map(|_| "function".to_string()), + function: tc_delta.function.as_ref().map(|f| FunctionCallDelta { + name: f.name.clone(), + arguments: f.arguments.clone(), + }), + }] + }), + reasoning_content: delta.analysis_delta.clone(), + }; + + // Build and emit chunk + let chunk = ChatCompletionStreamResponse { + id: dispatch.request_id.clone(), + object: "chat.completion.chunk".to_string(), + created: dispatch.created, + model: original_request.model.clone(), + system_fingerprint: dispatch.weight_version.clone(), + choices: vec![ChatStreamChoice { + index, + delta: chat_delta, + logprobs: None, + finish_reason: None, + matched_stop: None, + }], + usage: None, + }; + + let chunk_json = serde_json::to_string(&chunk) + .map_err(|e| format!("JSON serialization error: {}", e))?; + let sse_data = format!("data: {}\n\n", chunk_json); + + tx.send(Ok(Bytes::from(sse_data))) + .map_err(|_| "Failed to send chunk".to_string())?; + + Ok(()) + } + + /// Emit final chunk with finish_reason + fn emit_final_chunk( + index: u32, + finish_reason: &str, + matched_stop: Option<&serde_json::Value>, + dispatch: &context::DispatchMetadata, + original_request: &ChatCompletionRequest, + tx: &mpsc::UnboundedSender>, + ) -> Result<(), String> { + let chunk = ChatCompletionStreamResponse { + id: dispatch.request_id.clone(), + object: "chat.completion.chunk".to_string(), + created: dispatch.created, + model: original_request.model.clone(), + system_fingerprint: dispatch.weight_version.clone(), + choices: vec![ChatStreamChoice { + index, + delta: ChatMessageDelta { + role: None, + content: None, + tool_calls: None, + reasoning_content: None, + }, + logprobs: None, + finish_reason: Some(finish_reason.to_string()), + matched_stop: matched_stop.cloned(), + }], + usage: None, + }; + + let chunk_json = serde_json::to_string(&chunk) + .map_err(|e| format!("JSON serialization error: {}", e))?; + let sse_data = format!("data: {}\n\n", chunk_json); + + tx.send(Ok(Bytes::from(sse_data))) + .map_err(|_| "Failed to send final chunk".to_string())?; + + Ok(()) + } + + /// Emit usage chunk at the end + fn emit_usage_chunk( + prompt_tokens: u32, + completion_tokens: u32, + dispatch: &context::DispatchMetadata, + original_request: &ChatCompletionRequest, + tx: &mpsc::UnboundedSender>, + ) -> Result<(), String> { + let usage_chunk = ChatCompletionStreamResponse { + id: dispatch.request_id.clone(), + object: "chat.completion.chunk".to_string(), + created: dispatch.created, + model: original_request.model.clone(), + system_fingerprint: dispatch.weight_version.clone(), + choices: vec![], + usage: Some(Usage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + completion_tokens_details: None, + }), + }; + + let chunk_json = serde_json::to_string(&usage_chunk) + .map_err(|e| format!("JSON serialization error: {}", e))?; + let sse_data = format!("data: {}\n\n", chunk_json); + + tx.send(Ok(Bytes::from(sse_data))) + .map_err(|_| "Failed to send usage chunk".to_string())?; + + Ok(()) + } + + /// Build SSE response from receiver + fn build_sse_response(rx: mpsc::UnboundedReceiver>) -> Response { + let stream = UnboundedReceiverStream::new(rx); + let body = Body::from_stream(stream); + + Response::builder() + .status(StatusCode::OK) + .header( + CONTENT_TYPE, + HeaderValue::from_static("text/event-stream; charset=utf-8"), + ) + .header("Cache-Control", HeaderValue::from_static("no-cache")) + .header("Connection", HeaderValue::from_static("keep-alive")) + .body(body) + .unwrap() + } +} + +impl Default for HarmonyStreamingProcessor { + fn default() -> Self { + Self::new() + } +} diff --git a/sgl-router/src/routers/grpc/harmony/types.rs b/sgl-router/src/routers/grpc/harmony/types.rs new file mode 100644 index 000000000..8ce2a3602 --- /dev/null +++ b/sgl-router/src/routers/grpc/harmony/types.rs @@ -0,0 +1,137 @@ +//! Shared types for Harmony pipeline + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::protocols::common::ToolCall; + +/// Harmony message format +/// +/// Represents messages in the Harmony encoding format with role and content. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HarmonyMessage { + pub role: String, + pub content: String, +} + +impl HarmonyMessage { + pub fn new(role: impl Into, content: impl Into) -> Self { + Self { + role: role.into(), + content: content.into(), + } + } + + pub fn user(content: impl Into) -> Self { + Self::new("user", content) + } + + pub fn assistant(content: impl Into) -> Self { + Self::new("assistant", content) + } + + pub fn system(content: impl Into) -> Self { + Self::new("system", content) + } + + /// Convert from openai_harmony::chat::Message to our simplified HarmonyMessage + pub fn from_openai_harmony(msg: openai_harmony::chat::Message) -> Self { + use openai_harmony::chat::Content; + + // Extract role as string + let role = match msg.author.role { + openai_harmony::chat::Role::User => "user", + openai_harmony::chat::Role::Assistant => "assistant", + openai_harmony::chat::Role::System => "system", + openai_harmony::chat::Role::Developer => "developer", + openai_harmony::chat::Role::Tool => "tool", + } + .to_string(); + + // Extract text content from all Content::Text parts + let content = msg + .content + .iter() + .filter_map(|c| match c { + Content::Text(tc) => Some(tc.text.as_str()), + _ => None, + }) + .collect::>() + .join(""); + + Self { role, content } + } +} + +/// Output from Harmony encoding process +/// +/// Contains the encoded input_ids, stop tokens, selection text for worker routing, +/// and the Harmony message history. +#[derive(Debug, Clone)] +pub struct HarmonyBuildOutput { + /// Encoded token IDs to send to the model + pub input_ids: Vec, + + /// Stop token IDs for this model (injected into sampling params) + pub stop_token_ids: Vec, + + /// Selection text for worker routing (concise snippet from last user message) + pub selection_text: String, + + /// Harmony messages for this conversation (used for history tracking) + pub harmony_messages: Vec, +} + +/// Parsed output from all three Harmony channels +/// +/// Represents the complete response after parsing analysis, commentary, and final channels. +#[derive(Debug, Clone)] +pub struct HarmonyChannelOutput { + /// Analysis/reasoning content (from analysis channel) + pub analysis: Option, + + /// Tool calls (from commentary channel) + pub commentary: Option>, + + /// Final text content (from final channel) + pub final_text: String, + + /// Finish reason + pub finish_reason: String, + + /// Matched stop token (if any) + pub matched_stop: Option, +} + +/// Streaming delta for SSE responses +/// +/// Represents incremental updates as tokens are parsed from the stream. +#[derive(Debug, Clone)] +pub struct HarmonyChannelDelta { + /// Delta for analysis/reasoning content + pub analysis_delta: Option, + + /// Delta for tool calls + pub commentary_delta: Option, + + /// Delta for final text content + pub final_delta: Option, + + /// Whether this is the final delta + pub is_final: bool, +} + +/// Tool call delta for streaming +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCallDelta { + pub index: usize, + pub id: Option, + pub function: Option, +} + +/// Function call delta for streaming +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FunctionDelta { + pub name: Option, + pub arguments: Option, +} diff --git a/sgl-router/src/routers/grpc/mod.rs b/sgl-router/src/routers/grpc/mod.rs index 2483d6031..cd9cbc02b 100644 --- a/sgl-router/src/routers/grpc/mod.rs +++ b/sgl-router/src/routers/grpc/mod.rs @@ -3,11 +3,13 @@ use crate::{grpc_client::proto, protocols::common::StringOrArray}; pub mod context; +pub mod harmony; pub mod pd_router; pub mod pipeline; pub mod processing; pub mod responses; pub mod router; +pub mod stages; pub mod streaming; pub mod utils; diff --git a/sgl-router/src/routers/grpc/pd_router.rs b/sgl-router/src/routers/grpc/pd_router.rs index fe62b0dde..55cb653e5 100644 --- a/sgl-router/src/routers/grpc/pd_router.rs +++ b/sgl-router/src/routers/grpc/pd_router.rs @@ -14,9 +14,7 @@ use tracing::debug; use super::{context::SharedComponents, pipeline::RequestPipeline}; use crate::{ app_context::AppContext, - config::types::RetryConfig, core::{ConnectionMode, WorkerRegistry, WorkerType}, - policies::PolicyRegistry, protocols::{ chat::ChatCompletionRequest, classify::ClassifyRequest, @@ -26,26 +24,13 @@ use crate::{ rerank::RerankRequest, responses::{ResponsesGetParams, ResponsesRequest}, }, - reasoning_parser::ParserFactory as ReasoningParserFactory, routers::RouterTrait, - tokenizer::traits::Tokenizer, - tool_parser::ParserFactory as ToolParserFactory, }; /// gRPC PD (Prefill-Decode) router implementation for SGLang #[derive(Clone)] -#[allow(dead_code)] // Fields will be used once implementation is complete pub struct GrpcPDRouter { worker_registry: Arc, - policy_registry: Arc, - tokenizer: Arc, - reasoning_parser_factory: ReasoningParserFactory, - tool_parser_factory: ToolParserFactory, - dp_aware: bool, - api_key: Option, - retry_config: RetryConfig, - configured_reasoning_parser: Option, - configured_tool_parser: Option, pipeline: RequestPipeline, shared_components: Arc, } @@ -94,15 +79,6 @@ impl GrpcPDRouter { Ok(GrpcPDRouter { worker_registry, - policy_registry, - tokenizer, - reasoning_parser_factory, - tool_parser_factory, - dp_aware: ctx.router_config.dp_aware, - api_key: ctx.router_config.api_key.clone(), - retry_config: ctx.router_config.effective_retry_config(), - configured_reasoning_parser: ctx.configured_reasoning_parser.clone(), - configured_tool_parser: ctx.configured_tool_parser.clone(), pipeline, shared_components, }) @@ -174,7 +150,6 @@ impl std::fmt::Debug for GrpcPDRouter { f.debug_struct("GrpcPDRouter") .field("prefill_workers_count", &prefill_workers.len()) .field("decode_workers_count", &decode_workers.len()) - .field("dp_aware", &self.dp_aware) .finish() } } @@ -255,19 +230,19 @@ impl RouterTrait for GrpcPDRouter { (StatusCode::NOT_IMPLEMENTED).into_response() } - async fn route_classify( + async fn route_embeddings( &self, _headers: Option<&HeaderMap>, - _body: &ClassifyRequest, + _body: &EmbeddingRequest, _model_id: Option<&str>, ) -> Response { (StatusCode::NOT_IMPLEMENTED).into_response() } - async fn route_embeddings( + async fn route_classify( &self, _headers: Option<&HeaderMap>, - _body: &EmbeddingRequest, + _body: &ClassifyRequest, _model_id: Option<&str>, ) -> Response { (StatusCode::NOT_IMPLEMENTED).into_response() diff --git a/sgl-router/src/routers/grpc/pipeline.rs b/sgl-router/src/routers/grpc/pipeline.rs index 6dc487709..93a958e60 100644 --- a/sgl-router/src/routers/grpc/pipeline.rs +++ b/sgl-router/src/routers/grpc/pipeline.rs @@ -1,31 +1,22 @@ -//! Pipeline stages for gRPC router request processing +//! Pipeline orchestrator for gRPC router request processing //! -//! This module defines the core pipeline abstraction and individual processing stages -//! that transform a RequestContext through its lifecycle. +//! This module defines the RequestPipeline orchestrator that coordinates +//! the execution of pipeline stages from request preparation to response delivery. -use std::{ - borrow::Cow, - collections::HashMap, - sync::Arc, - time::{Instant, SystemTime, UNIX_EPOCH}, -}; +use std::{collections::HashMap, sync::Arc}; -use async_trait::async_trait; use axum::response::{IntoResponse, Response}; -use proto::DisaggregatedParams; -use rand::Rng; use tokio::sync::RwLock; -use tracing::{debug, error, warn}; -use uuid::Uuid; +use tracing::{debug, error}; -use super::{context::*, processing, responses::BackgroundTaskInfo, streaming, utils}; +// Import all stage types from the stages module +use super::stages::*; +use super::{context::*, harmony, processing, responses::BackgroundTaskInfo, streaming, utils}; use crate::{ - core::{ConnectionMode, Worker, WorkerRegistry, WorkerType}, - grpc_client::proto, + core::WorkerRegistry, policies::PolicyRegistry, protocols::{ chat::{ChatCompletionRequest, ChatCompletionResponse}, - common::InputIds, generate::GenerateRequest, }, reasoning_parser::ParserFactory as ReasoningParserFactory, @@ -33,875 +24,6 @@ use crate::{ tool_parser::ParserFactory as ToolParserFactory, }; -// ============================================================================ -// Pipeline Trait -// ============================================================================ - -/// Trait for pipeline stages that process requests -#[async_trait] -pub trait PipelineStage: Send + Sync { - /// Execute this stage, mutating the context - /// - /// Returns: - /// - `Ok(None)` - Continue to next stage - /// - `Ok(Some(response))` - Pipeline complete, return this response (e.g., streaming) - /// - `Err(response)` - Error occurred, return this error response - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response>; - - /// Stage name for logging - fn name(&self) -> &'static str; -} - -// ============================================================================ -// Stage 1: Preparation -// ============================================================================ - -/// Preparation stage: Filter tools, process messages, tokenize, build constraints -pub struct PreparationStage; - -#[async_trait] -impl PipelineStage for PreparationStage { - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - // Clone Arc before match to avoid borrow checker issues - // (matching borrows ctx, but prepare_* methods need mutable borrow) - // Arc clone is cheap (8 bytes) - avoids full request clone (15KB-200KB) - let is_chat = matches!(&ctx.input.request_type, RequestType::Chat(_)); - - if is_chat { - let request_arc = ctx.chat_request_arc(); - self.prepare_chat(ctx, &request_arc).await?; - } else { - let request_arc = ctx.generate_request_arc(); - self.prepare_generate(ctx, &request_arc).await?; - } - - Ok(None) - } - - fn name(&self) -> &'static str { - "Preparation" - } -} - -impl PreparationStage { - async fn prepare_chat( - &self, - ctx: &mut RequestContext, - request: &ChatCompletionRequest, - ) -> Result<(), Response> { - // Step 1: Filter tools if needed - let body_ref = utils::filter_tools_for_request(request); - - // Step 2: Process messages and apply chat template - let processed_messages = - match utils::process_chat_messages(&body_ref, &*ctx.components.tokenizer) { - Ok(msgs) => msgs, - Err(e) => { - return Err(utils::bad_request_error(e)); - } - }; - - // Step 3: Tokenize the processed text - let encoding = match ctx.components.tokenizer.encode(&processed_messages.text) { - Ok(encoding) => encoding, - Err(e) => { - return Err(utils::internal_error_message(format!( - "Tokenization failed: {}", - e - ))); - } - }; - - let token_ids = encoding.token_ids().to_vec(); - - // Step 4: Build tool constraints if needed - let tool_call_constraint = if let Some(tools) = body_ref.tools.as_ref() { - utils::generate_tool_constraints(tools, &request.tool_choice, &request.model).map_err( - |e| utils::bad_request_error(format!("Invalid tool configuration: {}", e)), - )? - } else { - None - }; - - // Step 5: Create stop sequence decoder (build once, reuse in non-stream) - let stop_decoder = utils::create_stop_decoder( - &ctx.components.tokenizer, - request.stop.as_ref(), - request.stop_token_ids.as_ref(), - request.skip_special_tokens, - request.no_stop_trim, - ); - - // Store results in context - ctx.state.preparation = Some(PreparationOutput { - original_text: Some(processed_messages.text.clone()), - token_ids, - processed_messages: Some(processed_messages), - tool_constraints: tool_call_constraint, - filtered_request: if matches!(body_ref, Cow::Owned(_)) { - Some(body_ref.into_owned()) - } else { - None - }, - }); - - // Store stop decoder for reuse in response processing - ctx.state.response.stop_decoder = Some(stop_decoder); - - Ok(()) - } - - async fn prepare_generate( - &self, - ctx: &mut RequestContext, - request: &GenerateRequest, - ) -> Result<(), Response> { - // Resolve input (text, prompt, or input_ids) - let (original_text, token_ids) = match self.resolve_generate_input(ctx, request) { - Ok(res) => res, - Err(msg) => { - return Err(utils::bad_request_error(msg)); - } - }; - - // Create stop sequence decoder for generate requests - let params = request.sampling_params.as_ref(); - let stop_decoder = utils::create_stop_decoder( - &ctx.components.tokenizer, - params.and_then(|p| p.stop.as_ref()), - params.and_then(|p| p.stop_token_ids.as_ref()), - params.and_then(|p| p.skip_special_tokens).unwrap_or(true), - params.and_then(|p| p.no_stop_trim).unwrap_or(false), - ); - - ctx.state.preparation = Some(PreparationOutput { - original_text, - token_ids, - processed_messages: None, - tool_constraints: None, - filtered_request: None, - }); - - // Store stop decoder - ctx.state.response.stop_decoder = Some(stop_decoder); - - Ok(()) - } - - fn resolve_generate_input( - &self, - ctx: &RequestContext, - request: &GenerateRequest, - ) -> Result<(Option, Vec), String> { - if let Some(text) = &request.text { - return self - .tokenize_single_text(&ctx.components.tokenizer, text) - .map(|(original, ids)| (Some(original), ids)); - } - - // Handle input_ids - validate and convert - if let Some(input_ids) = &request.input_ids { - return match input_ids { - InputIds::Single(ids) => ids - .iter() - .map(|&id| u32::try_from(id)) - .collect::, _>>() - .map(|converted| (None, converted)) - .map_err(|_| "input_ids must be non-negative".to_string()), - InputIds::Batch(_) => { - Err("Batch input_ids are not supported over gRPC generate yet".to_string()) - } - }; - } - - Err("Either `text` or `input_ids` must be provided".to_string()) - } - - fn tokenize_single_text( - &self, - tokenizer: &Arc, - text: &str, - ) -> Result<(String, Vec), String> { - let encoding = tokenizer - .encode(text) - .map_err(|e| format!("Tokenization failed: {}", e))?; - Ok((text.to_string(), encoding.token_ids().to_vec())) - } -} - -// ============================================================================ -// Stage 2: Worker Selection -// ============================================================================ - -/// Worker selection stage: Select appropriate worker(s) based on routing mode -pub struct WorkerSelectionStage { - worker_registry: Arc, - policy_registry: Arc, - mode: WorkerSelectionMode, -} - -pub enum WorkerSelectionMode { - /// Regular mode: select single worker - Regular, - /// PD mode: select prefill + decode workers - PrefillDecode, -} - -impl WorkerSelectionStage { - pub fn new( - worker_registry: Arc, - policy_registry: Arc, - mode: WorkerSelectionMode, - ) -> Self { - Self { - worker_registry, - policy_registry, - mode, - } - } -} - -#[async_trait] -impl PipelineStage for WorkerSelectionStage { - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - let prep = ctx - .state - .preparation - .as_ref() - .ok_or_else(|| utils::internal_error_static("Preparation stage not completed"))?; - - let text = prep.original_text.as_deref(); - - let workers = match self.mode { - WorkerSelectionMode::Regular => { - match self.select_single_worker(ctx.input.model_id.as_deref(), text) { - Some(w) => WorkerSelection::Single { worker: w }, - None => { - return Err(utils::service_unavailable_error(format!( - "No available workers for model: {:?}", - ctx.input.model_id - ))); - } - } - } - WorkerSelectionMode::PrefillDecode => { - match self.select_pd_pair(ctx.input.model_id.as_deref(), text) { - Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode }, - None => { - return Err(utils::service_unavailable_error(format!( - "No available PD worker pairs for model: {:?}", - ctx.input.model_id - ))); - } - } - } - }; - - ctx.state.workers = Some(workers); - Ok(None) - } - - fn name(&self) -> &'static str { - "WorkerSelection" - } -} - -impl WorkerSelectionStage { - fn select_single_worker( - &self, - model_id: Option<&str>, - text: Option<&str>, - ) -> Option> { - // Get workers for the specified model, filtered by connection mode - let workers = self.worker_registry.get_workers_filtered( - model_id, - Some(WorkerType::Regular), - Some(ConnectionMode::Grpc { port: None }), - false, // get all workers, we'll filter by is_available() next - ); - - // Filter by availability (health + circuit breaker) - let available: Vec> = workers - .iter() - .filter(|w| w.is_available()) - .cloned() - .collect(); - - if available.is_empty() { - return None; - } - - // Get the appropriate policy for this model - let policy = match model_id { - Some(model) => self.policy_registry.get_policy_or_default(model), - None => self.policy_registry.get_default_policy(), - }; - - // Select worker using the policy - let idx = policy.select_worker(&available, text)?; - Some(available[idx].clone()) - } - - fn select_pd_pair( - &self, - model_id: Option<&str>, - text: Option<&str>, - ) -> Option<(Arc, Arc)> { - let all_workers = self.worker_registry.get_workers_filtered( - model_id, - None, - Some(ConnectionMode::Grpc { port: None }), // Match any gRPC worker - false, - ); - - let (available_prefill, available_decode): (Vec<_>, Vec<_>) = - all_workers - .into_iter() - .fold((Vec::new(), Vec::new()), |mut acc, w| { - if w.is_available() { - match w.metadata().worker_type { - WorkerType::Prefill { .. } => acc.0.push(w), - WorkerType::Decode => acc.1.push(w), - _ => {} - } - } - acc - }); - - if available_prefill.is_empty() { - warn!("No available prefill workers"); - return None; - } - - if available_decode.is_empty() { - warn!("No available decode workers"); - return None; - } - - // Select using policies - let policy = match model_id { - Some(model) => self.policy_registry.get_policy_or_default(model), - None => self.policy_registry.get_default_policy(), - }; - - let prefill_idx = policy.select_worker(&available_prefill, text)?; - let decode_idx = policy.select_worker(&available_decode, text)?; - - Some(( - available_prefill[prefill_idx].clone(), - available_decode[decode_idx].clone(), - )) - } -} - -// ============================================================================ -// Stage 3: Client Acquisition -// ============================================================================ - -/// Client acquisition stage: Get gRPC clients from selected workers -pub struct ClientAcquisitionStage; - -#[async_trait] -impl PipelineStage for ClientAcquisitionStage { - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - let workers = ctx - .state - .workers - .as_ref() - .ok_or_else(|| utils::internal_error_static("Worker selection not completed"))?; - - let clients = match workers { - WorkerSelection::Single { worker } => { - let client = utils::get_grpc_client_from_worker(worker).await?; - ClientSelection::Single { client } - } - WorkerSelection::Dual { prefill, decode } => { - let prefill_client = utils::get_grpc_client_from_worker(prefill).await?; - let decode_client = utils::get_grpc_client_from_worker(decode).await?; - ClientSelection::Dual { - prefill: prefill_client, - decode: decode_client, - } - } - }; - - ctx.state.clients = Some(clients); - Ok(None) - } - - fn name(&self) -> &'static str { - "ClientAcquisition" - } -} - -// ============================================================================ -// Stage 4: Request Building -// ============================================================================ - -/// Request building stage: Build proto GenerateRequest -pub struct RequestBuildingStage { - inject_pd_metadata: bool, -} - -impl RequestBuildingStage { - pub fn new(inject_pd_metadata: bool) -> Self { - Self { inject_pd_metadata } - } -} - -#[async_trait] -impl PipelineStage for RequestBuildingStage { - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - let prep = ctx - .state - .preparation - .as_ref() - .ok_or_else(|| utils::internal_error_static("Preparation not completed"))?; - - let clients = ctx - .state - .clients - .as_ref() - .ok_or_else(|| utils::internal_error_static("Client acquisition not completed"))?; - - // Get client for building request (use prefill client if PD mode) - let builder_client = match clients { - ClientSelection::Single { client } => client, - ClientSelection::Dual { prefill, .. } => prefill, - }; - - let mut proto_request = match &ctx.input.request_type { - RequestType::Chat(request) => { - let request_id = format!("chatcmpl-{}", Uuid::new_v4()); - let body_ref = prep.filtered_request.as_ref().unwrap_or(request); - - builder_client - .build_generate_request( - request_id, - body_ref, - prep.processed_messages.as_ref().unwrap().text.clone(), - prep.token_ids.clone(), - prep.processed_messages - .as_ref() - .unwrap() - .multimodal_inputs - .clone(), - prep.tool_constraints.clone(), - ) - .map_err(|e| { - utils::bad_request_error(format!("Invalid request parameters: {}", e)) - })? - } - RequestType::Generate(request) => { - let request_id = request - .rid - .clone() - .unwrap_or_else(|| format!("gen-{}", Uuid::new_v4())); - - builder_client - .build_plain_generate_request( - request_id, - request, - prep.original_text.clone(), - prep.token_ids.clone(), - ) - .map_err(utils::bad_request_error)? - } - }; - - // Inject PD metadata if needed - if self.inject_pd_metadata { - if let WorkerSelection::Dual { prefill, .. } = ctx.state.workers.as_ref().unwrap() { - self.inject_bootstrap_metadata(&mut proto_request, prefill); - } - } - - ctx.state.proto_request = Some(proto_request); - Ok(None) - } - - fn name(&self) -> &'static str { - "RequestBuilding" - } -} - -impl RequestBuildingStage { - fn inject_bootstrap_metadata( - &self, - request: &mut proto::GenerateRequest, - prefill_worker: &Arc, - ) { - let hostname = prefill_worker.bootstrap_host(); - let bootstrap_port = prefill_worker.bootstrap_port().unwrap_or(8998); - - // Generate room ID for bootstrap - let room_id = rand::rng().random_range(0..i32::MAX); - - // Create DisaggregatedParams - let disagg_params = DisaggregatedParams { - bootstrap_host: hostname.to_string(), - bootstrap_port: bootstrap_port as i32, - bootstrap_room: room_id, - }; - - // Inject metadata directly into request - request.disaggregated_params = Some(disagg_params); - - debug!( - "Injected bootstrap metadata: host={}, port={}, room={}", - hostname, bootstrap_port, room_id - ); - } -} - -// ============================================================================ -// Stage 5: Dispatch Metadata -// ============================================================================ - -/// Dispatch metadata stage: Prepare metadata for dispatch -pub struct DispatchMetadataStage; - -#[async_trait] -impl PipelineStage for DispatchMetadataStage { - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - let proto_request = ctx - .state - .proto_request - .as_ref() - .ok_or_else(|| utils::internal_error_static("Proto request not built"))?; - - let request_id = proto_request.request_id.clone(); - let model = match &ctx.input.request_type { - RequestType::Chat(req) => req.model.clone(), - RequestType::Generate(_req) => { - // Generate requests don't have a model field - // Use model_id from input or default - ctx.input - .model_id - .clone() - .unwrap_or_else(|| "default".to_string()) - } - }; - - let weight_version = ctx - .state - .workers - .as_ref() - .map(|w| match w { - WorkerSelection::Single { worker } => worker, - WorkerSelection::Dual { decode, .. } => decode, - }) - .and_then(|w| w.metadata().labels.get("weight_version").cloned()) - .unwrap_or_else(|| "default".to_string()); - - let created = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs(); - - ctx.state.dispatch = Some(DispatchMetadata { - request_id, - model, - created, - weight_version: Some(weight_version), - is_streaming: ctx.is_streaming(), - }); - - Ok(None) - } - - fn name(&self) -> &'static str { - "DispatchMetadata" - } -} - -// ============================================================================ -// Stage 6: Request Execution -// ============================================================================ - -/// Request execution stage: Execute gRPC requests (single or dual dispatch) -pub struct RequestExecutionStage { - mode: ExecutionMode, -} - -pub enum ExecutionMode { - /// Regular mode: single worker execution - Single, - /// PD mode: dual dispatch to prefill + decode workers - DualDispatch, -} - -impl RequestExecutionStage { - pub fn new(mode: ExecutionMode) -> Self { - Self { mode } - } -} - -#[async_trait] -impl PipelineStage for RequestExecutionStage { - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - let proto_request = ctx - .state - .proto_request - .take() - .ok_or_else(|| utils::internal_error_static("Proto request not built"))?; - - let clients = ctx - .state - .clients - .as_mut() - .ok_or_else(|| utils::internal_error_static("Client acquisition not completed"))?; - - let result = match self.mode { - ExecutionMode::Single => self.execute_single(proto_request, clients).await?, - ExecutionMode::DualDispatch => { - self.execute_dual_dispatch(proto_request, clients).await? - } - }; - - // Store result in context for ResponseProcessingStage - ctx.state.response.execution_result = Some(result); - Ok(None) - } - - fn name(&self) -> &'static str { - "RequestExecution" - } -} - -impl RequestExecutionStage { - async fn execute_single( - &self, - proto_request: proto::GenerateRequest, - clients: &mut ClientSelection, - ) -> Result { - let client = clients - .single_mut() - .ok_or_else(|| utils::internal_error_static("Expected single client but got dual"))?; - - let stream = client.generate(proto_request).await.map_err(|e| { - utils::internal_error_message(format!("Failed to start generation: {}", e)) - })?; - - Ok(ExecutionResult::Single { stream }) - } - - async fn execute_dual_dispatch( - &self, - proto_request: proto::GenerateRequest, - clients: &mut ClientSelection, - ) -> Result { - let (prefill_client, decode_client) = clients - .dual_mut() - .ok_or_else(|| utils::internal_error_static("Expected dual clients but got single"))?; - - let prefill_request = proto_request.clone(); - let decode_request = proto_request; - - let (prefill_result, decode_result) = tokio::join!( - prefill_client.generate(prefill_request), - decode_client.generate(decode_request) - ); - - // Handle prefill result - let prefill_stream = match prefill_result { - Ok(s) => s, - Err(e) => { - return Err(utils::internal_error_message(format!( - "Prefill worker failed to start: {}", - e - ))); - } - }; - - // Handle decode result - let decode_stream = match decode_result { - Ok(s) => s, - Err(e) => { - return Err(utils::internal_error_message(format!( - "Decode worker failed to start: {}", - e - ))); - } - }; - - Ok(ExecutionResult::Dual { - prefill: prefill_stream, - decode: Box::new(decode_stream), - }) - } -} - -// ============================================================================ -// Stage 7: Response Processing -// ============================================================================ - -/// Response processing stage: Handles both streaming and non-streaming responses -/// -/// - For streaming: Spawns background task and returns SSE response (early exit) -/// - For non-streaming: Collects all responses and builds final ChatCompletionResponse -pub struct ResponseProcessingStage { - processor: processing::ResponseProcessor, - streaming_processor: Arc, -} - -impl ResponseProcessingStage { - pub fn new( - processor: processing::ResponseProcessor, - streaming_processor: Arc, - ) -> Self { - Self { - processor, - streaming_processor, - } - } -} - -#[async_trait] -impl PipelineStage for ResponseProcessingStage { - async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - // Delegate to request-type specific processing - match &ctx.input.request_type { - RequestType::Chat(_) => return self.process_chat_response(ctx).await, - RequestType::Generate(_) => return self.process_generate_response(ctx).await, - } - } - - fn name(&self) -> &'static str { - "ResponseProcessing" - } -} - -impl ResponseProcessingStage { - async fn process_chat_response( - &self, - ctx: &mut RequestContext, - ) -> Result, Response> { - let is_streaming = ctx.is_streaming(); - - // Extract execution result - let execution_result = ctx - .state - .response - .execution_result - .take() - .ok_or_else(|| utils::internal_error_static("No execution result"))?; - - // Get dispatch metadata (needed by both streaming and non-streaming) - let dispatch = ctx - .state - .dispatch - .as_ref() - .ok_or_else(|| utils::internal_error_static("Dispatch metadata not set"))? - .clone(); - - if is_streaming { - // Streaming: Use StreamingProcessor and return SSE response (done) - return Ok(Some( - self.streaming_processor.clone().process_streaming_response( - execution_result, - ctx.chat_request_arc(), // Cheap Arc clone (8 bytes) - dispatch, - ), - )); - } - - // Non-streaming: Delegate to ResponseProcessor - let request_logprobs = match &ctx.input.request_type { - RequestType::Chat(req) => req.logprobs, - _ => false, - }; - - let chat_request = ctx.chat_request_arc(); - - let stop_decoder = ctx - .state - .response - .stop_decoder - .as_mut() - .ok_or_else(|| utils::internal_error_static("Stop decoder not initialized"))?; - - let response = self - .processor - .process_non_streaming_chat_response( - execution_result, - chat_request, - dispatch, - stop_decoder, - request_logprobs, - ) - .await?; - - // Store the final response - ctx.state.response.final_response = Some(FinalResponse::Chat(response)); - - Ok(None) - } - - async fn process_generate_response( - &self, - ctx: &mut RequestContext, - ) -> Result, Response> { - let start_time = Instant::now(); - let is_streaming = ctx.is_streaming(); - - // Extract execution result - let execution_result = ctx - .state - .response - .execution_result - .take() - .ok_or_else(|| utils::internal_error_static("No execution result"))?; - - // Get dispatch metadata (needed by both streaming and non-streaming) - let dispatch = ctx - .state - .dispatch - .as_ref() - .ok_or_else(|| utils::internal_error_static("Dispatch metadata not set"))? - .clone(); - - if is_streaming { - // Streaming: Use StreamingProcessor and return SSE response (done) - return Ok(Some( - self.streaming_processor.clone().process_streaming_generate( - execution_result, - ctx.generate_request_arc(), // Cheap Arc clone (8 bytes) - dispatch, - ), - )); - } - - // Non-streaming: Delegate to ResponseProcessor - let request_logprobs = ctx.generate_request().return_logprob.unwrap_or(false); - let generate_request = ctx.generate_request_arc(); - - let stop_decoder = ctx - .state - .response - .stop_decoder - .as_mut() - .ok_or_else(|| utils::internal_error_static("Stop decoder not initialized"))?; - - let result_array = self - .processor - .process_non_streaming_generate_response( - execution_result, - generate_request, - dispatch, - stop_decoder, - request_logprobs, - start_time, - ) - .await?; - - // Store the final response - ctx.state.response.final_response = Some(FinalResponse::Generate(result_array)); - - Ok(None) - } -} - // ============================================================================ // Pipeline Orchestrator // ============================================================================ @@ -963,6 +85,64 @@ impl RequestPipeline { } } + /// Create a Harmony (single-worker) pipeline for Harmony-capable models + pub fn new_harmony( + worker_registry: Arc, + policy_registry: Arc, + _tokenizer: Arc, + _tool_parser_factory: ToolParserFactory, + _reasoning_parser_factory: ReasoningParserFactory, + _configured_tool_parser: Option, + _configured_reasoning_parser: Option, + ) -> Self { + let stages: Vec> = vec![ + Box::new(harmony::stages::HarmonyPreparationStage::new()), + Box::new(WorkerSelectionStage::new( + worker_registry, + policy_registry, + WorkerSelectionMode::Regular, + )), + Box::new(ClientAcquisitionStage), + Box::new(harmony::stages::HarmonyRequestBuildingStage::new(false)), + Box::new(DispatchMetadataStage), + Box::new(RequestExecutionStage::new(ExecutionMode::Single)), + Box::new(harmony::stages::HarmonyResponseProcessingStage::new()), + ]; + + Self { + stages: Arc::new(stages), + } + } + + /// Create a Harmony PD (prefill-decode) pipeline + pub fn new_harmony_pd( + worker_registry: Arc, + policy_registry: Arc, + _tokenizer: Arc, + _tool_parser_factory: ToolParserFactory, + _reasoning_parser_factory: ReasoningParserFactory, + _configured_tool_parser: Option, + _configured_reasoning_parser: Option, + ) -> Self { + let stages: Vec> = vec![ + Box::new(harmony::stages::HarmonyPreparationStage::new()), + Box::new(WorkerSelectionStage::new( + worker_registry, + policy_registry, + WorkerSelectionMode::PrefillDecode, + )), + Box::new(ClientAcquisitionStage), + Box::new(harmony::stages::HarmonyRequestBuildingStage::new(true)), + Box::new(DispatchMetadataStage), + Box::new(RequestExecutionStage::new(ExecutionMode::DualDispatch)), + Box::new(harmony::stages::HarmonyResponseProcessingStage::new()), + ]; + + Self { + stages: Arc::new(stages), + } + } + /// Create a PD (prefill-decode) pipeline pub fn new_pd( worker_registry: Arc, @@ -1186,4 +366,92 @@ impl RequestPipeline { None => Err(utils::internal_error_static("No response produced")), } } + + /// Execute Responses API pipeline + /// + /// TODO: Implement Responses API native execution + /// This is a stub to allow compilation. The actual implementation should: + /// 1. Support multi-turn MCP loop orchestration + /// 2. Handle tool call execution and result injection + /// 3. Emit proper SSE events for streaming mode + /// 4. Store responses in data connector + /// + /// For now, this returns an error indicating the feature is not implemented. + pub async fn execute_responses( + &self, + _request: Arc, + _headers: Option, + _model_id: Option, + _components: Arc, + ) -> Response { + utils::internal_error_static("Responses API execution not yet implemented") + } + + /// Execute Harmony Responses API request through all pipeline stages + /// + /// This method runs a single iteration of the Responses API request, + /// returning either ToolCallsFound (continue serving) or Completed (final response). + /// + /// Called by harmony::responses::serve_harmony_responses() for each iteration. + /// + /// # Arguments + /// + /// * `request` - Responses API request + /// * `ctx` - Harmony Responses context with MCP manager and components + /// + /// # Returns + /// + /// ResponsesIterationResult indicating whether to continue iteration or return + pub async fn execute_harmony_responses( + &self, + request: &crate::protocols::responses::ResponsesRequest, + harmony_ctx: &harmony::responses::HarmonyResponsesContext, + ) -> Result { + // Create RequestContext for this Responses request + let mut ctx = RequestContext::for_responses( + Arc::new(request.clone()), + None, // No headers needed for internal pipeline execution + None, // Model ID already set in request + harmony_ctx.components.clone(), + ); + + // Execute each pipeline stage in sequence + for (idx, stage) in self.stages.iter().enumerate() { + match stage.execute(&mut ctx).await { + Ok(Some(response)) => { + // Stage returned early response (e.g., streaming) - not expected for Responses iteration + error!( + "Stage {} ({}) returned unexpected response during Responses iteration", + idx + 1, + stage.name() + ); + return Err(response); + } + Ok(None) => { + // Stage completed successfully, continue to next stage + continue; + } + Err(response) => { + // Stage failed + error!( + "Stage {} ({}) failed with status {}", + idx + 1, + stage.name(), + response.status() + ); + return Err(response); + } + } + } + + // Extract ResponsesIterationResult from context + // This should have been set by HarmonyResponseProcessingStage + ctx.state + .response + .responses_iteration_result + .take() + .ok_or_else(|| { + utils::internal_error_static("No ResponsesIterationResult produced by pipeline") + }) + } } diff --git a/sgl-router/src/routers/grpc/processing.rs b/sgl-router/src/routers/grpc/processing.rs index 995bfbcd3..da228b17a 100644 --- a/sgl-router/src/routers/grpc/processing.rs +++ b/sgl-router/src/routers/grpc/processing.rs @@ -110,7 +110,7 @@ impl ResponseProcessor { Ok(all_responses) } - /// Process a single choice from GenerateComplete response (EXACT COPY from router.rs:1573-1725) + /// Process a single choice from GenerateComplete response #[allow(clippy::too_many_arguments)] pub async fn process_single_choice( &self, @@ -367,7 +367,7 @@ impl ResponseProcessor { Ok(response) } - /// Parse tool calls using model-specific parser (EXACT COPY from router.rs:296-361) + /// Parse tool calls using model-specific parser pub async fn parse_tool_calls( &self, processed_text: &str, diff --git a/sgl-router/src/routers/grpc/responses/conversions.rs b/sgl-router/src/routers/grpc/responses/conversions.rs index caa364fab..63ce79d16 100644 --- a/sgl-router/src/routers/grpc/responses/conversions.rs +++ b/sgl-router/src/routers/grpc/responses/conversions.rs @@ -12,7 +12,8 @@ use crate::protocols::{ common::{FunctionCallResponse, StreamOptions, ToolCall, UsageInfo}, responses::{ ResponseContentPart, ResponseInput, ResponseInputOutputItem, ResponseOutputItem, - ResponseStatus, ResponsesRequest, ResponsesResponse, ResponsesUsage, + ResponseReasoningContent::ReasoningText, ResponseStatus, ResponsesRequest, + ResponsesResponse, ResponsesUsage, StringOrContentParts, }, }; @@ -50,7 +51,6 @@ pub fn responses_to_chat(req: &ResponsesRequest) -> Result { // Convert SimpleInputMessage to chat message - use crate::protocols::responses::StringOrContentParts; let text = match content { StringOrContentParts::String(s) => s.clone(), StringOrContentParts::Array(parts) => { @@ -170,9 +170,7 @@ pub fn responses_to_chat(req: &ResponsesRequest) -> Result { - text.as_str() - } + ReasoningText { text } => text.as_str(), }) .collect::>() .join("\n"); @@ -184,6 +182,17 @@ pub fn responses_to_chat(req: &ResponsesRequest) -> Result { + // Function call output - add as tool message + // Note: The function name is looked up from prev_outputs in Harmony path + // For Chat path, we just use the call_id + messages.push(ChatMessage::Tool { + content: output.clone(), + tool_call_id: call_id.clone(), + }); + } } } } @@ -282,11 +291,9 @@ pub fn chat_to_responses( output.push(ResponseOutputItem::Reasoning { id: format!("reasoning_{}", chat_resp.id), summary: vec![], - content: vec![ - crate::protocols::responses::ResponseReasoningContent::ReasoningText { - text: reasoning.clone(), - }, - ], + content: vec![ReasoningText { + text: reasoning.clone(), + }], status: Some("completed".to_string()), }); } diff --git a/sgl-router/src/routers/grpc/router.rs b/sgl-router/src/routers/grpc/router.rs index f5a11cd5a..5c0ce4188 100644 --- a/sgl-router/src/routers/grpc/router.rs +++ b/sgl-router/src/routers/grpc/router.rs @@ -11,12 +11,15 @@ use axum::{ }; use tracing::debug; -use super::{context::SharedComponents, pipeline::RequestPipeline, responses}; +use super::{ + context::SharedComponents, + harmony::{serve_harmony_responses, HarmonyDetector, HarmonyResponsesContext}, + pipeline::RequestPipeline, + responses, +}; use crate::{ app_context::AppContext, - config::types::RetryConfig, core::WorkerRegistry, - policies::PolicyRegistry, protocols::{ chat::ChatCompletionRequest, classify::ClassifyRequest, @@ -26,10 +29,7 @@ use crate::{ rerank::RerankRequest, responses::{ResponsesGetParams, ResponsesRequest}, }, - reasoning_parser::ParserFactory as ReasoningParserFactory, routers::RouterTrait, - tokenizer::traits::Tokenizer, - tool_parser::ParserFactory as ToolParserFactory, }; /// gRPC router implementation for SGLang @@ -37,19 +37,13 @@ use crate::{ #[allow(dead_code)] pub struct GrpcRouter { worker_registry: Arc, - policy_registry: Arc, - tokenizer: Arc, - reasoning_parser_factory: ReasoningParserFactory, - tool_parser_factory: ToolParserFactory, - dp_aware: bool, - api_key: Option, - retry_config: RetryConfig, - configured_reasoning_parser: Option, - configured_tool_parser: Option, pipeline: RequestPipeline, + harmony_pipeline: RequestPipeline, shared_components: Arc, // Responses context (bundles all /v1/responses dependencies: storage, MCP, background_tasks) responses_context: responses::ResponsesContext, + // Harmony responses context (uses harmony pipeline) + harmony_responses_context: responses::ResponsesContext, } impl GrpcRouter { @@ -73,7 +67,7 @@ impl GrpcRouter { .clone(); let worker_registry = ctx.worker_registry.clone(); - let policy_registry = ctx.policy_registry.clone(); + let _policy_registry = ctx.policy_registry.clone(); // Create shared components for pipeline let shared_components = Arc::new(SharedComponents { @@ -82,10 +76,10 @@ impl GrpcRouter { reasoning_parser_factory: reasoning_parser_factory.clone(), }); - // Create pipeline + // Create regular pipeline let pipeline = RequestPipeline::new_regular( worker_registry.clone(), - policy_registry.clone(), + _policy_registry.clone(), tokenizer.clone(), tool_parser_factory.clone(), reasoning_parser_factory.clone(), @@ -93,34 +87,48 @@ impl GrpcRouter { ctx.configured_reasoning_parser.clone(), ); - // Create responses context with all dependencies - let responses_context = responses::ResponsesContext::new( - Arc::new(pipeline.clone()), - shared_components.clone(), + // Create Harmony pipelines + let harmony_pipeline = RequestPipeline::new_harmony( worker_registry.clone(), - ctx.response_storage.clone(), - ctx.conversation_storage.clone(), - ctx.conversation_item_storage.clone(), - ctx.mcp_manager - .get() - .ok_or_else(|| "gRPC router requires MCP manager".to_string())? - .clone(), + _policy_registry.clone(), + tokenizer.clone(), + tool_parser_factory.clone(), + reasoning_parser_factory.clone(), + ctx.configured_tool_parser.clone(), + ctx.configured_reasoning_parser.clone(), ); + // Extract shared dependencies for responses contexts + let mcp_manager = ctx + .mcp_manager + .get() + .ok_or_else(|| "gRPC router requires MCP manager".to_string())? + .clone(); + + // Helper closure to create responses context with a given pipeline + let create_responses_context = |pipeline: &RequestPipeline| { + responses::ResponsesContext::new( + Arc::new(pipeline.clone()), + shared_components.clone(), + worker_registry.clone(), + ctx.response_storage.clone(), + ctx.conversation_storage.clone(), + ctx.conversation_item_storage.clone(), + mcp_manager.clone(), + ) + }; + + // Create responses contexts for both pipelines + let responses_context = create_responses_context(&pipeline); + let harmony_responses_context = create_responses_context(&harmony_pipeline); + Ok(GrpcRouter { worker_registry, - policy_registry, - tokenizer, - reasoning_parser_factory, - tool_parser_factory, - dp_aware: ctx.router_config.dp_aware, - api_key: ctx.router_config.api_key.clone(), - retry_config: ctx.router_config.effective_retry_config(), - configured_reasoning_parser: ctx.configured_reasoning_parser.clone(), - configured_tool_parser: ctx.configured_tool_parser.clone(), pipeline, + harmony_pipeline, shared_components, responses_context, + harmony_responses_context, }) } @@ -131,13 +139,22 @@ impl GrpcRouter { body: &ChatCompletionRequest, model_id: Option<&str>, ) -> Response { + // Choose Harmony pipeline if model indicates Harmony + let is_harmony = HarmonyDetector::is_harmony_model(&body.model); + debug!( - "Processing chat completion request for model: {:?}", - model_id + "Processing chat completion request for model: {:?}, using_harmony={}", + model_id, is_harmony ); - // Use pipeline for ALL requests (streaming and non-streaming) - self.pipeline + let pipeline = if is_harmony { + &self.harmony_pipeline + } else { + &self.pipeline + }; + + // Use selected pipeline for ALL requests (streaming and non-streaming) + pipeline .execute_chat( Arc::new(body.clone()), headers.cloned(), @@ -166,6 +183,33 @@ impl GrpcRouter { ) .await } + + /// Main route_responses implementation (pipeline-based for Harmony) + async fn route_responses_impl( + &self, + _headers: Option<&HeaderMap>, + body: &ResponsesRequest, + model_id: Option<&str>, + ) -> Response { + debug!( + "Processing Harmony responses request for model: {:?}", + model_id + ); + + // Create HarmonyResponsesContext from existing responses context + let harmony_ctx = HarmonyResponsesContext::new( + Arc::new(self.harmony_pipeline.clone()), + self.shared_components.clone(), + self.harmony_responses_context.mcp_manager.clone(), + self.harmony_responses_context.response_storage.clone(), + ); + + // Use serve_harmony_responses for multi-turn MCP tool orchestration + match serve_harmony_responses(&harmony_ctx, body.clone()).await { + Ok(response) => axum::Json(response).into_response(), + Err(error_response) => error_response, + } + } } impl std::fmt::Debug for GrpcRouter { @@ -173,7 +217,6 @@ impl std::fmt::Debug for GrpcRouter { let stats = self.worker_registry.stats(); f.debug_struct("GrpcRouter") .field("workers_count", &stats.total_workers) - .field("dp_aware", &self.dp_aware) .finish() } } @@ -238,13 +281,27 @@ impl RouterTrait for GrpcRouter { body: &ResponsesRequest, model_id: Option<&str>, ) -> Response { - responses::route_responses( - &self.responses_context, - Arc::new(body.clone()), - headers.cloned(), - model_id.map(|s| s.to_string()), - ) - .await + // Choose implementation based on Harmony model detection + let is_harmony = HarmonyDetector::is_harmony_model(&body.model); + + debug!( + "Processing responses request for model: {:?}, using_harmony={}", + model_id, is_harmony + ); + + if is_harmony { + // Use pipeline-based implementation for Harmony models + self.route_responses_impl(headers, body, model_id).await + } else { + // Use legacy responses module for non-Harmony models + responses::route_responses( + &self.responses_context, + Arc::new(body.clone()), + headers.cloned(), + model_id.map(|s| s.to_string()), + ) + .await + } } async fn get_response( @@ -260,19 +317,19 @@ impl RouterTrait for GrpcRouter { responses::cancel_response_impl(&self.responses_context, response_id).await } - async fn route_classify( + async fn route_embeddings( &self, _headers: Option<&HeaderMap>, - _body: &ClassifyRequest, + _body: &EmbeddingRequest, _model_id: Option<&str>, ) -> Response { (StatusCode::NOT_IMPLEMENTED).into_response() } - async fn route_embeddings( + async fn route_classify( &self, _headers: Option<&HeaderMap>, - _body: &EmbeddingRequest, + _body: &ClassifyRequest, _model_id: Option<&str>, ) -> Response { (StatusCode::NOT_IMPLEMENTED).into_response() diff --git a/sgl-router/src/routers/grpc/stages/client_acquisition.rs b/sgl-router/src/routers/grpc/stages/client_acquisition.rs new file mode 100644 index 000000000..9915241a4 --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/client_acquisition.rs @@ -0,0 +1,46 @@ +//! Client acquisition stage: Get gRPC clients from selected workers + +use async_trait::async_trait; +use axum::response::Response; + +use super::PipelineStage; +use crate::routers::grpc::{ + context::{ClientSelection, RequestContext, WorkerSelection}, + utils, +}; + +/// Client acquisition stage: Get gRPC clients from selected workers +pub struct ClientAcquisitionStage; + +#[async_trait] +impl PipelineStage for ClientAcquisitionStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let workers = ctx + .state + .workers + .as_ref() + .ok_or_else(|| utils::internal_error_static("Worker selection not completed"))?; + + let clients = match workers { + WorkerSelection::Single { worker } => { + let client = utils::get_grpc_client_from_worker(worker).await?; + ClientSelection::Single { client } + } + WorkerSelection::Dual { prefill, decode } => { + let prefill_client = utils::get_grpc_client_from_worker(prefill).await?; + let decode_client = utils::get_grpc_client_from_worker(decode).await?; + ClientSelection::Dual { + prefill: prefill_client, + decode: decode_client, + } + } + }; + + ctx.state.clients = Some(clients); + Ok(None) + } + + fn name(&self) -> &'static str { + "ClientAcquisition" + } +} diff --git a/sgl-router/src/routers/grpc/stages/dispatch_metadata.rs b/sgl-router/src/routers/grpc/stages/dispatch_metadata.rs new file mode 100644 index 000000000..74b3f30ca --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/dispatch_metadata.rs @@ -0,0 +1,70 @@ +//! Dispatch metadata stage: Prepare metadata for dispatch + +use std::time::{SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use axum::response::Response; + +use super::PipelineStage; +use crate::routers::grpc::{ + context::{DispatchMetadata, RequestContext, RequestType, WorkerSelection}, + utils, +}; + +/// Dispatch metadata stage: Prepare metadata for dispatch +pub struct DispatchMetadataStage; + +#[async_trait] +impl PipelineStage for DispatchMetadataStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let proto_request = ctx + .state + .proto_request + .as_ref() + .ok_or_else(|| utils::internal_error_static("Proto request not built"))?; + + let request_id = proto_request.request_id.clone(); + let model = match &ctx.input.request_type { + RequestType::Chat(req) => req.model.clone(), + RequestType::Generate(_req) => { + // Generate requests don't have a model field + // Use model_id from input or default + ctx.input + .model_id + .clone() + .unwrap_or_else(|| "default".to_string()) + } + RequestType::Responses(req) => req.model.clone(), + }; + + let weight_version = ctx + .state + .workers + .as_ref() + .map(|w| match w { + WorkerSelection::Single { worker } => worker, + WorkerSelection::Dual { decode, .. } => decode, + }) + .and_then(|w| w.metadata().labels.get("weight_version").cloned()) + .unwrap_or_else(|| "default".to_string()); + + let created = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + + ctx.state.dispatch = Some(DispatchMetadata { + request_id, + model, + created, + weight_version: Some(weight_version), + is_streaming: ctx.is_streaming(), + }); + + Ok(None) + } + + fn name(&self) -> &'static str { + "DispatchMetadata" + } +} diff --git a/sgl-router/src/routers/grpc/stages/mod.rs b/sgl-router/src/routers/grpc/stages/mod.rs new file mode 100644 index 000000000..5c5d86dc9 --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/mod.rs @@ -0,0 +1,52 @@ +//! Pipeline stages for gRPC router request processing +//! +//! This module defines the core pipeline abstraction and individual processing stages +//! that transform a RequestContext through its lifecycle. + +use async_trait::async_trait; +use axum::response::Response; + +use crate::routers::grpc::context::RequestContext; + +// ============================================================================ +// Pipeline Trait +// ============================================================================ + +/// Trait for pipeline stages that process requests +#[async_trait] +pub trait PipelineStage: Send + Sync { + /// Execute this stage, mutating the context + /// + /// Returns: + /// - `Ok(None)` - Continue to next stage + /// - `Ok(Some(response))` - Pipeline complete, return this response (e.g., streaming) + /// - `Err(response)` - Error occurred, return this error response + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response>; + + /// Stage name for logging + fn name(&self) -> &'static str; +} + +// ============================================================================ +// Stage Modules +// ============================================================================ + +mod client_acquisition; +mod dispatch_metadata; +mod preparation; +mod request_building; +mod request_execution; +mod response_processing; +mod worker_selection; + +// ============================================================================ +// Public Exports +// ============================================================================ + +pub use client_acquisition::ClientAcquisitionStage; +pub use dispatch_metadata::DispatchMetadataStage; +pub use preparation::PreparationStage; +pub use request_building::RequestBuildingStage; +pub use request_execution::{ExecutionMode, RequestExecutionStage}; +pub use response_processing::ResponseProcessingStage; +pub use worker_selection::{WorkerSelectionMode, WorkerSelectionStage}; diff --git a/sgl-router/src/routers/grpc/stages/preparation.rs b/sgl-router/src/routers/grpc/stages/preparation.rs new file mode 100644 index 000000000..cd7c314e9 --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/preparation.rs @@ -0,0 +1,199 @@ +//! Preparation stage: Filter tools, process messages, tokenize, build constraints + +use std::{borrow::Cow, sync::Arc}; + +use async_trait::async_trait; +use axum::response::Response; + +use super::PipelineStage; +use crate::{ + protocols::{chat::ChatCompletionRequest, common::InputIds, generate::GenerateRequest}, + routers::grpc::{ + context::{PreparationOutput, RequestContext, RequestType}, + utils, + }, + tokenizer::traits::Tokenizer, +}; + +/// Preparation stage: Filter tools, process messages, tokenize, build constraints +pub struct PreparationStage; + +#[async_trait] +impl PipelineStage for PreparationStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + // Clone Arc before match to avoid borrow checker issues + // (matching borrows ctx, but prepare_* methods need mutable borrow) + // Arc clone is cheap (8 bytes) - avoids full request clone (15KB-200KB) + let is_chat = matches!(&ctx.input.request_type, RequestType::Chat(_)); + + if is_chat { + let request_arc = ctx.chat_request_arc(); + self.prepare_chat(ctx, &request_arc).await?; + } else { + let request_arc = ctx.generate_request_arc(); + self.prepare_generate(ctx, &request_arc).await?; + } + + Ok(None) + } + + fn name(&self) -> &'static str { + "Preparation" + } +} + +impl PreparationStage { + async fn prepare_chat( + &self, + ctx: &mut RequestContext, + request: &ChatCompletionRequest, + ) -> Result<(), Response> { + // Step 1: Filter tools if needed + let body_ref = utils::filter_tools_for_request(request); + + // Step 2: Process messages and apply chat template + let processed_messages = + match utils::process_chat_messages(&body_ref, &*ctx.components.tokenizer) { + Ok(msgs) => msgs, + Err(e) => { + return Err(utils::bad_request_error(e)); + } + }; + + // Step 3: Tokenize the processed text + let encoding = match ctx.components.tokenizer.encode(&processed_messages.text) { + Ok(encoding) => encoding, + Err(e) => { + return Err(utils::internal_error_message(format!( + "Tokenization failed: {}", + e + ))); + } + }; + + let token_ids = encoding.token_ids().to_vec(); + + // Step 4: Build tool constraints if needed + let tool_call_constraint = if let Some(tools) = body_ref.tools.as_ref() { + utils::generate_tool_constraints(tools, &request.tool_choice, &request.model).map_err( + |e| utils::bad_request_error(format!("Invalid tool configuration: {}", e)), + )? + } else { + None + }; + + // Step 5: Create stop sequence decoder (build once, reuse in non-stream) + let stop_decoder = utils::create_stop_decoder( + &ctx.components.tokenizer, + request.stop.as_ref(), + request.stop_token_ids.as_ref(), + request.skip_special_tokens, + request.no_stop_trim, + ); + + // Store results in context + ctx.state.preparation = Some(PreparationOutput { + original_text: Some(processed_messages.text.clone()), + token_ids, + processed_messages: Some(processed_messages), + tool_constraints: tool_call_constraint, + filtered_request: if matches!(body_ref, Cow::Owned(_)) { + Some(body_ref.into_owned()) + } else { + None + }, + // Harmony fields (not used for regular preparation) + harmony_mode: false, + selection_text: None, + harmony_messages: None, + harmony_stop_ids: None, + }); + + // Store stop decoder for reuse in response processing + ctx.state.response.stop_decoder = Some(stop_decoder); + + Ok(()) + } + + async fn prepare_generate( + &self, + ctx: &mut RequestContext, + request: &GenerateRequest, + ) -> Result<(), Response> { + // Resolve input (text, prompt, or input_ids) + let (original_text, token_ids) = match self.resolve_generate_input(ctx, request) { + Ok(res) => res, + Err(msg) => { + return Err(utils::bad_request_error(msg)); + } + }; + + // Create stop sequence decoder for generate requests + let params = request.sampling_params.as_ref(); + let stop_decoder = utils::create_stop_decoder( + &ctx.components.tokenizer, + params.and_then(|p| p.stop.as_ref()), + params.and_then(|p| p.stop_token_ids.as_ref()), + params.and_then(|p| p.skip_special_tokens).unwrap_or(true), + params.and_then(|p| p.no_stop_trim).unwrap_or(false), + ); + + ctx.state.preparation = Some(PreparationOutput { + original_text, + token_ids, + processed_messages: None, + tool_constraints: None, + filtered_request: None, + // Harmony fields (not used for generate requests) + harmony_mode: false, + selection_text: None, + harmony_messages: None, + harmony_stop_ids: None, + }); + + // Store stop decoder + ctx.state.response.stop_decoder = Some(stop_decoder); + + Ok(()) + } + + fn resolve_generate_input( + &self, + ctx: &RequestContext, + request: &GenerateRequest, + ) -> Result<(Option, Vec), String> { + if let Some(text) = &request.text { + return self + .tokenize_single_text(&ctx.components.tokenizer, text) + .map(|(original, ids)| (Some(original), ids)); + } + + // Handle input_ids - validate and convert + if let Some(input_ids) = &request.input_ids { + return match input_ids { + InputIds::Single(ids) => ids + .iter() + .map(|&id| u32::try_from(id)) + .collect::, _>>() + .map(|converted| (None, converted)) + .map_err(|_| "input_ids must be non-negative".to_string()), + InputIds::Batch(_) => { + Err("Batch input_ids are not supported over gRPC generate yet".to_string()) + } + }; + } + + Err("Either `text` or `input_ids` must be provided".to_string()) + } + + fn tokenize_single_text( + &self, + tokenizer: &Arc, + text: &str, + ) -> Result<(String, Vec), String> { + let encoding = tokenizer + .encode(text) + .map_err(|e| format!("Tokenization failed: {}", e))?; + Ok((text.to_string(), encoding.token_ids().to_vec())) + } +} diff --git a/sgl-router/src/routers/grpc/stages/request_building.rs b/sgl-router/src/routers/grpc/stages/request_building.rs new file mode 100644 index 000000000..7e10d47bc --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/request_building.rs @@ -0,0 +1,146 @@ +//! Request building stage: Build proto GenerateRequest + +use std::sync::Arc; + +use async_trait::async_trait; +use axum::response::Response; +use proto::DisaggregatedParams; +use rand::Rng; +use tracing::debug; +use uuid::Uuid; + +use super::PipelineStage; +use crate::{ + core::Worker, + grpc_client::proto, + routers::grpc::{ + context::{ClientSelection, RequestContext, RequestType, WorkerSelection}, + utils, + }, +}; + +/// Request building stage: Build proto GenerateRequest +pub struct RequestBuildingStage { + inject_pd_metadata: bool, +} + +impl RequestBuildingStage { + pub fn new(inject_pd_metadata: bool) -> Self { + Self { inject_pd_metadata } + } +} + +#[async_trait] +impl PipelineStage for RequestBuildingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let prep = ctx + .state + .preparation + .as_ref() + .ok_or_else(|| utils::internal_error_static("Preparation not completed"))?; + + let clients = ctx + .state + .clients + .as_ref() + .ok_or_else(|| utils::internal_error_static("Client acquisition not completed"))?; + + // Get client for building request (use prefill client if PD mode) + let builder_client = match clients { + ClientSelection::Single { client } => client, + ClientSelection::Dual { prefill, .. } => prefill, + }; + + let mut proto_request = match &ctx.input.request_type { + RequestType::Chat(request) => { + let request_id = format!("chatcmpl-{}", Uuid::new_v4()); + let body_ref = prep.filtered_request.as_ref().unwrap_or(request); + + builder_client + .build_generate_request( + request_id, + body_ref, + prep.processed_messages.as_ref().unwrap().text.clone(), + prep.token_ids.clone(), + prep.processed_messages + .as_ref() + .unwrap() + .multimodal_inputs + .clone(), + prep.tool_constraints.clone(), + ) + .map_err(|e| { + utils::bad_request_error(format!("Invalid request parameters: {}", e)) + })? + } + RequestType::Generate(request) => { + let request_id = request + .rid + .clone() + .unwrap_or_else(|| format!("gen-{}", Uuid::new_v4())); + + builder_client + .build_plain_generate_request( + request_id, + request, + prep.original_text.clone(), + prep.token_ids.clone(), + ) + .map_err(utils::bad_request_error)? + } + RequestType::Responses(_request) => { + // Responses API builds request during the MCP loop + // For now, create minimal request - responses handler will populate it + let request_id = format!("resp-{}", Uuid::new_v4()); + + proto::GenerateRequest { + request_id, + ..Default::default() + } + } + }; + + // Inject PD metadata if needed + if self.inject_pd_metadata { + if let WorkerSelection::Dual { prefill, .. } = ctx.state.workers.as_ref().unwrap() { + self.inject_bootstrap_metadata(&mut proto_request, prefill); + } + } + + ctx.state.proto_request = Some(proto_request); + Ok(None) + } + + fn name(&self) -> &'static str { + "RequestBuilding" + } +} + +impl RequestBuildingStage { + fn inject_bootstrap_metadata( + &self, + request: &mut proto::GenerateRequest, + prefill_worker: &Arc, + ) { + let hostname = prefill_worker.bootstrap_host(); + let bootstrap_port = prefill_worker.bootstrap_port().unwrap_or(8998); + + // Generate room ID for bootstrap + let room_id = rand::rng().random_range(0..i32::MAX); + + // Create DisaggregatedParams + let disagg_params = DisaggregatedParams { + bootstrap_host: hostname.to_string(), + bootstrap_port: bootstrap_port as i32, + bootstrap_room: room_id, + }; + + // Inject metadata directly into request + request.disaggregated_params = Some(disagg_params); + + debug!( + "Injected bootstrap metadata: host={}, port={}, room={}", + hostname, bootstrap_port, room_id + ); + } +} diff --git a/sgl-router/src/routers/grpc/stages/request_execution.rs b/sgl-router/src/routers/grpc/stages/request_execution.rs new file mode 100644 index 000000000..0bcf7862d --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/request_execution.rs @@ -0,0 +1,126 @@ +//! Request execution stage: Execute gRPC requests (single or dual dispatch) + +use async_trait::async_trait; +use axum::response::Response; + +use super::PipelineStage; +use crate::{ + grpc_client::proto, + routers::grpc::{ + context::{ClientSelection, ExecutionResult, RequestContext}, + utils, + }, +}; + +/// Request execution stage: Execute gRPC requests (single or dual dispatch) +pub struct RequestExecutionStage { + mode: ExecutionMode, +} + +pub enum ExecutionMode { + /// Regular mode: single worker execution + Single, + /// PD mode: dual dispatch to prefill + decode workers + DualDispatch, +} + +impl RequestExecutionStage { + pub fn new(mode: ExecutionMode) -> Self { + Self { mode } + } +} + +#[async_trait] +impl PipelineStage for RequestExecutionStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let proto_request = ctx + .state + .proto_request + .take() + .ok_or_else(|| utils::internal_error_static("Proto request not built"))?; + + let clients = ctx + .state + .clients + .as_mut() + .ok_or_else(|| utils::internal_error_static("Client acquisition not completed"))?; + + let result = match self.mode { + ExecutionMode::Single => self.execute_single(proto_request, clients).await?, + ExecutionMode::DualDispatch => { + self.execute_dual_dispatch(proto_request, clients).await? + } + }; + + // Store result in context for ResponseProcessingStage + ctx.state.response.execution_result = Some(result); + Ok(None) + } + + fn name(&self) -> &'static str { + "RequestExecution" + } +} + +impl RequestExecutionStage { + async fn execute_single( + &self, + proto_request: proto::GenerateRequest, + clients: &mut ClientSelection, + ) -> Result { + let client = clients + .single_mut() + .ok_or_else(|| utils::internal_error_static("Expected single client but got dual"))?; + + let stream = client.generate(proto_request).await.map_err(|e| { + utils::internal_error_message(format!("Failed to start generation: {}", e)) + })?; + + Ok(ExecutionResult::Single { stream }) + } + + async fn execute_dual_dispatch( + &self, + proto_request: proto::GenerateRequest, + clients: &mut ClientSelection, + ) -> Result { + let (prefill_client, decode_client) = clients + .dual_mut() + .ok_or_else(|| utils::internal_error_static("Expected dual clients but got single"))?; + + let prefill_request = proto_request.clone(); + let decode_request = proto_request; + + let (prefill_result, decode_result) = tokio::join!( + prefill_client.generate(prefill_request), + decode_client.generate(decode_request) + ); + + // Handle prefill result + let prefill_stream = match prefill_result { + Ok(s) => s, + Err(e) => { + return Err(utils::internal_error_message(format!( + "Prefill worker failed to start: {}", + e + ))); + } + }; + + // Handle decode result + let decode_stream = match decode_result { + Ok(s) => s, + Err(e) => { + return Err(utils::internal_error_message(format!( + "Decode worker failed to start: {}", + e + ))); + } + }; + + Ok(ExecutionResult::Dual { + prefill: prefill_stream, + decode: Box::new(decode_stream), + }) + } +} diff --git a/sgl-router/src/routers/grpc/stages/response_processing.rs b/sgl-router/src/routers/grpc/stages/response_processing.rs new file mode 100644 index 000000000..b7d8232b8 --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/response_processing.rs @@ -0,0 +1,184 @@ +//! Response processing stage: Handles both streaming and non-streaming responses +//! +//! - For streaming: Spawns background task and returns SSE response (early exit) +//! - For non-streaming: Collects all responses and builds final ChatCompletionResponse + +use std::{sync::Arc, time::Instant}; + +use async_trait::async_trait; +use axum::response::Response; + +use super::PipelineStage; +use crate::routers::grpc::{ + context::{FinalResponse, RequestContext, RequestType}, + processing, streaming, utils, +}; + +/// Response processing stage: Handles both streaming and non-streaming responses +/// +/// - For streaming: Spawns background task and returns SSE response (early exit) +/// - For non-streaming: Collects all responses and builds final ChatCompletionResponse +pub struct ResponseProcessingStage { + processor: processing::ResponseProcessor, + streaming_processor: Arc, +} + +impl ResponseProcessingStage { + pub fn new( + processor: processing::ResponseProcessor, + streaming_processor: Arc, + ) -> Self { + Self { + processor, + streaming_processor, + } + } +} + +#[async_trait] +impl PipelineStage for ResponseProcessingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + // Delegate to request-type specific processing + match &ctx.input.request_type { + RequestType::Chat(_) => self.process_chat_response(ctx).await, + RequestType::Generate(_) => self.process_generate_response(ctx).await, + RequestType::Responses(_) => Err(utils::bad_request_error( + "Responses API processing must be handled by responses handler".to_string(), + )), + } + } + + fn name(&self) -> &'static str { + "ResponseProcessing" + } +} + +impl ResponseProcessingStage { + async fn process_chat_response( + &self, + ctx: &mut RequestContext, + ) -> Result, Response> { + let is_streaming = ctx.is_streaming(); + + // Extract execution result + let execution_result = ctx + .state + .response + .execution_result + .take() + .ok_or_else(|| utils::internal_error_static("No execution result"))?; + + // Get dispatch metadata (needed by both streaming and non-streaming) + let dispatch = ctx + .state + .dispatch + .as_ref() + .ok_or_else(|| utils::internal_error_static("Dispatch metadata not set"))? + .clone(); + + if is_streaming { + // Streaming: Use StreamingProcessor and return SSE response (done) + return Ok(Some( + self.streaming_processor.clone().process_streaming_response( + execution_result, + ctx.chat_request_arc(), // Cheap Arc clone (8 bytes) + dispatch, + ), + )); + } + + // Non-streaming: Delegate to ResponseProcessor + let request_logprobs = match &ctx.input.request_type { + RequestType::Chat(req) => req.logprobs, + _ => false, + }; + + let chat_request = ctx.chat_request_arc(); + + let stop_decoder = ctx + .state + .response + .stop_decoder + .as_mut() + .ok_or_else(|| utils::internal_error_static("Stop decoder not initialized"))?; + + let response = self + .processor + .process_non_streaming_chat_response( + execution_result, + chat_request, + dispatch, + stop_decoder, + request_logprobs, + ) + .await?; + + // Store the final response + ctx.state.response.final_response = Some(FinalResponse::Chat(response)); + + Ok(None) + } + + async fn process_generate_response( + &self, + ctx: &mut RequestContext, + ) -> Result, Response> { + let start_time = Instant::now(); + let is_streaming = ctx.is_streaming(); + + // Extract execution result + let execution_result = ctx + .state + .response + .execution_result + .take() + .ok_or_else(|| utils::internal_error_static("No execution result"))?; + + // Get dispatch metadata (needed by both streaming and non-streaming) + let dispatch = ctx + .state + .dispatch + .as_ref() + .ok_or_else(|| utils::internal_error_static("Dispatch metadata not set"))? + .clone(); + + if is_streaming { + // Streaming: Use StreamingProcessor and return SSE response (done) + return Ok(Some( + self.streaming_processor.clone().process_streaming_generate( + execution_result, + ctx.generate_request_arc(), // Cheap Arc clone (8 bytes) + dispatch, + ), + )); + } + + // Non-streaming: Delegate to ResponseProcessor + let request_logprobs = ctx.generate_request().return_logprob.unwrap_or(false); + let generate_request = ctx.generate_request_arc(); + + let stop_decoder = ctx + .state + .response + .stop_decoder + .as_mut() + .ok_or_else(|| utils::internal_error_static("Stop decoder not initialized"))?; + + let result_array = self + .processor + .process_non_streaming_generate_response( + execution_result, + generate_request, + dispatch, + stop_decoder, + request_logprobs, + start_time, + ) + .await?; + + // Store the final response + ctx.state.response.final_response = Some(FinalResponse::Generate(result_array)); + + Ok(None) + } +} diff --git a/sgl-router/src/routers/grpc/stages/worker_selection.rs b/sgl-router/src/routers/grpc/stages/worker_selection.rs new file mode 100644 index 000000000..5014abf1f --- /dev/null +++ b/sgl-router/src/routers/grpc/stages/worker_selection.rs @@ -0,0 +1,184 @@ +//! Worker selection stage: Select appropriate worker(s) based on routing mode + +use std::sync::Arc; + +use async_trait::async_trait; +use axum::response::Response; +use tracing::warn; + +use super::PipelineStage; +use crate::{ + core::{ConnectionMode, Worker, WorkerRegistry, WorkerType}, + policies::PolicyRegistry, + routers::grpc::{ + context::{RequestContext, WorkerSelection}, + utils, + }, +}; + +/// Worker selection stage: Select appropriate worker(s) based on routing mode +pub struct WorkerSelectionStage { + worker_registry: Arc, + policy_registry: Arc, + mode: WorkerSelectionMode, +} + +pub enum WorkerSelectionMode { + /// Regular mode: select single worker + Regular, + /// PD mode: select prefill + decode workers + PrefillDecode, +} + +impl WorkerSelectionStage { + pub fn new( + worker_registry: Arc, + policy_registry: Arc, + mode: WorkerSelectionMode, + ) -> Self { + Self { + worker_registry, + policy_registry, + mode, + } + } +} + +#[async_trait] +impl PipelineStage for WorkerSelectionStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let prep = ctx + .state + .preparation + .as_ref() + .ok_or_else(|| utils::internal_error_static("Preparation stage not completed"))?; + + // For Harmony, use selection_text produced during Harmony encoding + // Otherwise, use original_text from regular preparation + let text = if prep.harmony_mode { + prep.selection_text.as_deref() + } else { + prep.original_text.as_deref() + }; + + let workers = match self.mode { + WorkerSelectionMode::Regular => { + match self.select_single_worker(ctx.input.model_id.as_deref(), text) { + Some(w) => WorkerSelection::Single { worker: w }, + None => { + return Err(utils::service_unavailable_error(format!( + "No available workers for model: {:?}", + ctx.input.model_id + ))); + } + } + } + WorkerSelectionMode::PrefillDecode => { + match self.select_pd_pair(ctx.input.model_id.as_deref(), text) { + Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode }, + None => { + return Err(utils::service_unavailable_error(format!( + "No available PD worker pairs for model: {:?}", + ctx.input.model_id + ))); + } + } + } + }; + + ctx.state.workers = Some(workers); + Ok(None) + } + + fn name(&self) -> &'static str { + "WorkerSelection" + } +} + +impl WorkerSelectionStage { + fn select_single_worker( + &self, + model_id: Option<&str>, + text: Option<&str>, + ) -> Option> { + // Get workers for the specified model, filtered by connection mode + let workers = self.worker_registry.get_workers_filtered( + model_id, + Some(WorkerType::Regular), + Some(ConnectionMode::Grpc { port: None }), + false, // get all workers, we'll filter by is_available() next + ); + + // Filter by availability (health + circuit breaker) + let available: Vec> = workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + + if available.is_empty() { + return None; + } + + // Get the appropriate policy for this model + let policy = match model_id { + Some(model) => self.policy_registry.get_policy_or_default(model), + None => self.policy_registry.get_default_policy(), + }; + + // Select worker using the policy + let idx = policy.select_worker(&available, text)?; + Some(available[idx].clone()) + } + + fn select_pd_pair( + &self, + model_id: Option<&str>, + text: Option<&str>, + ) -> Option<(Arc, Arc)> { + let all_workers = self.worker_registry.get_workers_filtered( + model_id, + None, + Some(ConnectionMode::Grpc { port: None }), // Match any gRPC worker + false, + ); + + let (available_prefill, available_decode): (Vec<_>, Vec<_>) = + all_workers + .into_iter() + .fold((Vec::new(), Vec::new()), |mut acc, w| { + if w.is_available() { + match w.metadata().worker_type { + WorkerType::Prefill { .. } => acc.0.push(w), + WorkerType::Decode => acc.1.push(w), + _ => {} + } + } + acc + }); + + if available_prefill.is_empty() { + warn!("No available prefill workers"); + return None; + } + + if available_decode.is_empty() { + warn!("No available decode workers"); + return None; + } + + // Select using policies + let policy = match model_id { + Some(model) => self.policy_registry.get_policy_or_default(model), + None => self.policy_registry.get_default_policy(), + }; + + let prefill_idx = policy.select_worker(&available_prefill, text)?; + let decode_idx = policy.select_worker(&available_decode, text)?; + + Some(( + available_prefill[prefill_idx].clone(), + available_decode[decode_idx].clone(), + )) + } +} diff --git a/sgl-router/tests/responses_api_test.rs b/sgl-router/tests/responses_api_test.rs index 92ced5ff7..9d10db992 100644 --- a/sgl-router/tests/responses_api_test.rs +++ b/sgl-router/tests/responses_api_test.rs @@ -80,6 +80,7 @@ async fn test_non_streaming_mcp_minimal_e2e_with_persistence() { tool_choice: Some(ToolChoice::default()), tools: Some(vec![ResponseTool { r#type: ResponseToolType::Mcp, + function: None, server_url: Some(mcp.url()), authorization: None, server_label: Some("mock".to_string()),