EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[router] responses api POST and GET with local storage (#10581)

Browse Source 
Co-authored-by: key4ng <rukeyang@gmail.com>
This commit is contained in:
Simo Lin
2025-09-23 12:12:02 -04:00
committed by GitHub
parent ddab4fc7c7
commit 98c3b04ff2
14 changed files with 1158 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
sgl-router/src/routers/factory.rs
View File

@@ -166,8 +166,12 @@ impl RouterFactory {
.cloned()
.ok_or_else(|| "OpenAI mode requires at least one worker URL".to_string())?;
let router =
OpenAIRouter::new(base_url, Some(ctx.router_config.circuit_breaker.clone())).await?;
let router = OpenAIRouter::new(
base_url,
Some(ctx.router_config.circuit_breaker.clone()),
ctx.response_storage.clone(),
)
.await?;
Ok(Box::new(router))
}

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

@@ -308,7 +308,12 @@ impl RouterTrait for GrpcPDRouter {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
async fn get_response(
&self,
_headers: Option<&HeaderMap>,
_response_id: &str,
_params: &crate::protocols::spec::ResponsesGetParams,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}

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

@@ -237,7 +237,12 @@ impl RouterTrait for GrpcRouter {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
async fn get_response(
&self,
_headers: Option<&HeaderMap>,
_response_id: &str,
_params: &crate::protocols::spec::ResponsesGetParams,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}

42
sgl-router/src/routers/header_utils.rs
View File

@@ -51,3 +51,45 @@ fn should_forward_header(name: &str) -> bool {
"host" // Should not forward the backend's host header
)
}
/// Apply headers to a reqwest request builder, filtering out headers that shouldn't be forwarded
/// or that will be set automatically by reqwest
pub fn apply_request_headers(
headers: &HeaderMap,
mut request_builder: reqwest::RequestBuilder,
skip_content_headers: bool,
) -> reqwest::RequestBuilder {
// Always forward Authorization header first if present
if let Some(auth) = headers
.get("authorization")
.or_else(|| headers.get("Authorization"))
{
request_builder = request_builder.header("Authorization", auth.clone());
}
// Forward other headers, filtering out problematic ones
for (key, value) in headers.iter() {
let key_str = key.as_str().to_lowercase();
// Skip headers that:
// - Are set automatically by reqwest (content-type, content-length for POST/PUT)
// - We already handled (authorization)
// - Are hop-by-hop headers (connection, transfer-encoding)
// - Should not be forwarded (host)
let should_skip = key_str == "authorization" || // Already handled above
key_str == "host" ||
key_str == "connection" ||
key_str == "transfer-encoding" ||
key_str == "keep-alive" ||
key_str == "te" ||
key_str == "trailers" ||
key_str == "upgrade" ||
(skip_content_headers && (key_str == "content-type" || key_str == "content-length"));
if !should_skip {
request_builder = request_builder.header(key.clone(), value.clone());
}
}
request_builder
}

583
sgl-router/src/routers/http/openai_router.rs
View File

@@ -1,10 +1,15 @@
//! OpenAI router implementation (reqwest-based)
//! OpenAI router implementation
use crate::config::CircuitBreakerConfig;
use crate::core::{CircuitBreaker, CircuitBreakerConfig as CoreCircuitBreakerConfig};
use crate::data_connector::{ResponseId, SharedResponseStorage, StoredResponse};
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, GenerateRequest, RerankRequest,
ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, RerankRequest,
ResponseContentPart, ResponseInput, ResponseInputOutputItem, ResponseOutputItem,
ResponseStatus, ResponseTextFormat, ResponsesGetParams, ResponsesRequest, ResponsesResponse,
TextFormatType,
};
use crate::routers::header_utils::{apply_request_headers, preserve_response_headers};
use async_trait::async_trait;
use axum::{
body::Body,
@@ -13,13 +18,17 @@ use axum::{
response::{IntoResponse, Response},
};
use futures_util::StreamExt;
use serde_json::{json, to_value, Value};
use std::{
any::Any,
collections::HashMap,
sync::atomic::{AtomicBool, Ordering},
};
use tokio::sync::mpsc;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tracing::{error, info, warn};
/// Router for OpenAI backend
#[derive(Debug)]
pub struct OpenAIRouter {
/// HTTP client for upstream OpenAI-compatible API
client: reqwest::Client,
@@ -29,6 +38,17 @@ pub struct OpenAIRouter {
circuit_breaker: CircuitBreaker,
/// Health status
healthy: AtomicBool,
/// Response storage for managing conversation history
response_storage: SharedResponseStorage,
}
impl std::fmt::Debug for OpenAIRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OpenAIRouter")
.field("base_url", &self.base_url)
.field("healthy", &self.healthy)
.finish()
}
}
impl OpenAIRouter {
@@ -36,6 +56,7 @@ impl OpenAIRouter {
pub async fn new(
base_url: String,
circuit_breaker_config: Option<CircuitBreakerConfig>,
response_storage: SharedResponseStorage,
) -> Result<Self, String> {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(300))
@@ -61,8 +82,246 @@ impl OpenAIRouter {
base_url,
circuit_breaker,
healthy: AtomicBool::new(true),
response_storage,
})
}
async fn handle_non_streaming_response(
&self,
url: String,
headers: Option<&HeaderMap>,
payload: Value,
original_body: &ResponsesRequest,
original_previous_response_id: Option<String>,
) -> Response {
let request_builder = self.client.post(&url).json(&payload);
// Apply headers with filtering
let request_builder = if let Some(headers) = headers {
apply_request_headers(headers, request_builder, true)
} else {
request_builder
};
match request_builder.send().await {
Ok(response) => {
let status = response.status();
if !status.is_success() {
let error_text = response
.text()
.await
.unwrap_or_else(|e| format!("Failed to get error body: {}", e));
return (status, error_text).into_response();
}
// Parse the response
match response.json::<Value>().await {
Ok(mut openai_response_json) => {
if let Some(prev_id) = original_previous_response_id {
if let Some(obj) = openai_response_json.as_object_mut() {
let should_insert = obj
.get("previous_response_id")
.map(|v| v.is_null())
.unwrap_or(true);
if should_insert {
obj.insert(
"previous_response_id".to_string(),
Value::String(prev_id),
);
}
}
}
if let Some(obj) = openai_response_json.as_object_mut() {
if !obj.contains_key("instructions") {
if let Some(instructions) = &original_body.instructions {
obj.insert(
"instructions".to_string(),
Value::String(instructions.clone()),
);
}
}
if !obj.contains_key("metadata") {
if let Some(metadata) = &original_body.metadata {
let metadata_map: serde_json::Map<String, Value> = metadata
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
obj.insert("metadata".to_string(), Value::Object(metadata_map));
}
}
// Reflect the client's requested store preference in the response body
obj.insert("store".to_string(), Value::Bool(original_body.store));
}
if original_body.store {
if let Err(e) = self
.store_response_internal(&openai_response_json, original_body)
.await
{
warn!("Failed to store response: {}", e);
}
}
match serde_json::to_string(&openai_response_json) {
Ok(json_str) => (
StatusCode::OK,
[("content-type", "application/json")],
json_str,
)
.into_response(),
Err(e) => {
error!("Failed to serialize response: {}", e);
(
StatusCode::INTERNAL_SERVER_ERROR,
json!({"error": {"message": "Failed to serialize response", "type": "internal_error"}}).to_string(),
)
.into_response()
}
}
}
Err(e) => {
error!("Failed to parse OpenAI response: {}", e);
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to parse response: {}", e),
)
.into_response()
}
}
}
Err(e) => (
StatusCode::BAD_GATEWAY,
format!("Failed to forward request to OpenAI: {}", e),
)
.into_response(),
}
}
async fn handle_streaming_response(
&self,
_url: String,
_headers: Option<&HeaderMap>,
_payload: Value,
_original_body: &ResponsesRequest,
_original_previous_response_id: Option<String>,
) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Streaming responses not yet implemented",
)
.into_response()
}
async fn store_response_internal(
&self,
response_json: &Value,
original_body: &ResponsesRequest,
) -> Result<(), String> {
if !original_body.store {
return Ok(());
}
match Self::store_response_impl(&self.response_storage, response_json, original_body).await
{
Ok(response_id) => {
info!(response_id = %response_id.0, "Stored response locally");
Ok(())
}
Err(e) => Err(e),
}
}
async fn store_response_impl(
response_storage: &SharedResponseStorage,
response_json: &Value,
original_body: &ResponsesRequest,
) -> Result<ResponseId, String> {
let input_text = match &original_body.input {
ResponseInput::Text(text) => text.clone(),
ResponseInput::Items(_) => "complex input".to_string(),
};
let output_text = Self::extract_primary_output_text(response_json).unwrap_or_default();
let mut stored_response = StoredResponse::new(input_text, output_text, None);
stored_response.instructions = response_json
.get("instructions")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.or_else(|| original_body.instructions.clone());
stored_response.model = response_json
.get("model")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.or_else(|| original_body.model.clone());
stored_response.user = response_json
.get("user")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.or_else(|| original_body.user.clone());
stored_response.metadata = response_json
.get("metadata")
.and_then(|v| v.as_object())
.map(|m| {
m.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect::<HashMap<_, _>>()
})
.unwrap_or_else(|| original_body.metadata.clone().unwrap_or_default());
stored_response.previous_response_id = response_json
.get("previous_response_id")
.and_then(|v| v.as_str())
.map(|s| ResponseId::from_string(s.to_string()))
.or_else(|| {
original_body
.previous_response_id
.as_ref()
.map(|id| ResponseId::from_string(id.clone()))
});
if let Some(id_str) = response_json.get("id").and_then(|v| v.as_str()) {
stored_response.id = ResponseId::from_string(id_str.to_string());
}
stored_response.raw_response = response_json.clone();
response_storage
.store_response(stored_response)
.await
.map_err(|e| format!("Failed to store response: {}", e))
}
fn extract_primary_output_text(response_json: &Value) -> Option<String> {
if let Some(items) = response_json.get("output").and_then(|v| v.as_array()) {
for item in items {
if let Some(content) = item.get("content").and_then(|v| v.as_array()) {
for part in content {
if part
.get("type")
.and_then(|v| v.as_str())
.map(|t| t == "output_text")
.unwrap_or(false)
{
if let Some(text) = part.get("text").and_then(|v| v.as_str()) {
return Some(text.to_string());
}
}
}
}
}
}
None
}
}
#[async_trait]
@@ -108,7 +367,7 @@ impl super::super::RouterTrait for OpenAIRouter {
}
async fn get_server_info(&self, _req: Request<Body>) -> Response {
let info = serde_json::json!({
let info = json!({
"router_type": "openai",
"workers": 1,
"base_url": &self.base_url
@@ -192,7 +451,7 @@ impl super::super::RouterTrait for OpenAIRouter {
}
// Serialize request body, removing SGLang-only fields
let mut payload = match serde_json::to_value(body) {
let mut payload = match to_value(body) {
Ok(v) => v,
Err(e) => {
return (
@@ -282,7 +541,7 @@ impl super::super::RouterTrait for OpenAIRouter {
} else {
// Stream SSE bytes to client
let stream = resp.bytes_stream();
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let (tx, rx) = mpsc::unbounded_channel();
tokio::spawn(async move {
let mut s = stream;
while let Some(chunk) = s.next().await {
@@ -299,9 +558,7 @@ impl super::super::RouterTrait for OpenAIRouter {
}
}
});
let mut response = Response::new(Body::from_stream(
tokio_stream::wrappers::UnboundedReceiverStream::new(rx),
));
let mut response = Response::new(Body::from_stream(UnboundedReceiverStream::new(rx)));
*response.status_mut() = status;
response
.headers_mut()
@@ -326,36 +583,294 @@ impl super::super::RouterTrait for OpenAIRouter {
async fn route_responses(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::ResponsesRequest,
_model_id: Option<&str>,
headers: Option<&HeaderMap>,
body: &ResponsesRequest,
model_id: Option<&str>,
) -> Response {
let url = format!("{}/v1/responses", self.base_url);
info!(
requested_store = body.store,
is_streaming = body.stream,
"openai_responses_request"
);
if body.stream {
return (
StatusCode::NOT_IMPLEMENTED,
"Streaming responses not yet implemented",
)
.into_response();
}
// Clone the body and override model if needed
let mut request_body = body.clone();
if let Some(model) = model_id {
request_body.model = Some(model.to_string());
}
// Store the original previous_response_id for the response
let original_previous_response_id = request_body.previous_response_id.clone();
// Handle previous_response_id by loading prior context
let mut conversation_items: Option<Vec<ResponseInputOutputItem>> = None;
if let Some(prev_id_str) = request_body.previous_response_id.clone() {
let prev_id = ResponseId::from_string(prev_id_str.clone());
match self
.response_storage
.get_response_chain(&prev_id, None)
.await
{
Ok(chain) => {
if !chain.responses.is_empty() {
let mut items = Vec::new();
for stored in chain.responses.iter() {
let trimmed_id = stored.id.0.trim_start_matches("resp_");
if !stored.input.is_empty() {
items.push(ResponseInputOutputItem::Message {
id: format!("msg_u_{}", trimmed_id),
role: "user".to_string(),
status: Some("completed".to_string()),
content: vec![ResponseContentPart::InputText {
text: stored.input.clone(),
}],
});
}
if !stored.output.is_empty() {
items.push(ResponseInputOutputItem::Message {
id: format!("msg_a_{}", trimmed_id),
role: "assistant".to_string(),
status: Some("completed".to_string()),
content: vec![ResponseContentPart::OutputText {
text: stored.output.clone(),
annotations: vec![],
logprobs: None,
}],
});
}
}
conversation_items = Some(items);
} else {
info!(previous_response_id = %prev_id_str, "previous chain empty");
}
}
Err(err) => {
warn!(previous_response_id = %prev_id_str, %err, "failed to fetch previous response chain");
}
}
// Clear previous_response_id from request since we're converting to conversation
request_body.previous_response_id = None;
}
if let Some(mut items) = conversation_items {
match &request_body.input {
ResponseInput::Text(text) => {
items.push(ResponseInputOutputItem::Message {
id: format!("msg_u_current_{}", items.len()),
role: "user".to_string(),
status: Some("completed".to_string()),
content: vec![ResponseContentPart::InputText { text: text.clone() }],
});
}
ResponseInput::Items(existing) => {
items.extend(existing.clone());
}
}
request_body.input = ResponseInput::Items(items);
}
// Always set store=false for OpenAI (we store internally)
request_body.store = false;
// Convert to JSON payload and strip SGLang-specific fields before forwarding
let mut payload = match to_value(&request_body) {
Ok(value) => value,
Err(err) => {
return (
StatusCode::BAD_REQUEST,
format!("Failed to serialize responses request: {}", err),
)
.into_response();
}
};
if let Some(obj) = payload.as_object_mut() {
for key in [
"request_id",
"priority",
"frequency_penalty",
"presence_penalty",
"stop",
"top_k",
"min_p",
"repetition_penalty",
] {
obj.remove(key);
}
}
// Check if streaming is requested
if body.stream {
// Handle streaming response
self.handle_streaming_response(
url,
headers,
payload,
body,
original_previous_response_id,
)
.await
} else {
// Handle non-streaming response
self.handle_non_streaming_response(
url,
headers,
payload,
body,
original_previous_response_id,
)
.await
}
}
async fn get_response(
&self,
_headers: Option<&HeaderMap>,
response_id: &str,
params: &ResponsesGetParams,
) -> Response {
let stored_id = ResponseId::from_string(response_id.to_string());
if let Ok(Some(stored_response)) = self.response_storage.get_response(&stored_id).await {
let stream_requested = params.stream.unwrap_or(false);
let raw_value = stored_response.raw_response.clone();
if !raw_value.is_null() {
if stream_requested {
return (
StatusCode::NOT_IMPLEMENTED,
"Streaming retrieval not yet implemented",
)
.into_response();
}
return (
StatusCode::OK,
[("content-type", "application/json")],
raw_value.to_string(),
)
.into_response();
}
let openai_response = ResponsesResponse {
id: stored_response.id.0.clone(),
object: "response".to_string(),
created_at: stored_response.created_at.timestamp(),
status: ResponseStatus::Completed,
error: None,
incomplete_details: None,
instructions: stored_response.instructions.clone(),
max_output_tokens: None,
model: stored_response
.model
.unwrap_or_else(|| "gpt-4o".to_string()),
output: vec![ResponseOutputItem::Message {
id: format!("msg_{}", stored_response.id.0),
role: "assistant".to_string(),
status: "completed".to_string(),
content: vec![ResponseContentPart::OutputText {
text: stored_response.output,
annotations: vec![],
logprobs: None,
}],
}],
parallel_tool_calls: true,
previous_response_id: stored_response.previous_response_id.map(|id| id.0),
reasoning: None,
store: true,
temperature: Some(1.0),
text: Some(ResponseTextFormat {
format: TextFormatType {
format_type: "text".to_string(),
},
}),
tool_choice: "auto".to_string(),
tools: vec![],
top_p: Some(1.0),
truncation: Some("disabled".to_string()),
usage: None,
user: stored_response.user.clone(),
metadata: stored_response.metadata.clone(),
};
if stream_requested {
return (
StatusCode::NOT_IMPLEMENTED,
"Streaming retrieval not yet implemented",
)
.into_response();
}
return (
StatusCode::OK,
[("content-type", "application/json")],
serde_json::to_string(&openai_response).unwrap_or_else(|e| {
format!("{{\"error\": \"Failed to serialize response: {}\"}}", e)
}),
)
.into_response();
}
(
StatusCode::NOT_IMPLEMENTED,
"Responses endpoint not implemented for OpenAI router",
StatusCode::NOT_FOUND,
format!(
"Response with id '{}' not found in local storage",
response_id
),
)
.into_response()
}
async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Responses retrieve endpoint not implemented for OpenAI router",
)
.into_response()
}
async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response {
// Forward to OpenAI's cancel endpoint
let url = format!("{}/v1/responses/{}/cancel", self.base_url, response_id);
async fn cancel_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Responses cancel endpoint not implemented for OpenAI router",
)
.into_response()
let request_builder = self.client.post(&url);
// Apply headers with filtering (skip content headers for POST without body)
let request_builder = if let Some(headers) = headers {
apply_request_headers(headers, request_builder, true)
} else {
request_builder
};
match request_builder.send().await {
Ok(response) => {
let status = response.status();
let headers = response.headers().clone();
match response.text().await {
Ok(body_text) => {
let mut response = (status, body_text).into_response();
*response.headers_mut() = preserve_response_headers(&headers);
response
}
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to read response body: {}", e),
)
.into_response(),
}
}
Err(e) => (
StatusCode::BAD_GATEWAY,
format!("Failed to cancel response on OpenAI: {}", e),
)
.into_response(),
}
}
async fn flush_cache(&self) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
StatusCode::FORBIDDEN,
"flush_cache not supported for OpenAI router",
)
.into_response()
@@ -363,7 +878,7 @@ impl super::super::RouterTrait for OpenAIRouter {
async fn get_worker_loads(&self) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
StatusCode::FORBIDDEN,
"get_worker_loads not supported for OpenAI router",
)
.into_response()
@@ -384,12 +899,12 @@ impl super::super::RouterTrait for OpenAIRouter {
async fn route_embeddings(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::EmbeddingRequest,
_body: &EmbeddingRequest,
_model_id: Option<&str>,
) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Embeddings endpoint not implemented for OpenAI backend",
StatusCode::FORBIDDEN,
"Embeddings endpoint not supported for OpenAI backend",
)
.into_response()
}
@@ -401,8 +916,8 @@ impl super::super::RouterTrait for OpenAIRouter {
_model_id: Option<&str>,
) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Rerank endpoint not implemented for OpenAI backend",
StatusCode::FORBIDDEN,
"Rerank endpoint not supported for OpenAI backend",
)
.into_response()
}

9
sgl-router/src/routers/http/pd_router.rs
View File

@@ -8,7 +8,7 @@ use crate::metrics::RouterMetrics;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
use crate::protocols::spec::{
ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateRequest, RerankRequest,
ResponsesRequest, StringOrArray, UserMessageContent,
ResponsesGetParams, ResponsesRequest, StringOrArray, UserMessageContent,
};
use crate::routers::header_utils;
use crate::routers::RouterTrait;
@@ -1424,7 +1424,12 @@ impl RouterTrait for PDRouter {
.into_response()
}
async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
async fn get_response(
&self,
_headers: Option<&HeaderMap>,
_response_id: &str,
_params: &ResponsesGetParams,
) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Responses retrieve endpoint not implemented for PD router",

9
sgl-router/src/routers/http/router.rs
View File

@@ -6,7 +6,7 @@ use crate::metrics::RouterMetrics;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, GenerationRequest,
RerankRequest, RerankResponse, RerankResult, ResponsesRequest,
RerankRequest, RerankResponse, RerankResult, ResponsesGetParams, ResponsesRequest,
};
use crate::routers::header_utils;
use crate::routers::RouterTrait;
@@ -903,7 +903,12 @@ impl RouterTrait for Router {
.await
}
async fn get_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response {
async fn get_response(
&self,
headers: Option<&HeaderMap>,
response_id: &str,
_params: &ResponsesGetParams,
) -> Response {
let endpoint = format!("v1/responses/{}", response_id);
self.route_get_request(headers, &endpoint).await
}

9
sgl-router/src/routers/mod.rs
View File

@@ -11,7 +11,7 @@ use std::fmt::Debug;
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, RerankRequest,
ResponsesRequest,
ResponsesGetParams, ResponsesRequest,
};
pub mod factory;
@@ -82,7 +82,12 @@ pub trait RouterTrait: Send + Sync + Debug {
) -> Response;
/// Retrieve a stored/background response by id
async fn get_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response;
async fn get_response(
&self,
headers: Option<&HeaderMap>,
response_id: &str,
params: &ResponsesGetParams,
) -> Response;
/// Cancel a background response by id
async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response;

66
sgl-router/src/routers/router_manager.rs
View File

@@ -8,7 +8,7 @@ use crate::config::{ConnectionMode, RoutingMode};
use crate::core::{WorkerRegistry, WorkerType};
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, RerankRequest,
ResponsesRequest,
ResponsesGetParams, ResponsesRequest,
};
use crate::routers::RouterTrait;
use crate::server::{AppContext, ServerConfig};
@@ -403,38 +403,19 @@ impl RouterTrait for RouterManager {
async fn route_responses(
&self,
_headers: Option<&HeaderMap>,
_body: &ResponsesRequest,
_model_id: Option<&str>,
headers: Option<&HeaderMap>,
body: &ResponsesRequest,
model_id: Option<&str>,
) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"responses api not yet implemented in inference gateway mode",
)
.into_response()
}
let selected_model = body.model.as_deref().or(model_id);
let router = self.select_router_for_request(headers, selected_model);
async fn get_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response {
let router = self.select_router_for_request(headers, None);
if let Some(router) = router {
router.get_response(headers, response_id).await
router.route_responses(headers, body, selected_model).await
} else {
(
StatusCode::NOT_FOUND,
format!("No router available to get response '{}'", response_id),
)
.into_response()
}
}
async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response {
let router = self.select_router_for_request(headers, None);
if let Some(router) = router {
router.cancel_response(headers, response_id).await
} else {
(
StatusCode::NOT_FOUND,
format!("No router available to cancel response '{}'", response_id),
"No router available to handle responses request",
)
.into_response()
}
@@ -460,6 +441,37 @@ impl RouterTrait for RouterManager {
.into_response()
}
async fn get_response(
&self,
headers: Option<&HeaderMap>,
response_id: &str,
params: &ResponsesGetParams,
) -> Response {
let router = self.select_router_for_request(headers, None);
if let Some(router) = router {
router.get_response(headers, response_id, params).await
} else {
(
StatusCode::NOT_FOUND,
format!("No router available to get response '{}'", response_id),
)
.into_response()
}
}
async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response {
let router = self.select_router_for_request(headers, None);
if let Some(router) = router {
router.cancel_response(headers, response_id).await
} else {
(
StatusCode::NOT_FOUND,
format!("No router available to cancel response '{}'", response_id),
)
.into_response()
}
}
async fn route_embeddings(
&self,
headers: Option<&HeaderMap>,