make arch-router cluster optional (#518)

This commit is contained in:
Adil Hafeez 2025-07-08 00:33:40 -07:00 committed by GitHub
parent a212dd79da
commit 147908ba7e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 44 additions and 30 deletions

View file

@ -23,6 +23,8 @@ use tracing::{debug, info, warn};
pub mod router;
const BIND_ADDRESS: &str = "0.0.0.0:9091";
const DEFAULT_ROUTING_LLM_PROVIDER: &str = "arch-router";
const DEFAULT_ROUTING_MODEL_NAME: &str = "Arch-Router";
// Utility function to extract the context from the incoming request headers
fn extract_context_from_request(req: &Request<Incoming>) -> Context {
@ -69,16 +71,23 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
info!("listening on http://{}", bind_address);
let listener = TcpListener::bind(bind_address).await?;
let model = arch_config
let routing_model_name: String = arch_config
.routing
.as_ref()
.map(|r| r.model.clone())
.unwrap_or_else(|| "none".to_string());
.and_then(|r| r.model.clone())
.unwrap_or_else(|| DEFAULT_ROUTING_MODEL_NAME.to_string());
let routing_llm_provider = arch_config
.routing
.as_ref()
.and_then(|r| r.llm_provider.clone())
.unwrap_or_else(|| DEFAULT_ROUTING_LLM_PROVIDER.to_string());
let router_service: Arc<RouterService> = Arc::new(RouterService::new(
arch_config.llm_providers.clone(),
llm_provider_endpoint.clone(),
model,
routing_model_name,
routing_llm_provider,
));
loop {

View file

@ -17,7 +17,7 @@ pub struct RouterService {
router_url: String,
client: reqwest::Client,
router_model: Arc<dyn RouterModel>,
routing_model_name: String,
routing_provider_name: String,
llm_usage_defined: bool,
llm_provider_map: HashMap<String, LlmProvider>,
}
@ -41,6 +41,7 @@ impl RouterService {
providers: Vec<LlmProvider>,
router_url: String,
routing_model_name: String,
routing_provider_name: String,
) -> Self {
let providers_with_usage = providers
.iter()
@ -65,7 +66,7 @@ impl RouterService {
router_url,
client: reqwest::Client::new(),
router_model,
routing_model_name,
routing_provider_name,
llm_usage_defined: !providers_with_usage.is_empty(),
llm_provider_map,
}
@ -104,7 +105,7 @@ impl RouterService {
llm_route_request_headers.insert(
header::HeaderName::from_static(ARCH_PROVIDER_HINT_HEADER),
header::HeaderValue::from_str(&self.routing_model_name).unwrap(),
header::HeaderValue::from_str(&self.routing_provider_name).unwrap(),
);
if let Some(trace_parent) = trace_parent {