[model-gateway] change sgl-router to sgl-model-gateway (#14312)
This commit is contained in:
Simo Lin
@@ -0,0 +1,174 @@
|
||||
// tests/common/mock_mcp_server.rs - Mock MCP server for testing
|
||||
use rmcp::{
|
||||
handler::server::{router::tool::ToolRouter, wrapper::Parameters},
|
||||
model::*,
|
||||
service::RequestContext,
|
||||
tool, tool_handler, tool_router,
|
||||
transport::streamable_http_server::{
|
||||
session::local::LocalSessionManager, StreamableHttpService,
|
||||
},
|
||||
ErrorData as McpError, RoleServer, ServerHandler,
|
||||
};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// Mock MCP server that returns hardcoded responses for testing
|
||||
pub struct MockMCPServer {
|
||||
pub port: u16,
|
||||
pub server_handle: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
/// Simple test server with mock search tools
|
||||
#[derive(Clone)]
|
||||
pub struct MockSearchServer {
|
||||
tool_router: ToolRouter<MockSearchServer>,
|
||||
}
|
||||
|
||||
#[tool_router]
|
||||
impl MockSearchServer {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tool_router: Self::tool_router(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tool(description = "Mock web search tool")]
|
||||
fn brave_web_search(
|
||||
&self,
|
||||
Parameters(params): Parameters<serde_json::Map<String, serde_json::Value>>,
|
||||
) -> Result<CallToolResult, McpError> {
|
||||
let query = params
|
||||
.get("query")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("test");
|
||||
Ok(CallToolResult::success(vec![Content::text(format!(
|
||||
"Mock search results for: {}",
|
||||
query
|
||||
))]))
|
||||
}
|
||||
|
||||
#[tool(description = "Mock local search tool")]
|
||||
fn brave_local_search(
|
||||
&self,
|
||||
Parameters(_params): Parameters<serde_json::Map<String, serde_json::Value>>,
|
||||
) -> Result<CallToolResult, McpError> {
|
||||
Ok(CallToolResult::success(vec![Content::text(
|
||||
"Mock local search results",
|
||||
)]))
|
||||
}
|
||||
}
|
||||
|
||||
#[tool_handler]
|
||||
impl ServerHandler for MockSearchServer {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo {
|
||||
protocol_version: ProtocolVersion::V_2024_11_05,
|
||||
capabilities: ServerCapabilities::builder().enable_tools().build(),
|
||||
server_info: Implementation::from_build_env(),
|
||||
instructions: Some("Mock server for testing".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn initialize(
|
||||
&self,
|
||||
_request: InitializeRequestParam,
|
||||
_context: RequestContext<RoleServer>,
|
||||
) -> Result<InitializeResult, McpError> {
|
||||
Ok(self.get_info())
|
||||
}
|
||||
}
|
||||
|
||||
impl MockMCPServer {
|
||||
/// Start a mock MCP server on an available port
|
||||
pub async fn start() -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
|
||||
// Find an available port
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await?;
|
||||
let port = listener.local_addr()?.port();
|
||||
|
||||
// Create the MCP service using rmcp's StreamableHttpService
|
||||
let service = StreamableHttpService::new(
|
||||
|| Ok(MockSearchServer::new()),
|
||||
LocalSessionManager::default().into(),
|
||||
Default::default(),
|
||||
);
|
||||
|
||||
let app = axum::Router::new().nest_service("/mcp", service);
|
||||
|
||||
let server_handle = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("Mock MCP server failed to start");
|
||||
});
|
||||
|
||||
// Give the server a moment to start
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
||||
|
||||
Ok(MockMCPServer {
|
||||
port,
|
||||
server_handle: Some(server_handle),
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the full URL for this mock server
|
||||
pub fn url(&self) -> String {
|
||||
format!("http://127.0.0.1:{}/mcp", self.port)
|
||||
}
|
||||
|
||||
/// Stop the mock server
|
||||
pub async fn stop(&mut self) {
|
||||
if let Some(handle) = self.server_handle.take() {
|
||||
handle.abort();
|
||||
// Wait a moment for cleanup
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MockMCPServer {
|
||||
fn drop(&mut self) {
|
||||
if let Some(handle) = self.server_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[allow(unused_imports)]
|
||||
use super::MockMCPServer;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mock_server_startup() {
|
||||
let mut server = MockMCPServer::start().await.unwrap();
|
||||
assert!(server.port > 0);
|
||||
assert!(server.url().contains(&server.port.to_string()));
|
||||
server.stop().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mock_server_with_rmcp_client() {
|
||||
let mut server = MockMCPServer::start().await.unwrap();
|
||||
|
||||
use rmcp::{transport::StreamableHttpClientTransport, ServiceExt};
|
||||
|
||||
let transport = StreamableHttpClientTransport::from_uri(server.url().as_str());
|
||||
let client = ().serve(transport).await;
|
||||
|
||||
assert!(client.is_ok(), "Should be able to connect to mock server");
|
||||
|
||||
if let Ok(client) = client {
|
||||
let tools = client.peer().list_all_tools().await;
|
||||
assert!(tools.is_ok(), "Should be able to list tools");
|
||||
|
||||
if let Ok(tools) = tools {
|
||||
assert_eq!(tools.len(), 2, "Should have 2 tools");
|
||||
assert!(tools.iter().any(|t| t.name == "brave_web_search"));
|
||||
assert!(tools.iter().any(|t| t.name == "brave_local_search"));
|
||||
}
|
||||
|
||||
// Shutdown by dropping the client
|
||||
drop(client);
|
||||
}
|
||||
|
||||
server.stop().await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
//! Mock servers for testing
|
||||
|
||||
#![allow(dead_code)]
|
||||
|
||||
use std::{net::SocketAddr, sync::Arc};
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Request, State},
|
||||
http::{HeaderValue, StatusCode},
|
||||
response::{
|
||||
sse::{Event, KeepAlive},
|
||||
IntoResponse, Response, Sse,
|
||||
},
|
||||
routing::post,
|
||||
Json, Router,
|
||||
};
|
||||
use futures_util::stream::{self, StreamExt};
|
||||
use serde_json::json;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// Mock OpenAI API server for testing
|
||||
pub struct MockOpenAIServer {
|
||||
addr: SocketAddr,
|
||||
_handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockServerState {
|
||||
require_auth: bool,
|
||||
expected_auth: Option<String>,
|
||||
}
|
||||
|
||||
impl MockOpenAIServer {
|
||||
/// Create and start a new mock OpenAI server
|
||||
pub async fn new() -> Self {
|
||||
Self::new_with_auth(None).await
|
||||
}
|
||||
|
||||
/// Create and start a new mock OpenAI server with optional auth requirement
|
||||
pub async fn new_with_auth(expected_auth: Option<String>) -> Self {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
|
||||
let state = Arc::new(MockServerState {
|
||||
require_auth: expected_auth.is_some(),
|
||||
expected_auth,
|
||||
});
|
||||
|
||||
let app = Router::new()
|
||||
.route("/v1/chat/completions", post(mock_chat_completions))
|
||||
.route("/v1/completions", post(mock_completions))
|
||||
.route("/v1/models", post(mock_models).get(mock_models))
|
||||
.with_state(state);
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
// Give the server a moment to start
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
|
||||
|
||||
Self {
|
||||
addr,
|
||||
_handle: handle,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the base URL for this mock server
|
||||
pub fn base_url(&self) -> String {
|
||||
format!("http://{}", self.addr)
|
||||
}
|
||||
}
|
||||
|
||||
/// Mock chat completions endpoint
|
||||
async fn mock_chat_completions(req: Request<Body>) -> Response {
|
||||
let (_, body) = req.into_parts();
|
||||
let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
|
||||
Ok(bytes) => bytes,
|
||||
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
|
||||
};
|
||||
|
||||
let request: serde_json::Value = match serde_json::from_slice(&body_bytes) {
|
||||
Ok(req) => req,
|
||||
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
|
||||
};
|
||||
|
||||
// Extract model from request or use default (owned String to satisfy 'static in stream)
|
||||
let model: String = request
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("gpt-3.5-turbo")
|
||||
.to_string();
|
||||
|
||||
// If stream requested, return SSE
|
||||
let is_stream = request
|
||||
.get("stream")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
if is_stream {
|
||||
let created = 1677652288u64;
|
||||
// Single chunk then [DONE]
|
||||
let model_chunk = model.clone();
|
||||
let event_stream = stream::once(async move {
|
||||
let chunk = json!({
|
||||
"id": "chatcmpl-123456789",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_chunk,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"content": "Hello!"
|
||||
},
|
||||
"finish_reason": null
|
||||
}]
|
||||
});
|
||||
Ok::<_, std::convert::Infallible>(Event::default().data(chunk.to_string()))
|
||||
})
|
||||
.chain(stream::once(async { Ok(Event::default().data("[DONE]")) }));
|
||||
|
||||
Sse::new(event_stream)
|
||||
.keep_alive(KeepAlive::default())
|
||||
.into_response()
|
||||
} else {
|
||||
// Create a mock non-streaming response
|
||||
let response = json!({
|
||||
"id": "chatcmpl-123456789",
|
||||
"object": "chat.completion",
|
||||
"created": 1677652288,
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello! I'm a mock OpenAI assistant. How can I help you today?"
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 9,
|
||||
"completion_tokens": 12,
|
||||
"total_tokens": 21
|
||||
}
|
||||
});
|
||||
|
||||
Json(response).into_response()
|
||||
}
|
||||
}
|
||||
|
||||
/// Mock completions endpoint (legacy)
|
||||
async fn mock_completions(req: Request<Body>) -> Response {
|
||||
let (_, body) = req.into_parts();
|
||||
let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
|
||||
Ok(bytes) => bytes,
|
||||
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
|
||||
};
|
||||
|
||||
let request: serde_json::Value = match serde_json::from_slice(&body_bytes) {
|
||||
Ok(req) => req,
|
||||
Err(_) => return StatusCode::BAD_REQUEST.into_response(),
|
||||
};
|
||||
|
||||
let model = request["model"].as_str().unwrap_or("text-davinci-003");
|
||||
|
||||
let response = json!({
|
||||
"id": "cmpl-123456789",
|
||||
"object": "text_completion",
|
||||
"created": 1677652288,
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"text": " This is a mock completion response.",
|
||||
"index": 0,
|
||||
"logprobs": null,
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"completion_tokens": 7,
|
||||
"total_tokens": 12
|
||||
}
|
||||
});
|
||||
|
||||
Json(response).into_response()
|
||||
}
|
||||
|
||||
/// Mock models endpoint
|
||||
async fn mock_models(State(state): State<Arc<MockServerState>>, req: Request<Body>) -> Response {
|
||||
// Optionally enforce Authorization header
|
||||
if state.require_auth {
|
||||
let auth = req
|
||||
.headers()
|
||||
.get("authorization")
|
||||
.or_else(|| req.headers().get("Authorization"))
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
let auth_ok = match (&state.expected_auth, auth) {
|
||||
(Some(expected), Some(got)) => &got == expected,
|
||||
(None, Some(_)) => true,
|
||||
_ => false,
|
||||
};
|
||||
if !auth_ok {
|
||||
let mut response = Response::new(Body::from(
|
||||
json!({
|
||||
"error": {
|
||||
"message": "Unauthorized",
|
||||
"type": "invalid_request_error"
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
));
|
||||
*response.status_mut() = StatusCode::UNAUTHORIZED;
|
||||
response
|
||||
.headers_mut()
|
||||
.insert("WWW-Authenticate", HeaderValue::from_static("Bearer"));
|
||||
return response;
|
||||
}
|
||||
}
|
||||
|
||||
let response = json!({
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "gpt-4",
|
||||
"object": "model",
|
||||
"created": 1677610602,
|
||||
"owned_by": "openai"
|
||||
},
|
||||
{
|
||||
"id": "gpt-3.5-turbo",
|
||||
"object": "model",
|
||||
"created": 1677610602,
|
||||
"owned_by": "openai"
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
Json(response).into_response()
|
||||
}
|
||||
+1271
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,659 @@
|
||||
// These modules are used by tests and benchmarks
|
||||
#![allow(dead_code)]
|
||||
|
||||
pub mod mock_mcp_server;
|
||||
pub mod mock_openai_server;
|
||||
pub mod mock_worker;
|
||||
pub mod streaming_helpers;
|
||||
pub mod test_app;
|
||||
|
||||
use std::{
|
||||
fs,
|
||||
path::PathBuf,
|
||||
sync::{Arc, Mutex, OnceLock},
|
||||
};
|
||||
|
||||
use serde_json::json;
|
||||
use sgl_model_gateway::{
|
||||
app_context::AppContext,
|
||||
config::{RouterConfig, RoutingMode},
|
||||
core::{
|
||||
BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType,
|
||||
},
|
||||
data_connector::{
|
||||
MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
|
||||
},
|
||||
middleware::TokenBucket,
|
||||
policies::PolicyRegistry,
|
||||
protocols::common::{Function, Tool},
|
||||
};
|
||||
|
||||
/// Helper function to create AppContext for tests
|
||||
pub async fn create_test_context(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());
|
||||
|
||||
let app_context = Arc::new(
|
||||
AppContext::builder()
|
||||
.router_config(config.clone())
|
||||
.client(client)
|
||||
.rate_limiter(rate_limiter)
|
||||
.tokenizer(None) // tokenizer
|
||||
.reasoning_parser_factory(None) // reasoning_parser_factory
|
||||
.tool_parser_factory(None) // 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::workflow::{
|
||||
create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
|
||||
};
|
||||
let engine = Arc::new(WorkflowEngine::new());
|
||||
engine.register_workflow(create_worker_registration_workflow(&config));
|
||||
engine.register_workflow(create_worker_removal_workflow());
|
||||
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,
|
||||
mcp_config_path: &str,
|
||||
) -> Arc<AppContext> {
|
||||
use sgl_model_gateway::mcp::{McpConfig, McpManager};
|
||||
|
||||
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());
|
||||
|
||||
let app_context = Arc::new(
|
||||
AppContext::builder()
|
||||
.router_config(config.clone())
|
||||
.client(client)
|
||||
.rate_limiter(rate_limiter)
|
||||
.tokenizer(None) // tokenizer
|
||||
.reasoning_parser_factory(None) // reasoning_parser_factory
|
||||
.tool_parser_factory(None) // 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::workflow::{
|
||||
create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
|
||||
};
|
||||
let engine = Arc::new(WorkflowEngine::new());
|
||||
engine.register_workflow(create_worker_registration_workflow(&config));
|
||||
engine.register_workflow(create_worker_removal_workflow());
|
||||
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 from config file
|
||||
let mcp_config = McpConfig::from_file(mcp_config_path)
|
||||
.await
|
||||
.expect("Failed to load MCP config from file");
|
||||
let mcp_manager = McpManager::with_defaults(mcp_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
|
||||
}
|
||||
|
||||
// Tokenizer download configuration
|
||||
const TINYLLAMA_TOKENIZER_URL: &str =
|
||||
"https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer.json";
|
||||
const CACHE_DIR: &str = ".tokenizer_cache";
|
||||
const TINYLLAMA_TOKENIZER_FILENAME: &str = "tinyllama_tokenizer.json";
|
||||
|
||||
// Global mutex to prevent concurrent downloads
|
||||
static DOWNLOAD_MUTEX: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
|
||||
/// Downloads the TinyLlama tokenizer from HuggingFace if not already cached.
|
||||
/// Returns the path to the cached tokenizer file.
|
||||
///
|
||||
/// This function is thread-safe and will only download the tokenizer once
|
||||
/// even if called from multiple threads concurrently.
|
||||
pub fn ensure_tokenizer_cached() -> PathBuf {
|
||||
// Get or initialize the mutex
|
||||
let mutex = DOWNLOAD_MUTEX.get_or_init(|| Mutex::new(()));
|
||||
|
||||
// Lock to ensure only one thread downloads at a time
|
||||
let _guard = mutex.lock().unwrap();
|
||||
|
||||
let cache_dir = PathBuf::from(CACHE_DIR);
|
||||
let tokenizer_path = cache_dir.join(TINYLLAMA_TOKENIZER_FILENAME);
|
||||
|
||||
// Create cache directory if it doesn't exist
|
||||
if !cache_dir.exists() {
|
||||
fs::create_dir_all(&cache_dir).expect("Failed to create cache directory");
|
||||
}
|
||||
|
||||
// Download tokenizer if not already cached
|
||||
if !tokenizer_path.exists() {
|
||||
println!("Downloading TinyLlama tokenizer from HuggingFace...");
|
||||
|
||||
// Use blocking reqwest client since we're in tests/benchmarks
|
||||
let client = reqwest::blocking::Client::new();
|
||||
let response = client
|
||||
.get(TINYLLAMA_TOKENIZER_URL)
|
||||
.send()
|
||||
.expect("Failed to download tokenizer");
|
||||
|
||||
if !response.status().is_success() {
|
||||
panic!("Failed to download tokenizer: HTTP {}", response.status());
|
||||
}
|
||||
|
||||
let content = response.bytes().expect("Failed to read tokenizer content");
|
||||
|
||||
if content.len() < 100 {
|
||||
panic!("Downloaded content too small: {} bytes", content.len());
|
||||
}
|
||||
|
||||
fs::write(&tokenizer_path, content).expect("Failed to write tokenizer to cache");
|
||||
println!(
|
||||
"Tokenizer downloaded and cached successfully ({} bytes)",
|
||||
tokenizer_path.metadata().unwrap().len()
|
||||
);
|
||||
}
|
||||
|
||||
tokenizer_path
|
||||
}
|
||||
|
||||
/// Common test prompts for consistency across tests
|
||||
pub const TEST_PROMPTS: [&str; 4] = [
|
||||
"deep learning is",
|
||||
"Deep learning is",
|
||||
"has anyone seen nemo lately",
|
||||
"another prompt",
|
||||
];
|
||||
|
||||
/// Pre-computed hashes for verification
|
||||
pub const EXPECTED_HASHES: [u64; 4] = [
|
||||
1209591529327510910,
|
||||
4181375434596349981,
|
||||
6245658446118930933,
|
||||
5097285695902185237,
|
||||
];
|
||||
|
||||
/// Create a comprehensive set of test tools covering all parser test scenarios
|
||||
#[allow(dead_code)]
|
||||
pub fn create_test_tools() -> Vec<Tool> {
|
||||
vec![
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "search".to_string(),
|
||||
description: Some("Search for information".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather information".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"location": {"type": "string"},
|
||||
"date": {"type": "string"},
|
||||
"units": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "calculate".to_string(),
|
||||
description: Some("Perform calculations".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"x": {"type": "number"},
|
||||
"y": {"type": "number"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "translate".to_string(),
|
||||
description: Some("Translate text".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {"type": "string"},
|
||||
"to": {"type": "string"},
|
||||
"target_lang": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "get_time".to_string(),
|
||||
description: Some("Get current time".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"timezone": {"type": "string"},
|
||||
"format": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "get_current_time".to_string(),
|
||||
description: Some("Get current time".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"timezone": {"type": "string"},
|
||||
"format": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "update_settings".to_string(),
|
||||
description: Some("Update settings".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"preferences": {"type": "object"},
|
||||
"notifications": {"type": "boolean"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "ping".to_string(),
|
||||
description: Some("Ping service".to_string()),
|
||||
parameters: json!({"type": "object", "properties": {}}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "test".to_string(),
|
||||
description: Some("Test function".to_string()),
|
||||
parameters: json!({"type": "object", "properties": {}}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "process".to_string(),
|
||||
description: Some("Process data".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": {"type": "number"},
|
||||
"rate": {"type": "number"},
|
||||
"enabled": {"type": "boolean"},
|
||||
"data": {"type": "object"},
|
||||
"text": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "web_search".to_string(),
|
||||
description: Some("Search the web".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string"},
|
||||
"num_results": {"type": "number"},
|
||||
"search_type": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "get_tourist_attractions".to_string(),
|
||||
description: Some("Get tourist attractions".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "config".to_string(),
|
||||
description: Some("Configuration function".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"debug": {"type": "boolean"},
|
||||
"verbose": {"type": "boolean"},
|
||||
"optional": {"type": "null"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "test_func".to_string(),
|
||||
description: Some("Test function".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bool_true": {"type": "boolean"},
|
||||
"bool_false": {"type": "boolean"},
|
||||
"none_val": {"type": "null"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "create".to_string(),
|
||||
description: Some("Create resource".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"email": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "add".to_string(),
|
||||
description: Some("Add operation".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"x": {"type": "number"},
|
||||
"y": {"type": "number"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "calc".to_string(),
|
||||
description: Some("Calculate".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"x": {"type": "number"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "func1".to_string(),
|
||||
description: Some("Function 1".to_string()),
|
||||
parameters: json!({"type": "object", "properties": {}}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "func2".to_string(),
|
||||
description: Some("Function 2".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"y": {"type": "number"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "tool1".to_string(),
|
||||
description: Some("Tool 1".to_string()),
|
||||
parameters: json!({"type": "object", "properties": {}}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "tool2".to_string(),
|
||||
description: Some("Tool 2".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"y": {"type": "number"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
//! Streaming Test Helpers
|
||||
//!
|
||||
//! Utilities for creating realistic streaming chunks that simulate
|
||||
//! how LLM tokens actually arrive (1-5 characters at a time).
|
||||
|
||||
/// Split input into realistic char-level chunks (2-3 chars each for determinism)
|
||||
pub fn create_realistic_chunks(input: &str) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
let chars: Vec<char> = input.chars().collect();
|
||||
let mut i = 0;
|
||||
|
||||
while i < chars.len() {
|
||||
// Take 2-3 characters at a time (deterministic for testing)
|
||||
let chunk_size = if i + 3 <= chars.len() && chars[i].is_ascii_alphanumeric() {
|
||||
3 // Longer chunks for alphanumeric sequences
|
||||
} else {
|
||||
2 // Shorter chunks for special characters
|
||||
};
|
||||
|
||||
let end = (i + chunk_size).min(chars.len());
|
||||
let chunk: String = chars[i..end].iter().collect();
|
||||
chunks.push(chunk);
|
||||
i = end;
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Split input at strategic positions to test edge cases
|
||||
/// This creates chunks that break at critical positions like after quotes, colons, etc.
|
||||
pub fn create_strategic_chunks(input: &str) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
let mut current = String::new();
|
||||
let chars: Vec<char> = input.chars().collect();
|
||||
|
||||
for (i, &ch) in chars.iter().enumerate() {
|
||||
current.push(ch);
|
||||
|
||||
// Break after strategic characters
|
||||
let should_break = matches!(ch, '"' | ':' | ',' | '{' | '}' | '[' | ']')
|
||||
|| (i > 0 && chars[i-1] == '"' && ch == ' ') // Space after quote
|
||||
|| current.len() >= 5; // Max 5 chars per chunk
|
||||
|
||||
if should_break && !current.is_empty() {
|
||||
chunks.push(current.clone());
|
||||
current.clear();
|
||||
}
|
||||
}
|
||||
|
||||
if !current.is_empty() {
|
||||
chunks.push(current);
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Create the bug scenario chunks: `{"name": "` arrives in parts
|
||||
pub fn create_bug_scenario_chunks() -> Vec<&'static str> {
|
||||
vec![
|
||||
r#"{"#,
|
||||
r#"""#,
|
||||
r#"name"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"""#, // Bug occurs here: parser has {"name": "
|
||||
r#"search"#, // Use valid tool name
|
||||
r#"""#,
|
||||
r#","#,
|
||||
r#" "#,
|
||||
r#"""#,
|
||||
r#"arguments"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"{"#,
|
||||
r#"""#,
|
||||
r#"query"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"""#,
|
||||
r#"test query"#,
|
||||
r#"""#,
|
||||
r#"}"#,
|
||||
r#"}"#,
|
||||
]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[allow(unused_imports)]
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_realistic_chunks() {
|
||||
let input = r#"{"name": "test"}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
// Should have multiple chunks
|
||||
assert!(chunks.len() > 3);
|
||||
|
||||
// Reconstructed should equal original
|
||||
let reconstructed: String = chunks.join("");
|
||||
assert_eq!(reconstructed, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strategic_chunks_breaks_after_quotes() {
|
||||
let input = r#"{"name": "value"}"#;
|
||||
let chunks = create_strategic_chunks(input);
|
||||
|
||||
// Should break after quotes and colons
|
||||
assert!(chunks.iter().any(|c| c.ends_with('"')));
|
||||
assert!(chunks.iter().any(|c| c.ends_with(':')));
|
||||
|
||||
// Reconstructed should equal original
|
||||
let reconstructed: String = chunks.join("");
|
||||
assert_eq!(reconstructed, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug_scenario_chunks() {
|
||||
let chunks = create_bug_scenario_chunks();
|
||||
let reconstructed: String = chunks.join("");
|
||||
|
||||
// Should reconstruct to valid JSON
|
||||
assert!(reconstructed.contains(r#"{"name": "search""#));
|
||||
|
||||
// The critical chunk sequence should be present (space after colon, then quote in next chunk)
|
||||
let joined = chunks.join("|");
|
||||
assert!(joined.contains(r#" |"#)); // The bug happens at {"name": " and then "
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,260 @@
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use axum::Router;
|
||||
use reqwest::Client;
|
||||
use sgl_model_gateway::{
|
||||
app_context::AppContext,
|
||||
config::RouterConfig,
|
||||
core::{
|
||||
BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType,
|
||||
},
|
||||
data_connector::{
|
||||
MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
|
||||
},
|
||||
mcp::{McpConfig, McpManager},
|
||||
middleware::{AuthConfig, TokenBucket},
|
||||
policies::PolicyRegistry,
|
||||
routers::RouterTrait,
|
||||
server::{build_app, AppState},
|
||||
};
|
||||
|
||||
/// Create a test Axum application using the actual server's build_app function
|
||||
#[allow(dead_code)]
|
||||
pub fn create_test_app(
|
||||
router: Arc<dyn RouterTrait>,
|
||||
client: Client,
|
||||
router_config: &RouterConfig,
|
||||
) -> Router {
|
||||
// Initialize rate limiter
|
||||
let rate_limiter = match router_config.max_concurrent_requests {
|
||||
n if n <= 0 => None,
|
||||
n => {
|
||||
let rate_limit_tokens = router_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(router_config.policy.clone()));
|
||||
|
||||
// Initialize storage backends
|
||||
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(),
|
||||
router_config.worker_startup_check_interval_secs,
|
||||
)));
|
||||
|
||||
// Create empty OnceLock for worker job queue and workflow engine
|
||||
let worker_job_queue = Arc::new(OnceLock::new());
|
||||
let workflow_engine = Arc::new(OnceLock::new());
|
||||
|
||||
// Create AppContext using builder pattern
|
||||
let app_context = Arc::new(
|
||||
AppContext::builder()
|
||||
.router_config(router_config.clone())
|
||||
.client(client)
|
||||
.rate_limiter(rate_limiter)
|
||||
.tokenizer(None) // tokenizer
|
||||
.reasoning_parser_factory(None) // reasoning_parser_factory
|
||||
.tool_parser_factory(None) // 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)
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
// Create AppState with the test router and context
|
||||
let app_state = Arc::new(AppState {
|
||||
router,
|
||||
context: app_context,
|
||||
concurrency_queue_tx: None,
|
||||
router_manager: None,
|
||||
});
|
||||
|
||||
// Configure request ID headers (use defaults if not specified)
|
||||
let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
|
||||
vec![
|
||||
"x-request-id".to_string(),
|
||||
"x-correlation-id".to_string(),
|
||||
"x-trace-id".to_string(),
|
||||
"request-id".to_string(),
|
||||
]
|
||||
});
|
||||
|
||||
// Create auth config from router config
|
||||
let auth_config = AuthConfig {
|
||||
api_key: router_config.api_key.clone(),
|
||||
};
|
||||
|
||||
// Use the actual server's build_app function
|
||||
build_app(
|
||||
app_state,
|
||||
auth_config,
|
||||
router_config.max_payload_size,
|
||||
request_id_headers,
|
||||
router_config.cors_allowed_origins.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Create a test Axum application with an existing AppContext
|
||||
#[allow(dead_code)]
|
||||
pub fn create_test_app_with_context(
|
||||
router: Arc<dyn RouterTrait>,
|
||||
app_context: Arc<AppContext>,
|
||||
) -> Router {
|
||||
// Create AppState with the test router and context
|
||||
let app_state = Arc::new(AppState {
|
||||
router,
|
||||
context: app_context.clone(),
|
||||
concurrency_queue_tx: None,
|
||||
router_manager: None,
|
||||
});
|
||||
|
||||
// Get config from the context
|
||||
let router_config = &app_context.router_config;
|
||||
|
||||
// Configure request ID headers (use defaults if not specified)
|
||||
let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| {
|
||||
vec![
|
||||
"x-request-id".to_string(),
|
||||
"x-correlation-id".to_string(),
|
||||
"x-trace-id".to_string(),
|
||||
"request-id".to_string(),
|
||||
]
|
||||
});
|
||||
|
||||
// Create auth config from router config
|
||||
let auth_config = AuthConfig {
|
||||
api_key: router_config.api_key.clone(),
|
||||
};
|
||||
|
||||
// Use the actual server's build_app function
|
||||
build_app(
|
||||
app_state,
|
||||
auth_config,
|
||||
router_config.max_payload_size,
|
||||
request_id_headers,
|
||||
router_config.cors_allowed_origins.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Create a minimal test AppContext for unit tests
|
||||
#[allow(dead_code)]
|
||||
pub async fn create_test_app_context() -> Arc<AppContext> {
|
||||
let router_config = RouterConfig::default();
|
||||
let client = Client::new();
|
||||
|
||||
// Initialize empty OnceLocks
|
||||
let worker_job_queue = Arc::new(OnceLock::new());
|
||||
let workflow_engine = Arc::new(OnceLock::new());
|
||||
|
||||
// Initialize MCP manager with empty config
|
||||
let mcp_manager_lock = Arc::new(OnceLock::new());
|
||||
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");
|
||||
mcp_manager_lock.set(Arc::new(mcp_manager)).ok();
|
||||
|
||||
// Initialize registries
|
||||
let worker_registry = Arc::new(WorkerRegistry::new());
|
||||
let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
|
||||
|
||||
// Initialize storage backends
|
||||
let response_storage = Arc::new(MemoryResponseStorage::new());
|
||||
let conversation_storage = Arc::new(MemoryConversationStorage::new());
|
||||
let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
|
||||
|
||||
Arc::new(
|
||||
AppContext::builder()
|
||||
.router_config(router_config)
|
||||
.client(client)
|
||||
.rate_limiter(None)
|
||||
.tokenizer(None)
|
||||
.reasoning_parser_factory(None)
|
||||
.tool_parser_factory(None)
|
||||
.worker_registry(worker_registry)
|
||||
.policy_registry(policy_registry)
|
||||
.response_storage(response_storage)
|
||||
.conversation_storage(conversation_storage)
|
||||
.conversation_item_storage(conversation_item_storage)
|
||||
.load_monitor(None)
|
||||
.worker_job_queue(worker_job_queue)
|
||||
.workflow_engine(workflow_engine)
|
||||
.mcp_manager(mcp_manager_lock)
|
||||
.build()
|
||||
.unwrap(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Register an external worker (OpenAI-compatible API endpoint) in the test AppContext.
|
||||
///
|
||||
/// This is used by tests that need to test the OpenAI router, which expects
|
||||
/// workers to be registered in the WorkerRegistry before routing requests.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `ctx` - The AppContext to register the worker in
|
||||
/// * `url` - The base URL of the external API endpoint
|
||||
/// * `models` - Optional list of model IDs this worker supports. If empty, uses "gpt-3.5-turbo" as default.
|
||||
#[allow(dead_code)]
|
||||
pub fn register_external_worker(ctx: &Arc<AppContext>, url: &str, models: Option<Vec<&str>>) {
|
||||
let model_list: Vec<ModelCard> = models
|
||||
.unwrap_or_else(|| vec!["gpt-3.5-turbo"])
|
||||
.into_iter()
|
||||
.map(ModelCard::new)
|
||||
.collect();
|
||||
|
||||
let worker: Arc<dyn Worker> = Arc::new(
|
||||
BasicWorkerBuilder::new(url)
|
||||
.worker_type(WorkerType::Regular)
|
||||
.runtime_type(RuntimeType::External)
|
||||
.models(model_list)
|
||||
.build(),
|
||||
);
|
||||
|
||||
ctx.worker_registry.register(worker);
|
||||
}
|
||||
|
||||
/// Register an external worker with a custom model card that has aliases.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `ctx` - The AppContext to register the worker in
|
||||
/// * `url` - The base URL of the external API endpoint
|
||||
/// * `model_card` - A fully configured ModelCard with aliases, provider, etc.
|
||||
#[allow(dead_code)]
|
||||
pub fn register_external_worker_with_card(ctx: &Arc<AppContext>, url: &str, model_card: ModelCard) {
|
||||
let worker: Arc<dyn Worker> = Arc::new(
|
||||
BasicWorkerBuilder::new(url)
|
||||
.worker_type(WorkerType::Regular)
|
||||
.runtime_type(RuntimeType::External)
|
||||
.model(model_card)
|
||||
.build(),
|
||||
);
|
||||
|
||||
ctx.worker_registry.register(worker);
|
||||
}
|
||||
Reference in New Issue
Block a user