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

[model-gateway] minor code clean up (#15578)

Browse Source
This commit is contained in:
Simo Lin
2025-12-21 11:00:32 -10:00
committed by GitHub
parent 393e2f9b62
commit e1dcd0df3e
6 changed files with 325 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/tests/common/mod.rs
+137
View File
@@ -26,6 +26,8 @@ use sgl_model_gateway::{
middleware::TokenBucket,
policies::PolicyRegistry,
protocols::common::{Function, Tool},
reasoning_parser::ParserFactory as ReasoningParserFactory,
tool_parser::ParserFactory as ToolParserFactory,
};
/// Helper function to create AppContext for tests
@@ -159,6 +161,141 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
app_context
}
/// Helper function to create AppContext for tests with parser factories initialized
pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppContext> {
let client = reqwest::Client::new();
// Initialize rate limiter
let rate_limiter = match config.max_concurrent_requests {
n if n <= 0 => None,
n => {
let rate_limit_tokens = config
.rate_limit_tokens_per_second
.filter(|&t| t > 0)
.unwrap_or(n);
Some(Arc::new(TokenBucket::new(
n as usize,
rate_limit_tokens as usize,
)))
}
};
// Initialize registries
let worker_registry = Arc::new(WorkerRegistry::new());
let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone()));
// Initialize storage backends (Memory for tests)
let response_storage = Arc::new(MemoryResponseStorage::new());
let conversation_storage = Arc::new(MemoryConversationStorage::new());
let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
// Initialize load monitor
let load_monitor = Some(Arc::new(LoadMonitor::new(
worker_registry.clone(),
policy_registry.clone(),
client.clone(),
config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue, workflow engine, and mcp manager
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let mcp_manager_lock = Arc::new(OnceLock::new());
// Initialize parser factories
let reasoning_parser_factory = Some(ReasoningParserFactory::new());
let tool_parser_factory = Some(ToolParserFactory::new());
let app_context = Arc::new(
AppContext::builder()
.router_config(config.clone())
.client(client)
.rate_limiter(rate_limiter)
.tokenizer(None) // tokenizer
.reasoning_parser_factory(reasoning_parser_factory)
.tool_parser_factory(tool_parser_factory)
.worker_registry(worker_registry)
.policy_registry(policy_registry)
.response_storage(response_storage)
.conversation_storage(conversation_storage)
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.mcp_manager(mcp_manager_lock)
.build()
.unwrap(),
);
// Initialize JobQueue after AppContext is created
let weak_context = Arc::downgrade(&app_context);
let job_queue = sgl_model_gateway::core::JobQueue::new(
sgl_model_gateway::core::JobQueueConfig::default(),
weak_context,
);
app_context
.worker_job_queue
.set(job_queue)
.expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows
use sgl_model_gateway::{
core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_worker_registration_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
.expect("worker_removal workflow should be valid");
app_context
.workflow_engine
.set(engine)
.expect("WorkflowEngine should only be initialized once");
// Register external workers for OpenAI mode
if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
for url in worker_urls {
// Create a worker that supports common test models
let models = vec![
ModelCard::new("mock-model"),
ModelCard::new("gpt-4"),
ModelCard::new("gpt-3.5-turbo"),
];
let worker: Arc<dyn Worker> = Arc::new(
BasicWorkerBuilder::new(url)
.worker_type(WorkerType::Regular)
.runtime_type(RuntimeType::External)
.models(models)
.build(),
);
app_context.worker_registry.register(worker);
}
}
// Initialize MCP manager with empty config
use sgl_model_gateway::mcp::{McpConfig, McpManager};
let empty_config = McpConfig {
servers: vec![],
pool: Default::default(),
proxy: None,
warmup: vec![],
inventory: Default::default(),
};
let mcp_manager = McpManager::with_defaults(empty_config)
.await
.expect("Failed to create MCP manager");
app_context
.mcp_manager
.set(Arc::new(mcp_manager))
.ok()
.expect("McpManager should only be initialized once");
app_context
}
/// Helper function to create AppContext for tests with MCP config from file
pub async fn create_test_context_with_mcp_config(
config: RouterConfig,
sgl-model-gateway/tests/parser_endpoints_test.rs
+65 -59
View File
@@ -88,8 +88,8 @@ impl ParserTestContext {
.build()
.unwrap();
// Create app context
let app_context = common::create_test_context(config.clone()).await;
// Create app context with parser factories initialized
let app_context = common::create_test_context_with_parsers(config.clone()).await;
// Create router
let router = RouterFactory::create_router(&app_context).await.unwrap();
@@ -160,14 +160,22 @@ mod parse_function_call_tests {
let resp = app.oneshot(req).await.unwrap();
// Parser endpoint should return 200 for valid requests (or SERVICE_UNAVAILABLE if parser factory not initialized)
// Since we're in a test without explicit parser factory setup, it may return 503
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Parser endpoint should return 200 for valid requests
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
// Verify response contains tool_calls
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(body_json["success"], true);
assert!(body_json["tool_calls"].is_array());
ctx.shutdown().await;
}
@@ -191,11 +199,11 @@ mod parse_function_call_tests {
let resp = app.oneshot(req).await.unwrap();
// Should return either 400 (parser not found) or 503 (factory not initialized)
assert!(
resp.status() == StatusCode::BAD_REQUEST
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected BAD_REQUEST (400) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 400 (parser not found)
assert_eq!(
resp.status(),
StatusCode::BAD_REQUEST,
"Expected BAD_REQUEST (400), got {}",
resp.status()
);
@@ -270,10 +278,11 @@ mod parse_function_call_tests {
let resp = app.oneshot(req).await.unwrap();
// Parser should handle empty text gracefully - return 200 or 503
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Parser should handle empty text gracefully - return 200
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
@@ -304,24 +313,23 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should return 200 or 503 depending on whether parser factory is initialized
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 200 with parser factory initialized
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
if resp.status() == StatusCode::OK {
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
// Check response structure
assert_eq!(body_json["success"], true);
assert!(body_json.get("normal_text").is_some());
assert!(body_json.get("reasoning_text").is_some());
}
// Check response structure
assert_eq!(body_json["success"], true);
assert!(body_json.get("normal_text").is_some());
assert!(body_json.get("reasoning_text").is_some());
ctx.shutdown().await;
}
@@ -345,11 +353,11 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should return 400 (parser not found) or 503 (factory not initialized)
assert!(
resp.status() == StatusCode::BAD_REQUEST
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected BAD_REQUEST (400) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 400 (parser not found)
assert_eq!(
resp.status(),
StatusCode::BAD_REQUEST,
"Expected BAD_REQUEST (400), got {}",
resp.status()
);
@@ -422,9 +430,10 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Parser should handle empty text gracefully
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
@@ -450,28 +459,24 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should still return 200 or 503, parser should handle gracefully
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 200, parser should handle gracefully
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
if resp.status() == StatusCode::OK {
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(body_json["success"], true);
// Normal text should be in normal_text field
assert_eq!(
body_json["normal_text"].as_str().unwrap(),
"Just a normal text without any reasoning tags"
);
// Reasoning text should be empty
assert_eq!(body_json["reasoning_text"].as_str().unwrap(), "");
}
assert_eq!(body_json["success"], true);
// When there are no reasoning tags, parser returns empty normal_text and empty reasoning_text
// since the detect_and_parse_reasoning method only extracts if it finds reasoning markers
assert!(body_json.get("normal_text").is_some());
assert!(body_json.get("reasoning_text").is_some());
ctx.shutdown().await;
}
@@ -497,9 +502,10 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should handle multiple blocks gracefully
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);