chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[router][grpc] Restructure modules and code clean up (#12598)

Browse Source
This commit is contained in:
Chang Su
2025-11-03 23:59:27 -08:00
committed by GitHub
parent d5fa019c36
commit 83804bc626
52 changed files with 1474 additions and 1453 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
sgl-router/src/protocols/responses.rs
View File

@@ -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 {

6
sgl-router/src/routers/grpc/common/mod.rs Normal file
View File

@@ -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;

83
sgl-router/src/routers/grpc/common/response_collection.rs Normal file
View File

@@ -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<Vec<proto::GenerateComplete>, 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());
}
}
}

65
sgl-router/src/routers/grpc/common/response_formatting.rs Normal file
View File

@@ -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<ChatChoice>,
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(),
}
}

196
sgl-router/src/routers/grpc/common/responses/handlers.rs Normal file
View File

@@ -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(),
}
}

7
sgl-router/src/routers/grpc/common/responses/mod.rs Normal file
View File

@@ -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};

14
sgl-router/src/routers/grpc/responses/streaming.rs → sgl-router/src/routers/grpc/common/responses/streaming.rs
View File

@@ -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<String, String>,
// Output item tracking (NEW)
// Output item tracking
output_items: Vec<OutputItemState>,
next_output_index: usize,
current_message_output_index: Option<usize>, // 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<Result<Bytes, std::io::Error>>,

0
sgl-router/src/routers/grpc/stages/client_acquisition.rs → sgl-router/src/routers/grpc/common/stages/client_acquisition.rs
View File

0
sgl-router/src/routers/grpc/stages/dispatch_metadata.rs → sgl-router/src/routers/grpc/common/stages/dispatch_metadata.rs
View File

38
sgl-router/src/routers/grpc/common/stages/helpers.rs Normal file
View File

@@ -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<dyn Worker>,
) {
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
);
}

29
sgl-router/src/routers/grpc/stages/mod.rs → sgl-router/src/routers/grpc/common/stages/mod.rs
View File

@@ -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};

0
sgl-router/src/routers/grpc/stages/request_execution.rs → sgl-router/src/routers/grpc/common/stages/request_execution.rs
View File

1
sgl-router/src/routers/grpc/stages/worker_selection.rs → sgl-router/src/routers/grpc/common/stages/worker_selection.rs
View File

@@ -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<Arc<dyn Worker>> = workers
.iter()
.filter(|w| w.is_available())

33
sgl-router/src/routers/grpc/context.rs
View File

@@ -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<SharedComponents>,
// === 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<u32, bool>,
}
// ============================================================================
// 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 {

69
sgl-router/src/routers/grpc/harmony/processor.rs
View File

@@ -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<Vec<proto::GenerateComplete>, 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<ChatCompletionResponse, Response> {
// 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<ResponsesIterationResult, Response> {
// 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,
})),

98
sgl-router/src/routers/grpc/harmony/responses.rs
View File

