Unverified Commit 4c9bcb9d authored by Keyang Ru's avatar Keyang Ru Committed by GitHub
Browse files

[Router] Refactor protocol definitions: split spec.rs into modular files (#11677)


Co-authored-by: default avatarChang Su <chang.s.su@oracle.com>
parent 86b04d25
...@@ -3,10 +3,11 @@ use serde_json::{from_str, to_string, to_value, to_vec}; ...@@ -3,10 +3,11 @@ use serde_json::{from_str, to_string, to_value, to_vec};
use std::time::Instant; use std::time::Instant;
use sglang_router_rs::core::{BasicWorker, BasicWorkerBuilder, Worker, WorkerType}; use sglang_router_rs::core::{BasicWorker, BasicWorkerBuilder, Worker, WorkerType};
use sglang_router_rs::protocols::spec::{ use sglang_router_rs::protocols::chat::{ChatCompletionRequest, ChatMessage, UserMessageContent};
ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateRequest, SamplingParams, use sglang_router_rs::protocols::common::StringOrArray;
StringOrArray, UserMessageContent, use sglang_router_rs::protocols::completion::CompletionRequest;
}; use sglang_router_rs::protocols::generate::GenerateRequest;
use sglang_router_rs::protocols::sampling_params::SamplingParams;
use sglang_router_rs::routers::http::pd_types::{generate_room_id, RequestWithBootstrap}; use sglang_router_rs::routers::http::pd_types::{generate_room_id, RequestWithBootstrap};
fn create_test_worker() -> BasicWorker { fn create_test_worker() -> BasicWorker {
......
...@@ -9,7 +9,7 @@ ...@@ -9,7 +9,7 @@
use criterion::{black_box, criterion_group, BenchmarkId, Criterion, Throughput}; use criterion::{black_box, criterion_group, BenchmarkId, Criterion, Throughput};
use serde_json::json; use serde_json::json;
use sglang_router_rs::protocols::spec::{Function, Tool}; use sglang_router_rs::protocols::common::{Function, Tool};
use sglang_router_rs::tool_parser::{JsonParser, ParserFactory as ToolParserFactory, ToolParser}; use sglang_router_rs::tool_parser::{JsonParser, ParserFactory as ToolParserFactory, ToolParser};
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
......
...@@ -7,10 +7,10 @@ use std::time::Duration; ...@@ -7,10 +7,10 @@ use std::time::Duration;
use tonic::{transport::Channel, Request, Streaming}; use tonic::{transport::Channel, Request, Streaming};
use tracing::{debug, warn}; use tracing::{debug, warn};
use crate::protocols::spec::{ use crate::protocols::chat::ChatCompletionRequest;
ChatCompletionRequest, GenerateRequest, ResponseFormat, use crate::protocols::common::{ResponseFormat, StringOrArray, ToolChoice, ToolChoiceValue};
SamplingParams as GenerateSamplingParams, StringOrArray, use crate::protocols::generate::GenerateRequest;
}; use crate::protocols::sampling_params::SamplingParams as GenerateSamplingParams;
// Include the generated protobuf code // Include the generated protobuf code
pub mod proto { pub mod proto {
...@@ -306,9 +306,7 @@ impl SglangSchedulerClient { ...@@ -306,9 +306,7 @@ impl SglangSchedulerClient {
// Handle skip_special_tokens: set to false if tools are present and tool_choice is not "none" // Handle skip_special_tokens: set to false if tools are present and tool_choice is not "none"
let skip_special_tokens = if request.tools.is_some() { let skip_special_tokens = if request.tools.is_some() {
match &request.tool_choice { match &request.tool_choice {
Some(crate::protocols::spec::ToolChoice::Value( Some(ToolChoice::Value(ToolChoiceValue::None)) => request.skip_special_tokens,
crate::protocols::spec::ToolChoiceValue::None,
)) => request.skip_special_tokens,
Some(_) => false, // tool_choice is not "none" Some(_) => false, // tool_choice is not "none"
None => false, // TODO: this assumes tool_choice defaults to "auto" when tools present None => false, // TODO: this assumes tool_choice defaults to "auto" when tools present
} }
......
This diff is collapsed.
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
// ============================================================================
// Default value helpers
// ============================================================================
/// Default model value when not specified
pub(crate) fn default_model() -> String {
"unknown".to_string()
}
/// Helper function for serde default value (returns true)
pub fn default_true() -> bool {
true
}
// ============================================================================
// GenerationRequest Trait
// ============================================================================
/// Trait for unified access to generation request properties
/// Implemented by ChatCompletionRequest, CompletionRequest, GenerateRequest,
/// EmbeddingRequest, RerankRequest, and ResponsesRequest
pub trait GenerationRequest: Send + Sync {
/// Check if the request is for streaming
fn is_stream(&self) -> bool;
/// Get the model name if specified
fn get_model(&self) -> Option<&str>;
/// Extract text content for routing decisions
fn extract_text_for_routing(&self) -> String;
}
// ============================================================================
// String/Array Utilities
// ============================================================================
/// A type that can be either a single string or an array of strings
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
#[serde(untagged)]
pub enum StringOrArray {
String(String),
Array(Vec<String>),
}
impl StringOrArray {
/// Get the number of items in the StringOrArray
pub fn len(&self) -> usize {
match self {
StringOrArray::String(_) => 1,
StringOrArray::Array(arr) => arr.len(),
}
}
/// Check if the StringOrArray is empty
pub fn is_empty(&self) -> bool {
match self {
StringOrArray::String(s) => s.is_empty(),
StringOrArray::Array(arr) => arr.is_empty(),
}
}
/// Convert to a vector of strings
pub fn to_vec(&self) -> Vec<String> {
match self {
StringOrArray::String(s) => vec![s.clone()],
StringOrArray::Array(arr) => arr.clone(),
}
}
}
// ============================================================================
// Content Parts (for multimodal messages)
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type")]
pub enum ContentPart {
#[serde(rename = "text")]
Text { text: String },
#[serde(rename = "image_url")]
ImageUrl { image_url: ImageUrl },
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ImageUrl {
pub url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub detail: Option<String>, // "auto", "low", or "high"
}
// ============================================================================
// Response Format (for structured outputs)
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type")]
pub enum ResponseFormat {
#[serde(rename = "text")]
Text,
#[serde(rename = "json_object")]
JsonObject,
#[serde(rename = "json_schema")]
JsonSchema { json_schema: JsonSchemaFormat },
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct JsonSchemaFormat {
pub name: String,
pub schema: Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub strict: Option<bool>,
}
// ============================================================================
// Streaming
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct StreamOptions {
#[serde(skip_serializing_if = "Option::is_none")]
pub include_usage: Option<bool>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ToolCallDelta {
pub index: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(rename = "type")]
pub tool_type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub function: Option<FunctionCallDelta>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct FunctionCallDelta {
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub arguments: Option<String>,
}
// ============================================================================
// Tools and Function Calling
// ============================================================================
/// Tool choice value for simple string options
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolChoiceValue {
Auto,
Required,
None,
}
/// Tool choice for both Chat Completion and Responses APIs
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub enum ToolChoice {
Value(ToolChoiceValue),
Function {
#[serde(rename = "type")]
tool_type: String, // "function"
function: FunctionChoice,
},
AllowedTools {
#[serde(rename = "type")]
tool_type: String, // "allowed_tools"
mode: String, // "auto" | "required" TODO: need validation
tools: Vec<ToolReference>,
},
}
impl Default for ToolChoice {
fn default() -> Self {
Self::Value(ToolChoiceValue::Auto)
}
}
/// Function choice specification for ToolChoice::Function
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct FunctionChoice {
pub name: String,
}
/// Tool reference for ToolChoice::AllowedTools
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ToolReference {
#[serde(rename = "type")]
pub tool_type: String, // "function"
pub name: String,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Tool {
#[serde(rename = "type")]
pub tool_type: String, // "function"
pub function: Function,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Function {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
pub parameters: Value, // JSON Schema
/// Whether to enable strict schema adherence (OpenAI structured outputs)
#[serde(skip_serializing_if = "Option::is_none")]
pub strict: Option<bool>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ToolCall {
pub id: String,
#[serde(rename = "type")]
pub tool_type: String, // "function"
pub function: FunctionCallResponse,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub enum FunctionCall {
None,
Auto,
Function { name: String },
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct FunctionCallResponse {
pub name: String,
#[serde(default)]
pub arguments: Option<String>, // JSON string
}
// ============================================================================
// Usage and Logging
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub completion_tokens_details: Option<CompletionTokensDetails>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CompletionTokensDetails {
pub reasoning_tokens: Option<u32>,
}
/// Usage information (used by rerank and other endpoints)
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct UsageInfo {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_tokens_details: Option<PromptTokenUsageInfo>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PromptTokenUsageInfo {
pub cached_tokens: u32,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct LogProbs {
pub tokens: Vec<String>,
pub token_logprobs: Vec<Option<f32>>,
pub top_logprobs: Vec<Option<HashMap<String, f32>>>,
pub text_offset: Vec<u32>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub enum ChatLogProbs {
Detailed {
#[serde(skip_serializing_if = "Option::is_none")]
content: Option<Vec<ChatLogProbsContent>>,
},
Raw(Value),
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ChatLogProbsContent {
pub token: String,
pub logprob: f32,
pub bytes: Option<Vec<u8>>,
pub top_logprobs: Vec<TopLogProb>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct TopLogProb {
pub token: String,
pub logprob: f32,
pub bytes: Option<Vec<u8>>,
}
// ============================================================================
// Error Types
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ErrorResponse {
pub error: ErrorDetail,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ErrorDetail {
pub message: String,
#[serde(rename = "type")]
pub error_type: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub param: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub code: Option<String>,
}
// ============================================================================
// Input Types
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub enum InputIds {
Single(Vec<i32>),
Batch(Vec<Vec<i32>>),
}
/// LoRA adapter path - can be single path or batch of paths (SGLang extension)
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub enum LoRAPath {
Single(Option<String>),
Batch(Vec<Option<String>>),
}
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use std::collections::HashMap;
use super::common::*;
// ============================================================================
// Completions API (v1/completions) - DEPRECATED but still supported
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CompletionRequest {
/// ID of the model to use (required for OpenAI, optional for some implementations, such as SGLang)
pub model: String,
/// The prompt(s) to generate completions for
pub prompt: StringOrArray,
/// The suffix that comes after a completion of inserted text
#[serde(skip_serializing_if = "Option::is_none")]
pub suffix: Option<String>,
/// The maximum number of tokens to generate
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
/// What sampling temperature to use, between 0 and 2
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
/// An alternative to sampling with temperature (nucleus sampling)
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
/// How many completions to generate for each prompt
#[serde(skip_serializing_if = "Option::is_none")]
pub n: Option<u32>,
/// Whether to stream back partial progress
#[serde(default)]
pub stream: bool,
/// Options for streaming response
#[serde(skip_serializing_if = "Option::is_none")]
pub stream_options: Option<StreamOptions>,
/// Include the log probabilities on the logprobs most likely tokens
#[serde(skip_serializing_if = "Option::is_none")]
pub logprobs: Option<u32>,
/// Echo back the prompt in addition to the completion
#[serde(default)]
pub echo: bool,
/// Up to 4 sequences where the API will stop generating further tokens
#[serde(skip_serializing_if = "Option::is_none")]
pub stop: Option<StringOrArray>,
/// Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far
#[serde(skip_serializing_if = "Option::is_none")]
pub presence_penalty: Option<f32>,
/// Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far
#[serde(skip_serializing_if = "Option::is_none")]
pub frequency_penalty: Option<f32>,
/// Generates best_of completions server-side and returns the "best"
#[serde(skip_serializing_if = "Option::is_none")]
pub best_of: Option<u32>,
/// Modify the likelihood of specified tokens appearing in the completion
#[serde(skip_serializing_if = "Option::is_none")]
pub logit_bias: Option<HashMap<String, f32>>,
/// A unique identifier representing your end-user
#[serde(skip_serializing_if = "Option::is_none")]
pub user: Option<String>,
/// If specified, our system will make a best effort to sample deterministically
#[serde(skip_serializing_if = "Option::is_none")]
pub seed: Option<i64>,
// -------- Engine Specific Sampling Parameters --------
/// Top-k sampling parameter (-1 to disable)
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<i32>,
/// Min-p nucleus sampling parameter
#[serde(skip_serializing_if = "Option::is_none")]
pub min_p: Option<f32>,
/// Minimum number of tokens to generate
#[serde(skip_serializing_if = "Option::is_none")]
pub min_tokens: Option<u32>,
/// Repetition penalty for reducing repetitive text
#[serde(skip_serializing_if = "Option::is_none")]
pub repetition_penalty: Option<f32>,
/// Regex constraint for output generation
#[serde(skip_serializing_if = "Option::is_none")]
pub regex: Option<String>,
/// EBNF grammar constraint for structured output
#[serde(skip_serializing_if = "Option::is_none")]
pub ebnf: Option<String>,
/// JSON schema constraint for structured output
#[serde(skip_serializing_if = "Option::is_none")]
pub json_schema: Option<String>,
/// Specific token IDs to use as stop conditions
#[serde(skip_serializing_if = "Option::is_none")]
pub stop_token_ids: Option<Vec<u32>>,
/// Skip trimming stop tokens from output
#[serde(default)]
pub no_stop_trim: bool,
/// Ignore end-of-sequence tokens during generation
#[serde(default)]
pub ignore_eos: bool,
/// Skip special tokens during detokenization
#[serde(default = "default_true")]
pub skip_special_tokens: bool,
/// Path to LoRA adapter(s) for model customization
#[serde(skip_serializing_if = "Option::is_none")]
pub lora_path: Option<String>,
/// Session parameters for continual prompting
#[serde(skip_serializing_if = "Option::is_none")]
pub session_params: Option<HashMap<String, Value>>,
/// Return model hidden states
#[serde(default)]
pub return_hidden_states: bool,
/// Sampling seed for deterministic outputs
#[serde(skip_serializing_if = "Option::is_none")]
pub sampling_seed: Option<u64>,
/// Additional fields including bootstrap info for PD routing
#[serde(flatten)]
pub other: Map<String, Value>,
}
impl GenerationRequest for CompletionRequest {
fn is_stream(&self) -> bool {
self.stream
}
fn get_model(&self) -> Option<&str> {
Some(&self.model)
}
fn extract_text_for_routing(&self) -> String {
match &self.prompt {
StringOrArray::String(s) => s.clone(),
StringOrArray::Array(v) => v.join(" "),
}
}
}
// ============================================================================
// Response Types
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CompletionResponse {
pub id: String,
pub object: String, // "text_completion"
pub created: u64,
pub model: String,
pub choices: Vec<CompletionChoice>,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage: Option<Usage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system_fingerprint: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CompletionChoice {
pub text: String,
pub index: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub logprobs: Option<LogProbs>,
pub finish_reason: Option<String>, // "stop", "length", "content_filter", etc.
/// Information about which stop condition was matched
#[serde(skip_serializing_if = "Option::is_none")]
pub matched_stop: Option<Value>, // Can be string or integer
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CompletionStreamResponse {
pub id: String,
pub object: String, // "text_completion"
pub created: u64,
pub choices: Vec<CompletionStreamChoice>,
pub model: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub system_fingerprint: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CompletionStreamChoice {
pub text: String,
pub index: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub logprobs: Option<LogProbs>,
pub finish_reason: Option<String>,
}
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::common::GenerationRequest;
// ============================================================================
// Embedding API
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct EmbeddingRequest {
/// ID of the model to use
pub model: String,
/// Input can be a string, array of strings, tokens, or batch inputs
pub input: Value,
/// Optional encoding format (e.g., "float", "base64")
#[serde(skip_serializing_if = "Option::is_none")]
pub encoding_format: Option<String>,
/// Optional user identifier
#[serde(skip_serializing_if = "Option::is_none")]
pub user: Option<String>,
/// Optional number of dimensions for the embedding
#[serde(skip_serializing_if = "Option::is_none")]
pub dimensions: Option<u32>,
/// SGLang extension: request id for tracking
#[serde(skip_serializing_if = "Option::is_none")]
pub rid: Option<String>,
}
impl GenerationRequest for EmbeddingRequest {
fn is_stream(&self) -> bool {
// Embeddings are non-streaming
false
}
fn get_model(&self) -> Option<&str> {
Some(&self.model)
}
fn extract_text_for_routing(&self) -> String {
// Best effort: extract text content for routing decisions
match &self.input {
Value::String(s) => s.clone(),
Value::Array(arr) => arr
.iter()
.filter_map(|v| v.as_str())
.collect::<Vec<_>>()
.join(" "),
_ => String::new(),
}
}
}
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use validator::Validate;
use super::common::{default_true, GenerationRequest, InputIds};
use super::sampling_params::SamplingParams;
use crate::protocols::validated::Normalizable;
// ============================================================================
// SGLang Generate API (native format)
// ============================================================================
#[derive(Clone, Debug, Serialize, Deserialize, Validate)]
#[validate(schema(function = "validate_generate_request"))]
pub struct GenerateRequest {
/// Text input - SGLang native format
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
/// Input IDs for tokenized input
#[serde(skip_serializing_if = "Option::is_none")]
pub input_ids: Option<InputIds>,
/// Input embeddings for direct embedding input
/// Can be a 2D array (single request) or 3D array (batch of requests)
/// Placeholder for future use
#[serde(skip_serializing_if = "Option::is_none")]
pub input_embeds: Option<Value>,
/// Image input data
/// Can be an image instance, file name, URL, or base64 encoded string
/// Supports single images, lists of images, or nested lists for batch processing
/// Placeholder for future use
#[serde(skip_serializing_if = "Option::is_none")]
pub image_data: Option<Value>,
/// Video input data
/// Can be a file name, URL, or base64 encoded string
/// Supports single videos, lists of videos, or nested lists for batch processing
/// Placeholder for future use
#[serde(skip_serializing_if = "Option::is_none")]
pub video_data: Option<Value>,
/// Audio input data
/// Can be a file name, URL, or base64 encoded string
/// Supports single audio files, lists of audio, or nested lists for batch processing
/// Placeholder for future use
#[serde(skip_serializing_if = "Option::is_none")]
pub audio_data: Option<Value>,
/// Sampling parameters (sglang style)
#[serde(skip_serializing_if = "Option::is_none")]
pub sampling_params: Option<SamplingParams>,
/// Whether to return logprobs
#[serde(skip_serializing_if = "Option::is_none")]
pub return_logprob: Option<bool>,
/// If return logprobs, the start location in the prompt for returning logprobs.
#[serde(skip_serializing_if = "Option::is_none")]
pub logprob_start_len: Option<i32>,
/// If return logprobs, the number of top logprobs to return at each position.
#[serde(skip_serializing_if = "Option::is_none")]
pub top_logprobs_num: Option<i32>,
/// If return logprobs, the token ids to return logprob for.
#[serde(skip_serializing_if = "Option::is_none")]
pub token_ids_logprob: Option<Vec<u32>>,
/// Whether to detokenize tokens in text in the returned logprobs.
#[serde(default)]
pub return_text_in_logprobs: bool,
/// Whether to stream the response
#[serde(default)]
pub stream: bool,
/// Whether to log metrics for this request (e.g. health_generate calls do not log metrics)
#[serde(default = "default_true")]
pub log_metrics: bool,
/// Return model hidden states
#[serde(default)]
pub return_hidden_states: bool,
/// The modalities of the image data [image, multi-images, video]
#[serde(skip_serializing_if = "Option::is_none")]
pub modalities: Option<Vec<String>>,
/// Session parameters for continual prompting
#[serde(skip_serializing_if = "Option::is_none")]
pub session_params: Option<HashMap<String, Value>>,
/// Path to LoRA adapter(s) for model customization
#[serde(skip_serializing_if = "Option::is_none")]
pub lora_path: Option<String>,
/// LoRA adapter ID (if pre-loaded)
#[serde(skip_serializing_if = "Option::is_none")]
pub lora_id: Option<String>,
/// Custom logit processor for advanced sampling control. Must be a serialized instance
/// of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py
/// Use the processor's `to_str()` method to generate the serialized string.
#[serde(skip_serializing_if = "Option::is_none")]
pub custom_logit_processor: Option<String>,
/// For disaggregated inference
#[serde(skip_serializing_if = "Option::is_none")]
pub bootstrap_host: Option<String>,
/// For disaggregated inference
#[serde(skip_serializing_if = "Option::is_none")]
pub bootstrap_port: Option<i32>,
/// For disaggregated inference
#[serde(skip_serializing_if = "Option::is_none")]
pub bootstrap_room: Option<i32>,
/// For disaggregated inference
#[serde(skip_serializing_if = "Option::is_none")]
pub bootstrap_pair_key: Option<String>,
/// Data parallel rank routing
#[serde(skip_serializing_if = "Option::is_none")]
pub data_parallel_rank: Option<i32>,
/// Background response
#[serde(default)]
pub background: bool,
/// Conversation ID for tracking
#[serde(skip_serializing_if = "Option::is_none")]
pub conversation_id: Option<String>,
/// Priority for the request
#[serde(skip_serializing_if = "Option::is_none")]
pub priority: Option<i32>,
/// Extra key for classifying the request (e.g. cache_salt)
#[serde(skip_serializing_if = "Option::is_none")]
pub extra_key: Option<String>,
/// Whether to disallow logging for this request (e.g. due to ZDR)
#[serde(default)]
pub no_logs: bool,
/// Custom metric labels
#[serde(skip_serializing_if = "Option::is_none")]
pub custom_labels: Option<HashMap<String, String>>,
/// Whether to return bytes for image generation
#[serde(default)]
pub return_bytes: bool,
/// Whether to return entropy
#[serde(default)]
pub return_entropy: bool,
/// Request ID for tracking (inherited from BaseReq in Python)
#[serde(skip_serializing_if = "Option::is_none")]
pub rid: Option<String>,
}
impl Normalizable for GenerateRequest {
// Use default no-op implementation - no normalization needed for GenerateRequest
}
/// Validation function for GenerateRequest - ensure exactly one input type is provided
fn validate_generate_request(req: &GenerateRequest) -> Result<(), validator::ValidationError> {
// Exactly one of text or input_ids must be provided
// Note: input_embeds not yet supported in Rust implementation
let has_text = req.text.is_some();
let has_input_ids = req.input_ids.is_some();
let count = [has_text, has_input_ids].iter().filter(|&&x| x).count();
if count == 0 {
return Err(validator::ValidationError::new(
"Either text or input_ids should be provided.",
));
}
if count > 1 {
return Err(validator::ValidationError::new(
"Either text or input_ids should be provided.",
));
}
Ok(())
}
impl GenerationRequest for GenerateRequest {
fn is_stream(&self) -> bool {
self.stream
}
fn get_model(&self) -> Option<&str> {
// Generate requests typically don't have a model field
None
}
fn extract_text_for_routing(&self) -> String {
// Check fields in priority order: text, input_ids
if let Some(ref text) = self.text {
return text.clone();
}
if let Some(ref input_ids) = self.input_ids {
return match input_ids {
InputIds::Single(ids) => ids
.iter()
.map(|&id| id.to_string())
.collect::<Vec<String>>()
.join(" "),
InputIds::Batch(batches) => batches
.iter()
.flat_map(|batch| batch.iter().map(|&id| id.to_string()))
.collect::<Vec<String>>()
.join(" "),
};
}
// No text input found
String::new()
}
}
// ============================================================================
// SGLang Generate Response Types
// ============================================================================
/// SGLang generate response (single completion or array for n>1)
///
/// Format for n=1:
/// ```json
/// {
/// "text": "...",
/// "output_ids": [...],
/// "meta_info": { ... }
/// }
/// ```
///
/// Format for n>1:
/// ```json
/// [
/// {"text": "...", "output_ids": [...], "meta_info": {...}},
/// {"text": "...", "output_ids": [...], "meta_info": {...}}
/// ]
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateResponse {
pub text: String,
pub output_ids: Vec<u32>,
pub meta_info: GenerateMetaInfo,
}
/// Metadata for a single generate completion
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateMetaInfo {
pub id: String,
pub finish_reason: GenerateFinishReason,
pub prompt_tokens: u32,
pub weight_version: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_token_logprobs: Option<Vec<Vec<Option<f64>>>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_token_logprobs: Option<Vec<Vec<Option<f64>>>>,
pub completion_tokens: u32,
pub cached_tokens: u32,
pub e2e_latency: f64,
#[serde(skip_serializing_if = "Option::is_none")]
pub matched_stop: Option<Value>,
}
/// Finish reason for generate endpoint
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum GenerateFinishReason {
Length {
length: u32,
},
Stop,
#[serde(untagged)]
Other(Value),
}
// Protocol definitions and validation for various LLM APIs // Protocol definitions and validation for various LLM APIs
// This module provides a structured approach to handling different API protocols // This module provides a structured approach to handling different API protocols
pub mod spec; pub mod chat;
pub mod common;
pub mod completion;
pub mod embedding;
pub mod generate;
pub mod rerank;
pub mod responses;
pub mod sampling_params;
pub mod validated; pub mod validated;
pub mod worker_spec; pub mod worker_spec;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use validator::Validate;
use super::common::{default_model, default_true, GenerationRequest, StringOrArray, UsageInfo};
fn default_rerank_object() -> String {
"rerank".to_string()
}
/// TODO: Create timestamp should not be in protocol layer
fn current_timestamp() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_else(|_| std::time::Duration::from_secs(0))
.as_secs() as i64
}
// ============================================================================
// Rerank API
// ============================================================================
#[derive(Debug, Clone, Deserialize, Serialize, Validate)]
#[validate(schema(function = "validate_rerank_request"))]
pub struct RerankRequest {
/// The query text to rank documents against
#[validate(custom(function = "validate_query"))]
pub query: String,
/// List of documents to be ranked
#[validate(custom(function = "validate_documents"))]
pub documents: Vec<String>,
/// Model to use for reranking
#[serde(default = "default_model")]
pub model: String,
/// Maximum number of documents to return (optional)
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1))]
pub top_k: Option<usize>,
/// Whether to return documents in addition to scores
#[serde(default = "default_true")]
pub return_documents: bool,
// SGLang specific extensions
/// Request ID for tracking
pub rid: Option<StringOrArray>,
/// User identifier
pub user: Option<String>,
}
impl GenerationRequest for RerankRequest {
fn get_model(&self) -> Option<&str> {
Some(&self.model)
}
fn is_stream(&self) -> bool {
false // Reranking doesn't support streaming
}
fn extract_text_for_routing(&self) -> String {
self.query.clone()
}
}
impl super::validated::Normalizable for RerankRequest {
// Use default no-op normalization
}
// ============================================================================
// Validation Functions
// ============================================================================
/// Validates that the query is not empty
fn validate_query(query: &str) -> Result<(), validator::ValidationError> {
if query.trim().is_empty() {
return Err(validator::ValidationError::new("query cannot be empty"));
}
Ok(())
}
/// Validates that the documents list is not empty
fn validate_documents(documents: &[String]) -> Result<(), validator::ValidationError> {
if documents.is_empty() {
return Err(validator::ValidationError::new(
"documents list cannot be empty",
));
}
Ok(())
}
/// Schema-level validation for cross-field dependencies
fn validate_rerank_request(req: &RerankRequest) -> Result<(), validator::ValidationError> {
// Validate top_k if specified
if let Some(k) = req.top_k {
if k > req.documents.len() {
// This is allowed but we log a warning
tracing::warn!(
"top_k ({}) is greater than number of documents ({})",
k,
req.documents.len()
);
}
}
Ok(())
}
impl RerankRequest {
/// Get the effective top_k value
pub fn effective_top_k(&self) -> usize {
self.top_k.unwrap_or(self.documents.len())
}
}
/// Individual rerank result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RerankResult {
/// Relevance score for the document
pub score: f32,
/// The document text (if return_documents was true)
#[serde(skip_serializing_if = "Option::is_none")]
pub document: Option<String>,
/// Original index of the document in the request
pub index: usize,
/// Additional metadata about the ranking
#[serde(skip_serializing_if = "Option::is_none")]
pub meta_info: Option<HashMap<String, Value>>,
}
/// Rerank response containing sorted results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RerankResponse {
/// Ranked results sorted by score (highest first)
pub results: Vec<RerankResult>,
/// Model used for reranking
pub model: String,
/// Usage information
pub usage: Option<UsageInfo>,
/// Response object type
#[serde(default = "default_rerank_object")]
pub object: String,
/// Response ID
pub id: Option<StringOrArray>,
/// Creation timestamp
pub created: i64,
}
impl RerankResponse {
/// Create a new RerankResponse with the given results and model
pub fn new(
results: Vec<RerankResult>,
model: String,
request_id: Option<StringOrArray>,
) -> Self {
RerankResponse {
results,
model,
usage: None,
object: default_rerank_object(),
id: request_id,
created: current_timestamp(),
}
}
/// Apply top_k limit to results
pub fn apply_top_k(&mut self, k: usize) {
self.results.truncate(k);
}
/// Drop documents from results (when return_documents is false)
pub fn drop_documents(&mut self) {
for result in &mut self.results {
result.document = None;
}
}
}
/// V1 API compatibility format for rerank requests
/// Matches Python's V1RerankReqInput
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct V1RerankReqInput {
pub query: String,
pub documents: Vec<String>,
}
/// Convert V1RerankReqInput to RerankRequest
impl From<V1RerankReqInput> for RerankRequest {
fn from(v1: V1RerankReqInput) -> Self {
RerankRequest {
query: v1.query,
documents: v1.documents,
model: default_model(),
top_k: None,
return_documents: true,
rid: None,
user: None,
}
}
}
This diff is collapsed.
use serde::{Deserialize, Serialize};
use validator::Validate;
use super::common::StringOrArray;
/// Sampling parameters for text generation
#[derive(Debug, Clone, Deserialize, Serialize, Default, Validate)]
#[validate(schema(function = "validate_sampling_params"))]
pub struct SamplingParams {
/// Temperature for sampling (must be >= 0.0, no upper limit)
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0.0))]
pub temperature: Option<f32>,
/// Maximum number of new tokens to generate (must be >= 0)
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0))]
pub max_new_tokens: Option<u32>,
/// Top-p nucleus sampling (0.0 < top_p <= 1.0)
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(custom(function = "validate_top_p_value"))]
pub top_p: Option<f32>,
/// Top-k sampling (-1 to disable, or >= 1)
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(custom(function = "validate_top_k_value"))]
pub top_k: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = -2.0, max = 2.0))]
pub frequency_penalty: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = -2.0, max = 2.0))]
pub presence_penalty: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0.0, max = 2.0))]
pub repetition_penalty: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop: Option<StringOrArray>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ignore_eos: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub skip_special_tokens: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub json_schema: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub regex: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ebnf: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0.0, max = 1.0))]
pub min_p: Option<f32>,
/// Minimum number of new tokens (validated in schema function for cross-field check with max_new_tokens)
#[serde(skip_serializing_if = "Option::is_none")]
pub min_new_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop_token_ids: Option<Vec<u32>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub no_stop_trim: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub n: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sampling_seed: Option<u64>,
}
// ============================================================================
// Shared Validation Functions
// ============================================================================
/// Validates top_p: 0.0 < top_p <= 1.0 (can't use range validator for open interval)
pub fn validate_top_p_value(top_p: f32) -> Result<(), validator::ValidationError> {
if !(top_p > 0.0 && top_p <= 1.0) {
return Err(validator::ValidationError::new(
"top_p must be in (0, 1] - greater than 0.0 and at most 1.0",
));
}
Ok(())
}
/// Validates top_k: -1 (disabled) or >= 1 (special -1 case - can't use range validator)
pub fn validate_top_k_value(top_k: i32) -> Result<(), validator::ValidationError> {
if top_k != -1 && top_k < 1 {
return Err(validator::ValidationError::new(
"top_k must be -1 (disabled) or at least 1",
));
}
Ok(())
}
// ============================================================================
// SamplingParams-Specific Validation
// ============================================================================
/// Validation function for SamplingParams - cross-field validation only
fn validate_sampling_params(params: &SamplingParams) -> Result<(), validator::ValidationError> {
// 1. Cross-field validation: min_new_tokens <= max_new_tokens
if let (Some(min), Some(max)) = (params.min_new_tokens, params.max_new_tokens) {
if min > max {
return Err(validator::ValidationError::new(
"min_new_tokens cannot exceed max_new_tokens",
));
}
}
// 2. Validate mutually exclusive structured output constraints
let constraint_count = [
params.regex.is_some(),
params.ebnf.is_some(),
params.json_schema.is_some(),
]
.iter()
.filter(|&&x| x)
.count();
if constraint_count > 1 {
return Err(validator::ValidationError::new(
"only one of regex, ebnf, or json_schema can be set",
));
}
Ok(())
}
This diff is collapsed.
...@@ -83,20 +83,11 @@ where ...@@ -83,20 +83,11 @@ where
// Then, automatically validate the data // Then, automatically validate the data
data.validate().map_err(|validation_errors| { data.validate().map_err(|validation_errors| {
// Extract the first error message from the validation errors
let error_message = validation_errors
.field_errors()
.values()
.flat_map(|errors| errors.iter())
.find_map(|e| e.message.as_ref())
.map(|m| m.to_string())
.unwrap_or_else(|| "Validation failed".to_string());
( (
StatusCode::BAD_REQUEST, StatusCode::BAD_REQUEST,
Json(json!({ Json(json!({
"error": { "error": {
"message": error_message, "message": validation_errors.to_string(),
"type": "invalid_request_error", "type": "invalid_request_error",
"code": 400 "code": 400
} }
......
...@@ -12,9 +12,8 @@ use serde_json::Value; ...@@ -12,9 +12,8 @@ use serde_json::Value;
use crate::core::Worker; use crate::core::Worker;
use crate::grpc_client::{proto, SglangSchedulerClient}; use crate::grpc_client::{proto, SglangSchedulerClient};
use crate::protocols::spec::{ use crate::protocols::chat::{ChatCompletionRequest, ChatCompletionResponse};
ChatCompletionRequest, ChatCompletionResponse, GenerateRequest, GenerateResponse, use crate::protocols::generate::{GenerateRequest, GenerateResponse};
};
use crate::reasoning_parser::ParserFactory as ReasoningParserFactory; use crate::reasoning_parser::ParserFactory as ReasoningParserFactory;
use crate::tokenizer::stop::StopSequenceDecoder; use crate::tokenizer::stop::StopSequenceDecoder;
use crate::tokenizer::traits::Tokenizer; use crate::tokenizer::traits::Tokenizer;
......
//! gRPC router implementations //! gRPC router implementations
use crate::grpc_client::proto; use crate::grpc_client::proto;
use crate::protocols::spec::StringOrArray; use crate::protocols::common::StringOrArray;
pub mod context; pub mod context;
pub mod pd_router; pub mod pd_router;
......
...@@ -3,10 +3,12 @@ ...@@ -3,10 +3,12 @@
use crate::config::types::RetryConfig; use crate::config::types::RetryConfig;
use crate::core::{ConnectionMode, WorkerRegistry, WorkerType}; use crate::core::{ConnectionMode, WorkerRegistry, WorkerType};
use crate::policies::PolicyRegistry; use crate::policies::PolicyRegistry;
use crate::protocols::spec::{ use crate::protocols::chat::ChatCompletionRequest;
ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, RerankRequest, use crate::protocols::completion::CompletionRequest;
ResponsesGetParams, ResponsesRequest, use crate::protocols::embedding::EmbeddingRequest;
}; use crate::protocols::generate::GenerateRequest;
use crate::protocols::rerank::RerankRequest;
use crate::protocols::responses::{ResponsesGetParams, ResponsesRequest};
use crate::reasoning_parser::ParserFactory as ReasoningParserFactory; use crate::reasoning_parser::ParserFactory as ReasoningParserFactory;
use crate::routers::RouterTrait; use crate::routers::RouterTrait;
use crate::server::AppContext; use crate::server::AppContext;
......
...@@ -14,7 +14,9 @@ use super::utils; ...@@ -14,7 +14,9 @@ use super::utils;
use crate::core::{ConnectionMode, Worker, WorkerRegistry, WorkerType}; use crate::core::{ConnectionMode, Worker, WorkerRegistry, WorkerType};
use crate::grpc_client::proto; use crate::grpc_client::proto;
use crate::policies::PolicyRegistry; use crate::policies::PolicyRegistry;
use crate::protocols::spec::{ChatCompletionRequest, GenerateRequest, InputIds}; use crate::protocols::chat::ChatCompletionRequest;
use crate::protocols::common::InputIds;
use crate::protocols::generate::GenerateRequest;
use crate::reasoning_parser::ParserFactory as ReasoningParserFactory; use crate::reasoning_parser::ParserFactory as ReasoningParserFactory;
use crate::tokenizer::traits::Tokenizer; use crate::tokenizer::traits::Tokenizer;
use crate::tool_parser::ParserFactory as ToolParserFactory; use crate::tool_parser::ParserFactory as ToolParserFactory;
......
...@@ -9,11 +9,13 @@ use serde_json::Value; ...@@ -9,11 +9,13 @@ use serde_json::Value;
use tracing::error; use tracing::error;
use crate::grpc_client::proto; use crate::grpc_client::proto;
use crate::protocols::spec::{ use crate::protocols::chat::{
ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse, ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse,
FunctionCallResponse, GenerateMetaInfo, GenerateRequest, GenerateResponse, ToolCall,
ToolChoice, ToolChoiceValue, Usage,
}; };
use crate::protocols::common::{
FunctionCallResponse, ToolCall, ToolChoice, ToolChoiceValue, Usage,
};
use crate::protocols::generate::{GenerateMetaInfo, GenerateRequest, GenerateResponse};
use crate::reasoning_parser::ParserFactory as ReasoningParserFactory; use crate::reasoning_parser::ParserFactory as ReasoningParserFactory;
use crate::tokenizer::stop::{SequenceDecoderOutput, StopSequenceDecoder}; use crate::tokenizer::stop::{SequenceDecoderOutput, StopSequenceDecoder};
use crate::tokenizer::traits::Tokenizer; use crate::tokenizer::traits::Tokenizer;
......
...@@ -14,10 +14,12 @@ use tracing::debug; ...@@ -14,10 +14,12 @@ use tracing::debug;
use crate::config::types::RetryConfig; use crate::config::types::RetryConfig;
use crate::core::WorkerRegistry; use crate::core::WorkerRegistry;
use crate::policies::PolicyRegistry; use crate::policies::PolicyRegistry;
use crate::protocols::spec::{ use crate::protocols::chat::ChatCompletionRequest;
ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, RerankRequest, use crate::protocols::completion::CompletionRequest;
ResponsesGetParams, ResponsesRequest, use crate::protocols::embedding::EmbeddingRequest;
}; use crate::protocols::generate::GenerateRequest;
use crate::protocols::rerank::RerankRequest;
use crate::protocols::responses::{ResponsesGetParams, ResponsesRequest};
use crate::reasoning_parser::ParserFactory as ReasoningParserFactory; use crate::reasoning_parser::ParserFactory as ReasoningParserFactory;
use crate::routers::RouterTrait; use crate::routers::RouterTrait;
use crate::server::AppContext; use crate::server::AppContext;
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment