diff --git a/.gitignore b/.gitignore index 0b188bc..e2aab4f 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,8 @@ target/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + + +# Added by cargo + +/target diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..c5407b9 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "llm-router" +version = "0.1.0" +edition = "2021" + +[dependencies] +axum = "0.7" +tokio = { version = "1.0", features = ["full"] } +reqwest = { version = "0.12", features = ["stream"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +toml = "0.8" +dirs = "5.0" +futures-util = "0.3" +ksni = "0.3" +clap = { version = "4.4", features = ["derive"] } diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..086721f --- /dev/null +++ b/src/main.rs @@ -0,0 +1,469 @@ +use axum::{ + body::Body, + extract::State, + http::{Request, StatusCode}, + response::{IntoResponse, Response}, + routing::any, + Router, +}; +use clap::Parser; +use futures_util::StreamExt; +use ksni::{Tray, TrayMethods}; +use serde::{Deserialize, Serialize}; +use std::{ + fs, + io::{self, Write}, + path::PathBuf, + process::Command, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, + }, + time::Duration, +}; + +#[derive(Parser, Debug)] +#[command( + name = "llm-router", + version, + about = "LLM Router Proxy with System Tray and Systemd Integration" +)] +struct Args { + /// Run the proxy service in daemon mode (non-interactive, used by systemd) + #[arg(long)] + daemon: bool, + + /// Start the installed systemd user service + #[arg(long)] + start_service: bool, + + /// Uninstall the service, remove configuration, and clean up systemd files + #[arg(long)] + uninstall: bool, +} + +#[derive(Serialize, Deserialize, Clone)] +struct Config { + listen_port: u16, + backend_urls: Vec, +} + +#[derive(Clone)] +struct Backend { + url: String, + healthy: Arc, +} + +#[derive(Clone)] +struct AppState { + backends: Vec, + client: reqwest::Client, + counter: Arc, +} +#[derive(Debug)] +struct LLMRouterTray { + backend_data: Vec<(String, bool)>, +} + +impl Tray for LLMRouterTray { + fn id(&self) -> String { + "llm-router".into() + } + + fn icon_name(&self) -> String { + "cpu".into() + } + + fn title(&self) -> String { + let online_count = self.backend_data.iter().filter(|(_, h)| *h).count(); + format!("LLM Router ({}/{})", online_count, self.backend_data.len()) + } + + fn tool_tip(&self) -> ksni::ToolTip { + let mut desc = String::from("Managed LLM Endpoints:\n"); + for (url, healthy) in &self.backend_data { + let status = if *healthy { + "🟢 Online" + } else { + "🔴 Offline" + }; + desc.push_str(&format!("• {}: {}\n", url, status)); + } + ksni::ToolTip { + title: "LLM Router Status".into(), + description: desc.trim().into(), + ..Default::default() + } + } + + fn menu(&self) -> Vec> { + use ksni::menu::*; + let mut items = vec![]; + items.push( + StandardItem { + label: "Live Endpoint Statuses:".into(), + enabled: false, + ..Default::default() + } + .into(), + ); + items.push(MenuItem::Separator); + for (url, healthy) in &self.backend_data { + let status_str = if *healthy { + "🟢 Online" + } else { + "🔴 Offline" + }; + items.push( + StandardItem { + label: format!("{} [{}]", url, status_str), + enabled: false, + ..Default::default() + } + .into(), + ); + } + items.push(MenuItem::Separator); + items.push( + StandardItem { + label: "Quit Router".into(), + icon_name: "application-exit".into(), + activate: Box::new(|_| std::process::exit(0)), + ..Default::default() + } + .into(), + ); + items + } +} + +fn get_config_path() -> PathBuf { + let mut path = dirs::config_dir().unwrap_or_else(|| PathBuf::from(".")); + path.push("llm-router"); + fs::create_dir_all(&path).unwrap(); + path.push("config.toml"); + path +} + +fn get_service_path() -> PathBuf { + let mut path = dirs::config_dir().unwrap_or_else(|| PathBuf::from(".")); + path.push("systemd"); + path.push("user"); + path.push("llm-router.service"); + path +} + +fn uninstall_service() { + println!("--- Uninstalling LLM Router ---"); + let _ = Command::new("systemctl") + .args(["--user", "stop", "llm-router.service"]) + .status(); + println!("Stopped systemd service (if running)."); + let _ = Command::new("systemctl") + .args(["--user", "disable", "llm-router.service"]) + .status(); + println!("Disabled systemd service."); + let service_path = get_service_path(); + if service_path.exists() { + if let Err(e) = fs::remove_file(&service_path) { + println!("Failed to remove service file: {}", e); + } else { + println!("Deleted service file: {}", service_path.display()); + } + } + let _ = Command::new("systemctl") + .args(["--user", "daemon-reload"]) + .status(); + let mut config_dir = dirs::config_dir().unwrap_or_else(|| PathBuf::from(".")); + config_dir.push("llm-router"); + if config_dir.exists() { + if let Err(e) = fs::remove_dir_all(&config_dir) { + println!("Failed to remove config folder: {}", e); + } else { + println!("Removed configuration folder: {}", config_dir.display()); + } + } + println!("\nUninstallation complete!"); + std::process::exit(0); +} + +fn start_systemd_service() { + match Command::new("systemctl") + .args(["--user", "start", "llm-router.service"]) + .status() + { + Ok(status) if status.success() => { + println!("Successfully started llm-router.service!"); + } + _ => { + println!("Failed to start systemd service. Ensure it has been installed first."); + } + } + std::process::exit(0); +} + +fn install_systemd_service() { + let exe_path = match std::env::current_exe() { + Ok(path) => path, + Err(e) => { + println!("Failed to get current executable path: {}", e); + return; + } + }; + let service_path = get_service_path(); + if let Some(parent) = service_path.parent() { + let _ = fs::create_dir_all(parent); + } + + let service_content = format!( + r#"[Unit] +Description=LLM Router Proxy Service +After=network.target + +[Service] +Type=simple +ExecStart={} --daemon +Restart=on-failure +RestartSec=5 + +[Install] +WantedBy=default.target +"#, + exe_path.display() + ); + + if let Err(e) = fs::write(&service_path, service_content) { + println!("Failed to write systemd service file: {}", e); + return; + } + println!("Created systemd service at: {}", service_path.display()); + + let _ = Command::new("systemctl") + .args(["--user", "daemon-reload"]) + .status(); + + match Command::new("systemctl") + .args(["--user", "enable", "--now", "llm-router.service"]) + .status() + { + Ok(status) if status.success() => { + println!("Successfully enabled and started llm-router.service as a user daemon!"); + } + _ => { + println!("Warning: Could not automatically start systemd service."); + } + } +} + +fn load_config(is_daemon: bool) -> Config { + let path = get_config_path(); + if path.exists() { + let content = fs::read_to_string(&path).expect("Failed to read config file"); + toml::from_str(&content).expect("Failed to parse config TOML") + } else { + if is_daemon { + eprintln!("Error: Configuration file not found at {}.", path.display()); + eprintln!("Please run the application interactively at least once without '--daemon' to complete setup."); + std::process::exit(1); + } + + println!("No configuration file found at {}.", path.display()); + println!("--- Interactive Setup Wizard ---"); + + print!("Enter port for proxy to listen on [default: 4000]: "); + io::stdout().flush().unwrap(); + let mut port_str = String::new(); + io::stdin().read_line(&mut port_str).unwrap(); + let listen_port = port_str.trim().parse().unwrap_or(4000); + + let mut backend_urls = Vec::new(); + println!("Enter your LLM server backend URLs (leave blank and press Enter when finished):"); + loop { + print!("Backend URL [e.g., http://192.168.1.50:11434]: "); + io::stdout().flush().unwrap(); + let mut url = String::new(); + io::stdin().read_line(&mut url).unwrap(); + let url = url.trim().to_string(); + + if url.is_empty() { + if backend_urls.is_empty() { + println!("Error: You must provide at least one backend server URL."); + continue; + } + break; + } + backend_urls.push(url); + } + + let config = Config { + listen_port, + backend_urls, + }; + + let toml_content = toml::to_string_pretty(&config).expect("Failed to serialize config"); + fs::write(&path, toml_content).expect("Failed to write config file"); + println!("\nConfiguration successfully saved to {}\n", path.display()); + + print!( + "Would you like to install and run this as a background systemd user service? [y/N]: " + ); + io::stdout().flush().unwrap(); + let mut install_ans = String::new(); + io::stdin().read_line(&mut install_ans).unwrap(); + if install_ans.trim().to_lowercase().starts_with('y') { + install_systemd_service(); + } + + config + } +} + +#[tokio::main] +async fn main() { + let args = Args::parse(); + if args.uninstall { + uninstall_service(); + } + if args.start_service { + start_systemd_service(); + } + let config = load_config(args.daemon); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(120)) + .build() + .unwrap(); + let mut backends = Vec::new(); + let mut initial_backend_data = Vec::new(); + for url in config.backend_urls { + backends.push(Backend { + url: url.clone(), + healthy: Arc::new(AtomicBool::new(false)), + }); + initial_backend_data.push((url, false)); + } + let tray = LLMRouterTray { + backend_data: initial_backend_data, + }; + let tray_handle = tray.spawn().await.expect("Failed to start system tray"); + + let state = AppState { + backends: backends.clone(), + client: client.clone(), + counter: Arc::new(AtomicUsize::new(0)), + }; + tokio::spawn(async move { + health_check_loop(backends, client, tray_handle).await; + }); + let app = Router::new() + .route("/*path", any(proxy_handler)) + .fallback(any(proxy_handler)) + .with_state(state); + + let bind_address = format!("0.0.0.0:{}", config.listen_port); + let listener = tokio::net::TcpListener::bind(&bind_address).await.unwrap(); + println!( + "LLM Router running on http://0.0.0.0:{}", + config.listen_port + ); + axum::serve(listener, app).await.unwrap(); +} + +async fn health_check_loop( + backends: Vec, + client: reqwest::Client, + tray_handle: ksni::Handle, +) { + loop { + for backend in &backends { + let health_url = format!("{}/v1/models", backend.url); + let res = client + .get(&health_url) + .timeout(Duration::from_secs(2)) + .send() + .await; + let is_healthy = res.map_or(false, |r| r.status().is_success()); + let prev = backend.healthy.load(Ordering::Relaxed); + + if prev != is_healthy { + println!( + "Backend [{}] status changed: {}", + backend.url, + if is_healthy { "ONLINE" } else { "OFFLINE" } + ); + } + backend.healthy.store(is_healthy, Ordering::Relaxed); + } + let snapshot: Vec<(String, bool)> = backends + .iter() + .map(|b| (b.url.clone(), b.healthy.load(Ordering::Relaxed))) + .collect(); + tray_handle + .update(move |tray: &mut LLMRouterTray| { + tray.backend_data = snapshot; + }) + .await; + tokio::time::sleep(Duration::from_secs(5)).await; + } +} + +fn select_backend(backends: &[Backend], counter: &AtomicUsize) -> Option { + let len = backends.len(); + if len == 0 { + return None; + } + let start_idx = counter.fetch_add(1, Ordering::Relaxed); + for i in 0..len { + let idx = (start_idx + i) % len; + if backends[idx].healthy.load(Ordering::Relaxed) { + return Some(backends[idx].url.clone()); + } + } + None +} + +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 path = req.uri().path().to_string(); + let query = req + .uri() + .query() + .map(|v| format!("?{}", v)) + .unwrap_or_default(); + let target_url = format!("{}{}{}", backend_url, path, query); + let method = req.method().clone(); + let headers = req.headers().clone(); + let body_bytes = match axum::body::to_bytes(req.into_body(), usize::MAX).await { + Ok(bytes) => bytes, + Err(_) => return (StatusCode::BAD_REQUEST, "Failed to read request body").into_response(), + }; + let mut forward_req = state.client.request(method, &target_url).body(body_bytes); + for (key, value) in headers.iter() { + forward_req = forward_req.header(key, value); + } + match forward_req.send().await { + Ok(resp) => { + let status = resp.status(); + let mut response_builder = Response::builder().status(status); + for (key, value) in resp.headers() { + response_builder = response_builder.header(key, value); + } + let stream = resp.bytes_stream().map(|result| { + result.map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err)) + }); + let body = Body::from_stream(stream); + response_builder.body(body).unwrap_or_else(|_| { + (StatusCode::INTERNAL_SERVER_ERROR, "Response build error").into_response() + }) + } + Err(_) => (StatusCode::BAD_GATEWAY, "Failed to reach LLM backend").into_response(), + } +}