@@ -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<Result<Bytes, io::Error>>, 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<Map> (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<String>,
partial_text: String,
) -> Result<ResponsesRequest, Box<Response>> {
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<ResponseTool> {
use serde_json::Value;
use crate::protocols::responses::ResponseToolType;
pub fn convert_mcp_tools_to_response_tools(mcp_tools: &[mcp::Tool]) -> Vec<ResponseTool> {
mcp_tools
.iter()
.map(|tool_info| ResponseTool {
@@ -1027,11 +1013,6 @@ fn inject_mcp_metadata(
tracking: &McpCallTracking,
mcp_manager: &Arc<McpManager>,
) {
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<McpToolInfo> = 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<ResponseInputOutputItem> {
arr.as_array()
.into_iter()
.flat_map(|items| items.iter())
.filter_map(|item| {
serde_json::from_value::<ResponseInputOutputItem>(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<ResponseInputOutputItem> {
arr.as_array()
.into_iter()
.flat_map(|items| items.iter())
.filter_map(|item| {
from_value::<ResponseInputOutputItem>(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"));

5
sgl-router/src/routers/grpc/harmony/stages/preparation.rs
View File

@@ -12,10 +12,9 @@ use crate::{
responses::ResponsesRequest,
},
routers::grpc::{
common::stages::PipelineStage,
context::{PreparationOutput, RequestContext, RequestType},
error,
stages::PipelineStage,
utils,
error, utils,
},
};

46
sgl-router/src/routers/grpc/harmony/stages/request_building.rs
View File

@@ -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<dyn Worker>,
) {
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);
}
}

2
sgl-router/src/routers/grpc/harmony/stages/response_processing.rs
View File

@@ -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

2
sgl-router/src/routers/grpc/harmony/streaming.rs
View File

@@ -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

3
sgl-router/src/routers/grpc/harmony/types.rs
View File

@@ -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",

6
sgl-router/src/routers/grpc/mod.rs
View File

@@ -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

3
sgl-router/src/routers/grpc/pd_router.rs
View File

@@ -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<Body>) -> Response {
// TODO: Implement actual generation test for gRPC PD mode
(
StatusCode::NOT_IMPLEMENTED,
"Health generate not yet implemented for gRPC PD",

104
sgl-router/src/routers/grpc/pipeline.rs
View File

@@ -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<String>,
configured_reasoning_parser: Option<String>,
) -> 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<Box<dyn PipelineStage>> = 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<String>,
configured_reasoning_parser: Option<String>,
) -> 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<Box<dyn PipelineStage>> = 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<ChatCompletionResponse, Response> 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<ChatCompletionRequest>,
headers: Option<http::HeaderMap>,
model_id: Option<String>,
components: Arc<SharedComponents>,
response_id: Option<String>,
background_tasks: Option<Arc<RwLock<HashMap<String, BackgroundTaskInfo>>>>,
) -> Result<ChatCompletionResponse, 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)) => {
@@ -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<crate::protocols::responses::ResponsesRequest>,
_headers: Option<http::HeaderMap>,
_model_id: Option<String>,
_components: Arc<SharedComponents>,
) -> 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)) => {

9
sgl-router/src/routers/grpc/regular/mod.rs Normal file
View File

@@ -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;

100
sgl-router/src/routers/grpc/processing.rs → sgl-router/src/routers/grpc/regular/processor.rs
View File

@@ -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<Vec<proto::GenerateComplete>, 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<String> = 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<ChatCompletionResponse, 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?;
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<Vec<GenerateResponse>, 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 =

0
sgl-router/src/routers/grpc/responses/context.rs → sgl-router/src/routers/grpc/regular/responses/context.rs
View File

0
sgl-router/src/routers/grpc/responses/conversions.rs → sgl-router/src/routers/grpc/regular/responses/conversions.rs
View File

398
sgl-router/src/routers/grpc/responses/handlers.rs → sgl-router/src/routers/grpc/regular/responses/handlers.rs
View File

@@ -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<ResponsesRequest>,
headers: Option<http::HeaderMap>,
model_id: Option<String>,
) -> 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<ResponsesRequest>,
@@ -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<crate::protocols::chat::ChatCompletionRequest>,
chat_request: Arc<chat::ChatCompletionRequest>,
headers: Option<http::HeaderMap>,
model_id: Option<String>,
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<crate::protocols::chat::ChatCompletionRequest>,
_chat_request: Arc<chat::ChatCompletionRequest>,
response_storage: Arc<dyn ResponseStorage>,
conversation_storage: Arc<dyn ConversationStorage>,
conversation_item_storage: Arc<dyn ConversationItemStorage>,
@@ -673,7 +537,7 @@ struct StreamingResponseAccumulator {
// Completion state
finish_reason: Option<String>,
usage: Option<crate::protocols::common::Usage>,
usage: Option<common::Usage>,
// 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(),
}
}

8
sgl-router/src/routers/grpc/responses/mod.rs → sgl-router/src/routers/grpc/regular/responses/mod.rs
View File

@@ -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;

82
sgl-router/src/routers/grpc/responses/tool_loop.rs → sgl-router/src/routers/grpc/regular/responses/tool_loop.rs
View File

@@ -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<crate::mcp::McpManager>,
server_label: &str,
) -> ResponseOutputItem {
fn build_mcp_list_tools_item(mcp: &Arc<McpManager>, server_label: &str) -> ResponseOutputItem {
let tools = mcp.list_tools();
let tools_info: Vec<McpToolInfo> = 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<Tool> {
fn convert_mcp_tools_to_chat_tools(mcp_tools: &[mcp::Tool]) -> Vec<Tool> {
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<Bytes, std::io::Error>>,
) -> Result<ChatCompletionResponse, String> {
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<usize, crate::protocols::common::ToolCall>,
tool_calls: HashMap<usize, ToolCall>,
finish_reason: Option<String>,
}
@@ -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

0
sgl-router/src/routers/grpc/responses/types.rs → sgl-router/src/routers/grpc/regular/responses/types.rs
View File

12
sgl-router/src/routers/grpc/regular/stages/chat/mod.rs Normal file
View File

@@ -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;

104
sgl-router/src/routers/grpc/regular/stages/chat/preparation.rs Normal file
View File

@@ -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<Option<Response>, 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(())
}
}

82
sgl-router/src/routers/grpc/regular/stages/chat/request_building.rs Normal file
View File

@@ -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<Option<Response>, 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"
}
}

111
sgl-router/src/routers/grpc/regular/stages/chat/response_processing.rs Normal file
View File

@@ -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<streaming::StreamingProcessor>,
}
impl ChatResponseProcessingStage {
pub fn new(
processor: processor::ResponseProcessor,
streaming_processor: Arc<streaming::StreamingProcessor>,
) -> Self {
Self {
processor,
streaming_processor,
}
}
}
#[async_trait]
impl PipelineStage for ChatResponseProcessingStage {
async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, 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<Option<Response>, 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)
}
}

