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.
This commit is contained in:
pyro committed 2026-10-28 10:36:50 -05:00
1 parent b8fd08985f
commit a6cf8d6a25
1 file changed
+57 -6
+57 -6
View File
@@ -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<String>,
#[serde(default)]
preferred_backend: Option<String>,
}
#[derive(Clone)]
@@ -57,12 +59,15 @@ struct Backend {
#[derive(Clone)]
struct AppState {
backends: Vec<Backend>,
preferred_backend: Option<String>,
client: reqwest::Client,
counter: Arc<AtomicUsize>,
}
#[derive(Debug)]
struct LLMRouterTray {
backend_data: Vec<(String, bool)>,
preferred_url: Option<String>,
}
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<Backend>,
preferred_url: Option<String>,
client: reqwest::Client,
tray_handle: ksni::Handle<LLMRouterTray>,
) {
@@ -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<String> {
fn select_backend(
backends: &[Backend],
preferred_backend: &Option<String>,
counter: &AtomicUsize,
) -> Option<String> {
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,7 +472,8 @@ fn select_backend(backends: &[Backend], counter: &AtomicUsize) -> Option<String>
}
async fn proxy_handler(State(state): State<AppState>, req: Request<Body>) -> Response {
let backend_url = match select_backend(&state.backends, &state.counter) {
let backend_url =
match select_backend(&state.backends, &state.preferred_backend, &state.counter) {
Some(url) => url,
None => {
return (