From a6cf8d6a2581918e3bb593bc6e4e3f3df354e17c Mon Sep 17 00:00:00 2001 From: pyro Date: Wed, 28 Oct 2026 10:36:50 -0500 Subject: [PATCH] updated the server so that you can set a preferred llm server. It will attempt to utilze that server for every request unless its not online then it will fallback to any other online serer. --- src/main.rs | 81 +++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 66 insertions(+), 15 deletions(-) diff --git a/src/main.rs b/src/main.rs index 086721f..750d48c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -26,7 +26,7 @@ use std::{ #[command( name = "llm-router", version, - about = "LLM Router Proxy with System Tray and Systemd Integration" + about = "LLM Router Proxy with System Tray, Systemd Integration, and Preferred Backend" )] struct Args { /// Run the proxy service in daemon mode (non-interactive, used by systemd) @@ -46,6 +46,8 @@ struct Args { struct Config { listen_port: u16, backend_urls: Vec, + #[serde(default)] + preferred_backend: Option, } #[derive(Clone)] @@ -57,12 +59,15 @@ struct Backend { #[derive(Clone)] struct AppState { backends: Vec, + preferred_backend: Option, client: reqwest::Client, counter: Arc, } + #[derive(Debug)] struct LLMRouterTray { backend_data: Vec<(String, bool)>, + preferred_url: Option, } impl Tray for LLMRouterTray { @@ -87,7 +92,9 @@ impl Tray for LLMRouterTray { } else { "🔴 Offline" }; - desc.push_str(&format!("• {}: {}\n", url, status)); + let is_pref = self.preferred_url.as_ref() == Some(url); + let pref_tag = if is_pref { " [Preferred]" } else { "" }; + desc.push_str(&format!("• {}{}: {}\n", url, pref_tag, status)); } ksni::ToolTip { title: "LLM Router Status".into(), @@ -114,9 +121,15 @@ impl Tray for LLMRouterTray { } else { "🔴 Offline" }; + let is_pref = self.preferred_url.as_ref() == Some(url); + let label = if is_pref { + format!("★ {} [{}] (Preferred)", url, status_str) + } else { + format!("{} [{}]", url, status_str) + }; items.push( StandardItem { - label: format!("{} [{}]", url, status_str), + label, enabled: false, ..Default::default() } @@ -295,9 +308,26 @@ fn load_config(is_daemon: bool) -> Config { backend_urls.push(url); } + let mut preferred_backend = None; + if backend_urls.len() > 1 { + print!("Enter a preferred backend URL from the list above (leave blank for none): "); + io::stdout().flush().unwrap(); + let mut pref = String::new(); + io::stdin().read_line(&mut pref).unwrap(); + let pref = pref.trim().to_string(); + if !pref.is_empty() { + if backend_urls.contains(&pref) { + preferred_backend = Some(pref); + } else { + println!("Warning: Provided preferred URL was not in the backend list. Ignoring preference."); + } + } + } + let config = Config { listen_port, backend_urls, + preferred_backend, }; let toml_content = toml::to_string_pretty(&config).expect("Failed to serialize config"); @@ -343,17 +373,23 @@ async fn main() { } let tray = LLMRouterTray { backend_data: initial_backend_data, + preferred_url: config.preferred_backend.clone(), }; let tray_handle = tray.spawn().await.expect("Failed to start system tray"); let state = AppState { backends: backends.clone(), + preferred_backend: config.preferred_backend.clone(), client: client.clone(), counter: Arc::new(AtomicUsize::new(0)), }; + + let tray_handle_clone = tray_handle.clone(); + let pref_url_for_health = config.preferred_backend.clone(); tokio::spawn(async move { - health_check_loop(backends, client, tray_handle).await; + health_check_loop(backends, pref_url_for_health, client, tray_handle_clone).await; }); + let app = Router::new() .route("/*path", any(proxy_handler)) .fallback(any(proxy_handler)) @@ -370,6 +406,7 @@ async fn main() { async fn health_check_loop( backends: Vec, + preferred_url: Option, client: reqwest::Client, tray_handle: ksni::Handle, ) { @@ -397,20 +434,33 @@ async fn health_check_loop( .iter() .map(|b| (b.url.clone(), b.healthy.load(Ordering::Relaxed))) .collect(); + let pref_url_clone = preferred_url.clone(); tray_handle .update(move |tray: &mut LLMRouterTray| { tray.backend_data = snapshot; + tray.preferred_url = pref_url_clone; }) .await; tokio::time::sleep(Duration::from_secs(5)).await; } } -fn select_backend(backends: &[Backend], counter: &AtomicUsize) -> Option { +fn select_backend( + backends: &[Backend], + preferred_backend: &Option, + counter: &AtomicUsize, +) -> Option { let len = backends.len(); if len == 0 { return None; } + if let Some(pref) = preferred_backend { + if let Some(backend) = backends.iter().find(|b| &b.url == pref) { + if backend.healthy.load(Ordering::Relaxed) { + return Some(backend.url.clone()); + } + } + } let start_idx = counter.fetch_add(1, Ordering::Relaxed); for i in 0..len { let idx = (start_idx + i) % len; @@ -422,16 +472,17 @@ fn select_backend(backends: &[Backend], counter: &AtomicUsize) -> Option } async fn proxy_handler(State(state): State, req: Request) -> Response { - let backend_url = match select_backend(&state.backends, &state.counter) { - Some(url) => url, - None => { - return ( - StatusCode::SERVICE_UNAVAILABLE, - "All LLM backends are offline", - ) - .into_response() - } - }; + let backend_url = + match select_backend(&state.backends, &state.preferred_backend, &state.counter) { + Some(url) => url, + None => { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "All LLM backends are offline", + ) + .into_response() + } + }; let path = req.uri().path().to_string(); let query = req .uri()