12
sgl-router/src/routers/grpc/regular/stages/generate/mod.rs Normal file
View File

@@ -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;

119
sgl-router/src/routers/grpc/regular/stages/generate/preparation.rs Normal file
View File

@@ -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<Option<Response>, 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<String>, Vec<u32>), 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::<Result<Vec<u32>, _>>()
.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<dyn Tokenizer>,
text: &str,
) -> Result<(String, Vec<u32>), String> {
let encoding = tokenizer
.encode(text)
.map_err(|e| format!("Tokenization failed: {}", e))?;
Ok((text.to_string(), encoding.token_ids().to_vec()))
}
}

78
sgl-router/src/routers/grpc/regular/stages/generate/request_building.rs Normal file
View File

@@ -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<Option<Response>, 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"
}
}

109
sgl-router/src/routers/grpc/regular/stages/generate/response_processing.rs Normal file
View File

@@ -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<streaming::StreamingProcessor>,
}
impl GenerateResponseProcessingStage {
pub fn new(
processor: processor::ResponseProcessor,
streaming_processor: Arc<streaming::StreamingProcessor>,
) -> Self {
Self {
processor,
streaming_processor,
}
}
}
#[async_trait]
impl PipelineStage for GenerateResponseProcessingStage {
async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, 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<Option<Response>, 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)
}
}

17
sgl-router/src/routers/grpc/regular/stages/mod.rs Normal file
View File

@@ -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;

52
sgl-router/src/routers/grpc/regular/stages/preparation.rs Normal file
View File

@@ -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<Option<Response>, 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"
}
}

54
sgl-router/src/routers/grpc/regular/stages/request_building.rs Normal file
View File

@@ -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<Option<Response>, 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"
}
}

52
sgl-router/src/routers/grpc/regular/stages/response_processing.rs Normal file
View File

@@ -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<streaming::StreamingProcessor>,
) -> 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<Option<Response>, 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"
}
}

37
sgl-router/src/routers/grpc/streaming.rs → sgl-router/src/routers/grpc/regular/streaming.rs
View File

@@ -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<dyn Tokenizer>,
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<String>,
configured_reasoning_parser: Option<String>,
}
@@ -51,8 +50,8 @@ pub struct StreamingProcessor {
impl StreamingProcessor {
pub fn new(
tokenizer: Arc<dyn Tokenizer>,
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<String>,
configured_reasoning_parser: Option<String>,
) -> 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<StringOrArray>, Option<Vec<u32>>, bool, bool),
original_request: Arc<ChatCompletionRequest>,
@@ -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<StringOrArray>, Option<Vec<u32>>, bool, bool),
original_request: Arc<ChatCompletionRequest>,
@@ -696,7 +695,7 @@ impl StreamingProcessor {
/// Process streaming chunks for generate endpoint (no tool/reasoning parsing)
async fn process_generate_streaming(
tokenizer: Arc<dyn Tokenizer>,
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<dyn Tokenizer>,
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<dyn Tokenizer>,
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 {

90
sgl-router/src/routers/grpc/router.rs
View File

@@ -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<SharedComponents>,
// 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<Body>) -> 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(

195
sgl-router/src/routers/grpc/stages/preparation.rs
View File

@@ -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<Option<Response>, 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<String>, Vec<u32>), 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::<Result<Vec<u32>, _>>()
.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<dyn Tokenizer>,
text: &str,
) -> Result<(String, Vec<u32>), String> {
let encoding = tokenizer
.encode(text)
.map_err(|e| format!("Tokenization failed: {}", e))?;
Ok((text.to_string(), encoding.token_ids().to_vec()))
}
}

144
sgl-router/src/routers/grpc/stages/request_building.rs
View File

@@ -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<Option<Response>, 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<dyn Worker>,
) {
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
);
}
}

184
sgl-router/src/routers/grpc/stages/response_processing.rs
View File

@@ -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<streaming::StreamingProcessor>,
}
impl ResponseProcessingStage {
pub fn new(
processor: processing::ResponseProcessor,
streaming_processor: Arc<streaming::StreamingProcessor>,
) -> Self {
Self {
processor,
streaming_processor,
}
}
}
#[async_trait]
impl PipelineStage for ResponseProcessingStage {
async fn execute(&self, ctx: &mut RequestContext) -> Result<Option<Response>, 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<Option<Response>, 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<Option<Response>, 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)
}
}

39
sgl-router/src/routers/grpc/utils.rs
View File

@@ -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<String, String> {
fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
// 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<Box<dyn crate::reasoning_parser::ReasoningParser>> {
) -> Option<Box<dyn ReasoningParser>> {
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<Box<dyn crate::tool_parser::ToolParser>> {
) -> Option<Box<dyn ToolParser>> {
if let Some(parser_name) = configured_parser {
// Use configured parser if specified
tool_parser_factory

9
sgl-router/src/tokenizer/chat_template.rs
View File

@@ -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 {

2
sgl-router/tests/chat_template_integration.rs
View File

@@ -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,