diff --git a/sgl-router/src/protocols/responses.rs b/sgl-router/src/protocols/responses.rs index d09f96790..cb5d52d85 100644 --- a/sgl-router/src/protocols/responses.rs +++ b/sgl-router/src/protocols/responses.rs @@ -251,9 +251,10 @@ pub enum ResponseOutputItem { // Configuration Enums // ============================================================================ -#[derive(Debug, Clone, Deserialize, Serialize)] +#[derive(Debug, Clone, Deserialize, Serialize, Default)] #[serde(rename_all = "snake_case")] pub enum ServiceTier { + #[default] Auto, Default, Flex, @@ -261,25 +262,14 @@ pub enum ServiceTier { Priority, } -impl Default for ServiceTier { - fn default() -> Self { - Self::Auto - } -} - -#[derive(Debug, Clone, Deserialize, Serialize)] +#[derive(Debug, Clone, Deserialize, Serialize, Default)] #[serde(rename_all = "snake_case")] pub enum Truncation { Auto, + #[default] Disabled, } -impl Default for Truncation { - fn default() -> Self { - Self::Disabled - } -} - #[derive(Debug, Clone, Deserialize, Serialize)] #[serde(rename_all = "snake_case")] pub enum ResponseStatus { diff --git a/sgl-router/src/routers/grpc/common/mod.rs b/sgl-router/src/routers/grpc/common/mod.rs new file mode 100644 index 000000000..c42ea3b05 --- /dev/null +++ b/sgl-router/src/routers/grpc/common/mod.rs @@ -0,0 +1,6 @@ +//! Shared code for both regular and harmony routers + +pub mod response_collection; +pub mod response_formatting; +pub mod responses; +pub mod stages; diff --git a/sgl-router/src/routers/grpc/common/response_collection.rs b/sgl-router/src/routers/grpc/common/response_collection.rs new file mode 100644 index 000000000..46db1894d --- /dev/null +++ b/sgl-router/src/routers/grpc/common/response_collection.rs @@ -0,0 +1,83 @@ +//! Shared response collection logic +//! +//! This module contains common logic for collecting responses from execution results. +//! Both regular and harmony processors use these functions to avoid duplication. + +use axum::response::Response; + +use crate::{ + grpc_client::proto, + routers::grpc::{context::ExecutionResult, error, utils}, +}; + +/// Collect and merge responses from execution result +/// +/// Handles both Single and Dual (prefill-decode) execution modes. +/// For Dual mode, merges prefill input_logprobs into decode responses if requested. +/// +/// # Arguments +/// * `execution_result` - The execution result containing stream(s) +/// * `merge_logprobs` - Whether to merge prefill input_logprobs (for chat with logprobs=true) +/// +/// # Returns +/// Vector of GenerateComplete responses, one per index (n parameter) +pub async fn collect_responses( + execution_result: ExecutionResult, + merge_logprobs: bool, +) -> Result, Response> { + let all_responses = match execution_result { + ExecutionResult::Single { mut stream } => { + let responses = utils::collect_stream_responses(&mut stream, "Single").await?; + stream.mark_completed(); + responses + } + ExecutionResult::Dual { + mut prefill, + decode, + } => { + // Collect prefill for input_logprobs (don't mark completed yet) + let prefill_responses = + utils::collect_stream_responses(&mut prefill, "Prefill").await?; + + // Collect decode for actual output (don't mark completed yet) + let mut decode_stream = *decode; + let mut decode_responses = + utils::collect_stream_responses(&mut decode_stream, "Decode").await?; + + // Mark both streams as completed now that both succeeded + prefill.mark_completed(); + decode_stream.mark_completed(); + + // Merge prefill input_logprobs if requested + if merge_logprobs { + merge_prefill_logprobs(&prefill_responses, &mut decode_responses); + } + + decode_responses + } + }; + + if all_responses.is_empty() { + return Err(error::internal_error("No responses from server")); + } + + Ok(all_responses) +} + +/// Merge prefill input_logprobs into decode responses +/// +/// Takes input_logprobs from the first prefill response and copies them +/// into all decode responses. This is used in PD mode when logprobs are requested. +fn merge_prefill_logprobs( + prefill_responses: &[proto::GenerateComplete], + decode_responses: &mut [proto::GenerateComplete], +) { + if let Some(prefill_input_logprobs) = prefill_responses + .first() + .and_then(|r| r.input_logprobs.clone()) + { + for response in decode_responses.iter_mut() { + response.input_logprobs = Some(prefill_input_logprobs.clone()); + } + } +} diff --git a/sgl-router/src/routers/grpc/common/response_formatting.rs b/sgl-router/src/routers/grpc/common/response_formatting.rs new file mode 100644 index 000000000..8cd66520d --- /dev/null +++ b/sgl-router/src/routers/grpc/common/response_formatting.rs @@ -0,0 +1,65 @@ +//! Shared response formatting logic +//! +//! This module contains common logic for formatting responses, including: +//! - Usage calculation from gRPC responses +//! - ChatCompletionResponse construction + +use crate::{ + grpc_client::proto, + protocols::{ + chat::{ChatChoice, ChatCompletionResponse}, + common::Usage, + }, + routers::grpc::context::DispatchMetadata, +}; + +/// Build usage information from collected gRPC responses +/// +/// Sums prompt_tokens and completion_tokens across all responses. +/// Typically used with n>1 parameter where multiple completions are generated. +/// +/// # Arguments +/// * `responses` - Vector of GenerateComplete responses from the backend +/// +/// # Returns +/// Usage object with aggregated token counts +pub fn build_usage(responses: &[proto::GenerateComplete]) -> Usage { + let total_prompt_tokens: u32 = responses.iter().map(|r| r.prompt_tokens as u32).sum(); + let total_completion_tokens: u32 = responses.iter().map(|r| r.completion_tokens as u32).sum(); + + Usage { + prompt_tokens: total_prompt_tokens, + completion_tokens: total_completion_tokens, + total_tokens: total_prompt_tokens + total_completion_tokens, + completion_tokens_details: None, + } +} + +/// Build final ChatCompletionResponse from processed choices +/// +/// Constructs the OpenAI-compatible response object with all metadata. +/// +/// # Arguments +/// * `choices` - Processed chat choices (after parsing, logprobs, etc.) +/// * `dispatch` - Dispatch metadata (request_id, created timestamp, etc.) +/// * `model` - Model name to include in response +/// * `usage` - Token usage information +/// +/// # Returns +/// Complete ChatCompletionResponse ready to send to client +pub fn build_chat_response( + choices: Vec, + dispatch: &DispatchMetadata, + model: String, + usage: Usage, +) -> ChatCompletionResponse { + ChatCompletionResponse { + id: dispatch.request_id.clone(), + object: "chat.completion".to_string(), + created: dispatch.created, + model, + choices, + usage: Some(usage), + system_fingerprint: dispatch.weight_version.clone(), + } +} diff --git a/sgl-router/src/routers/grpc/common/responses/handlers.rs b/sgl-router/src/routers/grpc/common/responses/handlers.rs new file mode 100644 index 000000000..4b7b8c090 --- /dev/null +++ b/sgl-router/src/routers/grpc/common/responses/handlers.rs @@ -0,0 +1,196 @@ +//! Shared response handlers for both regular and harmony implementations +//! +//! These handlers are used by both pipelines for retrieving and cancelling responses. + +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, +}; +use serde_json::json; +use tracing::{debug, error, warn}; + +use crate::{ + data_connector::ResponseId, routers::grpc::regular::responses::context::ResponsesContext, +}; + +/// Implementation for GET /v1/responses/{response_id} +/// +/// Retrieves a stored response from the database. +/// Used by both regular and harmony implementations. +pub async fn get_response_impl(ctx: &ResponsesContext, response_id: &str) -> Response { + let resp_id = ResponseId::from(response_id); + + // Retrieve response from storage + match ctx.response_storage.get_response(&resp_id).await { + Ok(Some(stored_response)) => axum::Json(stored_response.raw_response).into_response(), + Ok(None) => ( + StatusCode::NOT_FOUND, + axum::Json(json!({ + "error": { + "message": format!("Response with id '{}' not found", response_id), + "type": "not_found_error", + "code": "response_not_found" + } + })), + ) + .into_response(), + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + axum::Json(json!({ + "error": { + "message": format!("Failed to retrieve response: {}", e), + "type": "internal_error" + } + })), + ) + .into_response(), + } +} + +/// Implementation for POST /v1/responses/{response_id}/cancel +/// +/// Cancels a background response if it's still in progress. +pub async fn cancel_response_impl(ctx: &ResponsesContext, response_id: &str) -> Response { + let resp_id = ResponseId::from(response_id); + + // Retrieve response from storage to check if it exists and get current status + match ctx.response_storage.get_response(&resp_id).await { + Ok(Some(stored_response)) => { + // Check current status - only queued or in_progress responses can be cancelled + let current_status = stored_response + .raw_response + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + + match current_status { + "queued" | "in_progress" => { + // Attempt to abort the background task + let mut tasks = ctx.background_tasks.write().await; + if let Some(task_info) = tasks.remove(response_id) { + // Abort the Rust task immediately + task_info.handle.abort(); + + // Abort the Python/scheduler request via gRPC (if client is available) + let client_opt = task_info.client.read().await; + if let Some(ref client) = *client_opt { + if let Err(e) = client + .abort_request( + task_info.grpc_request_id.clone(), + "User cancelled via API".to_string(), + ) + .await + { + warn!( + "Failed to abort Python request {}: {}", + task_info.grpc_request_id, e + ); + } else { + debug!( + "Successfully aborted Python request: {}", + task_info.grpc_request_id + ); + } + } else { + debug!("Client not yet available for abort, request may not have started yet"); + } + + // Task was found and aborted + ( + StatusCode::OK, + axum::Json(json!({ + "id": response_id, + "status": "cancelled", + "message": "Background task has been cancelled" + })), + ) + .into_response() + } else { + // Task handle not found but status is queued/in_progress + // This can happen if: (1) task crashed, or (2) storage persistence failed + error!( + "Response {} has status '{}' but task handle is missing. Task may have crashed or storage update failed.", + response_id, current_status + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + axum::Json(json!({ + "error": { + "message": "Internal error: background task completed but failed to update status in storage", + "type": "internal_error", + "code": "status_update_failed" + } + })), + ) + .into_response() + } + } + "completed" => ( + StatusCode::BAD_REQUEST, + axum::Json(json!({ + "error": { + "message": "Cannot cancel completed response", + "type": "invalid_request_error", + "code": "response_already_completed" + } + })), + ) + .into_response(), + "failed" => ( + StatusCode::BAD_REQUEST, + axum::Json(json!({ + "error": { + "message": "Cannot cancel failed response", + "type": "invalid_request_error", + "code": "response_already_failed" + } + })), + ) + .into_response(), + "cancelled" => ( + StatusCode::OK, + axum::Json(json!({ + "id": response_id, + "status": "cancelled", + "message": "Response was already cancelled" + })), + ) + .into_response(), + _ => { + // Unknown status + ( + StatusCode::INTERNAL_SERVER_ERROR, + axum::Json(json!({ + "error": { + "message": format!("Unknown response status: {}", current_status), + "type": "internal_error" + } + })), + ) + .into_response() + } + } + } + Ok(None) => ( + StatusCode::NOT_FOUND, + axum::Json(json!({ + "error": { + "message": format!("Response with id '{}' not found", response_id), + "type": "not_found_error", + "code": "response_not_found" + } + })), + ) + .into_response(), + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + axum::Json(json!({ + "error": { + "message": format!("Failed to retrieve response: {}", e), + "type": "internal_error" + } + })), + ) + .into_response(), + } +} diff --git a/sgl-router/src/routers/grpc/common/responses/mod.rs b/sgl-router/src/routers/grpc/common/responses/mod.rs new file mode 100644 index 000000000..bb39a2575 --- /dev/null +++ b/sgl-router/src/routers/grpc/common/responses/mod.rs @@ -0,0 +1,7 @@ +//! Shared response functionality used by both regular and harmony implementations + +pub mod handlers; +pub mod streaming; + +pub use handlers::{cancel_response_impl, get_response_impl}; +pub use streaming::{OutputItemType, ResponseStreamEventEmitter}; diff --git a/sgl-router/src/routers/grpc/responses/streaming.rs b/sgl-router/src/routers/grpc/common/responses/streaming.rs similarity index 98% rename from sgl-router/src/routers/grpc/responses/streaming.rs rename to sgl-router/src/routers/grpc/common/responses/streaming.rs index 62062adce..219f5fe46 100644 --- a/sgl-router/src/routers/grpc/responses/streaming.rs +++ b/sgl-router/src/routers/grpc/common/responses/streaming.rs @@ -7,7 +7,7 @@ use serde_json::json; use tokio::sync::mpsc; use uuid::Uuid; -use crate::protocols::chat::ChatCompletionStreamResponse; +use crate::{mcp, protocols::chat::ChatCompletionStreamResponse}; pub enum OutputItemType { Message, @@ -30,10 +30,6 @@ struct OutputItemState { status: ItemStatus, } -// ============================================================================ -// Streaming Event Emitter -// ============================================================================ - /// OpenAI-compatible event emitter for /v1/responses streaming /// /// Manages state and sequence numbers to emit proper event types: @@ -66,7 +62,7 @@ pub struct ResponseStreamEventEmitter { has_emitted_content_part_added: bool, // MCP call tracking mcp_call_accumulated_args: HashMap, - // Output item tracking (NEW) + // Output item tracking output_items: Vec, next_output_index: usize, current_message_output_index: Option, // Tracks output_index of current message @@ -248,7 +244,7 @@ impl ResponseStreamEventEmitter { pub fn emit_mcp_list_tools_completed( &mut self, output_index: usize, - tools: &[crate::mcp::Tool], + tools: &[mcp::Tool], ) -> serde_json::Value { let tool_items: Vec<_> = tools .iter() @@ -331,7 +327,7 @@ impl ResponseStreamEventEmitter { }) } - pub(super) fn emit_mcp_call_failed( + pub fn emit_mcp_call_failed( &mut self, output_index: usize, item_id: &str, @@ -453,7 +449,7 @@ impl ResponseStreamEventEmitter { } /// Process a chunk and emit appropriate events - pub(super) fn process_chunk( + pub fn process_chunk( &mut self, chunk: &ChatCompletionStreamResponse, tx: &mpsc::UnboundedSender>, diff --git a/sgl-router/src/routers/grpc/stages/client_acquisition.rs b/sgl-router/src/routers/grpc/common/stages/client_acquisition.rs similarity index 100% rename from sgl-router/src/routers/grpc/stages/client_acquisition.rs rename to sgl-router/src/routers/grpc/common/stages/client_acquisition.rs diff --git a/sgl-router/src/routers/grpc/stages/dispatch_metadata.rs b/sgl-router/src/routers/grpc/common/stages/dispatch_metadata.rs similarity index 100% rename from sgl-router/src/routers/grpc/stages/dispatch_metadata.rs rename to sgl-router/src/routers/grpc/common/stages/dispatch_metadata.rs diff --git a/sgl-router/src/routers/grpc/common/stages/helpers.rs b/sgl-router/src/routers/grpc/common/stages/helpers.rs new file mode 100644 index 000000000..6fa8788b7 --- /dev/null +++ b/sgl-router/src/routers/grpc/common/stages/helpers.rs @@ -0,0 +1,38 @@ +//! Common helper functions shared across stages + +use std::sync::Arc; + +use proto::DisaggregatedParams; +use rand::Rng; +use tracing::debug; + +use crate::{core::Worker, grpc_client::proto}; + +/// Inject PD bootstrap metadata into a gRPC request +/// +/// Used by both chat and generate request building stages when in PD mode. +pub fn inject_bootstrap_metadata( + 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/mod.rs b/sgl-router/src/routers/grpc/common/stages/mod.rs similarity index 51% rename from sgl-router/src/routers/grpc/stages/mod.rs rename to sgl-router/src/routers/grpc/common/stages/mod.rs index 5c5d86dc9..eafc2261b 100644 --- a/sgl-router/src/routers/grpc/stages/mod.rs +++ b/sgl-router/src/routers/grpc/common/stages/mod.rs @@ -1,17 +1,16 @@ -//! Pipeline stages for gRPC router request processing +//! Common pipeline stages shared across all endpoints and model types //! -//! This module defines the core pipeline abstraction and individual processing stages -//! that transform a RequestContext through its lifecycle. +//! These stages are endpoint-agnostic and model-agnostic: +//! - Worker selection +//! - Client acquisition +//! - Dispatch metadata generation +//! - Request execution 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 { @@ -27,26 +26,14 @@ pub trait PipelineStage: Send + Sync { fn name(&self) -> &'static str; } -// ============================================================================ -// Stage Modules -// ============================================================================ - mod client_acquisition; mod dispatch_metadata; -mod preparation; -mod request_building; +pub mod helpers; mod request_execution; -mod response_processing; mod worker_selection; -// ============================================================================ -// Public Exports -// ============================================================================ - +// Export stage implementations 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/request_execution.rs b/sgl-router/src/routers/grpc/common/stages/request_execution.rs similarity index 100% rename from sgl-router/src/routers/grpc/stages/request_execution.rs rename to sgl-router/src/routers/grpc/common/stages/request_execution.rs diff --git a/sgl-router/src/routers/grpc/stages/worker_selection.rs b/sgl-router/src/routers/grpc/common/stages/worker_selection.rs similarity index 98% rename from sgl-router/src/routers/grpc/stages/worker_selection.rs rename to sgl-router/src/routers/grpc/common/stages/worker_selection.rs index 2362164a9..554ccc47f 100644 --- a/sgl-router/src/routers/grpc/stages/worker_selection.rs +++ b/sgl-router/src/routers/grpc/common/stages/worker_selection.rs @@ -109,7 +109,6 @@ impl WorkerSelectionStage { 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()) diff --git a/sgl-router/src/routers/grpc/context.rs b/sgl-router/src/routers/grpc/context.rs index e6f9c1725..6fc8f308b 100644 --- a/sgl-router/src/routers/grpc/context.rs +++ b/sgl-router/src/routers/grpc/context.rs @@ -11,7 +11,7 @@ use serde_json::Value; use crate::{ core::Worker, - grpc_client::{proto, SglangSchedulerClient}, + grpc_client::{proto, sglang_scheduler::AbortOnDropStream, SglangSchedulerClient}, protocols::{ chat::{ChatCompletionRequest, ChatCompletionResponse}, generate::{GenerateRequest, GenerateResponse}, @@ -22,23 +22,14 @@ use crate::{ tool_parser::ParserFactory as ToolParserFactory, }; -// ============================================================================ -// Core Context Types -// ============================================================================ - /// Main request processing context /// /// This is the single source of truth for all request state as it flows /// through the pipeline stages. Uses Rust's type system to enforce proper /// stage ordering at compile time. pub struct RequestContext { - // === Input (Immutable) === pub input: RequestInput, - - // === Shared Components (Immutable References) === pub components: Arc, - - // === Processing State (Mutable, evolves through pipeline) === pub state: ProcessingState, } @@ -86,10 +77,6 @@ pub struct ProcessingState { pub response: ResponseState, } -// ============================================================================ -// Stage-Specific Output Types -// ============================================================================ - /// Output from preparation stage (Step 1) pub struct PreparationOutput { /// Original text (for chat) or resolved text (for generate) @@ -201,10 +188,6 @@ pub struct StreamingState { pub has_tool_calls: HashMap, } -// ============================================================================ -// Context Builders -// ============================================================================ - impl RequestContext { /// Create context for chat completion request pub fn for_chat( @@ -323,14 +306,6 @@ impl RequestContext { } } -// ============================================================================ -// Default Implementations -// ============================================================================ - -// ============================================================================ -// Helper Methods -// ============================================================================ - impl WorkerSelection { pub fn is_dual(&self) -> bool { matches!(self, Self::Dual { .. }) @@ -428,12 +403,6 @@ impl ClientSelection { } } -// ============================================================================ -// Execution and Response Types -// ============================================================================ - -use crate::grpc_client::sglang_scheduler::AbortOnDropStream; - /// Result of request execution (streams from workers) /// Uses AbortOnDropStream to automatically abort on cancellation pub enum ExecutionResult { diff --git a/sgl-router/src/routers/grpc/harmony/processor.rs b/sgl-router/src/routers/grpc/harmony/processor.rs index c27e1c43f..31678fb99 100644 --- a/sgl-router/src/routers/grpc/harmony/processor.rs +++ b/sgl-router/src/routers/grpc/harmony/processor.rs @@ -17,8 +17,9 @@ use crate::{ }, }, routers::grpc::{ + common::{response_collection, response_formatting}, context::{DispatchMetadata, ExecutionResult}, - error, utils, + error, }, }; @@ -34,28 +35,6 @@ impl HarmonyResponseProcessor { 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, @@ -64,7 +43,7 @@ impl HarmonyResponseProcessor { dispatch: DispatchMetadata, ) -> Result { // Collect all completed responses (one per choice) - let all_responses = Self::collect_responses(execution_result).await?; + let all_responses = response_collection::collect_responses(execution_result, false).await?; if all_responses.is_empty() { return Err(error::internal_error("No responses from server")); } @@ -117,28 +96,15 @@ impl HarmonyResponseProcessor { } // 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, - }; + let usage = response_formatting::build_usage(&all_responses); // Final ChatCompletionResponse - let response = ChatCompletionResponse { - id: dispatch.request_id.clone(), - object: "chat.completion".to_string(), - created: dispatch.created, - model: chat_request.model.clone(), + let response = response_formatting::build_chat_response( choices, - usage: Some(usage), - system_fingerprint: dispatch.weight_version.clone(), - }; + &dispatch, + chat_request.model.clone(), + usage, + ); Ok(response) } @@ -191,7 +157,7 @@ impl HarmonyResponseProcessor { dispatch: DispatchMetadata, ) -> Result { // Collect all completed responses - let all_responses = Self::collect_responses(execution_result).await?; + let all_responses = response_collection::collect_responses(execution_result, false).await?; if all_responses.is_empty() { return Err(error::internal_error("No responses from server")); } @@ -280,14 +246,7 @@ impl HarmonyResponseProcessor { } // 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, - }; + let usage = response_formatting::build_usage(std::slice::from_ref(complete)); // Build ResponsesResponse with all required fields let response = ResponsesResponse { @@ -316,9 +275,9 @@ impl HarmonyResponseProcessor { 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: usage.prompt_tokens, + output_tokens: usage.completion_tokens, + total_tokens: usage.total_tokens, input_tokens_details: None, output_tokens_details: None, })), diff --git a/sgl-router/src/routers/grpc/harmony/responses.rs b/sgl-router/src/routers/grpc/harmony/responses.rs index 6309b9cb2..b0c558b43 100644 --- a/sgl-router/src/routers/grpc/harmony/responses.rs +++ b/sgl-router/src/routers/grpc/harmony/responses.rs @@ -37,12 +37,14 @@ //! for complete architecture, rationale, and implementation details. use std::{ + io, sync::Arc, time::{SystemTime, UNIX_EPOCH}, }; use axum::{body::Body, http::StatusCode, response::Response}; -use serde_json::Value as JsonValue; +use bytes::Bytes; +use serde_json::{from_str, from_value, json, to_string, to_value, Value}; use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::{debug, warn}; @@ -50,21 +52,22 @@ use uuid::Uuid; use crate::{ data_connector::{ResponseId, ResponseStorage}, - mcp::McpManager, + mcp::{self, McpManager}, protocols::{ common::{Function, ToolCall}, responses::{ - ResponseInput, ResponseInputOutputItem, ResponseTool, ResponseToolType, + McpToolInfo, ResponseContentPart, ResponseInput, ResponseInputOutputItem, + ResponseOutputItem, ResponseReasoningContent, ResponseTool, ResponseToolType, ResponsesRequest, ResponsesResponse, StringOrContentParts, }, }, routers::{ grpc::{ + common::responses::streaming::{OutputItemType, ResponseStreamEventEmitter}, context::SharedComponents, error, harmony::processor::ResponsesIterationResult, pipeline::RequestPipeline, - responses::streaming::{OutputItemType, ResponseStreamEventEmitter}, }, openai::mcp::ensure_request_mcp_client, }, @@ -414,10 +417,6 @@ pub async fn serve_harmony_responses_stream( Err(err_response) => return err_response, }; - use std::io; - - use bytes::Bytes; - // Create SSE channel let (tx, rx) = mpsc::unbounded_channel(); let stream = UnboundedReceiverStream::new(rx); @@ -444,7 +443,7 @@ pub async fn serve_harmony_responses_stream( // Helper to emit error and return let emit_error = |tx: &mpsc::UnboundedSender>, error_msg: &str| { // Create error event manually since emit_failed doesn't exist - let event = serde_json::json!({ + let event = json!({ "type": "response.failed", "response_id": response_id_for_error, "error": { @@ -452,7 +451,7 @@ pub async fn serve_harmony_responses_stream( "type": "internal_error" } }); - let sse_data = format!("data: {}\n\n", serde_json::to_string(&event).unwrap()); + let sse_data = format!("data: {}\n\n", to_string(&event).unwrap()); let _ = tx.send(Ok(Bytes::from(sse_data))); }; @@ -517,7 +516,6 @@ pub async fn serve_harmony_responses_stream( let tool_items: Vec<_> = mcp_tools .iter() .map(|t| { - use serde_json::{json, Value}; json!({ "name": t.name, "description": t.description, @@ -527,7 +525,7 @@ pub async fn serve_harmony_responses_stream( .collect(); // Emit output_item.added - let item = serde_json::json!({ + let item = json!({ "id": item_id, "type": "mcp_list_tools", "server_label": "sglang-mcp", @@ -552,7 +550,7 @@ pub async fn serve_harmony_responses_stream( } // Emit output_item.done - let item_done = serde_json::json!({ + let item_done = json!({ "id": item_id, "type": "mcp_list_tools", "server_label": "sglang-mcp", @@ -677,7 +675,7 @@ pub async fn serve_harmony_responses_stream( ); // Emit response.completed with usage - let usage_json = serde_json::json!({ + let usage_json = json!({ "prompt_tokens": usage.prompt_tokens, "completion_tokens": usage.completion_tokens, "total_tokens": usage.total_tokens, @@ -733,7 +731,7 @@ async fn execute_mcp_tools( // 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| { + let args: Value = from_str(args_str).map_err(|e| { error::internal_error(format!( "Invalid tool arguments JSON for tool '{}': {}", tool_call.function.name, e @@ -741,8 +739,7 @@ async fn execute_mcp_tools( })?; // Execute tool via MCP manager - // Convert JsonValue to ToolArgs via Option (MCP manager expects this) - let args_map = if let JsonValue::Object(map) = args { + let args_map = if let Value::Object(map) = args { Some(map) } else { None @@ -763,15 +760,14 @@ async fn execute_mcp_tools( 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"}), - ) + to_value(content) + .unwrap_or_else(|_| json!({"error": "Failed to serialize tool result"})) } else { - serde_json::json!({"result": "success"}) + json!({"result": "success"}) }; let is_error = mcp_result.is_error.unwrap_or(false); - let output_str = serde_json::to_string(&output) + let output_str = to_string(&output) .unwrap_or_else(|_| r#"{"error": "Failed to serialize output"}"#.to_string()); // Record this call in tracking @@ -804,10 +800,10 @@ async fn execute_mcp_tools( ); let error_msg = format!("Tool execution failed: {}", e); - let error_output = serde_json::json!({ + let error_output = json!({ "error": error_msg.clone() }); - let error_output_str = serde_json::to_string(&error_output) + let error_output_str = to_string(&error_output) .unwrap_or_else(|_| format!(r#"{{"error": "{}"}}"#, error_msg)); // Record failed call in tracking @@ -859,12 +855,6 @@ fn build_next_request_with_tools( 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, @@ -926,7 +916,7 @@ fn build_next_request_with_tools( // 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| { + let output_str = to_string(&tool_result.output).unwrap_or_else(|e| { format!("{{\"error\": \"Failed to serialize tool output: {}\"}}", e) }); @@ -967,7 +957,7 @@ struct ToolResult { tool_name: String, /// Tool output (JSON value) - output: JsonValue, + output: Value, /// Whether this is an error result is_error: bool, @@ -985,11 +975,7 @@ struct ToolResult { /// # Returns /// /// Vector of ResponseTool entries in MCP format -pub fn convert_mcp_tools_to_response_tools(mcp_tools: &[crate::mcp::Tool]) -> Vec { - use serde_json::Value; - - use crate::protocols::responses::ResponseToolType; - +pub fn convert_mcp_tools_to_response_tools(mcp_tools: &[mcp::Tool]) -> Vec { mcp_tools .iter() .map(|tool_info| ResponseTool { @@ -1027,11 +1013,6 @@ fn inject_mcp_metadata( 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 @@ -1121,23 +1102,22 @@ async fn load_previous_messages( 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| { - warn!( - "Failed to deserialize stored {} item: {}. Item: {}", - item_type, e, item - ); - }) - .ok() - }) - .collect() - }; + let deserialize_items = |arr: &Value, item_type: &str| -> Vec { + arr.as_array() + .into_iter() + .flat_map(|items| items.iter()) + .filter_map(|item| { + from_value::(item.clone()) + .map_err(|e| { + 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")); diff --git a/sgl-router/src/routers/grpc/harmony/stages/preparation.rs b/sgl-router/src/routers/grpc/harmony/stages/preparation.rs index 2f36c4482..9d00680ea 100644 --- a/sgl-router/src/routers/grpc/harmony/stages/preparation.rs +++ b/sgl-router/src/routers/grpc/harmony/stages/preparation.rs @@ -12,10 +12,9 @@ use crate::{ responses::ResponsesRequest, }, routers::grpc::{ + common::stages::PipelineStage, context::{PreparationOutput, RequestContext, RequestType}, - error, - stages::PipelineStage, - utils, + error, utils, }, }; diff --git a/sgl-router/src/routers/grpc/harmony/stages/request_building.rs b/sgl-router/src/routers/grpc/harmony/stages/request_building.rs index d6011af95..3396607c0 100644 --- a/sgl-router/src/routers/grpc/harmony/stages/request_building.rs +++ b/sgl-router/src/routers/grpc/harmony/stages/request_building.rs @@ -1,21 +1,14 @@ //! 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}, - error, - stages::PipelineStage, - }, +use crate::routers::grpc::{ + common::stages::{helpers, PipelineStage}, + context::{ClientSelection, RequestContext, RequestType, WorkerSelection}, + error, }; /// Harmony Request Building stage: Convert Harmony tokens to gRPC request @@ -31,34 +24,6 @@ impl HarmonyRequestBuildingStage { 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] @@ -94,7 +59,6 @@ impl PipelineStage for HarmonyRequestBuildingStage { }; // 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 { @@ -141,7 +105,7 @@ impl PipelineStage for HarmonyRequestBuildingStage { // 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); + helpers::inject_bootstrap_metadata(&mut proto_request, prefill); } } diff --git a/sgl-router/src/routers/grpc/harmony/stages/response_processing.rs b/sgl-router/src/routers/grpc/harmony/stages/response_processing.rs index 55f855be1..bc640efb5 100644 --- a/sgl-router/src/routers/grpc/harmony/stages/response_processing.rs +++ b/sgl-router/src/routers/grpc/harmony/stages/response_processing.rs @@ -7,9 +7,9 @@ use axum::response::Response; use super::super::{HarmonyResponseProcessor, HarmonyStreamingProcessor}; use crate::routers::grpc::{ + common::stages::PipelineStage, context::{FinalResponse, RequestContext, RequestType}, error, - stages::PipelineStage, }; /// Harmony Response Processing stage: Parse and format Harmony responses diff --git a/sgl-router/src/routers/grpc/harmony/streaming.rs b/sgl-router/src/routers/grpc/harmony/streaming.rs index ba174a549..25210d40a 100644 --- a/sgl-router/src/routers/grpc/harmony/streaming.rs +++ b/sgl-router/src/routers/grpc/harmony/streaming.rs @@ -31,8 +31,8 @@ use crate::{ responses::{ResponseStatus, ResponseUsage, ResponsesResponse, ResponsesUsage}, }, routers::grpc::{ + common::responses::streaming::{OutputItemType, ResponseStreamEventEmitter}, context, - responses::streaming::{OutputItemType, ResponseStreamEventEmitter}, }, }; /// Processor for streaming Harmony responses diff --git a/sgl-router/src/routers/grpc/harmony/types.rs b/sgl-router/src/routers/grpc/harmony/types.rs index 8ce2a3602..bfeeac475 100644 --- a/sgl-router/src/routers/grpc/harmony/types.rs +++ b/sgl-router/src/routers/grpc/harmony/types.rs @@ -1,5 +1,6 @@ //! Shared types for Harmony pipeline +use openai_harmony::chat::Content; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -36,8 +37,6 @@ impl HarmonyMessage { /// 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", diff --git a/sgl-router/src/routers/grpc/mod.rs b/sgl-router/src/routers/grpc/mod.rs index 57384c683..213d5399d 100644 --- a/sgl-router/src/routers/grpc/mod.rs +++ b/sgl-router/src/routers/grpc/mod.rs @@ -2,16 +2,14 @@ use crate::{grpc_client::proto, protocols::common::StringOrArray}; +pub mod common; pub mod context; pub mod error; pub mod harmony; pub mod pd_router; pub mod pipeline; -pub mod processing; -pub mod responses; +pub mod regular; pub mod router; -pub mod stages; -pub mod streaming; pub mod utils; /// Processed chat messages ready for gRPC generation diff --git a/sgl-router/src/routers/grpc/pd_router.rs b/sgl-router/src/routers/grpc/pd_router.rs index 55cb653e5..953804c01 100644 --- a/sgl-router/src/routers/grpc/pd_router.rs +++ b/sgl-router/src/routers/grpc/pd_router.rs @@ -1,5 +1,3 @@ -// PD (Prefill-Decode) gRPC Router Implementation - use std::sync::Arc; use async_trait::async_trait; @@ -161,7 +159,6 @@ impl RouterTrait for GrpcPDRouter { } async fn health_generate(&self, _req: Request) -> Response { - // TODO: Implement actual generation test for gRPC PD mode ( StatusCode::NOT_IMPLEMENTED, "Health generate not yet implemented for gRPC PD", diff --git a/sgl-router/src/routers/grpc/pipeline.rs b/sgl-router/src/routers/grpc/pipeline.rs index 475233c0d..a05d6b81a 100644 --- a/sgl-router/src/routers/grpc/pipeline.rs +++ b/sgl-router/src/routers/grpc/pipeline.rs @@ -3,15 +3,17 @@ //! This module defines the RequestPipeline orchestrator that coordinates //! the execution of pipeline stages from request preparation to response delivery. -use std::{collections::HashMap, sync::Arc}; +use std::sync::Arc; use axum::response::{IntoResponse, Response}; -use tokio::sync::RwLock; -use tracing::{debug, error}; +use tracing::error; -// Import all stage types from the stages module -use super::stages::*; -use super::{context::*, error, harmony, processing, responses::BackgroundTaskInfo, streaming}; +use super::{ + common::stages::*, + context::*, + error, harmony, + regular::{processor, stages::*, streaming}, +}; use crate::{ core::WorkerRegistry, policies::PolicyRegistry, @@ -24,10 +26,6 @@ use crate::{ tool_parser::ParserFactory as ToolParserFactory, }; -// ============================================================================ -// Pipeline Orchestrator -// ============================================================================ - /// Generic request pipeline for all request types /// /// Orchestrates all stages from request preparation to response delivery. @@ -48,8 +46,7 @@ impl RequestPipeline { configured_tool_parser: Option, configured_reasoning_parser: Option, ) -> Self { - // Create response processor - let processor = processing::ResponseProcessor::new( + let processor = processor::ResponseProcessor::new( tokenizer.clone(), tool_parser_factory.clone(), reasoning_parser_factory.clone(), @@ -57,7 +54,6 @@ impl RequestPipeline { configured_reasoning_parser.clone(), ); - // Create streaming processor let streaming_processor = Arc::new(streaming::StreamingProcessor::new( tokenizer, tool_parser_factory, @@ -67,7 +63,7 @@ impl RequestPipeline { )); let stages: Vec> = vec![ - Box::new(PreparationStage), + Box::new(PreparationStage::new()), Box::new(WorkerSelectionStage::new( worker_registry, policy_registry, @@ -153,8 +149,7 @@ impl RequestPipeline { configured_tool_parser: Option, configured_reasoning_parser: Option, ) -> Self { - // Create response processor - let processor = processing::ResponseProcessor::new( + let processor = processor::ResponseProcessor::new( tokenizer.clone(), tool_parser_factory.clone(), reasoning_parser_factory.clone(), @@ -162,7 +157,6 @@ impl RequestPipeline { configured_reasoning_parser.clone(), ); - // Create streaming processor let streaming_processor = Arc::new(streaming::StreamingProcessor::new( tokenizer, tool_parser_factory, @@ -172,7 +166,7 @@ impl RequestPipeline { )); let stages: Vec> = vec![ - Box::new(PreparationStage), + Box::new(PreparationStage::new()), Box::new(WorkerSelectionStage::new( worker_registry, policy_registry, @@ -200,7 +194,6 @@ impl RequestPipeline { ) -> Response { let mut ctx = RequestContext::for_chat(request, headers, model_id, components); - // Execute each stage in sequence for (idx, stage) in self.stages.iter().enumerate() { match stage.execute(&mut ctx).await { Ok(Some(response)) => { @@ -208,7 +201,6 @@ impl RequestPipeline { return response; } Ok(None) => { - // Continue to next stage continue; } Err(response) => { @@ -224,7 +216,6 @@ impl RequestPipeline { } } - // Extract final response match ctx.state.response.final_response { Some(FinalResponse::Chat(response)) => axum::Json(response).into_response(), Some(FinalResponse::Generate(_)) => { @@ -244,7 +235,6 @@ impl RequestPipeline { ) -> Response { let mut ctx = RequestContext::for_generate(request, headers, model_id, components); - // Execute each stage in sequence for (idx, stage) in self.stages.iter().enumerate() { match stage.execute(&mut ctx).await { Ok(Some(response)) => { @@ -252,7 +242,6 @@ impl RequestPipeline { return response; } Ok(None) => { - // Continue to next stage continue; } Err(response) => { @@ -268,7 +257,6 @@ impl RequestPipeline { } } - // Extract final response match ctx.state.response.final_response { Some(FinalResponse::Generate(response)) => axum::Json(response).into_response(), Some(FinalResponse::Chat(_)) => { @@ -280,25 +268,19 @@ impl RequestPipeline { /// Execute chat pipeline for responses endpoint /// - /// TODO: The support for background tasks is not scalable. Consider replacing this with - /// a better design in the future. - /// Used by ALL non-streaming /v1/responses requests (both sync and background modes). - /// Uses the same 7 pipeline stages as execute_chat(), with three differences: + /// Used by ALL non-streaming /v1/responses requests. + /// Uses the same 7 pipeline stages as execute_chat(), with two differences: /// 1. Returns Result for tool_loop composition /// 2. Disallows streaming (responses endpoint uses different SSE format) - /// 3. Injects hooks for background task cancellation (only active when response_id provided) pub async fn execute_chat_for_responses( &self, request: Arc, headers: Option, model_id: Option, components: Arc, - response_id: Option, - background_tasks: Option>>>, ) -> Result { let mut ctx = RequestContext::for_chat(request, headers, model_id, components); - // Execute each stage in sequence for (idx, stage) in self.stages.iter().enumerate() { match stage.execute(&mut ctx).await { Ok(Some(_response)) => { @@ -308,40 +290,6 @@ impl RequestPipeline { )); } Ok(None) => { - let stage_name = stage.name(); - - // After ClientAcquisitionStage, store client for background task cancellation - if stage_name == "ClientAcquisition" { - if let (Some(ref clients), Some(ref resp_id), Some(ref tasks)) = - (&ctx.state.clients, &response_id, &background_tasks) - { - let client_to_store = match clients { - ClientSelection::Single { client } => client.clone(), - ClientSelection::Dual { decode, .. } => decode.clone(), - }; - - if let Some(task_info) = tasks.write().await.get_mut(resp_id.as_str()) { - *task_info.client.write().await = Some(client_to_store); - debug!("Stored client for response_id: {}", resp_id); - } - } - } - - // After DispatchMetadataStage, store grpc_request_id for background task cancellation - if stage_name == "DispatchMetadata" { - if let (Some(ref dispatch), Some(ref resp_id), Some(ref tasks)) = - (&ctx.state.dispatch, &response_id, &background_tasks) - { - let grpc_request_id = dispatch.request_id.clone(); - - if let Some(task_info) = tasks.write().await.get_mut(resp_id.as_str()) { - task_info.grpc_request_id = grpc_request_id.clone(); - debug!("Stored grpc_request_id for response_id: {}", resp_id); - } - } - } - - // Continue to next stage continue; } Err(response) => { @@ -357,7 +305,6 @@ impl RequestPipeline { } } - // Extract final response match ctx.state.response.final_response { Some(FinalResponse::Chat(response)) => Ok(response), Some(FinalResponse::Generate(_)) => { @@ -367,26 +314,6 @@ impl RequestPipeline { } } - /// 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 { - error::internal_error("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, @@ -415,7 +342,6 @@ impl RequestPipeline { 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)) => { @@ -428,7 +354,6 @@ impl RequestPipeline { return Err(response); } Ok(None) => { - // Stage completed successfully, continue to next stage continue; } Err(response) => { @@ -472,7 +397,6 @@ impl RequestPipeline { harmony_ctx.components.clone(), ); - // Execute pipeline stages up to dispatch (which creates the stream) for (idx, stage) in self.stages.iter().enumerate() { match stage.execute(&mut ctx).await { Ok(Some(response)) => { diff --git a/sgl-router/src/routers/grpc/regular/mod.rs b/sgl-router/src/routers/grpc/regular/mod.rs new file mode 100644 index 000000000..2303a2348 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/mod.rs @@ -0,0 +1,9 @@ +//! Regular (non-harmony) model processing +//! +//! This module contains all code specific to regular tokenizer-based models, +//! including pipeline stages, response processing, and streaming. + +pub mod processor; +pub mod responses; +pub mod stages; +pub mod streaming; diff --git a/sgl-router/src/routers/grpc/processing.rs b/sgl-router/src/routers/grpc/regular/processor.rs similarity index 81% rename from sgl-router/src/routers/grpc/processing.rs rename to sgl-router/src/routers/grpc/regular/processor.rs index f7aff5d6f..59e95084a 100644 --- a/sgl-router/src/routers/grpc/processing.rs +++ b/sgl-router/src/routers/grpc/regular/processor.rs @@ -1,7 +1,7 @@ //! Shared response processing logic for gRPC routers //! //! This module contains response processing functions that are shared between -//! the regular router and PD router, eliminating ~1,200 lines of exact duplicates. +//! the regular router and PD router. use std::{sync::Arc, time::Instant}; @@ -9,18 +9,19 @@ use proto::generate_complete::MatchedStop; use serde_json::Value; use tracing::error; -use super::{ - context::{DispatchMetadata, ExecutionResult}, - error, utils, -}; use crate::{ grpc_client::proto, protocols::{ chat::{ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse}, - common::{FunctionCallResponse, ToolCall, ToolChoice, ToolChoiceValue, Usage}, + common::{FunctionCallResponse, ToolCall, ToolChoice, ToolChoiceValue}, generate::{GenerateMetaInfo, GenerateRequest, GenerateResponse}, }, reasoning_parser::ParserFactory as ReasoningParserFactory, + routers::grpc::{ + common::{response_collection, response_formatting}, + context::{DispatchMetadata, ExecutionResult}, + error, utils, + }, tokenizer::{ stop::{SequenceDecoderOutput, StopSequenceDecoder}, traits::Tokenizer, @@ -28,10 +29,6 @@ use crate::{ tool_parser::ParserFactory as ToolParserFactory, }; -// ============================================================================ -// Response Processor - Main Entry Point -// ============================================================================ - /// Unified response processor for both routers #[derive(Clone)] pub struct ResponseProcessor { @@ -59,57 +56,6 @@ impl ResponseProcessor { } } - /// Helper to collect responses from execution result and merge logprobs if needed - async fn collect_and_merge_responses( - execution_result: ExecutionResult, - request_logprobs: bool, - ) -> Result, axum::response::Response> { - let all_responses = match execution_result { - ExecutionResult::Single { mut stream } => { - let responses = utils::collect_stream_responses(&mut stream, "Single").await?; - stream.mark_completed(); - responses - } - ExecutionResult::Dual { - mut prefill, - decode, - } => { - // Collect prefill for input_logprobs (don't mark completed yet) - let prefill_responses = - utils::collect_stream_responses(&mut prefill, "Prefill").await?; - - // Collect decode for actual output (don't mark completed yet) - let mut decode_stream = *decode; - let mut decode_responses = - utils::collect_stream_responses(&mut decode_stream, "Decode").await?; - - // Mark both streams as completed now that both succeeded - prefill.mark_completed(); - decode_stream.mark_completed(); - - // Merge prefill input_logprobs if requested - if request_logprobs { - if let Some(prefill_input_logprobs) = prefill_responses - .first() - .and_then(|r| r.input_logprobs.clone()) - { - for response in &mut decode_responses { - response.input_logprobs = Some(prefill_input_logprobs.clone()); - } - } - } - - decode_responses - } - }; - - if all_responses.is_empty() { - return Err(error::internal_error("No responses from server")); - } - - Ok(all_responses) - } - /// Process a single choice from GenerateComplete response #[allow(clippy::too_many_arguments)] pub async fn process_single_choice( @@ -151,7 +97,6 @@ impl ResponseProcessor { let mut reasoning_text: Option = None; let mut processed_text = final_text; - // Check if reasoning parsing is enabled and parser is available if original_request.separate_reasoning && reasoning_parser_available { let pooled_parser = utils::get_reasoning_parser( &self.reasoning_parser_factory, @@ -275,7 +220,7 @@ impl ResponseProcessor { ) -> Result { // Collect all responses from the execution result let all_responses = - Self::collect_and_merge_responses(execution_result, request_logprobs).await?; + response_collection::collect_responses(execution_result, request_logprobs).await?; let history_tool_calls_count = utils::get_history_tool_calls_count(&chat_request); @@ -341,28 +286,15 @@ impl ResponseProcessor { } // Build usage - let total_prompt_tokens: u32 = all_responses.iter().map(|r| r.prompt_tokens as u32).sum(); - let total_completion_tokens: u32 = all_responses - .iter() - .map(|r| r.completion_tokens as u32) - .sum(); - let usage = Usage { - prompt_tokens: total_prompt_tokens, - completion_tokens: total_completion_tokens, - total_tokens: total_prompt_tokens + total_completion_tokens, - completion_tokens_details: None, - }; + let usage = response_formatting::build_usage(&all_responses); // Build final ChatCompletionResponse - let response = ChatCompletionResponse { - id: dispatch.request_id.clone(), - object: "chat.completion".to_string(), - created: dispatch.created, - model: dispatch.model.clone(), + let response = response_formatting::build_chat_response( choices, - usage: Some(usage), - system_fingerprint: dispatch.weight_version.clone(), - }; + &dispatch, + dispatch.model.clone(), + usage, + ); Ok(response) } @@ -436,7 +368,7 @@ impl ResponseProcessor { ) -> Result, axum::response::Response> { // Collect all responses from the execution result let all_responses = - Self::collect_and_merge_responses(execution_result, request_logprobs).await?; + response_collection::collect_responses(execution_result, request_logprobs).await?; // Process each completion let mut result_array = Vec::new(); @@ -474,7 +406,7 @@ impl ResponseProcessor { } let output_ids = std::mem::take(&mut complete.output_ids); - let finish_reason_str = std::mem::take(&mut complete.finish_reason); + let finish_reason_str = complete.finish_reason.to_string(); // Parse finish_reason from string to proper type let finish_reason = diff --git a/sgl-router/src/routers/grpc/responses/context.rs b/sgl-router/src/routers/grpc/regular/responses/context.rs similarity index 100% rename from sgl-router/src/routers/grpc/responses/context.rs rename to sgl-router/src/routers/grpc/regular/responses/context.rs diff --git a/sgl-router/src/routers/grpc/responses/conversions.rs b/sgl-router/src/routers/grpc/regular/responses/conversions.rs similarity index 100% rename from sgl-router/src/routers/grpc/responses/conversions.rs rename to sgl-router/src/routers/grpc/regular/responses/conversions.rs diff --git a/sgl-router/src/routers/grpc/responses/handlers.rs b/sgl-router/src/routers/grpc/regular/responses/handlers.rs similarity index 69% rename from sgl-router/src/routers/grpc/responses/handlers.rs rename to sgl-router/src/routers/grpc/regular/responses/handlers.rs index 47e86f8b6..be299e4f5 100644 --- a/sgl-router/src/routers/grpc/responses/handlers.rs +++ b/sgl-router/src/routers/grpc/regular/responses/handlers.rs @@ -9,18 +9,19 @@ //! # Architecture //! //! This module orchestrates all request handling for the /v1/responses endpoint. -//! It supports three execution modes: +//! It supports two execution modes: //! //! 1. **Synchronous** - Returns complete response immediately -//! 2. **Background** - Returns queued response, executes in background task -//! 3. **Streaming** - Returns SSE stream with real-time events +//! 2. **Streaming** - Returns SSE stream with real-time events +//! +//! Note: Background mode is no longer supported. Requests with background=true +//! will be rejected with a 400 error. //! //! # Request Flow //! //! ```text //! route_responses() //! ├─► route_responses_sync() → route_responses_internal() -//! ├─► route_responses_background() → spawn(route_responses_internal()) //! └─► route_responses_streaming() → convert_chat_stream_to_responses_stream() //! //! route_responses_internal() @@ -31,10 +32,7 @@ //! └─► pipeline.execute_chat_for_responses() //! ``` -use std::{ - sync::Arc, - time::{SystemTime, UNIX_EPOCH}, -}; +use std::sync::Arc; use axum::{ body::Body, @@ -44,39 +42,36 @@ use axum::{ use bytes::Bytes; use futures_util::StreamExt; use serde_json::json; -use tokio::sync::{mpsc, RwLock}; +use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; -use tracing::{debug, error, warn}; +use tracing::{debug, warn}; use uuid::Uuid; use validator::Validate; use super::{ conversions, - streaming::ResponseStreamEventEmitter, tool_loop::{execute_tool_loop, execute_tool_loop_streaming}, - types::BackgroundTaskInfo, }; use crate::{ data_connector::{ - ConversationId, ConversationItemStorage, ConversationStorage, ResponseId, ResponseStorage, + self, ConversationId, ConversationItemStorage, ConversationStorage, ResponseId, + ResponseStorage, }, protocols::{ - chat::ChatCompletionStreamResponse, + chat::{self, ChatCompletionStreamResponse}, + common, responses::{ - ResponseContentPart, ResponseInput, ResponseInputOutputItem, ResponseOutputItem, - ResponseStatus, ResponsesRequest, ResponsesResponse, ResponsesUsage, + self, ResponseContentPart, ResponseInput, ResponseInputOutputItem, ResponseOutputItem, + ResponseReasoningContent, ResponseStatus, ResponsesRequest, ResponsesResponse, + ResponsesUsage, }, }, routers::{ - grpc::error, + grpc::{common::responses::streaming::ResponseStreamEventEmitter, error}, openai::{conversations::persist_conversation_items, mcp::ensure_request_mcp_client}, }, }; -// ============================================================================ -// Main Request Handler -// ============================================================================ - /// Main handler for POST /v1/responses /// /// Validates request, determines execution mode (sync/async/streaming), and delegates @@ -154,19 +149,17 @@ pub async fn route_responses( .into_response(); } - // 3. Check for incompatible parameter combinations - let is_streaming = request.stream.unwrap_or(false); + // 3. Reject background mode (no longer supported) let is_background = request.background.unwrap_or(false); - - if is_streaming && is_background { + if is_background { return ( StatusCode::BAD_REQUEST, axum::Json(json!({ "error": { - "message": "Cannot use streaming with background mode. Please set either 'stream' or 'background' to false.", + "message": "Background mode is not supported. Please set 'background' to false or omit it.", "type": "invalid_request_error", - "param": serde_json::Value::Null, - "code": "incompatible_parameters" + "param": "background", + "code": "unsupported_parameter" } })), ) @@ -174,10 +167,9 @@ pub async fn route_responses( } // 4. Route based on execution mode + let is_streaming = request.stream.unwrap_or(false); if is_streaming { route_responses_streaming(ctx, request, headers, model_id).await - } else if is_background { - route_responses_background(ctx, request, headers, model_id).await } else { route_responses_sync(ctx, request, headers, model_id, None).await } @@ -284,134 +276,7 @@ async fn route_responses_internal( Ok(responses_response) } -// ============================================================================ -// Background Mode Execution -// ============================================================================ - -/// Execute responses request in background mode -#[allow(clippy::too_many_arguments)] -async fn route_responses_background( - ctx: &super::context::ResponsesContext, - request: Arc, - headers: Option, - model_id: Option, -) -> Response { - // Generate response_id for background tracking - let response_id = format!("resp_{}", Uuid::new_v4()); - - // Get current timestamp - let created_at = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() as i64; - - // Create queued response - let queued_response = ResponsesResponse { - id: response_id.clone(), - object: "response".to_string(), - created_at, - status: ResponseStatus::Queued, - error: None, - incomplete_details: None, - instructions: request.instructions.clone(), - max_output_tokens: request.max_output_tokens, - model: request.model.clone(), - output: Vec::new(), - parallel_tool_calls: request.parallel_tool_calls.unwrap_or(true), - previous_response_id: request.previous_response_id.clone(), - reasoning: None, - store: request.store.unwrap_or(true), - temperature: request.temperature, - text: None, - tool_choice: "auto".to_string(), - tools: request.tools.clone().unwrap_or_default(), - top_p: request.top_p, - truncation: None, - usage: None, - user: None, - safety_identifier: request.user.clone(), - metadata: request.metadata.clone().unwrap_or_default(), - }; - - // Persist queued response to storage - if let Ok(response_json) = serde_json::to_value(&queued_response) { - if let Err(e) = persist_conversation_items( - ctx.conversation_storage.clone(), - ctx.conversation_item_storage.clone(), - ctx.response_storage.clone(), - &response_json, - &request, - ) - .await - { - warn!("Failed to persist queued response: {}", e); - } - } - - // Spawn background task - let ctx_clone = ctx.clone(); - let request_clone = request.clone(); - let headers_clone = headers.clone(); - let model_id_clone = model_id.clone(); - let response_id_clone = response_id.clone(); - - let handle = tokio::task::spawn(async move { - // Execute synchronously (set background=false to prevent recursion) - let mut background_request = (*request_clone).clone(); - background_request.background = Some(false); - - match route_responses_internal( - &ctx_clone, - Arc::new(background_request), - headers_clone, - model_id_clone, - Some(response_id_clone.clone()), - ) - .await - { - Ok(_) => { - debug!( - "Background response {} completed successfully", - response_id_clone - ); - } - Err(response) => { - warn!( - "Background response {} failed with status {}", - response_id_clone, - response.status() - ); - } - } - - // Clean up task handle when done - ctx_clone - .background_tasks - .write() - .await - .remove(&response_id_clone); - }); - - // Store task info for cancellation support - ctx.background_tasks.write().await.insert( - response_id.clone(), - BackgroundTaskInfo { - handle, - grpc_request_id: String::new(), // Will be populated by pipeline at DispatchMetadataStage - client: Arc::new(RwLock::new(None)), - }, - ); - - // Return queued response immediately - axum::Json(queued_response).into_response() -} - -// ============================================================================ -// Streaming Mode Execution -// ============================================================================ - /// Execute streaming responses request -#[allow(clippy::too_many_arguments)] async fn route_responses_streaming( ctx: &super::context::ResponsesContext, request: Arc, @@ -467,10 +332,9 @@ async fn route_responses_streaming( /// 3. Converts ChatCompletionStreamResponse → ResponsesResponse delta /// 4. Accumulates response state for final persistence /// 5. Emits transformed SSE events in responses format -#[allow(clippy::too_many_arguments)] async fn convert_chat_stream_to_responses_stream( ctx: &super::context::ResponsesContext, - chat_request: Arc, + chat_request: Arc, headers: Option, model_id: Option, original_request: &ResponsesRequest, @@ -557,7 +421,7 @@ async fn convert_chat_stream_to_responses_stream( async fn process_and_transform_sse_stream( body: Body, original_request: ResponsesRequest, - _chat_request: Arc, + _chat_request: Arc, response_storage: Arc, conversation_storage: Arc, conversation_item_storage: Arc, @@ -673,7 +537,7 @@ struct StreamingResponseAccumulator { // Completion state finish_reason: Option, - usage: Option, + usage: Option, // Original request for final response construction original_request: ResponsesRequest, @@ -789,11 +653,9 @@ impl StreamingResponseAccumulator { output.push(ResponseOutputItem::Reasoning { id: format!("reasoning_{}", self.response_id), summary: vec![], - content: vec![ - crate::protocols::responses::ResponseReasoningContent::ReasoningText { - text: self.reasoning_buffer, - }, - ], + content: vec![ResponseReasoningContent::ReasoningText { + text: self.reasoning_buffer, + }], status: Some("completed".to_string()), }); } @@ -811,7 +673,7 @@ impl StreamingResponseAccumulator { // Convert usage let usage = self.usage.as_ref().map(|u| { - let usage_info = crate::protocols::common::UsageInfo { + let usage_info = common::UsageInfo { prompt_tokens: u.prompt_tokens, completion_tokens: u.completion_tokens, total_tokens: u.total_tokens, @@ -878,8 +740,6 @@ async fn execute_without_mcp( headers, model_id, ctx.components.clone(), - response_id.clone(), - Some(ctx.background_tasks.clone()), ) .await?; // Preserve the Response error as-is @@ -973,9 +833,9 @@ async fn load_conversation_history( // Load conversation history const MAX_CONVERSATION_HISTORY_ITEMS: usize = 100; - let params = crate::data_connector::ListParams { + let params = data_connector::ListParams { limit: MAX_CONVERSATION_HISTORY_ITEMS, - order: crate::data_connector::SortOrder::Asc, + order: data_connector::SortOrder::Asc, after: None, }; @@ -1014,8 +874,7 @@ async fn load_conversation_history( ResponseInput::Items(current_items) => { // Process all item types, converting SimpleInputMessage to Message for item in current_items.iter() { - let normalized = - crate::protocols::responses::normalize_input_item(item); + let normalized = responses::normalize_input_item(item); items.push(normalized); } } @@ -1050,7 +909,7 @@ async fn load_conversation_history( ResponseInput::Items(current_items) => { // Process all item types, converting SimpleInputMessage to Message for item in current_items.iter() { - let normalized = crate::protocols::responses::normalize_input_item(item); + let normalized = responses::normalize_input_item(item); items.push(normalized); } } @@ -1061,194 +920,3 @@ async fn load_conversation_history( Ok(modified_request) } - -// ============================================================================ -// GET Response Implementation -// ============================================================================ - -/// Implementation for GET /v1/responses/{response_id} -pub async fn get_response_impl( - ctx: &super::context::ResponsesContext, - response_id: &str, -) -> Response { - let resp_id = ResponseId::from(response_id); - - // Retrieve response from storage - match ctx.response_storage.get_response(&resp_id).await { - Ok(Some(stored_response)) => axum::Json(stored_response.raw_response).into_response(), - Ok(None) => ( - StatusCode::NOT_FOUND, - axum::Json(json!({ - "error": { - "message": format!("Response with id '{}' not found", response_id), - "type": "not_found_error", - "code": "response_not_found" - } - })), - ) - .into_response(), - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - axum::Json(json!({ - "error": { - "message": format!("Failed to retrieve response: {}", e), - "type": "internal_error" - } - })), - ) - .into_response(), - } -} - -// ============================================================================ -// CANCEL Response Implementation -// ============================================================================ - -/// Implementation for POST /v1/responses/{response_id}/cancel -pub async fn cancel_response_impl( - ctx: &super::context::ResponsesContext, - response_id: &str, -) -> Response { - let resp_id = ResponseId::from(response_id); - - // Retrieve response from storage to check if it exists and get current status - match ctx.response_storage.get_response(&resp_id).await { - Ok(Some(stored_response)) => { - // Check current status - only queued or in_progress responses can be cancelled - let current_status = stored_response - .raw_response - .get("status") - .and_then(|v| v.as_str()) - .unwrap_or("unknown"); - - match current_status { - "queued" | "in_progress" => { - // Attempt to abort the background task - let mut tasks = ctx.background_tasks.write().await; - if let Some(task_info) = tasks.remove(response_id) { - // Abort the Rust task immediately - task_info.handle.abort(); - - // Abort the Python/scheduler request via gRPC (if client is available) - let client_opt = task_info.client.read().await; - if let Some(ref client) = *client_opt { - if let Err(e) = client - .abort_request( - task_info.grpc_request_id.clone(), - "User cancelled via API".to_string(), - ) - .await - { - warn!( - "Failed to abort Python request {}: {}", - task_info.grpc_request_id, e - ); - } else { - debug!( - "Successfully aborted Python request: {}", - task_info.grpc_request_id - ); - } - } else { - debug!("Client not yet available for abort, request may not have started yet"); - } - - // Task was found and aborted - ( - StatusCode::OK, - axum::Json(json!({ - "id": response_id, - "status": "cancelled", - "message": "Background task has been cancelled" - })), - ) - .into_response() - } else { - // Task handle not found but status is queued/in_progress - // This can happen if: (1) task crashed, or (2) storage persistence failed - error!( - "Response {} has status '{}' but task handle is missing. Task may have crashed or storage update failed.", - response_id, current_status - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - axum::Json(json!({ - "error": { - "message": "Internal error: background task completed but failed to update status in storage", - "type": "internal_error", - "code": "status_update_failed" - } - })), - ) - .into_response() - } - } - "completed" => ( - StatusCode::BAD_REQUEST, - axum::Json(json!({ - "error": { - "message": "Cannot cancel completed response", - "type": "invalid_request_error", - "code": "response_already_completed" - } - })), - ) - .into_response(), - "failed" => ( - StatusCode::BAD_REQUEST, - axum::Json(json!({ - "error": { - "message": "Cannot cancel failed response", - "type": "invalid_request_error", - "code": "response_already_failed" - } - })), - ) - .into_response(), - "cancelled" => ( - StatusCode::OK, - axum::Json(json!({ - "id": response_id, - "status": "cancelled", - "message": "Response was already cancelled" - })), - ) - .into_response(), - _ => { - // Unknown status - ( - StatusCode::INTERNAL_SERVER_ERROR, - axum::Json(json!({ - "error": { - "message": format!("Unknown response status: {}", current_status), - "type": "internal_error" - } - })), - ) - .into_response() - } - } - } - Ok(None) => ( - StatusCode::NOT_FOUND, - axum::Json(json!({ - "error": { - "message": format!("Response with id '{}' not found", response_id), - "type": "not_found_error", - "code": "response_not_found" - } - })), - ) - .into_response(), - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - axum::Json(json!({ - "error": { - "message": format!("Failed to retrieve response: {}", e), - "type": "internal_error" - } - })), - ) - .into_response(), - } -} diff --git a/sgl-router/src/routers/grpc/responses/mod.rs b/sgl-router/src/routers/grpc/regular/responses/mod.rs similarity index 58% rename from sgl-router/src/routers/grpc/responses/mod.rs rename to sgl-router/src/routers/grpc/regular/responses/mod.rs index a73a04d9b..ce67c0191 100644 --- a/sgl-router/src/routers/grpc/responses/mod.rs +++ b/sgl-router/src/routers/grpc/regular/responses/mod.rs @@ -1,9 +1,8 @@ -//! gRPC Router `/v1/responses` endpoint implementation +//! Regular gRPC Router `/v1/responses` endpoint implementation //! -//! This module handles all responses-specific logic including: +//! This module handles all responses-specific logic for the regular (non-Harmony) pipeline including: //! - Request validation //! - Conversation history and response chain loading -//! - Background mode execution //! - Streaming support //! - MCP tool loop wrapper //! - Response persistence @@ -12,11 +11,10 @@ pub mod context; mod conversions; mod handlers; -pub mod streaming; pub mod tool_loop; pub mod types; // Public exports pub use context::ResponsesContext; -pub use handlers::{cancel_response_impl, get_response_impl, route_responses}; +pub use handlers::route_responses; pub use types::BackgroundTaskInfo; diff --git a/sgl-router/src/routers/grpc/responses/tool_loop.rs b/sgl-router/src/routers/grpc/regular/responses/tool_loop.rs similarity index 95% rename from sgl-router/src/routers/grpc/responses/tool_loop.rs rename to sgl-router/src/routers/grpc/regular/responses/tool_loop.rs index 1b2b9e38b..9300dc962 100644 --- a/sgl-router/src/routers/grpc/responses/tool_loop.rs +++ b/sgl-router/src/routers/grpc/regular/responses/tool_loop.rs @@ -12,23 +12,30 @@ use axum::{ response::Response, }; use bytes::Bytes; +use futures_util::StreamExt; use serde_json::json; use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::{debug, warn}; use uuid::Uuid; -use super::{ - super::error, - conversions, - streaming::{OutputItemType, ResponseStreamEventEmitter}, -}; -use crate::protocols::{ - chat::ChatCompletionResponse, - common::{Tool, ToolChoice, ToolChoiceValue}, - responses::{ - McpToolInfo, ResponseContentPart, ResponseInput, ResponseInputOutputItem, - ResponseOutputItem, ResponseStatus, ResponseToolType, ResponsesRequest, ResponsesResponse, +use super::conversions; +use crate::{ + mcp::{self, McpManager}, + protocols::{ + chat::{ + ChatChoice, ChatCompletionMessage, ChatCompletionResponse, ChatCompletionStreamResponse, + }, + common::{Function, FunctionCallResponse, Tool, ToolCall, ToolChoice, ToolChoiceValue}, + responses::{ + self, McpToolInfo, ResponseContentPart, ResponseInput, ResponseInputOutputItem, + ResponseOutputItem, ResponseStatus, ResponseToolType, ResponsesRequest, + ResponsesResponse, + }, + }, + routers::grpc::{ + common::responses::streaming::{OutputItemType, ResponseStreamEventEmitter}, + error, }, }; @@ -155,10 +162,7 @@ fn generate_mcp_id(prefix: &str) -> String { } /// Build mcp_list_tools output item -fn build_mcp_list_tools_item( - mcp: &Arc, - server_label: &str, -) -> ResponseOutputItem { +fn build_mcp_list_tools_item(mcp: &Arc, server_label: &str) -> ResponseOutputItem { let tools = mcp.list_tools(); let tools_info: Vec = tools .iter() @@ -263,8 +267,6 @@ pub(super) async fn execute_tool_loop( headers.clone(), model_id.clone(), ctx.components.clone(), - response_id.clone(), - Some(ctx.background_tasks.clone()), ) .await?; @@ -358,10 +360,9 @@ pub(super) async fn execute_tool_loop( content: vec![ResponseContentPart::InputText { text: text.clone() }], status: Some("completed".to_string()), }], - ResponseInput::Items(items) => items - .iter() - .map(crate::protocols::responses::normalize_input_item) - .collect(), + ResponseInput::Items(items) => { + items.iter().map(responses::normalize_input_item).collect() + } }; // Append all conversation history (function calls and outputs) @@ -830,10 +831,9 @@ async fn execute_tool_loop_streaming_internal( content: vec![ResponseContentPart::InputText { text: text.clone() }], status: Some("completed".to_string()), }], - ResponseInput::Items(items) => items - .iter() - .map(crate::protocols::responses::normalize_input_item) - .collect(), + ResponseInput::Items(items) => { + items.iter().map(responses::normalize_input_item).collect() + } }; input_items.extend_from_slice(&state.conversation_history); @@ -911,13 +911,13 @@ async fn execute_tool_loop_streaming_internal( } /// Convert MCP tools to Chat API tool format -fn convert_mcp_tools_to_chat_tools(mcp_tools: &[crate::mcp::Tool]) -> Vec { +fn convert_mcp_tools_to_chat_tools(mcp_tools: &[mcp::Tool]) -> Vec { use serde_json::Value; mcp_tools .iter() .map(|tool_info| Tool { tool_type: "function".to_string(), - function: crate::protocols::common::Function { + function: 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()), @@ -933,10 +933,6 @@ async fn convert_and_accumulate_stream( emitter: &mut ResponseStreamEventEmitter, tx: &mpsc::UnboundedSender>, ) -> Result { - use futures_util::StreamExt; - - use crate::protocols::chat::ChatCompletionStreamResponse; - let mut accumulator = ChatResponseAccumulator::new(); let mut stream = body.into_data_stream(); @@ -971,7 +967,7 @@ struct ChatResponseAccumulator { id: String, model: String, content: String, - tool_calls: HashMap, + tool_calls: HashMap, finish_reason: Option, } @@ -986,7 +982,7 @@ impl ChatResponseAccumulator { } } - fn process_chunk(&mut self, chunk: &crate::protocols::chat::ChatCompletionStreamResponse) { + fn process_chunk(&mut self, chunk: &ChatCompletionStreamResponse) { if !chunk.id.is_empty() { self.id = chunk.id.clone(); } @@ -1004,15 +1000,13 @@ impl ChatResponseAccumulator { if let Some(tool_call_deltas) = &choice.delta.tool_calls { for delta in tool_call_deltas { let index = delta.index as usize; - let entry = self.tool_calls.entry(index).or_insert_with(|| { - crate::protocols::common::ToolCall { - id: String::new(), - tool_type: "function".to_string(), - function: crate::protocols::common::FunctionCallResponse { - name: String::new(), - arguments: Some(String::new()), - }, - } + let entry = self.tool_calls.entry(index).or_insert_with(|| ToolCall { + id: String::new(), + tool_type: "function".to_string(), + function: FunctionCallResponse { + name: String::new(), + arguments: Some(String::new()), + }, }); if let Some(id) = &delta.id { @@ -1048,9 +1042,9 @@ impl ChatResponseAccumulator { object: "chat.completion".to_string(), created: chrono::Utc::now().timestamp() as u64, model: self.model, - choices: vec![crate::protocols::chat::ChatChoice { + choices: vec![ChatChoice { index: 0, - message: crate::protocols::chat::ChatCompletionMessage { + message: ChatCompletionMessage { role: "assistant".to_string(), content: if self.content.is_empty() { None diff --git a/sgl-router/src/routers/grpc/responses/types.rs b/sgl-router/src/routers/grpc/regular/responses/types.rs similarity index 100% rename from sgl-router/src/routers/grpc/responses/types.rs rename to sgl-router/src/routers/grpc/regular/responses/types.rs diff --git a/sgl-router/src/routers/grpc/regular/stages/chat/mod.rs b/sgl-router/src/routers/grpc/regular/stages/chat/mod.rs new file mode 100644 index 000000000..606849c78 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/chat/mod.rs @@ -0,0 +1,12 @@ +//! Chat endpoint pipeline stages +//! +//! These stages handle chat-specific preprocessing, request building, and response processing. +//! They work with any model type by using injected model adapters. + +mod preparation; +mod request_building; +mod response_processing; + +pub use preparation::ChatPreparationStage; +pub use request_building::ChatRequestBuildingStage; +pub use response_processing::ChatResponseProcessingStage; diff --git a/sgl-router/src/routers/grpc/regular/stages/chat/preparation.rs b/sgl-router/src/routers/grpc/regular/stages/chat/preparation.rs new file mode 100644 index 000000000..49a7a41ef --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/chat/preparation.rs @@ -0,0 +1,104 @@ +//! Chat preparation stage: Filter tools, process messages, tokenize, build constraints + +use std::borrow::Cow; + +use async_trait::async_trait; +use axum::response::Response; + +use crate::{ + protocols::chat::ChatCompletionRequest, + routers::grpc::{ + common::stages::PipelineStage, + context::{PreparationOutput, RequestContext}, + error, utils, + }, +}; + +/// Chat preparation stage +/// +/// Extracts chat-specific preparation logic from the old unified PreparationStage. +/// This is a direct extraction without architectural changes. +pub struct ChatPreparationStage; + +#[async_trait] +impl PipelineStage for ChatPreparationStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let request = ctx.chat_request_arc(); + self.prepare_chat(ctx, &request).await?; + Ok(None) + } + + fn name(&self) -> &'static str { + "ChatPreparation" + } +} + +impl ChatPreparationStage { + 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(error::bad_request(e)); + } + }; + + // Step 3: Tokenize the processed text + let encoding = match ctx.components.tokenizer.encode(&processed_messages.text) { + Ok(encoding) => encoding, + Err(e) => { + return Err(error::internal_error(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| error::bad_request(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(()) + } +} diff --git a/sgl-router/src/routers/grpc/regular/stages/chat/request_building.rs b/sgl-router/src/routers/grpc/regular/stages/chat/request_building.rs new file mode 100644 index 000000000..2d69b7de1 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/chat/request_building.rs @@ -0,0 +1,82 @@ +//! Chat request building stage: Build proto GenerateRequest for chat requests + +use async_trait::async_trait; +use axum::response::Response; +use uuid::Uuid; + +use crate::routers::grpc::{ + common::stages::{helpers, PipelineStage}, + context::{ClientSelection, RequestContext, WorkerSelection}, + error, +}; + +/// Chat request building stage +/// +/// Extracts chat-specific request building logic from the old unified RequestBuildingStage. +pub struct ChatRequestBuildingStage { + inject_pd_metadata: bool, +} + +impl ChatRequestBuildingStage { + pub fn new(inject_pd_metadata: bool) -> Self { + Self { inject_pd_metadata } + } +} + +#[async_trait] +impl PipelineStage for ChatRequestBuildingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let prep = ctx + .state + .preparation + .as_ref() + .ok_or_else(|| error::internal_error("Preparation not completed"))?; + + let clients = ctx + .state + .clients + .as_ref() + .ok_or_else(|| error::internal_error("Client acquisition not completed"))?; + + let chat_request = ctx.chat_request_arc(); + + // Get client for building request (use prefill client if PD mode) + let builder_client = match clients { + ClientSelection::Single { client } => client, + ClientSelection::Dual { prefill, .. } => prefill, + }; + + // Build chat request + let request_id = format!("chatcmpl-{}", Uuid::new_v4()); + let body_ref = prep.filtered_request.as_ref().unwrap_or(&chat_request); + + let mut proto_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| error::bad_request(format!("Invalid request parameters: {}", e)))?; + + // Inject PD metadata if needed + if self.inject_pd_metadata { + if let WorkerSelection::Dual { prefill, .. } = ctx.state.workers.as_ref().unwrap() { + helpers::inject_bootstrap_metadata(&mut proto_request, prefill); + } + } + + ctx.state.proto_request = Some(proto_request); + Ok(None) + } + + fn name(&self) -> &'static str { + "ChatRequestBuilding" + } +} diff --git a/sgl-router/src/routers/grpc/regular/stages/chat/response_processing.rs b/sgl-router/src/routers/grpc/regular/stages/chat/response_processing.rs new file mode 100644 index 000000000..1d3a20220 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/chat/response_processing.rs @@ -0,0 +1,111 @@ +//! Chat 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; + +use async_trait::async_trait; +use axum::response::Response; + +use crate::routers::grpc::{ + common::stages::PipelineStage, + context::{FinalResponse, RequestContext}, + error, + regular::{processor, streaming}, +}; + +/// Chat response processing stage +/// +/// Extracts chat-specific response processing logic from the old unified ResponseProcessingStage. +pub struct ChatResponseProcessingStage { + processor: processor::ResponseProcessor, + streaming_processor: Arc, +} + +impl ChatResponseProcessingStage { + pub fn new( + processor: processor::ResponseProcessor, + streaming_processor: Arc, + ) -> Self { + Self { + processor, + streaming_processor, + } + } +} + +#[async_trait] +impl PipelineStage for ChatResponseProcessingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + self.process_chat_response(ctx).await + } + + fn name(&self) -> &'static str { + "ChatResponseProcessing" + } +} + +impl ChatResponseProcessingStage { + 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(|| error::internal_error("No execution result"))?; + + // Get dispatch metadata (needed by both streaming and non-streaming) + let dispatch = ctx + .state + .dispatch + .as_ref() + .ok_or_else(|| error::internal_error("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 = ctx.chat_request().logprobs; + + let chat_request = ctx.chat_request_arc(); + + let stop_decoder = ctx + .state + .response + .stop_decoder + .as_mut() + .ok_or_else(|| error::internal_error("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) + } +} diff --git a/sgl-router/src/routers/grpc/regular/stages/generate/mod.rs b/sgl-router/src/routers/grpc/regular/stages/generate/mod.rs new file mode 100644 index 000000000..0d3d33e97 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/generate/mod.rs @@ -0,0 +1,12 @@ +//! Generate endpoint pipeline stages +//! +//! These stages handle generate-specific preprocessing, request building, and response processing. +//! They work with any model type by using injected model adapters. + +mod preparation; +mod request_building; +mod response_processing; + +pub use preparation::GeneratePreparationStage; +pub use request_building::GenerateRequestBuildingStage; +pub use response_processing::GenerateResponseProcessingStage; diff --git a/sgl-router/src/routers/grpc/regular/stages/generate/preparation.rs b/sgl-router/src/routers/grpc/regular/stages/generate/preparation.rs new file mode 100644 index 000000000..95f2b2ef8 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/generate/preparation.rs @@ -0,0 +1,119 @@ +//! Generate preparation stage: Resolve input, tokenize, create stop decoder + +use std::sync::Arc; + +use async_trait::async_trait; +use axum::response::Response; + +use crate::{ + protocols::{common::InputIds, generate::GenerateRequest}, + routers::grpc::{ + common::stages::PipelineStage, + context::{PreparationOutput, RequestContext}, + error, utils, + }, + tokenizer::traits::Tokenizer, +}; + +/// Generate preparation stage +/// +/// Extracts generate-specific preparation logic from the old unified PreparationStage. +/// This is a direct extraction without architectural changes. +pub struct GeneratePreparationStage; + +#[async_trait] +impl PipelineStage for GeneratePreparationStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let request = ctx.generate_request_arc(); + self.prepare_generate(ctx, &request).await?; + Ok(None) + } + + fn name(&self) -> &'static str { + "GeneratePreparation" + } +} + +impl GeneratePreparationStage { + 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(error::bad_request(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/regular/stages/generate/request_building.rs b/sgl-router/src/routers/grpc/regular/stages/generate/request_building.rs new file mode 100644 index 000000000..956a21ec9 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/generate/request_building.rs @@ -0,0 +1,78 @@ +//! Generate request building stage: Build proto GenerateRequest for generate requests + +use async_trait::async_trait; +use axum::response::Response; +use uuid::Uuid; + +use crate::routers::grpc::{ + common::stages::{helpers, PipelineStage}, + context::{ClientSelection, RequestContext, WorkerSelection}, + error, +}; + +/// Generate request building stage +/// +/// Extracts generate-specific request building logic from the old unified RequestBuildingStage. +pub struct GenerateRequestBuildingStage { + inject_pd_metadata: bool, +} + +impl GenerateRequestBuildingStage { + pub fn new(inject_pd_metadata: bool) -> Self { + Self { inject_pd_metadata } + } +} + +#[async_trait] +impl PipelineStage for GenerateRequestBuildingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let prep = ctx + .state + .preparation + .as_ref() + .ok_or_else(|| error::internal_error("Preparation not completed"))?; + + let clients = ctx + .state + .clients + .as_ref() + .ok_or_else(|| error::internal_error("Client acquisition not completed"))?; + + let generate_request = ctx.generate_request_arc(); + + // Get client for building request (use prefill client if PD mode) + let builder_client = match clients { + ClientSelection::Single { client } => client, + ClientSelection::Dual { prefill, .. } => prefill, + }; + + // Build generate request + let request_id = generate_request + .rid + .clone() + .unwrap_or_else(|| format!("gen-{}", Uuid::new_v4())); + + let mut proto_request = builder_client + .build_plain_generate_request( + request_id, + &generate_request, + prep.original_text.clone(), + prep.token_ids.clone(), + ) + .map_err(error::bad_request)?; + + // Inject PD metadata if needed + if self.inject_pd_metadata { + if let WorkerSelection::Dual { prefill, .. } = ctx.state.workers.as_ref().unwrap() { + helpers::inject_bootstrap_metadata(&mut proto_request, prefill); + } + } + + ctx.state.proto_request = Some(proto_request); + Ok(None) + } + + fn name(&self) -> &'static str { + "GenerateRequestBuilding" + } +} diff --git a/sgl-router/src/routers/grpc/regular/stages/generate/response_processing.rs b/sgl-router/src/routers/grpc/regular/stages/generate/response_processing.rs new file mode 100644 index 000000000..29bb69eda --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/generate/response_processing.rs @@ -0,0 +1,109 @@ +//! Generate response processing stage: Handles both streaming and non-streaming responses + +use std::{sync::Arc, time::Instant}; + +use async_trait::async_trait; +use axum::response::Response; + +use crate::routers::grpc::{ + common::stages::PipelineStage, + context::{FinalResponse, RequestContext}, + error, + regular::{processor, streaming}, +}; + +/// Generate response processing stage +/// +/// Extracts generate-specific response processing logic from the old unified ResponseProcessingStage. +pub struct GenerateResponseProcessingStage { + processor: processor::ResponseProcessor, + streaming_processor: Arc, +} + +impl GenerateResponseProcessingStage { + pub fn new( + processor: processor::ResponseProcessor, + streaming_processor: Arc, + ) -> Self { + Self { + processor, + streaming_processor, + } + } +} + +#[async_trait] +impl PipelineStage for GenerateResponseProcessingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + self.process_generate_response(ctx).await + } + + fn name(&self) -> &'static str { + "GenerateResponseProcessing" + } +} + +impl GenerateResponseProcessingStage { + 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(|| error::internal_error("No execution result"))?; + + // Get dispatch metadata (needed by both streaming and non-streaming) + let dispatch = ctx + .state + .dispatch + .as_ref() + .ok_or_else(|| error::internal_error("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(|| error::internal_error("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/regular/stages/mod.rs b/sgl-router/src/routers/grpc/regular/stages/mod.rs new file mode 100644 index 000000000..a448a82ea --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/mod.rs @@ -0,0 +1,17 @@ +//! Pipeline stages for regular (non-harmony) model processing +//! +//! This module defines stages specific to regular tokenizer-based models. + +pub mod chat; +pub mod generate; +mod preparation; +mod request_building; +mod response_processing; + +pub use chat::{ChatPreparationStage, ChatRequestBuildingStage, ChatResponseProcessingStage}; +pub use generate::{ + GeneratePreparationStage, GenerateRequestBuildingStage, GenerateResponseProcessingStage, +}; +pub use preparation::PreparationStage; +pub use request_building::RequestBuildingStage; +pub use response_processing::ResponseProcessingStage; diff --git a/sgl-router/src/routers/grpc/regular/stages/preparation.rs b/sgl-router/src/routers/grpc/regular/stages/preparation.rs new file mode 100644 index 000000000..0436cf936 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/preparation.rs @@ -0,0 +1,52 @@ +//! Preparation stage that delegates to endpoint-specific implementations +//! +//! This stage checks RequestType at runtime and delegates to the appropriate +//! endpoint-specific stage (ChatPreparationStage or GeneratePreparationStage). + +use async_trait::async_trait; +use axum::response::Response; + +use super::{chat::ChatPreparationStage, generate::GeneratePreparationStage}; +use crate::routers::grpc::{ + common::stages::PipelineStage, + context::{RequestContext, RequestType}, +}; + +/// Preparation stage (delegates to endpoint-specific implementations) +pub struct PreparationStage { + chat_stage: ChatPreparationStage, + generate_stage: GeneratePreparationStage, +} + +impl PreparationStage { + pub fn new() -> Self { + Self { + chat_stage: ChatPreparationStage, + generate_stage: GeneratePreparationStage, + } + } +} + +impl Default for PreparationStage { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl PipelineStage for PreparationStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + match &ctx.input.request_type { + RequestType::Chat(_) => self.chat_stage.execute(ctx).await, + RequestType::Generate(_) => self.generate_stage.execute(ctx).await, + RequestType::Responses(_) => { + // Responses API has its own preparation handled elsewhere + Ok(None) + } + } + } + + fn name(&self) -> &'static str { + "Preparation" + } +} diff --git a/sgl-router/src/routers/grpc/regular/stages/request_building.rs b/sgl-router/src/routers/grpc/regular/stages/request_building.rs new file mode 100644 index 000000000..db9401728 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/request_building.rs @@ -0,0 +1,54 @@ +//! Request building stage that delegates to endpoint-specific implementations + +use async_trait::async_trait; +use axum::response::Response; +use uuid::Uuid; + +use super::{chat::ChatRequestBuildingStage, generate::GenerateRequestBuildingStage}; +use crate::{ + grpc_client::proto, + routers::grpc::{ + common::stages::PipelineStage, + context::{RequestContext, RequestType}, + }, +}; + +/// Request building stage (delegates to endpoint-specific implementations) +pub struct RequestBuildingStage { + chat_stage: ChatRequestBuildingStage, + generate_stage: GenerateRequestBuildingStage, +} + +impl RequestBuildingStage { + pub fn new(inject_pd_metadata: bool) -> Self { + Self { + chat_stage: ChatRequestBuildingStage::new(inject_pd_metadata), + generate_stage: GenerateRequestBuildingStage::new(inject_pd_metadata), + } + } +} + +#[async_trait] +impl PipelineStage for RequestBuildingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + match &ctx.input.request_type { + RequestType::Chat(_) => self.chat_stage.execute(ctx).await, + RequestType::Generate(_) => self.generate_stage.execute(ctx).await, + 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()); + + ctx.state.proto_request = Some(proto::GenerateRequest { + request_id, + ..Default::default() + }); + Ok(None) + } + } + } + + fn name(&self) -> &'static str { + "RequestBuilding" + } +} diff --git a/sgl-router/src/routers/grpc/regular/stages/response_processing.rs b/sgl-router/src/routers/grpc/regular/stages/response_processing.rs new file mode 100644 index 000000000..71e885be4 --- /dev/null +++ b/sgl-router/src/routers/grpc/regular/stages/response_processing.rs @@ -0,0 +1,52 @@ +//! Response processing stage that delegates to endpoint-specific implementations + +use std::sync::Arc; + +use async_trait::async_trait; +use axum::response::Response; + +use super::{chat::ChatResponseProcessingStage, generate::GenerateResponseProcessingStage}; +use crate::routers::grpc::{ + common::stages::PipelineStage, + context::{RequestContext, RequestType}, + error, + regular::{processor, streaming}, +}; + +/// Response processing stage (delegates to endpoint-specific implementations) +pub struct ResponseProcessingStage { + chat_stage: ChatResponseProcessingStage, + generate_stage: GenerateResponseProcessingStage, +} + +impl ResponseProcessingStage { + pub fn new( + processor: processor::ResponseProcessor, + streaming_processor: Arc, + ) -> Self { + Self { + chat_stage: ChatResponseProcessingStage::new( + processor.clone(), + streaming_processor.clone(), + ), + generate_stage: GenerateResponseProcessingStage::new(processor, streaming_processor), + } + } +} + +#[async_trait] +impl PipelineStage for ResponseProcessingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + match &ctx.input.request_type { + RequestType::Chat(_) => self.chat_stage.execute(ctx).await, + RequestType::Generate(_) => self.generate_stage.execute(ctx).await, + RequestType::Responses(_) => Err(error::bad_request( + "Responses API processing must be handled by responses handler".to_string(), + )), + } + } + + fn name(&self) -> &'static str { + "ResponseProcessing" + } +} diff --git a/sgl-router/src/routers/grpc/streaming.rs b/sgl-router/src/routers/grpc/regular/streaming.rs similarity index 97% rename from sgl-router/src/routers/grpc/streaming.rs rename to sgl-router/src/routers/grpc/regular/streaming.rs index 1202ed69b..93cfbbd72 100644 --- a/sgl-router/src/routers/grpc/streaming.rs +++ b/sgl-router/src/routers/grpc/regular/streaming.rs @@ -1,7 +1,6 @@ //! Streaming response processor for gRPC routers //! -//! This module contains shared streaming logic for both Regular and PD routers, -//! eliminating ~600 lines of duplication. +//! This module contains shared streaming logic for both Regular and PD router. use std::{collections::HashMap, io, sync::Arc, time::Instant}; @@ -17,9 +16,8 @@ use tokio::sync::{mpsc, mpsc::UnboundedSender}; use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt}; use tracing::{debug, error, warn}; -use super::{context, utils}; use crate::{ - grpc_client::proto, + grpc_client::{proto, sglang_scheduler::AbortOnDropStream}, protocols::{ chat::{ ChatCompletionRequest, ChatCompletionStreamResponse, ChatMessageDelta, ChatStreamChoice, @@ -30,20 +28,21 @@ use crate::{ }, generate::GenerateRequest, }, - reasoning_parser::ReasoningParser, + reasoning_parser::{ParserFactory as ReasoningParserFactory, ParserResult, ReasoningParser}, + routers::grpc::{context, utils}, tokenizer::{ stop::{SequenceDecoderOutput, StopSequenceDecoder}, traits::Tokenizer, }, - tool_parser::ToolParser, + tool_parser::{ParserFactory as ToolParserFactory, StreamingParseResult, ToolParser}, }; /// Shared streaming processor for both single and dual dispatch modes #[derive(Clone)] pub struct StreamingProcessor { tokenizer: Arc, - tool_parser_factory: crate::tool_parser::ParserFactory, - reasoning_parser_factory: crate::reasoning_parser::ParserFactory, + tool_parser_factory: ToolParserFactory, + reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, configured_reasoning_parser: Option, } @@ -51,8 +50,8 @@ pub struct StreamingProcessor { impl StreamingProcessor { pub fn new( tokenizer: Arc, - tool_parser_factory: crate::tool_parser::ParserFactory, - reasoning_parser_factory: crate::reasoning_parser::ParserFactory, + tool_parser_factory: ToolParserFactory, + reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, configured_reasoning_parser: Option, ) -> Self { @@ -161,7 +160,7 @@ impl StreamingProcessor { /// Process streaming chunks from a single stream (Regular mode) pub async fn process_streaming_chunks( &self, - mut grpc_stream: crate::grpc_client::sglang_scheduler::AbortOnDropStream, + mut grpc_stream: AbortOnDropStream, dispatch: context::DispatchMetadata, stop_params: (Option, Option>, bool, bool), original_request: Arc, @@ -576,8 +575,8 @@ impl StreamingProcessor { /// Process dual streaming chunks (prefill + decode) - PD mode pub async fn process_dual_streaming_chunks( &self, - mut prefill_stream: crate::grpc_client::sglang_scheduler::AbortOnDropStream, - decode_stream: crate::grpc_client::sglang_scheduler::AbortOnDropStream, + mut prefill_stream: AbortOnDropStream, + decode_stream: AbortOnDropStream, dispatch: context::DispatchMetadata, stop_params: (Option, Option>, bool, bool), original_request: Arc, @@ -696,7 +695,7 @@ impl StreamingProcessor { /// Process streaming chunks for generate endpoint (no tool/reasoning parsing) async fn process_generate_streaming( tokenizer: Arc, - mut stream: crate::grpc_client::sglang_scheduler::AbortOnDropStream, + mut stream: AbortOnDropStream, request_id: String, weight_version: String, _include_logprobs: bool, @@ -800,8 +799,8 @@ impl StreamingProcessor { /// Process dual streaming for generate endpoint (PD mode with logprobs support) async fn process_generate_streaming_dual( tokenizer: Arc, - mut prefill_stream: crate::grpc_client::sglang_scheduler::AbortOnDropStream, - decode_stream: crate::grpc_client::sglang_scheduler::AbortOnDropStream, + mut prefill_stream: AbortOnDropStream, + decode_stream: AbortOnDropStream, request_id: String, weight_version: String, return_logprob: bool, @@ -857,7 +856,7 @@ impl StreamingProcessor { /// Process generate streaming with optional input_logprobs async fn process_generate_streaming_with_input_logprobs( tokenizer: Arc, - mut stream: crate::grpc_client::sglang_scheduler::AbortOnDropStream, + mut stream: AbortOnDropStream, request_id: String, weight_version: String, _include_logprobs: bool, @@ -1051,7 +1050,7 @@ impl StreamingProcessor { }; match parse_result { - Ok(crate::reasoning_parser::ParserResult { + Ok(ParserResult { reasoning_text, normal_text, }) => { @@ -1122,7 +1121,7 @@ impl StreamingProcessor { let mut parser = pooled_parser.lock().await; match parser.parse_incremental(delta, tools).await { - Ok(crate::tool_parser::StreamingParseResult { normal_text, calls }) => { + Ok(StreamingParseResult { normal_text, calls }) => { // Emit normal text if present if !normal_text.is_empty() { chunks.push(ChatCompletionStreamResponse { diff --git a/sgl-router/src/routers/grpc/router.rs b/sgl-router/src/routers/grpc/router.rs index c317ed1d3..ca323eec9 100644 --- a/sgl-router/src/routers/grpc/router.rs +++ b/sgl-router/src/routers/grpc/router.rs @@ -1,5 +1,3 @@ -// gRPC Router Implementation - use std::sync::Arc; use async_trait::async_trait; @@ -12,13 +10,14 @@ use axum::{ use tracing::debug; use super::{ + common::responses::handlers::{cancel_response_impl, get_response_impl}, context::SharedComponents, harmony::{ serve_harmony_responses, serve_harmony_responses_stream, HarmonyDetector, HarmonyResponsesContext, }, pipeline::RequestPipeline, - responses, + regular::responses, }; use crate::{ app_context::AppContext, @@ -43,9 +42,7 @@ pub struct GrpcRouter { 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, } @@ -156,7 +153,6 @@ impl GrpcRouter { &self.pipeline }; - // Use selected pipeline for ALL requests (streaming and non-streaming) pipeline .execute_chat( Arc::new(body.clone()), @@ -176,7 +172,6 @@ impl GrpcRouter { ) -> Response { debug!("Processing generate request for model: {:?}", model_id); - // Use pipeline for ALL requests (streaming and non-streaming) self.pipeline .execute_generate( Arc::new(body.clone()), @@ -187,35 +182,51 @@ impl GrpcRouter { .await } - /// Main route_responses implementation (pipeline-based for Harmony) + /// Main route_responses implementation + /// + /// Routes to either Harmony or regular responses implementation based on model detection async fn route_responses_impl( &self, - _headers: Option<&HeaderMap>, + headers: Option<&HeaderMap>, body: &ResponsesRequest, model_id: Option<&str>, ) -> Response { + // Choose implementation based on Harmony model detection + let is_harmony = HarmonyDetector::is_harmony_model(&body.model); + debug!( - "Processing Harmony responses request for model: {:?}, streaming: {:?}", - model_id, body.stream + "Processing responses request for model: {:?}, using_harmony={}", + model_id, is_harmony ); - // 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(), - ); + if is_harmony { + debug!( + "Processing Harmony responses request for model: {:?}, streaming: {:?}", + model_id, body.stream + ); + 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(), + ); - // Check if streaming is requested - if body.stream.unwrap_or(false) { - serve_harmony_responses_stream(&harmony_ctx, body.clone()).await - } else { - // Use non-streaming version for standard JSON responses - match serve_harmony_responses(&harmony_ctx, body.clone()).await { - Ok(response) => axum::Json(response).into_response(), - Err(error_response) => error_response, + if body.stream.unwrap_or(false) { + serve_harmony_responses_stream(&harmony_ctx, body.clone()).await + } else { + match serve_harmony_responses(&harmony_ctx, body.clone()).await { + Ok(response) => axum::Json(response).into_response(), + Err(error_response) => error_response, + } } + } else { + responses::route_responses( + &self.responses_context, + Arc::new(body.clone()), + headers.cloned(), + model_id.map(|s| s.to_string()), + ) + .await } } } @@ -236,7 +247,6 @@ impl RouterTrait for GrpcRouter { } async fn health_generate(&self, _req: Request) -> Response { - // TODO: Implement actual generation test for gRPC ( StatusCode::NOT_IMPLEMENTED, "Health generate not yet implemented for gRPC", @@ -289,27 +299,7 @@ impl RouterTrait for GrpcRouter { body: &ResponsesRequest, model_id: Option<&str>, ) -> Response { - // 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 - } + self.route_responses_impl(headers, body, model_id).await } async fn get_response( @@ -318,11 +308,11 @@ impl RouterTrait for GrpcRouter { response_id: &str, _params: &ResponsesGetParams, ) -> Response { - responses::get_response_impl(&self.responses_context, response_id).await + get_response_impl(&self.responses_context, response_id).await } async fn cancel_response(&self, _headers: Option<&HeaderMap>, response_id: &str) -> Response { - responses::cancel_response_impl(&self.responses_context, response_id).await + cancel_response_impl(&self.responses_context, response_id).await } async fn route_embeddings( diff --git a/sgl-router/src/routers/grpc/stages/preparation.rs b/sgl-router/src/routers/grpc/stages/preparation.rs deleted file mode 100644 index fe7f77e56..000000000 --- a/sgl-router/src/routers/grpc/stages/preparation.rs +++ /dev/null @@ -1,195 +0,0 @@ -//! 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}, - error, 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(error::bad_request(e)); - } - }; - - // Step 3: Tokenize the processed text - let encoding = match ctx.components.tokenizer.encode(&processed_messages.text) { - Ok(encoding) => encoding, - Err(e) => { - return Err(error::internal_error(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| error::bad_request(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(error::bad_request(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 deleted file mode 100644 index 2879073b2..000000000 --- a/sgl-router/src/routers/grpc/stages/request_building.rs +++ /dev/null @@ -1,144 +0,0 @@ -//! 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}, - error, - }, -}; - -/// 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(|| error::internal_error("Preparation not completed"))?; - - let clients = ctx - .state - .clients - .as_ref() - .ok_or_else(|| error::internal_error("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| error::bad_request(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(error::bad_request)? - } - 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/response_processing.rs b/sgl-router/src/routers/grpc/stages/response_processing.rs deleted file mode 100644 index 5a9ff9b12..000000000 --- a/sgl-router/src/routers/grpc/stages/response_processing.rs +++ /dev/null @@ -1,184 +0,0 @@ -//! 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}, - error, processing, streaming, -}; - -/// 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(error::bad_request( - "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(|| error::internal_error("No execution result"))?; - - // Get dispatch metadata (needed by both streaming and non-streaming) - let dispatch = ctx - .state - .dispatch - .as_ref() - .ok_or_else(|| error::internal_error("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(|| error::internal_error("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(|| error::internal_error("No execution result"))?; - - // Get dispatch metadata (needed by both streaming and non-streaming) - let dispatch = ctx - .state - .dispatch - .as_ref() - .ok_or_else(|| error::internal_error("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(|| error::internal_error("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/utils.rs b/sgl-router/src/routers/grpc/utils.rs index 0d5c217d9..0ad45e96a 100644 --- a/sgl-router/src/routers/grpc/utils.rs +++ b/sgl-router/src/routers/grpc/utils.rs @@ -21,12 +21,19 @@ use crate::{ }, generate::GenerateFinishReason, }, + reasoning_parser::{ + ParserFactory as ReasoningParserFactory, PooledParser as ReasoningPooledParser, + ReasoningParser, + }, tokenizer::{ cache::CachedTokenizer, chat_template::{ChatTemplateContentFormat, ChatTemplateParams}, traits::Tokenizer, HuggingFaceTokenizer, }, + tool_parser::{ + ParserFactory as ToolParserFactory, PooledParser as ToolPooledParser, ToolParser, + }, }; /// Get gRPC client from worker, returning appropriate error response on failure @@ -44,20 +51,17 @@ pub async fn get_grpc_client_from_worker( /// Process tool call arguments in messages /// Per Transformers docs, tool call arguments in assistant messages should be dicts -pub fn process_tool_call_arguments(messages: &mut [Value]) -> Result<(), String> { +fn process_tool_call_arguments(messages: &mut [Value]) -> Result<(), String> { for msg in messages { - // Early return if not assistant message let role = msg.get("role").and_then(|v| v.as_str()); if role != Some("assistant") { continue; } - // Early return if no tool_calls let Some(tool_calls) = msg.get_mut("tool_calls").and_then(|tc| tc.as_array_mut()) else { continue; }; - // Process each tool call's arguments for call in tool_calls { let Some(function) = call.get_mut("function") else { continue; @@ -107,10 +111,7 @@ pub fn process_content_format( } /// Transform a single content field based on content format -pub fn transform_content_field( - content_value: &mut Value, - content_format: ChatTemplateContentFormat, -) { +fn transform_content_field(content_value: &mut Value, content_format: ChatTemplateContentFormat) { let Some(content_array) = content_value.as_array() else { return; // Not multimodal, keep as-is }; @@ -209,7 +210,7 @@ pub fn generate_tool_constraints( /// Build JSON schema for required tool calls (array with minItems: 1) /// Includes $defs consolidation from all tools (matching Python's behavior) -pub fn build_required_array_schema(tools: &[Tool]) -> Result { +fn build_required_array_schema(tools: &[Tool]) -> Result { // Build anyOf schemas for each tool let mut any_of_schemas = Vec::new(); for tool in tools { @@ -651,7 +652,7 @@ pub fn generate_tool_call_id( /// Check if a reasoning parser is available for the given model pub fn check_reasoning_parser_availability( - reasoning_parser_factory: &crate::reasoning_parser::ParserFactory, + reasoning_parser_factory: &ReasoningParserFactory, configured_parser: Option<&String>, model: &str, ) -> bool { @@ -666,7 +667,7 @@ pub fn check_reasoning_parser_availability( /// Check if a tool parser is available for the given model pub fn check_tool_parser_availability( - tool_parser_factory: &crate::tool_parser::ParserFactory, + tool_parser_factory: &ToolParserFactory, configured_parser: Option<&String>, model: &str, ) -> bool { @@ -683,10 +684,10 @@ pub fn check_tool_parser_availability( /// Otherwise, auto-detect based on the model name. /// Get a pooled reasoning parser (for non-streaming where state doesn't matter) pub fn get_reasoning_parser( - reasoning_parser_factory: &crate::reasoning_parser::ParserFactory, + reasoning_parser_factory: &ReasoningParserFactory, configured_parser: Option<&String>, model: &str, -) -> crate::reasoning_parser::PooledParser { +) -> ReasoningPooledParser { if let Some(parser_name) = configured_parser { // Use configured parser if specified reasoning_parser_factory @@ -707,10 +708,10 @@ pub fn get_reasoning_parser( /// Create a fresh reasoning parser instance (for streaming where state isolation is needed) pub fn create_reasoning_parser( - reasoning_parser_factory: &crate::reasoning_parser::ParserFactory, + reasoning_parser_factory: &ReasoningParserFactory, configured_parser: Option<&String>, model: &str, -) -> Option> { +) -> Option> { if let Some(parser_name) = configured_parser { // Use configured parser if specified reasoning_parser_factory @@ -735,10 +736,10 @@ pub fn create_reasoning_parser( /// Otherwise, auto-detect based on the model name. /// Get a pooled tool parser (for non-streaming where state doesn't matter) pub fn get_tool_parser( - tool_parser_factory: &crate::tool_parser::ParserFactory, + tool_parser_factory: &ToolParserFactory, configured_parser: Option<&String>, model: &str, -) -> crate::tool_parser::PooledParser { +) -> ToolPooledParser { if let Some(parser_name) = configured_parser { // Use configured parser if specified tool_parser_factory @@ -759,10 +760,10 @@ pub fn get_tool_parser( /// Create a fresh tool parser instance (for streaming where state isolation is needed) pub fn create_tool_parser( - tool_parser_factory: &crate::tool_parser::ParserFactory, + tool_parser_factory: &ToolParserFactory, configured_parser: Option<&String>, model: &str, -) -> Option> { +) -> Option> { if let Some(parser_name) = configured_parser { // Use configured parser if specified tool_parser_factory diff --git a/sgl-router/src/tokenizer/chat_template.rs b/sgl-router/src/tokenizer/chat_template.rs index 387974f59..9f33921c5 100644 --- a/sgl-router/src/tokenizer/chat_template.rs +++ b/sgl-router/src/tokenizer/chat_template.rs @@ -18,20 +18,15 @@ use minijinja::{ use serde_json; /// Chat template content format -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum ChatTemplateContentFormat { /// Content is a simple string + #[default] String, /// Content is a list of structured parts (OpenAI format) OpenAI, } -impl Default for ChatTemplateContentFormat { - fn default() -> Self { - Self::String - } -} - impl std::fmt::Display for ChatTemplateContentFormat { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { diff --git a/sgl-router/tests/chat_template_integration.rs b/sgl-router/tests/chat_template_integration.rs index 2f345a4f1..4fdc61eb5 100644 --- a/sgl-router/tests/chat_template_integration.rs +++ b/sgl-router/tests/chat_template_integration.rs @@ -171,7 +171,7 @@ fn test_chatml_template() { let processor = ChatTemplateProcessor::new(template.to_string()); - let messages = vec![ + let messages = [ ChatMessage::User { content: UserMessageContent::Text("Hello".to_string()), name: None,