wrote the tool
This commit is contained in:
pyro committed
1 parent
2711c895c0
commit
f87bf2f43e
3 files changed
+396
No files matched your search
+378
@@ -0,0 +1,378 @@
|
||||
use clap::Parser;
|
||||
use ksni::{Tray, TrayMethods};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::{
|
||||
fs,
|
||||
io::{self, Write},
|
||||
path::PathBuf,
|
||||
process::Command,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(
|
||||
name = "llm-router-tray",
|
||||
version,
|
||||
about = "LLM Server Status Tray Monitor with Systemd Integration and Preferred Backend"
|
||||
)]
|
||||
struct Args {
|
||||
/// Interactively setup configuration and install the systemd user service
|
||||
#[arg(long)]
|
||||
install: 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 {
|
||||
backend_urls: Vec<String>,
|
||||
#[serde(default)]
|
||||
preferred_backend: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct Backend {
|
||||
url: String,
|
||||
healthy: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct LLMTray {
|
||||
backend_data: Vec<(String, bool)>,
|
||||
preferred_url: Option<String>,
|
||||
}
|
||||
|
||||
impl Tray for LLMTray {
|
||||
fn id(&self) -> String {
|
||||
"llm-router-tray".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 Monitor ({}/{})", online_count, self.backend_data.len())
|
||||
}
|
||||
|
||||
fn tool_tip(&self) -> ksni::ToolTip {
|
||||
let mut desc = String::from("Monitored LLM Endpoints:\n");
|
||||
for (url, healthy) in &self.backend_data {
|
||||
let status = if *healthy {
|
||||
"🟢 Online"
|
||||
} else {
|
||||
"🔴 Offline"
|
||||
};
|
||||
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 Server Status".into(),
|
||||
description: desc.trim().into(),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn menu(&self) -> Vec<ksni::menu::MenuItem<Self>> {
|
||||
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"
|
||||
};
|
||||
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,
|
||||
enabled: false,
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
);
|
||||
}
|
||||
|
||||
items.push(MenuItem::Separator);
|
||||
items.push(
|
||||
StandardItem {
|
||||
label: "Quit Monitor".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-tray");
|
||||
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-tray.service");
|
||||
path
|
||||
}
|
||||
|
||||
fn uninstall_service() {
|
||||
println!("--- Uninstalling LLM Router Tray ---");
|
||||
let _ = Command::new("systemctl")
|
||||
.args(["--user", "stop", "llm-router-tray.service"])
|
||||
.status();
|
||||
println!("Stopped systemd service (if running).");
|
||||
|
||||
let _ = Command::new("systemctl")
|
||||
.args(["--user", "disable", "llm-router-tray.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-tray");
|
||||
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-tray.service"])
|
||||
.status()
|
||||
{
|
||||
Ok(status) if status.success() => {
|
||||
println!("Successfully started llm-router-tray.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 Tray Monitor Service
|
||||
After=network.target graphical-session.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
ExecStart={}
|
||||
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-tray.service"])
|
||||
.status()
|
||||
{
|
||||
Ok(status) if status.success() => {
|
||||
println!("Successfully enabled and started llm-router-tray.service as a user daemon!");
|
||||
}
|
||||
_ => {
|
||||
println!("Warning: Could not automatically start systemd service.");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn install_interactive() {
|
||||
let path = get_config_path();
|
||||
println!("--- LLM Router Tray Interactive Setup Wizard ---");
|
||||
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 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 {
|
||||
backend_urls,
|
||||
preferred_backend,
|
||||
};
|
||||
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());
|
||||
install_systemd_service();
|
||||
std::process::exit(0);
|
||||
}
|
||||
|
||||
fn load_config() -> 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 {
|
||||
println!("No configuration file found at {}.", path.display());
|
||||
println!("Running interactive setup wizard...");
|
||||
install_interactive();
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let args = Args::parse();
|
||||
if args.uninstall {
|
||||
uninstall_service();
|
||||
}
|
||||
if args.start_service {
|
||||
start_systemd_service();
|
||||
}
|
||||
if args.install {
|
||||
install_interactive();
|
||||
}
|
||||
let config = load_config();
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(2))
|
||||
.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 = LLMTray {
|
||||
backend_data: initial_backend_data,
|
||||
preferred_url: config.preferred_backend.clone(),
|
||||
};
|
||||
let tray_handle = tray.spawn().await.expect("Failed to start system tray");
|
||||
let preferred_url_clone = config.preferred_backend.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
for backend in &backends {
|
||||
let health_url = format!("{}/v1/models", backend.url);
|
||||
let res = client.get(&health_url).send().await;
|
||||
let is_healthy = res.map_or(false, |r| r.status().is_success());
|
||||
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();
|
||||
let pref_url_ref = preferred_url_clone.clone();
|
||||
tray_handle
|
||||
.update(move |tray: &mut LLMTray| {
|
||||
tray.backend_data = snapshot;
|
||||
tray.preferred_url = pref_url_ref;
|
||||
})
|
||||
.await;
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
}
|
||||
});
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
Reference in new issue
Block a user