From bf3193d8909b433a64142642deefea879486d217 Mon Sep 17 00:00:00 2001 From: jh-block Date: Tue, 10 Feb 2026 21:14:19 +0100 Subject: [PATCH] Switched local inference to use llama.cpp and added HF model download --- Cargo.lock | 253 +- crates/goose-cli/src/cli.rs | 205 ++ crates/goose-server/src/openapi.rs | 15 + .../src/routes/local_inference.rs | 367 ++- crates/goose-server/src/state.rs | 3 + crates/goose/Cargo.toml | 11 +- crates/goose/examples/candle_quantized.rs | 702 ----- crates/goose/examples/test_candle_minimal.rs | 54 - crates/goose/src/agents/extension_manager.rs | 6 +- .../goose/src/dictation/download_manager.rs | 37 +- crates/goose/src/prompts/tiny_model_system.md | 27 +- crates/goose/src/providers/local_inference.rs | 2261 ++++++++++++----- .../providers/local_inference/hf_models.rs | 565 ++++ .../local_inference/local_model_registry.rs | 297 +++ ui/desktop/openapi.json | 597 ++++- ui/desktop/src/api/index.ts | 4 +- ui/desktop/src/api/sdk.gen.ts | 26 +- ui/desktop/src/api/types.gen.ts | 277 +- .../src/components/ToolCallWithResponse.tsx | 2 +- .../src/components/settings/SettingsView.tsx | 19 +- .../localInference/HuggingFaceModelSearch.tsx | 331 +++ .../localInference/LocalInferenceSection.tsx | 9 + .../localInference/LocalInferenceSettings.tsx | 487 ++-- .../localInference/ModelSettingsPanel.tsx | 417 +++ .../settings/models/ModelsSection.tsx | 6 - 25 files changed, 5218 insertions(+), 1760 deletions(-) delete mode 100644 crates/goose/examples/candle_quantized.rs delete mode 100644 crates/goose/examples/test_candle_minimal.rs create mode 100644 crates/goose/src/providers/local_inference/hf_models.rs create mode 100644 crates/goose/src/providers/local_inference/local_model_registry.rs create mode 100644 ui/desktop/src/components/settings/localInference/HuggingFaceModelSearch.tsx create mode 100644 ui/desktop/src/components/settings/localInference/LocalInferenceSection.tsx create mode 100644 ui/desktop/src/components/settings/localInference/ModelSettingsPanel.tsx diff --git a/Cargo.lock b/Cargo.lock index d76750d1ff..975fa80b9d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -140,7 +140,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -151,7 +151,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1568,7 +1568,7 @@ dependencies = [ "candle-kernels", "candle-metal-kernels", "candle-ug", - "cudarc 0.19.1", + "cudarc 0.19.2", "float8 0.6.1", "gemm 0.19.0", "half", @@ -2285,9 +2285,9 @@ dependencies = [ [[package]] name = "cudarc" -version = "0.19.1" +version = "0.19.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60d4882b3e023670c25b8a3e1ac97349070f8e5400807c091b70e0bdad6e9b93" +checksum = "aed81f178e780f3d5d354d12b4c5c5a484c4a9c329ecd037ac57f2a0e0648397" dependencies = [ "float8 0.7.0", "half", @@ -3020,7 +3020,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3233,6 +3233,26 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "enumflags2" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1027f7680c853e056ebcec683615fb6fbbc07dbaa13b4d5d9442b146ded4ecef" +dependencies = [ + "enumflags2_derive", +] + +[[package]] +name = "enumflags2_derive" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67c78a4d8fdf9953a5c9d458f9efe940fd97a0cab0941c075a813ac594733827" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.114", +] + [[package]] name = "env-lock" version = "1.0.2" @@ -3292,7 +3312,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3459,6 +3479,15 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "find_cuda_helper" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f9e65c593dd01ac77daad909ea4ad17f0d6d1776193fc8ea766356177abdad" +dependencies = [ + "glob", +] + [[package]] name = "fixedbitset" version = "0.5.7" @@ -3481,7 +3510,7 @@ version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "719a903cc23e4a89e87962c2a80fdb45cdaad0983a89bd150bb57b4c8571a7d5" dependencies = [ - "cudarc 0.19.1", + "cudarc 0.19.2", "half", "num-traits", "rand 0.9.2", @@ -4065,6 +4094,19 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", + "wasip3", +] + [[package]] name = "gif" version = "0.13.3" @@ -4135,6 +4177,7 @@ dependencies = [ "dashmap 6.1.0", "dirs 5.0.1", "dotenvy", + "encoding_rs", "env-lock", "etcetera 0.11.0", "fs2", @@ -4151,6 +4194,7 @@ dependencies = [ "jsonwebtoken", "keyring", "lazy_static", + "llama-cpp-2", "lru", "minijinja", "mockall", @@ -4973,6 +5017,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + [[package]] name = "ident_case" version = "1.0.1" @@ -5275,7 +5325,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5476,6 +5526,12 @@ dependencies = [ "spin", ] +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + [[package]] name = "lebe" version = "0.5.3" @@ -5484,9 +5540,9 @@ checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8" [[package]] name = "libc" -version = "0.2.180" +version = "0.2.181" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" +checksum = "459427e2af2b9c839b132acb702a1c654d95e10f8c326bfc2ad11310e458b1c5" [[package]] name = "libdbus-sys" @@ -5580,6 +5636,34 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" +[[package]] +name = "llama-cpp-2" +version = "0.1.133" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "888c8805527f4c35ec16f26003d54a318cde1629e7439da8e9ef2d6d3883e106" +dependencies = [ + "encoding_rs", + "enumflags2", + "llama-cpp-sys-2", + "thiserror 1.0.69", + "tracing", + "tracing-core", +] + +[[package]] +name = "llama-cpp-sys-2" +version = "0.1.133" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a180dfa6d6f9d1df1e031bcdf0464bbad4f9b326395bfd28f2fa539d8cbc9c2b" +dependencies = [ + "bindgen", + "cc", + "cmake", + "find_cuda_helper", + "glob", + "walkdir", +] + [[package]] name = "lock_api" version = "0.4.14" @@ -5996,7 +6080,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -8146,7 +8230,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.11.0", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -9971,15 +10055,15 @@ dependencies = [ [[package]] name = "tempfile" -version = "3.24.0" +version = "3.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "655da9c7eb6305c55742045d5a8d2037996d61d8de95806335c7c86ce0f82e9c" +checksum = "0136791f7c95b1f6dd99f9cc786b91bb81c3800b639b3478e561ddb7be95e5f1" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.1", "once_cell", "rustix 1.1.3", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -10430,9 +10514,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.9.11+spec-1.1.0" +version = "0.9.12+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3afc9a848309fe1aaffaed6e1546a7a14de1f935dc9d89d32afd9a44bab7c46" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" dependencies = [ "indexmap 2.13.0", "serde_core", @@ -10454,9 +10538,9 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.0.6+spec-1.1.0" +version = "1.0.7+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3198b4b0a8e11f09dd03e133c0280504d0801269e9afa46362ffde1cbeebf44" +checksum = "247eaa3197818b831697600aadf81514e577e0cba5eab10f7e064e78ae154df1" dependencies = [ "winnow", ] @@ -11285,6 +11369,15 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasite" version = "0.1.0" @@ -11350,6 +11443,28 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap 2.13.0", + "wasm-encoder", + "wasmparser", +] + [[package]] name = "wasm-streams" version = "0.4.2" @@ -11373,6 +11488,18 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags 2.10.0", + "hashbrown 0.15.5", + "indexmap 2.13.0", + "semver", +] + [[package]] name = "web-sys" version = "0.3.85" @@ -11500,7 +11627,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -12101,6 +12228,88 @@ name = "wit-bindgen" version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap 2.13.0", + "prettyplease", + "syn 2.0.114", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.114", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags 2.10.0", + "indexmap 2.13.0", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap 2.13.0", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] [[package]] name = "writeable" diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index 327d945bc8..ee8350f5ab 100644 --- a/crates/goose-cli/src/cli.rs +++ b/crates/goose-cli/src/cli.rs @@ -854,6 +854,13 @@ enum Command { #[command(subcommand)] command: TermCommand, }, + /// Manage local inference models + #[command(about = "Manage local inference models", visible_alias = "lm")] + LocalModels { + #[command(subcommand)] + command: LocalModelsCommand, + }, + /// Generate completions for various shells #[command(about = "Generate the autocompletion script for the specified shell")] Completion { @@ -865,6 +872,38 @@ enum Command { }, } +#[derive(Subcommand)] +enum LocalModelsCommand { + /// Search HuggingFace for GGUF models + #[command(about = "Search HuggingFace for GGUF models")] + Search { + /// Search query + query: String, + + /// Maximum number of results + #[arg(short, long, default_value = "10")] + limit: usize, + }, + + /// Download a model from HuggingFace + #[command(about = "Download a GGUF model (e.g. bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M)")] + Download { + /// Model spec in user/repo:quantization format + spec: String, + }, + + /// List downloaded local models + #[command(about = "List downloaded local models")] + List, + + /// Delete a downloaded model + #[command(about = "Delete a downloaded local model")] + Delete { + /// Model ID to delete + id: String, + }, +} + #[derive(Subcommand)] enum TermCommand { /// Print shell initialization script @@ -954,6 +993,7 @@ fn get_command_name(command: &Option) -> &'static str { Some(Command::Recipe { .. }) => "recipe", Some(Command::Web { .. }) => "web", Some(Command::Term { .. }) => "term", + Some(Command::LocalModels { .. }) => "local-models", Some(Command::Completion { .. }) => "completion", None => "default_session", } @@ -1389,6 +1429,170 @@ async fn handle_term_subcommand(command: TermCommand) -> Result<()> { } } +async fn handle_local_models_command(command: LocalModelsCommand) -> Result<()> { + use goose::providers::local_inference::hf_models; + use goose::providers::local_inference::local_model_registry::{ + display_name_from_repo, get_registry, model_id_from_repo, LocalModelEntry, + }; + + match command { + LocalModelsCommand::Search { query, limit } => { + println!("Searching HuggingFace for '{}'...", query); + let results = hf_models::search_gguf_models(&query, limit).await?; + + if results.is_empty() { + println!("No GGUF models found."); + return Ok(()); + } + + for model in &results { + println!( + "\n{} (by {}) — {} downloads", + model.model_name, model.author, model.downloads + ); + for file in &model.gguf_files { + let size = if file.size_bytes > 0 { + format!( + "{:.1}GB", + file.size_bytes as f64 / (1024.0 * 1024.0 * 1024.0) + ) + } else { + "unknown".to_string() + }; + println!(" {} — {}", file.quantization, size); + } + println!( + " Download: goose local-models download {}:", + model.repo_id + ); + } + } + LocalModelsCommand::Download { spec } => { + println!("Resolving {}...", spec); + let (repo_id, file) = hf_models::resolve_model_spec(&spec).await?; + let model_id = model_id_from_repo(&repo_id, &file.quantization); + let display_name = display_name_from_repo(&repo_id, &file.quantization); + let local_path = + goose::config::paths::Paths::in_data_dir("models").join(&file.filename); + + println!( + "Downloading {} ({})...", + display_name, + if file.size_bytes > 0 { + format!( + "{:.1}GB", + file.size_bytes as f64 / (1024.0 * 1024.0 * 1024.0) + ) + } else { + "unknown size".to_string() + } + ); + + // Register + let entry = LocalModelEntry { + id: model_id.clone(), + display_name: display_name.clone(), + repo_id: repo_id.clone(), + filename: file.filename.clone(), + quantization: file.quantization.clone(), + local_path: local_path.clone(), + source_url: file.download_url.clone(), + settings: Default::default(), + }; + + { + let mut registry = get_registry() + .lock() + .map_err(|_| anyhow::anyhow!("Failed to acquire registry lock"))?; + registry.add_model(entry)?; + } + + // Download + let manager = goose::dictation::download_manager::get_download_manager(); + manager + .download_model( + format!("{}-model", model_id), + file.download_url, + local_path, + None, + None, + ) + .await?; + + // Poll progress + loop { + if let Some(progress) = manager.get_progress(&format!("{}-model", model_id)) { + match progress.status { + goose::dictation::download_manager::DownloadStatus::Downloading => { + print!( + "\r {:.1}% ({:.0}MB / {:.0}MB)", + progress.progress_percent, + progress.bytes_downloaded as f64 / (1024.0 * 1024.0), + progress.total_bytes as f64 / (1024.0 * 1024.0), + ); + use std::io::Write; + std::io::stdout().flush().ok(); + } + goose::dictation::download_manager::DownloadStatus::Completed => { + println!("\nDownloaded: {} (id: {})", display_name, model_id); + break; + } + goose::dictation::download_manager::DownloadStatus::Failed => { + let err = progress.error.unwrap_or_default(); + anyhow::bail!("Download failed: {}", err); + } + goose::dictation::download_manager::DownloadStatus::Cancelled => { + println!("\nDownload cancelled."); + break; + } + } + } + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + } + } + LocalModelsCommand::List => { + let registry = get_registry() + .lock() + .map_err(|_| anyhow::anyhow!("Failed to acquire registry lock"))?; + let models = registry.list_models(); + + if models.is_empty() { + println!("No local models downloaded."); + return Ok(()); + } + + println!("{:<40} {:<20} {:<10} Downloaded", "ID", "Name", "Quant"); + println!("{}", "-".repeat(80)); + for m in models { + println!( + "{:<40} {:<20} {:<10} {}", + m.id, + m.display_name, + m.quantization, + if m.is_downloaded() { "✓" } else { "✗" } + ); + } + } + LocalModelsCommand::Delete { id } => { + let mut registry = get_registry() + .lock() + .map_err(|_| anyhow::anyhow!("Failed to acquire registry lock"))?; + + if let Some(entry) = registry.get_model(&id) { + if entry.local_path.exists() { + std::fs::remove_file(&entry.local_path)?; + } + registry.remove_model(&id)?; + println!("Deleted model: {}", id); + } else { + println!("Model not found: {}", id); + } + } + } + + Ok(()) +} + async fn handle_default_session() -> Result<()> { if !Config::global().exists() { return handle_configure().await; @@ -1519,6 +1723,7 @@ pub async fn cli() -> anyhow::Result<()> { no_auth, }) => crate::commands::web::handle_web(port, host, open, auth_token, no_auth).await, Some(Command::Term { command }) => handle_term_subcommand(command).await, + Some(Command::LocalModels { command }) => handle_local_models_command(command).await, None => handle_default_session().await, } } diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index 99fd20e079..3848cf7d65 100644 --- a/crates/goose-server/src/openapi.rs +++ b/crates/goose-server/src/openapi.rs @@ -422,10 +422,15 @@ derive_utoipa!(Icon as IconSchema); super::routes::dictation::cancel_download, super::routes::dictation::delete_model, super::routes::local_inference::list_local_models, + super::routes::local_inference::search_hf_models, + super::routes::local_inference::get_repo_files, + super::routes::local_inference::download_hf_model, super::routes::local_inference::download_local_model, super::routes::local_inference::get_local_model_download_progress, super::routes::local_inference::cancel_local_model_download, super::routes::local_inference::delete_local_model, + super::routes::local_inference::get_model_settings, + super::routes::local_inference::update_model_settings, ), components(schemas( super::routes::config_management::UpsertConfigQuery, @@ -592,8 +597,18 @@ derive_utoipa!(Icon as IconSchema); super::routes::dictation::DictationProviderStatus, super::routes::dictation::WhisperModelResponse, super::routes::local_inference::LocalModelResponse, + super::routes::local_inference::RegistryModelResponse, + super::routes::local_inference::ModelListItem, + super::routes::local_inference::HfDownloadRequest, + super::routes::local_inference::HfDownloadResponse, goose::providers::local_inference::LocalLlmModel, goose::providers::local_inference::ModelTier, + goose::providers::local_inference::hf_models::HfModelInfo, + goose::providers::local_inference::hf_models::HfGgufFile, + goose::providers::local_inference::hf_models::HfQuantVariant, + super::routes::local_inference::RepoVariantsResponse, + goose::providers::local_inference::local_model_registry::ModelSettings, + goose::providers::local_inference::local_model_registry::SamplingConfig, DownloadProgress, DownloadStatus, )) diff --git a/crates/goose-server/src/routes/local_inference.rs b/crates/goose-server/src/routes/local_inference.rs index 099319607e..7a58cfc8d5 100644 --- a/crates/goose-server/src/routes/local_inference.rs +++ b/crates/goose-server/src/routes/local_inference.rs @@ -1,17 +1,22 @@ use crate::routes::errors::ErrorResponse; use crate::state::AppState; use axum::{ - extract::Path, + extract::{Path, Query}, http::StatusCode, routing::{delete, get, post}, Json, Router, }; +use goose::config::paths::Paths; use goose::dictation::download_manager::{get_download_manager, DownloadProgress}; -use goose::providers::local_inference::{ - available_local_models, get_local_model, recommend_local_model, LocalLlmModel, - LOCAL_LLM_MODEL_CONFIG_KEY, +use goose::providers::local_inference::hf_models::{self, HfModelInfo, HfQuantVariant}; +use goose::providers::local_inference::local_model_registry::{ + display_name_from_repo, get_registry, model_id_from_repo, LocalModelEntry, ModelSettings, }; -use serde::Serialize; +use goose::providers::local_inference::{ + available_inference_memory_bytes, available_local_models, get_local_model, + recommend_local_model, LocalLlmModel, LOCAL_LLM_MODEL_CONFIG_KEY, +}; +use serde::{Deserialize, Serialize}; use std::sync::Arc; use utoipa::ToSchema; @@ -22,6 +27,27 @@ pub struct LocalModelResponse { pub model: &'static LocalLlmModel, pub downloaded: bool, pub recommended: bool, + pub featured: bool, +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct RegistryModelResponse { + pub id: String, + pub display_name: String, + pub repo_id: String, + pub filename: String, + pub quantization: String, + pub size_bytes: u64, + pub downloaded: bool, + pub featured: bool, + pub settings: ModelSettings, +} + +#[derive(Debug, Serialize, ToSchema)] +#[serde(untagged)] +pub enum ModelListItem { + Featured(LocalModelResponse), + Registry(Box), } fn convert_error(e: anyhow::Error) -> ErrorResponse { @@ -45,21 +71,194 @@ fn convert_error(e: anyhow::Error) -> ErrorResponse { get, path = "/local-inference/models", responses( - (status = 200, description = "List of available local LLM models", body = Vec) + (status = 200, description = "List of available local LLM models", body = Vec) ) )] -pub async fn list_local_models() -> Result>, ErrorResponse> { - let recommended_id = recommend_local_model(); - let models = available_local_models() +pub async fn list_local_models( + axum::extract::State(state): axum::extract::State>, +) -> Result>, ErrorResponse> { + let recommended_id = recommend_local_model(&state.inference_runtime); + let featured_ids: Vec<&str> = available_local_models().iter().map(|m| m.id).collect(); + + let mut items: Vec = available_local_models() .iter() - .map(|m| LocalModelResponse { - model: m, - downloaded: m.is_downloaded(), - recommended: m.id == recommended_id, + .map(|m| { + ModelListItem::Featured(LocalModelResponse { + model: m, + downloaded: m.is_downloaded(), + recommended: m.id == recommended_id, + featured: true, + }) }) .collect(); - Ok(Json(models)) + // Add registry models that aren't in the featured list + if let Ok(registry) = get_registry().lock() { + for entry in registry.list_models() { + if !featured_ids.contains(&entry.id.as_str()) { + items.push(ModelListItem::Registry(Box::new(RegistryModelResponse { + id: entry.id.clone(), + display_name: entry.display_name.clone(), + repo_id: entry.repo_id.clone(), + filename: entry.filename.clone(), + quantization: entry.quantization.clone(), + size_bytes: entry.file_size(), + downloaded: entry.is_downloaded(), + featured: false, + settings: entry.settings.clone(), + }))); + } + } + } + + Ok(Json(items)) +} + +#[derive(Debug, Deserialize)] +pub struct SearchQuery { + pub q: String, + pub limit: Option, +} + +#[utoipa::path( + get, + path = "/local-inference/search", + params( + ("q" = String, Query, description = "Search query"), + ("limit" = Option, Query, description = "Max results") + ), + responses( + (status = 200, description = "Search results", body = Vec), + (status = 500, description = "Search failed") + ) +)] +pub async fn search_hf_models( + Query(params): Query, +) -> Result>, ErrorResponse> { + let limit = params.limit.unwrap_or(20).min(50); + let results = hf_models::search_gguf_models(¶ms.q, limit) + .await + .map_err(|e| ErrorResponse::internal(format!("Search failed: {}", e)))?; + Ok(Json(results)) +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct RepoVariantsResponse { + pub variants: Vec, + pub recommended_index: Option, +} + +#[utoipa::path( + get, + path = "/local-inference/repo/{author}/{repo}/files", + responses( + (status = 200, description = "GGUF quantization variants in repo", body = RepoVariantsResponse), + (status = 500, description = "Failed to fetch repo files") + ) +)] +pub async fn get_repo_files( + axum::extract::State(state): axum::extract::State>, + Path((author, repo)): Path<(String, String)>, +) -> Result, ErrorResponse> { + let repo_id = format!("{}/{}", author, repo); + let variants = hf_models::get_repo_gguf_variants(&repo_id) + .await + .map_err(|e| ErrorResponse::internal(format!("Failed to fetch repo files: {}", e)))?; + + let available_memory = available_inference_memory_bytes(&state.inference_runtime); + let recommended_index = hf_models::recommend_variant(&variants, available_memory); + + Ok(Json(RepoVariantsResponse { + variants, + recommended_index, + })) +} + +#[derive(Debug, Deserialize, ToSchema)] +pub struct HfDownloadRequest { + pub repo_id: Option, + pub filename: Option, + pub spec: Option, +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct HfDownloadResponse { + pub model_id: String, +} + +#[utoipa::path( + post, + path = "/local-inference/download", + request_body = HfDownloadRequest, + responses( + (status = 202, description = "Download started", body = HfDownloadResponse), + (status = 400, description = "Invalid request"), + (status = 500, description = "Internal server error") + ) +)] +pub async fn download_hf_model( + Json(req): Json, +) -> Result<(StatusCode, Json), ErrorResponse> { + let (repo_id, filename, quantization, download_url) = if let Some(spec) = &req.spec { + let (repo_id, file) = hf_models::resolve_model_spec(spec) + .await + .map_err(|e| ErrorResponse::bad_request(format!("{}", e)))?; + (repo_id, file.filename, file.quantization, file.download_url) + } else if let (Some(repo_id), Some(filename)) = (&req.repo_id, &req.filename) { + let quantization = hf_models::parse_quantization_from_filename(filename); + let download_url = format!( + "https://huggingface.co/{}/resolve/main/{}", + repo_id, filename + ); + ( + repo_id.clone(), + filename.clone(), + quantization, + download_url, + ) + } else { + return Err(ErrorResponse::bad_request( + "Provide either 'spec' or both 'repo_id' and 'filename'", + )); + }; + + let model_id = model_id_from_repo(&repo_id, &quantization); + let display_name = display_name_from_repo(&repo_id, &quantization); + let local_path = Paths::in_data_dir("models").join(&filename); + + let entry = LocalModelEntry { + id: model_id.clone(), + display_name, + repo_id, + filename: filename.clone(), + quantization, + local_path: local_path.clone(), + source_url: download_url.clone(), + settings: ModelSettings::default(), + }; + + { + let mut registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + registry + .add_model(entry) + .map_err(|e| ErrorResponse::internal(format!("{}", e)))?; + } + + let manager = get_download_manager(); + manager + .download_model( + format!("{}-model", model_id), + download_url, + local_path, + Some(LOCAL_LLM_MODEL_CONFIG_KEY.to_string()), + Some(model_id.clone()), + ) + .await + .map_err(convert_error)?; + + Ok((StatusCode::ACCEPTED, Json(HfDownloadResponse { model_id }))) } #[utoipa::path( @@ -79,26 +278,11 @@ pub async fn download_local_model( let manager = get_download_manager(); - // Download model file (don't set config yet - wait for tokenizer) manager .download_model( format!("{}-model", model.id), model.url.to_string(), model.local_path(), - None, - None, - ) - .await - .map_err(convert_error)?; - - // Download tokenizer file (set config and provider when this completes) - // We'll set GOOSE_PROVIDER to "local" after the tokenizer download completes - // This is handled in the download_manager callback - manager - .download_model( - format!("{}-tokenizer", model.id), - model.tokenizer_url.to_string(), - model.tokenizer_path(), Some(LOCAL_LLM_MODEL_CONFIG_KEY.to_string()), Some(model.id.to_string()), ) @@ -121,26 +305,10 @@ pub async fn get_local_model_download_progress( ) -> Result, ErrorResponse> { let manager = get_download_manager(); - // Check both model and tokenizer progress let model_progress = manager .get_progress(&format!("{}-model", model_id)) .ok_or_else(|| ErrorResponse::not_found("Download not found"))?; - let tokenizer_progress = manager.get_progress(&format!("{}-tokenizer", model_id)); - - // If tokenizer failed, return that error - if let Some(tok_prog) = tokenizer_progress { - if tok_prog.status == goose::dictation::download_manager::DownloadStatus::Failed { - return Ok(Json(tok_prog)); - } - } - - // If model failed, return that error - if model_progress.status == goose::dictation::download_manager::DownloadStatus::Failed { - return Ok(Json(model_progress)); - } - - // Otherwise return model progress (which shows overall download progress) Ok(Json(model_progress)) } @@ -159,9 +327,6 @@ pub async fn cancel_local_model_download( manager .cancel_download(&format!("{}-model", model_id)) .map_err(convert_error)?; - manager - .cancel_download(&format!("{}-tokenizer", model_id)) - .map_err(convert_error)?; Ok(StatusCode::OK) } @@ -176,34 +341,104 @@ pub async fn cancel_local_model_download( ) )] pub async fn delete_local_model(Path(model_id): Path) -> Result { - let model = - get_local_model(&model_id).ok_or_else(|| ErrorResponse::not_found("Model not found"))?; - - let model_path = model.local_path(); - let tokenizer_path = model.tokenizer_path(); - - if !model_path.exists() && !tokenizer_path.exists() { - return Err(ErrorResponse::not_found("Model not downloaded")); + // Try featured model first + if let Some(model) = get_local_model(&model_id) { + let model_path = model.local_path(); + if !model_path.exists() { + return Err(ErrorResponse::not_found("Model not downloaded")); + } + tokio::fs::remove_file(&model_path) + .await + .map_err(|e| ErrorResponse::internal(format!("Failed to delete model: {}", e)))?; + return Ok(StatusCode::OK); } - // Delete both files + // Try registry model + let model_path = { + let registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + let entry = registry + .get_model(&model_id) + .ok_or_else(|| ErrorResponse::not_found("Model not found"))?; + entry.local_path.clone() + }; + if model_path.exists() { tokio::fs::remove_file(&model_path) .await .map_err(|e| ErrorResponse::internal(format!("Failed to delete model: {}", e)))?; } - if tokenizer_path.exists() { - tokio::fs::remove_file(&tokenizer_path) - .await - .map_err(|e| ErrorResponse::internal(format!("Failed to delete tokenizer: {}", e)))?; - } + + // Remove from registry + let mut registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + registry + .remove_model(&model_id) + .map_err(|e| ErrorResponse::internal(format!("{}", e)))?; Ok(StatusCode::OK) } +#[utoipa::path( + get, + path = "/local-inference/models/{model_id}/settings", + responses( + (status = 200, description = "Model settings", body = ModelSettings), + (status = 404, description = "Model not found") + ) +)] +pub async fn get_model_settings( + Path(model_id): Path, +) -> Result, ErrorResponse> { + let registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + + if let Some(settings) = registry.get_model_settings(&model_id) { + return Ok(Json(settings.clone())); + } + + Err(ErrorResponse::not_found("Model not found")) +} + +#[utoipa::path( + put, + path = "/local-inference/models/{model_id}/settings", + request_body = ModelSettings, + responses( + (status = 200, description = "Settings updated", body = ModelSettings), + (status = 404, description = "Model not found"), + (status = 500, description = "Failed to save settings") + ) +)] +pub async fn update_model_settings( + Path(model_id): Path, + Json(settings): Json, +) -> Result, ErrorResponse> { + let mut registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + + registry + .update_model_settings(&model_id, settings.clone()) + .map_err(|e| ErrorResponse::not_found(format!("{}", e)))?; + + Ok(Json(settings)) +} + pub fn routes(state: Arc) -> Router { + goose::dictation::download_manager::cleanup_partial_downloads(&Paths::in_data_dir("models")); + Router::new() .route("/local-inference/models", get(list_local_models)) + .route("/local-inference/search", get(search_hf_models)) + .route( + "/local-inference/repo/{author}/{repo}/files", + get(get_repo_files), + ) + .route("/local-inference/download", post(download_hf_model)) .route( "/local-inference/models/{model_id}/download", post(download_local_model), @@ -220,5 +455,13 @@ pub fn routes(state: Arc) -> Router { "/local-inference/models/{model_id}", delete(delete_local_model), ) + .route( + "/local-inference/models/{model_id}/settings", + get(get_model_settings), + ) + .route( + "/local-inference/models/{model_id}/settings", + axum::routing::put(update_model_settings), + ) .with_state(state) } diff --git a/crates/goose-server/src/state.rs b/crates/goose-server/src/state.rs index 85f4ab76d2..025a25b376 100644 --- a/crates/goose-server/src/state.rs +++ b/crates/goose-server/src/state.rs @@ -11,6 +11,7 @@ use tokio::task::JoinHandle; use crate::tunnel::TunnelManager; use goose::agents::ExtensionLoadResult; +use goose::providers::local_inference::InferenceRuntime; type ExtensionLoadingTasks = Arc>>>>>>>; @@ -23,6 +24,7 @@ pub struct AppState { recipe_session_tracker: Arc>>, pub tunnel_manager: Arc, pub extension_loading_tasks: ExtensionLoadingTasks, + pub inference_runtime: Arc, } impl AppState { @@ -38,6 +40,7 @@ impl AppState { recipe_session_tracker: Arc::new(Mutex::new(HashSet::new())), tunnel_manager, extension_loading_tasks: Arc::new(Mutex::new(HashMap::new())), + inference_runtime: InferenceRuntime::get_or_init(), })) } diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 3f29f108b3..44f5537a79 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -9,7 +9,7 @@ description.workspace = true [features] default = [] -cuda = ["candle-core/cuda", "candle-nn/cuda"] +cuda = ["candle-core/cuda", "candle-nn/cuda", "llama-cpp-2/cuda"] [lints] workspace = true @@ -128,14 +128,17 @@ ignore = { workspace = true } which = { workspace = true } pctx_code_mode = "^0.2.3" unbinder = "0.1.7" +llama-cpp-2 = { version = "0.1.133", features = ["sampler"] } +encoding_rs = "0.8.35" [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3", features = ["wincred"] } -# Platform-specific GPU acceleration for Whisper +# Platform-specific GPU acceleration for Whisper and local inference [target.'cfg(target_os = "macos")'.dependencies] candle-core = { version = "0.9", default-features = false, features = ["metal"] } candle-nn = { version = "0.9", default-features = false, features = ["metal"] } +llama-cpp-2 = { version = "0.1.133", features = ["sampler", "metal"] } [dev-dependencies] serial_test = { workspace = true } @@ -158,10 +161,6 @@ path = "examples/agent.rs" name = "databricks_oauth" path = "examples/databricks_oauth.rs" -[[example]] -name = "candle_quantized" -path = "examples/candle_quantized.rs" - [[example]] name = "test_local_provider" path = "examples/test_local_provider.rs" diff --git a/crates/goose/examples/candle_quantized.rs b/crates/goose/examples/candle_quantized.rs deleted file mode 100644 index eb7e348a05..0000000000 --- a/crates/goose/examples/candle_quantized.rs +++ /dev/null @@ -1,702 +0,0 @@ -#[cfg(feature = "mkl")] -extern crate intel_mkl_src; - -#[cfg(feature = "accelerate")] -extern crate accelerate_src; - -use clap::{Parser, ValueEnum}; -use std::io::Write; -use tokenizers::Tokenizer; - -use candle::quantized::{ggml_file, gguf_file}; -use candle::Tensor; -use candle_transformers::generation::{LogitsProcessor, Sampling}; - -use candle_examples::token_output_stream::TokenOutputStream; -use candle_transformers::models::quantized_llama as model; -use model::ModelWeights; - -const DEFAULT_PROMPT: &str = "My favorite theorem is "; - -#[derive(Debug)] -enum Prompt { - Interactive, - Chat, - One(String), -} - -#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)] -enum Which { - #[value(name = "7b")] - L7b, - #[value(name = "13b")] - L13b, - #[value(name = "70b")] - L70b, - #[value(name = "7b-chat")] - L7bChat, - #[value(name = "13b-chat")] - L13bChat, - #[value(name = "70b-chat")] - L70bChat, - #[value(name = "7b-code")] - L7bCode, - #[value(name = "13b-code")] - L13bCode, - #[value(name = "32b-code")] - L34bCode, - #[value(name = "7b-leo")] - Leo7b, - #[value(name = "13b-leo")] - Leo13b, - #[value(name = "7b-mistral")] - Mistral7b, - #[value(name = "7b-mistral-instruct")] - Mistral7bInstruct, - #[value(name = "7b-mistral-instruct-v0.2")] - Mistral7bInstructV02, - #[value(name = "7b-zephyr-a")] - Zephyr7bAlpha, - #[value(name = "7b-zephyr-b")] - Zephyr7bBeta, - #[value(name = "7b-open-chat-3.5")] - OpenChat35, - #[value(name = "7b-starling-a")] - Starling7bAlpha, - #[value(name = "mixtral")] - Mixtral, - #[value(name = "mixtral-instruct")] - MixtralInstruct, - #[value(name = "llama3-8b")] - L8b, - #[value(name = "phi3")] - Phi3, - #[value(name = "SmoLM2-360M-Instruct")] - SmolLM2_360MInstruct, - #[value(name = "SmoLM2-1.7B-Instruct")] - SmolLM2_1BInstruct, - #[value(name = "deepseekr1-llama8b")] - DeepseekR1Llama8b, -} - -impl Which { - fn is_mistral(&self) -> bool { - match self { - Self::L7b - | Self::L13b - | Self::L70b - | Self::L7bChat - | Self::L13bChat - | Self::L70bChat - | Self::L7bCode - | Self::L13bCode - | Self::L34bCode - | Self::Leo7b - | Self::Leo13b - | Self::L8b - | Self::Phi3 - | Self::SmolLM2_1BInstruct - | Self::SmolLM2_360MInstruct - | Self::DeepseekR1Llama8b => false, - // Zephyr and OpenChat are fine tuned versions of mistral and should be treated in the - // same way. Starling is a fine tuned version of OpenChat. - Self::OpenChat35 - | Self::Starling7bAlpha - | Self::Zephyr7bAlpha - | Self::Zephyr7bBeta - | Self::Mixtral - | Self::MixtralInstruct - | Self::Mistral7b - | Self::Mistral7bInstruct - | Self::Mistral7bInstructV02 => true, - } - } - - fn is_zephyr(&self) -> bool { - match self { - Self::L7b - | Self::L13b - | Self::L70b - | Self::L7bChat - | Self::L13bChat - | Self::L70bChat - | Self::L7bCode - | Self::L13bCode - | Self::L34bCode - | Self::Leo7b - | Self::Leo13b - | Self::Mixtral - | Self::MixtralInstruct - | Self::Mistral7b - | Self::Mistral7bInstruct - | Self::Mistral7bInstructV02 - | Self::OpenChat35 - | Self::Starling7bAlpha - | Self::L8b - | Self::SmolLM2_1BInstruct - | Self::SmolLM2_360MInstruct - | Self::Phi3 - | Self::DeepseekR1Llama8b => false, - Self::Zephyr7bAlpha | Self::Zephyr7bBeta => true, - } - } - - fn is_open_chat(&self) -> bool { - match self { - Self::L7b - | Self::L13b - | Self::L70b - | Self::L7bChat - | Self::L13bChat - | Self::L70bChat - | Self::L7bCode - | Self::L13bCode - | Self::L34bCode - | Self::Leo7b - | Self::Leo13b - | Self::Mixtral - | Self::MixtralInstruct - | Self::Mistral7b - | Self::Mistral7bInstruct - | Self::Mistral7bInstructV02 - | Self::Zephyr7bAlpha - | Self::Zephyr7bBeta - | Self::L8b - | Self::SmolLM2_1BInstruct - | Self::SmolLM2_360MInstruct - | Self::Phi3 - | Self::DeepseekR1Llama8b => false, - Self::OpenChat35 | Self::Starling7bAlpha => true, - } - } - - fn is_deepseek(&self) -> bool { - match self { - Self::L7b - | Self::L13b - | Self::L70b - | Self::L7bChat - | Self::L13bChat - | Self::L70bChat - | Self::L7bCode - | Self::L13bCode - | Self::L34bCode - | Self::Leo7b - | Self::Leo13b - | Self::Mixtral - | Self::MixtralInstruct - | Self::Mistral7b - | Self::Mistral7bInstruct - | Self::Mistral7bInstructV02 - | Self::Zephyr7bAlpha - | Self::Zephyr7bBeta - | Self::L8b - | Self::SmolLM2_1BInstruct - | Self::SmolLM2_360MInstruct - | Self::Phi3 - | Self::OpenChat35 - | Self::Starling7bAlpha => false, - Self::DeepseekR1Llama8b => true, - } - } - fn tokenizer_repo(&self) -> &'static str { - match self { - Self::L7b - | Self::L13b - | Self::L70b - | Self::L7bChat - | Self::L13bChat - | Self::L70bChat - | Self::L7bCode - | Self::L13bCode - | Self::L34bCode => "hf-internal-testing/llama-tokenizer", - Self::Leo7b => "LeoLM/leo-hessianai-7b", - Self::Leo13b => "LeoLM/leo-hessianai-13b", - Self::Mixtral => "mistralai/Mixtral-8x7B-v0.1", - Self::MixtralInstruct => "mistralai/Mixtral-8x7B-Instruct-v0.1", - Self::Mistral7b - | Self::Mistral7bInstruct - | Self::Mistral7bInstructV02 - | Self::Zephyr7bAlpha - | Self::Zephyr7bBeta => "mistralai/Mistral-7B-v0.1", - Self::OpenChat35 => "openchat/openchat_3.5", - Self::Starling7bAlpha => "berkeley-nest/Starling-LM-7B-alpha", - Self::L8b => "meta-llama/Meta-Llama-3-8B", - Self::Phi3 => "microsoft/Phi-3-mini-4k-instruct", - Self::SmolLM2_360MInstruct => "HuggingFaceTB/SmolLM2-360M-Instruct", - Self::SmolLM2_1BInstruct => "HuggingFaceTB/SmolLM2-1.7B-Instruct", - Self::DeepseekR1Llama8b => "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", - } - } -} - -#[derive(Parser, Debug)] -#[command(author, version, about, long_about = None)] -struct Args { - /// GGML/GGUF file to load, typically a .bin/.gguf file generated by the quantize command from llama.cpp - #[arg(long)] - model: Option, - - /// The initial prompt, use 'interactive' for entering multiple prompts in an interactive way - /// and 'chat' for an interactive model where history of previous prompts and generated tokens - /// is preserved. - #[arg(long)] - prompt: Option, - - /// The length of the sample to generate (in tokens). - #[arg(short = 'n', long, default_value_t = 1000)] - sample_len: usize, - - /// The tokenizer config in json format. - #[arg(long)] - tokenizer: Option, - - /// The temperature used to generate samples, use 0 for greedy sampling. - #[arg(long, default_value_t = 0.8)] - temperature: f64, - - /// Nucleus sampling probability cutoff. - #[arg(long)] - top_p: Option, - - /// Only sample among the top K samples. - #[arg(long)] - top_k: Option, - - /// The seed to use when generating random samples. - #[arg(long, default_value_t = 299792458)] - seed: u64, - - /// Enable tracing (generates a trace-timestamp.json file). - #[arg(long)] - tracing: bool, - - /// Display the token for the specified prompt. - #[arg(long)] - verbose_prompt: bool, - - /// Process prompt elements separately. - #[arg(long)] - split_prompt: bool, - - /// Run on CPU rather than GPU even if a GPU is available. - #[arg(long)] - cpu: bool, - - /// Penalty to be applied for repeating tokens, 1. means no penalty. - #[arg(long, default_value_t = 1.1)] - repeat_penalty: f32, - - /// The context size to consider for the repeat penalty. - #[arg(long, default_value_t = 64)] - repeat_last_n: usize, - - /// The model size to use. - #[arg(long, default_value = "7b")] - which: Which, - - /// Group-Query Attention, use 8 for the 70B version of LLaMAv2. - #[arg(long)] - gqa: Option, - - /// Use the slower dmmv cuda kernel. - #[arg(long)] - force_dmmv: bool, -} - -impl Args { - fn tokenizer(&self) -> anyhow::Result { - let tokenizer_path = match &self.tokenizer { - Some(config) => std::path::PathBuf::from(config), - None => { - let api = hf_hub::api::sync::Api::new()?; - let repo = self.which.tokenizer_repo(); - let api = api.model(repo.to_string()); - api.get("tokenizer.json")? - } - }; - Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg) - } - - fn model(&self) -> anyhow::Result { - let model_path = match &self.model { - Some(config) => std::path::PathBuf::from(config), - None => { - let (repo, filename) = match self.which { - Which::L7b => ("TheBloke/Llama-2-7B-GGML", "llama-2-7b.ggmlv3.q4_0.bin"), - Which::L13b => ("TheBloke/Llama-2-13B-GGML", "llama-2-13b.ggmlv3.q4_0.bin"), - Which::L70b => ("TheBloke/Llama-2-70B-GGML", "llama-2-70b.ggmlv3.q4_0.bin"), - Which::L7bChat => ( - "TheBloke/Llama-2-7B-Chat-GGML", - "llama-2-7b-chat.ggmlv3.q4_0.bin", - ), - Which::L13bChat => ( - "TheBloke/Llama-2-13B-Chat-GGML", - "llama-2-13b-chat.ggmlv3.q4_0.bin", - ), - Which::L70bChat => ( - "TheBloke/Llama-2-70B-Chat-GGML", - "llama-2-70b-chat.ggmlv3.q4_0.bin", - ), - Which::L7bCode => ("TheBloke/CodeLlama-7B-GGUF", "codellama-7b.Q8_0.gguf"), - Which::L13bCode => ("TheBloke/CodeLlama-13B-GGUF", "codellama-13b.Q8_0.gguf"), - Which::L34bCode => ("TheBloke/CodeLlama-34B-GGUF", "codellama-34b.Q8_0.gguf"), - Which::Leo7b => ( - "TheBloke/leo-hessianai-7B-GGUF", - "leo-hessianai-7b.Q4_K_M.gguf", - ), - Which::Leo13b => ( - "TheBloke/leo-hessianai-13B-GGUF", - "leo-hessianai-13b.Q4_K_M.gguf", - ), - Which::Mixtral => ( - "TheBloke/Mixtral-8x7B-v0.1-GGUF", - "mixtral-8x7b-v0.1.Q4_K_M.gguf", - ), - Which::MixtralInstruct => ( - "TheBloke/Mixtral-8x7B-Instruct-v0.1-GGUF", - "mixtral-8x7b-instruct-v0.1.Q4_K_M.gguf", - ), - Which::Mistral7b => ( - "TheBloke/Mistral-7B-v0.1-GGUF", - "mistral-7b-v0.1.Q4_K_S.gguf", - ), - Which::Mistral7bInstruct => ( - "TheBloke/Mistral-7B-Instruct-v0.1-GGUF", - "mistral-7b-instruct-v0.1.Q4_K_S.gguf", - ), - Which::Mistral7bInstructV02 => ( - "TheBloke/Mistral-7B-Instruct-v0.2-GGUF", - "mistral-7b-instruct-v0.2.Q4_K_S.gguf", - ), - Which::Zephyr7bAlpha => ( - "TheBloke/zephyr-7B-alpha-GGUF", - "zephyr-7b-alpha.Q4_K_M.gguf", - ), - Which::Zephyr7bBeta => { - ("TheBloke/zephyr-7B-beta-GGUF", "zephyr-7b-beta.Q4_K_M.gguf") - } - Which::OpenChat35 => ("TheBloke/openchat_3.5-GGUF", "openchat_3.5.Q4_K_M.gguf"), - Which::Starling7bAlpha => ( - "TheBloke/Starling-LM-7B-alpha-GGUF", - "starling-lm-7b-alpha.Q4_K_M.gguf", - ), - // TODO: swap to TheBloke model when available - Which::L8b => ( - "QuantFactory/Meta-Llama-3-8B-GGUF", - "Meta-Llama-3-8B.Q4_K_S.gguf", - ), - Which::Phi3 => ( - "microsoft/Phi-3-mini-4k-instruct-gguf", - "Phi-3-mini-4k-instruct-q4.gguf", - ), - Which::SmolLM2_360MInstruct => ( - "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF", - "smollm2-360m-instruct-q8_0.gguf", - ), - Which::SmolLM2_1BInstruct => ( - "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF", - "smollm2-1.7b-instruct-q4_k_m.gguf", - ), - Which::DeepseekR1Llama8b => ( - "unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF", - "DeepSeek-R1-Distill-Llama-8B-Q4_K_M.gguf", - ), - }; - let revision = if self.which == Which::Phi3 { - "5eef2ce24766d31909c0b269fe90c817a8f263fb" - } else { - "main" - }; - let api = hf_hub::api::sync::Api::new()?; - api.repo(hf_hub::Repo::with_revision( - repo.to_string(), - hf_hub::RepoType::Model, - revision.to_string(), - )) - .get(filename)? - } - }; - Ok(model_path) - } -} - -fn format_size(size_in_bytes: usize) -> String { - if size_in_bytes < 1_000 { - format!("{size_in_bytes}B") - } else if size_in_bytes < 1_000_000 { - format!("{:.2}KB", size_in_bytes as f64 / 1e3) - } else if size_in_bytes < 1_000_000_000 { - format!("{:.2}MB", size_in_bytes as f64 / 1e6) - } else { - format!("{:.2}GB", size_in_bytes as f64 / 1e9) - } -} - -fn main() -> anyhow::Result<()> { - use tracing_chrome::ChromeLayerBuilder; - use tracing_subscriber::prelude::*; - - let args = Args::parse(); - - #[cfg(feature = "cuda")] - candle::quantized::cuda::set_force_dmmv(args.force_dmmv); - - candle::cuda::set_gemm_reduced_precision_f16(true); - candle::cuda::set_gemm_reduced_precision_bf16(true); - - let _guard = if args.tracing { - let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); - tracing_subscriber::registry().with(chrome_layer).init(); - Some(guard) - } else { - None - }; - - println!( - "avx: {}, neon: {}, simd128: {}, f16c: {}", - candle::utils::with_avx(), - candle::utils::with_neon(), - candle::utils::with_simd128(), - candle::utils::with_f16c() - ); - println!( - "temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}", - args.temperature, args.repeat_penalty, args.repeat_last_n - ); - - let model_path = args.model()?; - let mut file = std::fs::File::open(&model_path)?; - let start = std::time::Instant::now(); - let device = candle_examples::device(args.cpu)?; - - let mut model = match model_path.extension().and_then(|v| v.to_str()) { - Some("gguf") => { - let model = gguf_file::Content::read(&mut file).map_err(|e| e.with_path(model_path))?; - let mut total_size_in_bytes = 0; - for (_, tensor) in model.tensor_infos.iter() { - let elem_count = tensor.shape.elem_count(); - total_size_in_bytes += - elem_count * tensor.ggml_dtype.type_size() / tensor.ggml_dtype.block_size(); - } - println!( - "loaded {:?} tensors ({}) in {:.2}s", - model.tensor_infos.len(), - &format_size(total_size_in_bytes), - start.elapsed().as_secs_f32(), - ); - ModelWeights::from_gguf(model, &mut file, &device)? - } - Some("ggml" | "bin") | Some(_) | None => { - let model = ggml_file::Content::read(&mut file, &device) - .map_err(|e| e.with_path(model_path))?; - let mut total_size_in_bytes = 0; - for (_, tensor) in model.tensors.iter() { - let elem_count = tensor.shape().elem_count(); - total_size_in_bytes += - elem_count * tensor.dtype().type_size() / tensor.dtype().block_size(); - } - println!( - "loaded {:?} tensors ({}) in {:.2}s", - model.tensors.len(), - &format_size(total_size_in_bytes), - start.elapsed().as_secs_f32(), - ); - println!("params: {:?}", model.hparams); - let default_gqa = match args.which { - Which::L7b - | Which::L13b - | Which::L7bChat - | Which::L13bChat - | Which::L7bCode - | Which::L13bCode - | Which::L34bCode - | Which::Leo7b - | Which::Leo13b - | Which::L8b - | Which::SmolLM2_1BInstruct - | Which::SmolLM2_360MInstruct - | Which::DeepseekR1Llama8b - | Which::Phi3 => 1, - Which::Mixtral - | Which::MixtralInstruct - | Which::Mistral7b - | Which::Mistral7bInstruct - | Which::Mistral7bInstructV02 - | Which::Zephyr7bAlpha - | Which::Zephyr7bBeta - | Which::L70b - | Which::L70bChat - | Which::OpenChat35 - | Which::Starling7bAlpha => 8, - }; - ModelWeights::from_ggml(model, args.gqa.unwrap_or(default_gqa))? - } - }; - println!("model built"); - - let tokenizer = args.tokenizer()?; - let mut tos = TokenOutputStream::new(tokenizer); - let prompt = match args.prompt.as_deref() { - Some("chat") => Prompt::Chat, - Some("interactive") => Prompt::Interactive, - Some(s) => Prompt::One(s.to_string()), - None => Prompt::One(DEFAULT_PROMPT.to_string()), - }; - - let mut pre_prompt_tokens = vec![]; - for prompt_index in 0.. { - let prompt_str = match &prompt { - Prompt::One(prompt) => prompt.clone(), - Prompt::Interactive | Prompt::Chat => { - let is_interactive = matches!(prompt, Prompt::Interactive); - print!("> "); - std::io::stdout().flush()?; - let mut prompt = String::new(); - std::io::stdin().read_line(&mut prompt)?; - if prompt.ends_with('\n') { - prompt.pop(); - if prompt.ends_with('\r') { - prompt.pop(); - } - } - if args.which.is_open_chat() { - format!("GPT4 Correct User: {prompt}<|end_of_turn|>GPT4 Correct Assistant:") - } else if args.which.is_zephyr() { - if prompt_index == 0 || is_interactive { - format!("<|system|>\n\n<|user|>\n{prompt}\n<|assistant|>",) - } else { - format!("<|user|>\n{prompt}\n<|assistant|>") - } - } else if args.which.is_mistral() { - format!("[INST] {prompt} [/INST]") - } else if args.which.is_deepseek() { - format!("<｜User｜>{prompt}<｜Assistant｜>") - } else { - prompt - } - } - }; - print!("{}", &prompt_str); - let tokens = tos - .tokenizer() - .encode(prompt_str, true) - .map_err(anyhow::Error::msg)?; - if args.verbose_prompt { - for (token, id) in tokens.get_tokens().iter().zip(tokens.get_ids().iter()) { - let token = token.replace('▁', " ").replace("<0x0A>", "\n"); - println!("{id:7} -> '{token}'"); - } - } - - let prompt_tokens = [&pre_prompt_tokens, tokens.get_ids()].concat(); - let to_sample = args.sample_len.saturating_sub(1); - let prompt_tokens = if prompt_tokens.len() + to_sample > model::MAX_SEQ_LEN - 10 { - let to_remove = prompt_tokens.len() + to_sample + 10 - model::MAX_SEQ_LEN; - prompt_tokens[prompt_tokens.len().saturating_sub(to_remove)..].to_vec() - } else { - prompt_tokens - }; - let mut all_tokens = vec![]; - let mut logits_processor = { - let temperature = args.temperature; - let sampling = if temperature <= 0. { - Sampling::ArgMax - } else { - match (args.top_k, args.top_p) { - (None, None) => Sampling::All { temperature }, - (Some(k), None) => Sampling::TopK { k, temperature }, - (None, Some(p)) => Sampling::TopP { p, temperature }, - (Some(k), Some(p)) => Sampling::TopKThenTopP { k, p, temperature }, - } - }; - LogitsProcessor::from_sampling(args.seed, sampling) - }; - - let start_prompt_processing = std::time::Instant::now(); - let mut next_token = if !args.split_prompt { - let input = Tensor::new(prompt_tokens.as_slice(), &device)?.unsqueeze(0)?; - let logits = model.forward(&input, 0)?; - let logits = logits.squeeze(0)?; - logits_processor.sample(&logits)? - } else { - let mut next_token = 0; - for (pos, token) in prompt_tokens.iter().enumerate() { - let input = Tensor::new(&[*token], &device)?.unsqueeze(0)?; - let logits = model.forward(&input, pos)?; - let logits = logits.squeeze(0)?; - next_token = logits_processor.sample(&logits)? - } - next_token - }; - let prompt_dt = start_prompt_processing.elapsed(); - all_tokens.push(next_token); - if let Some(t) = tos.next_token(next_token)? { - print!("{t}"); - std::io::stdout().flush()?; - } - - let eos_token = match args.which { - Which::SmolLM2_360MInstruct | Which::SmolLM2_1BInstruct => "<|endoftext|>", - Which::L8b => "<|end_of_text|>", - Which::DeepseekR1Llama8b => "<｜end▁of▁sentence｜>", - _ => match args.which.is_open_chat() { - true => "<|end_of_turn|>", - false => "", - }, - }; - - let eos_token = *tos.tokenizer().get_vocab(true).get(eos_token).unwrap(); - let start_post_prompt = std::time::Instant::now(); - let mut sampled = 0; - for index in 0..to_sample { - let input = Tensor::new(&[next_token], &device)?.unsqueeze(0)?; - let logits = model.forward(&input, prompt_tokens.len() + index)?; - let logits = logits.squeeze(0)?; - let logits = if args.repeat_penalty == 1. { - logits - } else { - let start_at = all_tokens.len().saturating_sub(args.repeat_last_n); - candle_transformers::utils::apply_repeat_penalty( - &logits, - args.repeat_penalty, - &all_tokens[start_at..], - )? - }; - next_token = logits_processor.sample(&logits)?; - all_tokens.push(next_token); - if let Some(t) = tos.next_token(next_token)? { - print!("{t}"); - std::io::stdout().flush()?; - } - sampled += 1; - if next_token == eos_token { - break; - }; - } - if let Some(rest) = tos.decode_rest().map_err(candle::Error::msg)? { - print!("{rest}"); - } - std::io::stdout().flush()?; - let dt = start_post_prompt.elapsed(); - println!( - "\n\n{:4} prompt tokens processed: {:.2} token/s", - prompt_tokens.len(), - prompt_tokens.len() as f64 / prompt_dt.as_secs_f64(), - ); - println!( - "{sampled:4} tokens generated: {:.2} token/s", - sampled as f64 / dt.as_secs_f64(), - ); - - match prompt { - Prompt::One(_) => break, - Prompt::Interactive => {} - Prompt::Chat => { - pre_prompt_tokens = [prompt_tokens.as_slice(), all_tokens.as_slice()].concat() - } - } - } - - Ok(()) -} diff --git a/crates/goose/examples/test_candle_minimal.rs b/crates/goose/examples/test_candle_minimal.rs deleted file mode 100644 index e15e08b5f1..0000000000 --- a/crates/goose/examples/test_candle_minimal.rs +++ /dev/null @@ -1,54 +0,0 @@ -// Minimal test to see if candle works the same in goose repo as in candle repo -use anyhow::Result; -use candle_core::{Device, Tensor}; -use candle_transformers::models::quantized_llama::ModelWeights; -use candle_core::quantized::gguf_file; -use tokenizers::Tokenizer; - -fn main() -> Result<()> { - let home = std::env::var("HOME").map_err(anyhow::Error::msg)?; - let model_path = std::path::PathBuf::from(format!("{}/.local/share/goose/models/llama-3.2-3b.gguf", home)); - let tokenizer_path = std::path::PathBuf::from(format!("{}/.local/share/goose/models/llama-3.2-3b_tokenizer.json", home)); - - let prompt = std::fs::read_to_string("/tmp/goose_prompt_stream.txt")?; - - // Device - let device = if let Ok(device) = Device::new_metal(0) { - device - } else { - Device::Cpu - }; - - // Load model - let mut file = std::fs::File::open(&model_path)?; - let content = gguf_file::Content::read(&mut file)?; - let mut model = ModelWeights::from_gguf(content, &mut file, &device)?; - - // Load tokenizer - let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg)?; - - // Tokenize - use false since prompt already has <|begin_of_text|> - let tokens = tokenizer.encode(prompt.as_str(), false).map_err(anyhow::Error::msg)?; - let prompt_tokens = tokens.get_ids().to_vec(); - - println!("Prompt tokens: {}", prompt_tokens.len()); - - // Split-prompt prefill - let mut next_token = 0u32; - for (pos, &token) in prompt_tokens.iter().enumerate() { - let input = Tensor::new(&[token], &device)?.unsqueeze(0)?; - let logits = model.forward(&input, pos)?; - let logits = logits.squeeze(0)?; - next_token = logits.argmax(0)?.to_scalar::()?; - - if pos >= prompt_tokens.len().saturating_sub(5) { - println!("pos={}, input_token={}, next_token={}", pos, token, next_token); - } - } - - // Decode first token - let decoded = tokenizer.decode(&[next_token], false).map_err(anyhow::Error::msg)?; - println!("\nFirst generated token: ID={}, text='{}'", next_token, decoded); - - Ok(()) -} diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index b40f46d393..cd1a3d81d7 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -1495,9 +1495,9 @@ impl ExtensionManager { ) -> Option { if let Ok(provider_guard) = self.provider.try_lock() { if let Some(provider) = provider_guard.as_ref() { - if provider.get_model_config().context_limit() < 9 * 1024 * 1024 { - return None; - } + if provider.get_model_config().context_limit() < 9 * 1024 * 1024 { + return None; + } } } diff --git a/crates/goose/src/dictation/download_manager.rs b/crates/goose/src/dictation/download_manager.rs index 7b7878da9b..0f2f80f25a 100644 --- a/crates/goose/src/dictation/download_manager.rs +++ b/crates/goose/src/dictation/download_manager.rs @@ -1,11 +1,32 @@ use anyhow::Result; use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use tokio::io::AsyncWriteExt; use utoipa::ToSchema; +fn partial_path_for(destination: &Path) -> PathBuf { + destination.with_extension( + destination + .extension() + .map(|e| format!("{}.part", e.to_string_lossy())) + .unwrap_or_else(|| "part".to_string()), + ) +} + +/// Remove any leftover `.part` files in the given directory. +pub fn cleanup_partial_downloads(dir: &Path) { + if let Ok(entries) = std::fs::read_dir(dir) { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().is_some_and(|e| e == "part") { + let _ = std::fs::remove_file(&path); + } + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct DownloadProgress { /// Model ID being downloaded @@ -116,6 +137,8 @@ impl DownloadManager { let downloads = self.downloads.clone(); let model_id_clone = model_id.clone(); + let destination_for_cleanup = destination.clone(); + // Download in background task tokio::spawn(async move { match Self::download_file(&url, &destination, &downloads, &model_id_clone).await { @@ -133,6 +156,10 @@ impl DownloadManager { } } Err(e) => { + // Clean up partial file on failure + let partial = partial_path_for(&destination_for_cleanup); + let _ = tokio::fs::remove_file(&partial).await; + if let Ok(mut downloads) = downloads.lock() { if let Some(progress) = downloads.get_mut(&model_id_clone) { progress.status = DownloadStatus::Failed; @@ -169,7 +196,8 @@ impl DownloadManager { } } - let mut file = tokio::fs::File::create(destination).await?; + let partial_path = partial_path_for(destination); + let mut file = tokio::fs::File::create(&partial_path).await?; let mut bytes_downloaded = 0u64; let start_time = std::time::Instant::now(); @@ -188,8 +216,7 @@ impl DownloadManager { }; if should_cancel { - // Clean up partial download - let _ = tokio::fs::remove_file(destination).await; + let _ = tokio::fs::remove_file(&partial_path).await; return Ok(()); } @@ -229,6 +256,8 @@ impl DownloadManager { } file.flush().await?; + drop(file); + tokio::fs::rename(&partial_path, destination).await?; Ok(()) } diff --git a/crates/goose/src/prompts/tiny_model_system.md b/crates/goose/src/prompts/tiny_model_system.md index a2bc97cb67..08b1f116c0 100644 --- a/crates/goose/src/prompts/tiny_model_system.md +++ b/crates/goose/src/prompts/tiny_model_system.md @@ -1,21 +1,22 @@ -You are goose (lowercase), an AI assistant created by Block +You are goose, an autonomous AI agent created by Block. You act on the user's +behalf — you do not explain how to do things, you DO them directly. -Help the user using your knowledge, your reasoning and by executing commands -in the {{shell}} shell. +The OS is {{os}}, the shell is {{shell}}, and the working directory is {{working_directory}} -The OS is {{os}} and the current directory is {{working_directory}} +When the user asks you to do something, take action immediately. Do not describe +what you would do or give instructions — execute the commands yourself. -If you need to execute a shell command, you can do so by starting a new line with $, for example -to look at the files in the current folder, just end your message on +To run a shell command, start a new line with $: $ ls -Other useful commands are: `rg` to search for text, `cat` to read or write files -or `head` to just see part of it. use `echo "content" > file` for small files, -`cat` for longer. +Keep your responses brief. State what you are doing, then do it. For example: -# Important +User: how many files are in /tmp? +You: Let me check. +$ ls -1 /tmp | wc -l -Only execute shell commands when you need to read a file you know exists or when -you need to create a file or execute a command. Do not use shell commands if you -know the answer. Do not assume files or folders exists until you check. \ No newline at end of file +After a command runs, you will see its output. Use the output to answer the user +or take the next step. Do not repeat commands you have already run. + +Do not use shell commands if you already know the answer. \ No newline at end of file diff --git a/crates/goose/src/providers/local_inference.rs b/crates/goose/src/providers/local_inference.rs index 47cc7b100f..260d8ed80b 100644 --- a/crates/goose/src/providers/local_inference.rs +++ b/crates/goose/src/providers/local_inference.rs @@ -1,3 +1,6 @@ +pub mod hf_models; +pub mod local_model_registry; + use crate::config::paths::Paths; use crate::conversation::message::{Message, MessageContent}; use crate::model::ModelConfig; @@ -5,29 +8,94 @@ use crate::providers::base::{ MessageStream, Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage, }; use crate::providers::errors::ProviderError; +use crate::providers::formats::openai::format_tools; use anyhow::Result; use async_stream::try_stream; use async_trait::async_trait; -use candle_core::{Device, Tensor}; -use candle_transformers::models::{quantized_llama, quantized_phi, quantized_phi3}; use futures::future::BoxFuture; +use llama_cpp_2::context::params::LlamaContextParams; +use llama_cpp_2::llama_backend::LlamaBackend; +use llama_cpp_2::llama_batch::LlamaBatch; +use llama_cpp_2::model::params::LlamaModelParams; +use llama_cpp_2::model::{AddBos, LlamaChatMessage, LlamaChatTemplate, LlamaModel}; +use llama_cpp_2::sampling::LlamaSampler; +use llama_cpp_2::{list_llama_ggml_backend_devices, LlamaBackendDeviceType}; use rmcp::model::{CallToolRequestParams, Role, Tool}; use serde::{Deserialize, Serialize}; -use serde_json::json; +use serde_json::{json, Value}; use std::borrow::Cow; +use std::collections::HashMap; +use std::num::NonZeroU32; use std::path::PathBuf; -use std::sync::Arc; -use tokenizers::Tokenizer; +use std::sync::{Arc, Mutex as StdMutex, Weak}; use tokio::sync::Mutex; use utoipa::ToSchema; use uuid::Uuid; +type ModelSlot = Arc>>; + +/// Owns the llama backend and all cached models. Field order matters: +/// `models` is declared before `backend` so Rust drops all loaded models +/// (and their Metal/GPU resources) before the backend calls +/// `llama_backend_free()`, avoiding the ggml-metal assertion on shutdown. +pub struct InferenceRuntime { + models: StdMutex>, + backend: LlamaBackend, +} + +/// Global weak reference used to share a single `InferenceRuntime` across +/// all providers and server routes. Only a `Weak` is stored — strong `Arc`s +/// live in providers and `AppState`. When all strong refs drop (normal +/// shutdown), the runtime is deallocated and the backend freed. The `Weak` +/// left behind is inert during `__cxa_finalize`, so no ggml statics race. +static RUNTIME: StdMutex> = StdMutex::new(Weak::new()); + +impl InferenceRuntime { + pub fn get_or_init() -> Arc { + let mut guard = RUNTIME.lock().expect("runtime lock poisoned"); + if let Some(runtime) = guard.upgrade() { + return runtime; + } + let backend = match LlamaBackend::init() { + Ok(b) => b, + Err(llama_cpp_2::LlamaCppError::BackendAlreadyInitialized) => { + panic!("LlamaBackend already initialized but Weak was dead — should be impossible") + } + Err(e) => panic!("Failed to init llama backend: {}", e), + }; + let runtime = Arc::new(Self { + models: StdMutex::new(HashMap::new()), + backend, + }); + *guard = Arc::downgrade(&runtime); + runtime + } + + pub fn backend(&self) -> &LlamaBackend { + &self.backend + } + + fn get_or_create_model_slot(&self, model_id: &str) -> ModelSlot { + let mut map = self.models.lock().expect("model cache lock poisoned"); + map.entry(model_id.to_string()) + .or_insert_with(|| Arc::new(Mutex::new(None))) + .clone() + } + + fn other_model_slots(&self, keep_model_id: &str) -> Vec { + let map = self.models.lock().expect("model cache lock poisoned"); + map.iter() + .filter(|(id, _)| id.as_str() != keep_model_id) + .map(|(_, slot)| slot.clone()) + .collect() + } +} + const PROVIDER_NAME: &str = "local"; const DEFAULT_MODEL: &str = "llama-3.2-1b"; pub const LOCAL_LLM_MODEL_CONFIG_KEY: &str = "LOCAL_LLM_MODEL"; -// Load tiny model system prompt with environment context fn load_tiny_model_prompt() -> String { use std::env; @@ -54,7 +122,6 @@ fn load_tiny_model_prompt() -> String { }); crate::prompt_template::render_template("tiny_model_system.md", &context).unwrap_or_else(|e| { - // Fallback if template fails to load eprintln!("WARNING: Failed to load tiny_model_system.md: {:?}", e); "You are Goose, an AI assistant. You can execute shell commands by starting lines with $." .to_string() @@ -70,51 +137,15 @@ pub enum ModelTier { Large, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ChatTemplate { - Llama3, - ChatML, - Mistral, -} - -impl Default for ChatTemplate { - fn default() -> Self { - ChatTemplate::Llama3 - } -} - -impl ChatTemplate { - /// Get EOS token strings to strip from output - fn eos_strings(&self) -> &[&str] { - match self { - ChatTemplate::Llama3 => &["<|eot_id|>", "<|end_of_text|>"], - ChatTemplate::ChatML => &["<|im_end|>"], - ChatTemplate::Mistral => &[""], - } - } -} - #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct LocalLlmModel { - /// Model identifier (e.g., "llama-3.2-1b") pub id: &'static str, - /// Display name pub name: &'static str, - /// Model file size in MB pub size_mb: u32, - /// Maximum context window in tokens pub context_limit: usize, - /// Download URL for the model GGUF file pub url: &'static str, - /// Download URL for the tokenizer JSON - pub tokenizer_url: &'static str, - /// Description and use case pub description: &'static str, - /// Model tier/category pub tier: ModelTier, - /// Chat template format - #[serde(skip)] - pub chat_template: ChatTemplate, } const LOCAL_LLM_MODELS: &[LocalLlmModel] = &[ @@ -124,10 +155,8 @@ const LOCAL_LLM_MODELS: &[LocalLlmModel] = &[ size_mb: 700, context_limit: 4096, url: "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_K_M.gguf", - tokenizer_url: "https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B/resolve/main/tokenizer.json", description: "Fastest, CPU-optimized for quick responses", tier: ModelTier::Tiny, - chat_template: ChatTemplate::Llama3, }, LocalLlmModel { id: "llama-3.2-3b", @@ -135,10 +164,8 @@ const LOCAL_LLM_MODELS: &[LocalLlmModel] = &[ size_mb: 2000, context_limit: 8192, url: "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q4_K_M.gguf", - tokenizer_url: "https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B/resolve/main/tokenizer.json", description: "Good balance of speed and quality for laptops", tier: ModelTier::Small, - chat_template: ChatTemplate::Llama3, }, LocalLlmModel { id: "hermes-2-pro-7b", @@ -146,10 +173,8 @@ const LOCAL_LLM_MODELS: &[LocalLlmModel] = &[ size_mb: 4500, context_limit: 8192, url: "https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B-GGUF/resolve/main/Hermes-2-Pro-Llama-3-8B-Q4_K_M.gguf", - tokenizer_url: "https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B/resolve/main/tokenizer.json", description: "High quality for desktops with GPU", tier: ModelTier::Medium, - chat_template: ChatTemplate::ChatML, }, LocalLlmModel { id: "mistral-small-22b", @@ -157,10 +182,8 @@ const LOCAL_LLM_MODELS: &[LocalLlmModel] = &[ size_mb: 13000, context_limit: 32768, url: "https://huggingface.co/bartowski/Mistral-Small-Instruct-2409-GGUF/resolve/main/Mistral-Small-Instruct-2409-Q4_K_M.gguf", - tokenizer_url: "https://huggingface.co/mistralai/Mistral-Small-Instruct-2409/resolve/main/tokenizer.json", description: "Highest quality with long context support", tier: ModelTier::Large, - chat_template: ChatTemplate::Mistral, }, ]; @@ -169,12 +192,8 @@ impl LocalLlmModel { Paths::in_data_dir("models").join(format!("{}.gguf", self.id)) } - pub fn tokenizer_path(&self) -> PathBuf { - Paths::in_data_dir("models").join(format!("{}_tokenizer.json", self.id)) - } - pub fn is_downloaded(&self) -> bool { - self.local_path().exists() && self.tokenizer_path().exists() + self.local_path().exists() } } @@ -186,120 +205,200 @@ pub fn get_local_model(id: &str) -> Option<&'static LocalLlmModel> { LOCAL_LLM_MODELS.iter().find(|m| m.id == id) } -pub fn recommend_local_model() -> &'static str { - let has_gpu = Device::new_cuda(0).is_ok() || Device::new_metal(0).is_ok(); - let mem_mb = sys_info::mem_info().map(|m| m.avail / 1024).unwrap_or(0); +/// Resolve model path, context limit, and settings for any model ID — checks registry first, +/// then falls back to the hardcoded featured list. +pub fn resolve_model_path( + model_id: &str, +) -> Option<( + PathBuf, + usize, + crate::providers::local_inference::local_model_registry::ModelSettings, +)> { + use crate::providers::local_inference::local_model_registry::get_registry; - if has_gpu && mem_mb >= 16_000 { - "hermes-2-pro-7b" // Medium tier - GPU with lots of memory - } else if mem_mb >= 4_000 { - "llama-3.2-3b" // Small tier - decent memory + // Check registry first (covers both HF-downloaded and migrated legacy models) + if let Ok(registry) = get_registry().lock() { + if let Some(entry) = registry.get_model(model_id) { + let ctx = entry.settings.context_size.unwrap_or(0) as usize; + return Some((entry.local_path.clone(), ctx, entry.settings.clone())); + } + } + + // Fall back to hardcoded featured list + get_local_model(model_id).map(|m| { + ( + m.local_path(), + m.context_limit, + crate::providers::local_inference::local_model_registry::ModelSettings::default(), + ) + }) +} + +pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 { + let _ = &runtime.backend; + let devices = list_llama_ggml_backend_devices(); + + let accel_memory = devices + .iter() + .filter(|d| { + matches!( + d.device_type, + LlamaBackendDeviceType::Gpu + | LlamaBackendDeviceType::IntegratedGpu + | LlamaBackendDeviceType::Accelerator + ) + }) + .map(|d| d.memory_free as u64) + .max() + .unwrap_or(0); + + if accel_memory > 0 { + accel_memory } else { - "llama-3.2-1b" // Tiny tier - low memory + devices + .iter() + .filter(|d| d.device_type == LlamaBackendDeviceType::Cpu) + .map(|d| d.memory_free as u64) + .max() + .unwrap_or(0) } } -enum ModelWeights { - Llama(quantized_llama::ModelWeights), - Phi(quantized_phi::ModelWeights), - Phi3(quantized_phi3::ModelWeights), -} +pub fn recommend_local_model(runtime: &InferenceRuntime) -> &'static str { + let effective_memory_mb = available_inference_memory_bytes(runtime) / (1024 * 1024); -impl ModelWeights { - fn forward(&mut self, input: &Tensor, pos: usize) -> candle_core::Result { - match self { - ModelWeights::Llama(m) => m.forward(input, pos), - ModelWeights::Phi(m) => m.forward(input, pos), - ModelWeights::Phi3(m) => m.forward(input, pos), - } + if effective_memory_mb >= 16_000 { + "mistral-small-22b" + } else if effective_memory_mb >= 6_000 { + "hermes-2-pro-7b" + } else if effective_memory_mb >= 3_000 { + "llama-3.2-3b" + } else { + "llama-3.2-1b" } } struct LoadedModel { - model: ModelWeights, - tokenizer: Tokenizer, - device: Device, - eos_token_ids: Vec, + model: LlamaModel, + template: LlamaChatTemplate, +} + +enum EmulatorAction { + Text(String), + ShellCommand(String), + ExecuteCode(String), +} + +enum ParserState { + Normal, + InCommand, + InExecuteBlock, } -/// Streaming parser for emulator commands -/// Accumulates chunks and emits complete text or commands struct StreamingEmulatorParser { buffer: String, - in_command: bool, - command_start_pos: usize, + state: ParserState, + code_mode_enabled: bool, } impl StreamingEmulatorParser { - fn new() -> Self { + fn new(code_mode_enabled: bool) -> Self { Self { buffer: String::new(), - in_command: false, - command_start_pos: 0, + state: ParserState::Normal, + code_mode_enabled, } } - /// Process a chunk and return any complete items (text or commands) - /// Returns (optional_text, optional_command) - fn process_chunk(&mut self, chunk: &str) -> Vec<(Option, Option)> { + fn process_chunk(&mut self, chunk: &str) -> Vec { self.buffer.push_str(chunk); let mut results = Vec::new(); loop { - if self.in_command { - // Look for newline to end the command - if let Some(newline_pos) = self.buffer[self.command_start_pos..].find('\n') { - let absolute_pos = self.command_start_pos + newline_pos; - // Extract command from "$ command" - let command_line = &self.buffer[self.command_start_pos..absolute_pos]; - if let Some(command) = command_line.strip_prefix('$') { - let command = command.trim(); - if !command.is_empty() { - results.push((None, Some(command.to_string()))); + match self.state { + ParserState::InCommand => { + if let Some((command_line, rest)) = self.buffer.split_once('\n') { + if let Some(command) = command_line.strip_prefix('$') { + let command = command.trim(); + if !command.is_empty() { + results.push(EmulatorAction::ShellCommand(command.to_string())); + } } + self.buffer = rest.to_string(); + self.state = ParserState::Normal; + } else { + break; } - // Remove processed part from buffer - self.buffer = self.buffer[absolute_pos + 1..].to_string(); - self.in_command = false; - self.command_start_pos = 0; - } else { - // Command not complete yet, wait for more chunks - break; } - } else { - // Look for command start: "\n$" or "$" at beginning - if let Some(pos) = self.buffer.find("\n$") { - // Emit text before the command - let text = self.buffer[..pos + 1].to_string(); // Include the \n - if !text.trim().is_empty() { - results.push((Some(text), None)); - } - // Remove text from buffer, start command parsing - self.buffer = self.buffer[pos + 1..].to_string(); // Buffer now starts with "$" - self.in_command = true; - self.command_start_pos = 0; - } else if self.buffer.starts_with('$') && self.buffer.len() == chunk.len() { - // Command at very start of response (first chunk) - self.in_command = true; - self.command_start_pos = 0; - } else { - // No command found, but keep last few chars in case of split pattern - // E.g., chunk ends with "\n" and next starts with "$" - if self.buffer.chars().count() > 2 && !self.buffer.ends_with('\n') { - // Emit all but last 2 characters as safe text (use char boundaries) - let mut chars = self.buffer.chars(); - let keep_count = 2; - let emit_count = self.buffer.chars().count() - keep_count; - - let emit_text: String = chars.by_ref().take(emit_count).collect(); - let keep_text: String = chars.collect(); - - if !emit_text.is_empty() { - results.push((Some(emit_text), None)); + ParserState::InExecuteBlock => { + // Look for closing ``` to end the execute block + if let Some(end_idx) = self.buffer.find("\n```") { + #[allow(clippy::string_slice)] + let code = self.buffer[..end_idx].to_string(); + // Skip past the closing ``` and any trailing newline + #[allow(clippy::string_slice)] + let rest = &self.buffer[end_idx + 4..]; + let rest = rest.strip_prefix('\n').unwrap_or(rest); + self.buffer = rest.to_string(); + self.state = ParserState::Normal; + if !code.trim().is_empty() { + results.push(EmulatorAction::ExecuteCode(code)); } - self.buffer = keep_text; + } else { + // Still accumulating code — wait for closing fence + break; + } + } + ParserState::Normal => { + // Check for ```execute block (code mode) + if self.code_mode_enabled { + if let Some((before, after)) = self.buffer.split_once("```execute\n") { + if !before.trim().is_empty() { + results.push(EmulatorAction::Text(before.to_string())); + } + self.buffer = after.to_string(); + self.state = ParserState::InExecuteBlock; + continue; + } + // Also handle without newline after tag (accumulating) + if self.buffer.ends_with("```execute") { + let before = self.buffer.trim_end_matches("```execute"); + if !before.trim().is_empty() { + results.push(EmulatorAction::Text(before.to_string())); + } + self.buffer.clear(); + self.state = ParserState::InExecuteBlock; + continue; + } + } + + // Check for $ command + if let Some((before_dollar, from_dollar)) = self.buffer.split_once("\n$") { + let text = format!("{}\n", before_dollar); + if !text.trim().is_empty() { + results.push(EmulatorAction::Text(text)); + } + self.buffer = format!("${}", from_dollar); + self.state = ParserState::InCommand; + } else if self.buffer.starts_with('$') && self.buffer.len() == chunk.len() { + self.state = ParserState::InCommand; + } else { + // Hold back a small tail in case it's the start of + // a ``` fence or a \n$ command prefix. + let hold_back = if self.code_mode_enabled { 12 } else { 2 }; + let char_count = self.buffer.chars().count(); + if char_count > hold_back && !self.buffer.ends_with('\n') { + let mut chars = self.buffer.chars(); + let emit_count = char_count - hold_back; + let emit_text: String = chars.by_ref().take(emit_count).collect(); + let keep_text: String = chars.collect(); + if !emit_text.is_empty() { + results.push(EmulatorAction::Text(emit_text)); + } + self.buffer = keep_text; + } + break; } - break; } } } @@ -307,376 +406,579 @@ impl StreamingEmulatorParser { results } - /// Flush any remaining buffer content, handling incomplete commands - fn flush(&mut self) -> Vec<(Option, Option)> { + fn flush(&mut self) -> Vec { let mut results = Vec::new(); if !self.buffer.is_empty() { - if self.in_command { - // We're in the middle of parsing a command - complete it - let command_line = self.buffer.trim(); - if let Some(command) = command_line.strip_prefix('$') { - let command = command.trim(); - if !command.is_empty() { - results.push((None, Some(command.to_string()))); + match self.state { + ParserState::InCommand => { + let command_line = self.buffer.trim(); + if let Some(command) = command_line.strip_prefix('$') { + let command = command.trim(); + if !command.is_empty() { + results.push(EmulatorAction::ShellCommand(command.to_string())); + } + } else if !command_line.is_empty() { + results.push(EmulatorAction::Text(self.buffer.clone())); } - } else if !command_line.is_empty() { - // Malformed command, just emit as text - results.push((Some(self.buffer.clone()), None)); } - } else { - // Just regular text remaining - results.push((Some(self.buffer.clone()), None)); + ParserState::InExecuteBlock => { + // Unterminated code block — execute what we have + let code = self.buffer.trim(); + if !code.is_empty() { + results.push(EmulatorAction::ExecuteCode(code.to_string())); + } + } + ParserState::Normal => { + results.push(EmulatorAction::Text(self.buffer.clone())); + } } self.buffer.clear(); - self.in_command = false; + self.state = ParserState::Normal; } results } } +fn extract_text_content(msg: &Message) -> String { + let mut parts = Vec::new(); + + for content in &msg.content { + match content { + MessageContent::Text(text) => { + parts.push(text.text.clone()); + } + MessageContent::ToolRequest(req) => { + if let Ok(call) = &req.tool_call { + if let Some(cmd) = call + .arguments + .as_ref() + .and_then(|a| a.get("command")) + .and_then(|v| v.as_str()) + { + parts.push(format!("$ {}", cmd)); + } else if let Some(code) = call + .arguments + .as_ref() + .and_then(|a| a.get("code")) + .and_then(|v| v.as_str()) + { + parts.push(format!("```execute\n{}\n```", code)); + } + } + } + MessageContent::ToolResponse(response) => match &response.tool_result { + Ok(result) => { + let mut output_parts = Vec::new(); + for content_item in &result.content { + if let Some(text_content) = content_item.as_text() { + output_parts.push(text_content.text.to_string()); + } + } + if !output_parts.is_empty() { + parts.push(format!("Command output:\n{}", output_parts.join("\n"))); + } + } + Err(e) => { + parts.push(format!("Command error: {}", e)); + } + }, + _ => {} + } + } + + parts.join("\n") +} + +/// Build a compact JSON string of tools with only name and description +/// (no parameter schemas) to reduce token count for small context windows. +fn compact_tools_json(tools: &[Tool]) -> Option { + let compact: Vec = tools + .iter() + .map(|t| { + json!({ + "type": "function", + "function": { + "name": t.name, + "description": t.description.as_ref().map(|d| d.as_ref()).unwrap_or(""), + } + }) + }) + .collect(); + serde_json::to_string(&compact).ok() +} + +/// Estimate the maximum context length that can fit in available accelerator/CPU +/// memory based on the model's KV cache requirements. +/// +/// Returns `None` if the model architecture values are unavailable. +fn estimate_max_context_for_memory( + model: &LlamaModel, + runtime: &InferenceRuntime, +) -> Option { + let available = available_inference_memory_bytes(runtime); + if available == 0 { + return None; + } + + // Reserve 20% of available memory for computation scratch buffers and overhead. + let usable = (available as f64 * 0.8) as u64; + + let n_layer = model.n_layer() as u64; + let n_head_kv = model.n_head_kv() as u64; + let n_head = model.n_head() as u64; + let n_embd = model.n_embd() as u64; + + if n_head == 0 || n_layer == 0 || n_head_kv == 0 || n_embd == 0 { + return None; + } + + let head_dim = n_embd / n_head; + // KV cache: 2 (K+V) * n_layer * n_head_kv * head_dim * 2 bytes (f16) per context token + let bytes_per_token = 2 * n_layer * n_head_kv * head_dim * 2; + + if bytes_per_token == 0 { + return None; + } + + Some((usable / bytes_per_token) as usize) +} + +fn effective_context_size( + prompt_token_count: usize, + context_limit: usize, + n_ctx_train: usize, + memory_max_ctx: Option, +) -> usize { + let limit = if context_limit > 0 { + context_limit + } else { + n_ctx_train + }; + + // Cap by estimated memory capacity when available. + let limit = match memory_max_ctx { + Some(mem_max) if mem_max < limit => { + tracing::info!( + "Capping context from {} to {} based on available memory", + limit, + mem_max, + ); + mem_max + } + _ => limit, + }; + + let min_generation_headroom = 512; + let needed = prompt_token_count + min_generation_headroom; + needed.max(limit) +} + +/// Split generated text into (content, tool_calls_json). +/// Looks for the last top-level JSON object containing `"tool_calls"`. +/// Returns the text before it as content, and the JSON string if found. +#[allow(clippy::string_slice)] +fn split_content_and_tool_calls(text: &str) -> (String, Option) { + let trimmed = text.trim_end(); + if !trimmed.ends_with('}') { + return (text.to_string(), None); + } + + // Scan backwards for the matching '{' of the final '}'. + // We only match on ASCII braces so `start` is always a char boundary. + let bytes = trimmed.as_bytes(); + let mut depth = 0i32; + let mut json_start = None; + for i in (0..bytes.len()).rev() { + match bytes[i] { + b'}' => depth += 1, + b'{' => { + depth -= 1; + if depth == 0 { + json_start = Some(i); + break; + } + } + _ => {} + } + } + + let Some(start) = json_start else { + return (text.to_string(), None); + }; + + let json_str = &trimmed[start..]; + let parsed: Value = match serde_json::from_str(json_str) { + Ok(v) => v, + Err(_) => return (text.to_string(), None), + }; + + if parsed + .get("tool_calls") + .and_then(|v| v.as_array()) + .is_none() + { + return (text.to_string(), None); + } + + let content = trimmed[..start].trim_end().to_string(); + (content, Some(json_str.to_string())) +} + +/// Return the byte length of text that is safe to stream. +/// Everything before the last unmatched top-level `{` is safe — the `{` could +/// be the start of a tool-call JSON block still being generated. +/// If all braces are balanced the entire text is safe. +fn safe_stream_end(text: &str) -> usize { + // Hold back from the start of any incomplete tag. + // If we find an unmatched opening, nothing from that point should be streamed. + let xml_hold = text.find("").unwrap_or(text.len()); + + let bytes = text.as_bytes(); + let mut safe_end = bytes.len(); + let mut depth = 0i32; + for (i, &b) in bytes.iter().enumerate() { + match b { + b'{' => { + if depth == 0 { + safe_end = i; + } + depth += 1; + } + b'}' => { + depth -= 1; + if depth == 0 { + safe_end = i + 1; + } + } + _ => { + if depth == 0 { + safe_end = i + 1; + } + } + } + } + + // Also hold back a partial ``, hold them. + let tag = b""; + let tail_hold = { + let mut hold = safe_end; + let check_len = tag.len().min(bytes.len()); + for start in (safe_end.saturating_sub(check_len))..safe_end { + let tail = &bytes[start..safe_end]; + if tag.starts_with(tail) { + hold = start; + break; + } + } + hold + }; + + safe_end.min(xml_hold).min(tail_hold) +} + +/// Extract tool call messages from a JSON object containing "tool_calls". +/// Handles both the model's native format (name/arguments at top level) +/// and the OpenAI format (function.name/function.arguments). +fn extract_tool_call_messages(tool_calls_json: &str, message_id: &str) -> Vec { + let parsed: Value = match serde_json::from_str(tool_calls_json) { + Ok(v) => v, + Err(_) => return vec![], + }; + + let Some(tool_calls) = parsed.get("tool_calls").and_then(|v| v.as_array()) else { + return vec![]; + }; + + let mut messages = Vec::new(); + for tc in tool_calls { + // Try OpenAI format first: {"function": {"name": ..., "arguments": ...}, "id": ...} + // Then model's native format: {"name": ..., "arguments": {...}, "id": ...} + let (name, arguments) = if let Some(func) = tc.get("function") { + let n = func.get("name").and_then(|v| v.as_str()).unwrap_or(""); + let args_str = func + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("{}"); + let args: Option> = serde_json::from_str(args_str).ok(); + (n.to_string(), args) + } else { + let n = tc.get("name").and_then(|v| v.as_str()).unwrap_or(""); + // Arguments may be an object directly (model format) or a string (OAI format) + let args = if let Some(obj) = tc.get("arguments").and_then(|v| v.as_object()) { + Some(obj.clone()) + } else if let Some(s) = tc.get("arguments").and_then(|v| v.as_str()) { + serde_json::from_str(s).ok() + } else { + None + }; + (n.to_string(), args) + }; + + if name.is_empty() { + continue; + } + + let id = tc + .get("id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); + + let tool_call = CallToolRequestParams { + meta: None, + task: None, + name: Cow::Owned(name), + arguments, + }; + + let mut msg = Message::assistant(); + msg.content + .push(MessageContent::tool_request(id, Ok(tool_call))); + msg.id = Some(message_id.to_string()); + messages.push(msg); + } + + messages +} + +/// Parse XML-style tool calls used by models like qwen3-coder. +/// Format: +/// ```text +/// +/// +/// value1 +/// value2 +/// +/// +/// ``` +/// Returns (content_before_tool_calls, vec_of_tool_calls) or None if no XML tool calls found. +#[allow(clippy::type_complexity)] +fn split_content_and_xml_tool_calls( + text: &str, +) -> Option<(String, Vec<(String, serde_json::Map)>)> { + let (content, first_block_and_rest) = text.split_once("")?; + let content = content.trim_end().to_string(); + let mut tool_calls = Vec::new(); + + // Process the first block, then keep splitting on subsequent tags + let mut remaining = first_block_and_rest; + loop { + // Split off the block up to (or take the rest if unclosed) + let (block, after_close) = remaining + .split_once("") + .unwrap_or((remaining, "")); + + if let Some(tool_call) = parse_single_xml_tool_call(block) { + tool_calls.push(tool_call); + } + + // Try to find the next in what remains + match after_close.split_once("") { + Some((_between, next_remaining)) => remaining = next_remaining, + None => break, + } + } + + if tool_calls.is_empty() { + None + } else { + Some((content, tool_calls)) + } +} + +fn parse_single_xml_tool_call(block: &str) -> Option<(String, serde_json::Map)> { + // Extract function name from ... + let (_, after_func_eq) = block.split_once("')?; + let func_name = func_name.trim().to_string(); + + let mut args = serde_json::Map::new(); + let mut rest = func_body; + + // Extract each VALUE + while let Some((_, after_param_eq)) = rest.split_once("') else { + break; + }; + let param_name = param_name.trim().to_string(); + + let (value, after_value) = after_name_close + .split_once("") + .unwrap_or((after_name_close, "")); + let value = value.trim(); + + let json_value = + serde_json::from_str(value).unwrap_or_else(|_| Value::String(value.to_string())); + args.insert(param_name, json_value); + + rest = after_value; + } + + Some((func_name, args)) +} + +fn extract_xml_tool_call_messages( + tool_calls: Vec<(String, serde_json::Map)>, + message_id: &str, +) -> Vec { + tool_calls + .into_iter() + .map(|(name, args)| { + let tool_call = CallToolRequestParams { + meta: None, + task: None, + name: Cow::Owned(name), + arguments: if args.is_empty() { None } else { Some(args) }, + }; + let mut msg = Message::assistant(); + msg.content.push(MessageContent::tool_request( + Uuid::new_v4().to_string(), + Ok(tool_call), + )); + msg.id = Some(message_id.to_string()); + msg + }) + .collect() +} + +fn build_context_params( + ctx_size: u32, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, +) -> LlamaContextParams { + let mut params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(ctx_size)); + + if let Some(n_batch) = settings.n_batch { + params = params.with_n_batch(n_batch); + } + if let Some(n_threads) = settings.n_threads { + params = params.with_n_threads(n_threads); + params = params.with_n_threads_batch(n_threads); + } + if let Some(flash_attn) = settings.flash_attention { + // llama_flash_attn_type: 0 = disabled, 1 = enabled + let policy = if flash_attn { 1 } else { 0 }; + params = params.with_flash_attention_policy(policy); + } + + params +} + +fn build_sampler( + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, +) -> LlamaSampler { + use crate::providers::local_inference::local_model_registry::SamplingConfig; + + let has_penalties = settings.repeat_penalty != 1.0 + || settings.frequency_penalty != 0.0 + || settings.presence_penalty != 0.0; + + let mut samplers: Vec = Vec::new(); + + if has_penalties { + samplers.push(LlamaSampler::penalties( + settings.repeat_last_n, + settings.repeat_penalty, + settings.frequency_penalty, + settings.presence_penalty, + )); + } + + match &settings.sampling { + SamplingConfig::Greedy => { + samplers.push(LlamaSampler::greedy()); + } + SamplingConfig::Temperature { + temperature, + top_k, + top_p, + min_p, + seed, + } => { + samplers.push(LlamaSampler::top_k(*top_k)); + samplers.push(LlamaSampler::top_p(*top_p, 1)); + samplers.push(LlamaSampler::min_p(*min_p, 1)); + samplers.push(LlamaSampler::temp(*temperature)); + samplers.push(LlamaSampler::dist(seed.unwrap_or(0))); + } + SamplingConfig::MirostatV2 { tau, eta, seed } => { + samplers.push(LlamaSampler::mirostat_v2(seed.unwrap_or(0), *tau, *eta)); + } + } + + if samplers.len() == 1 { + samplers.pop().unwrap() + } else { + LlamaSampler::chain_simple(samplers) + } +} + pub struct LocalInferenceProvider { - model: Arc>>, + runtime: Arc, + model: ModelSlot, model_config: ModelConfig, name: String, } impl LocalInferenceProvider { pub async fn from_env(model: ModelConfig) -> Result { + let runtime = InferenceRuntime::get_or_init(); + let model_slot = runtime.get_or_create_model_slot(&model.model_name); Ok(Self { - model: Arc::new(Mutex::new(None)), + runtime, + model: model_slot, model_config: model, name: PROVIDER_NAME.to_string(), }) } - async fn load_model(model_id: &str) -> Result { - // Get model definition - let model_info = get_local_model(model_id) + fn load_model_sync( + runtime: &InferenceRuntime, + model_id: &str, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + ) -> Result { + let (model_path, _context_limit, _) = resolve_model_path(model_id) .ok_or_else(|| ProviderError::ExecutionError(format!("Unknown model: {}", model_id)))?; - let model_path = model_info.local_path(); - let tokenizer_path = model_info.tokenizer_path(); - if !model_path.exists() { return Err(ProviderError::ExecutionError(format!( "Model not downloaded: {}. Please download it from Settings > Local Inference.", - model_info.name + model_id ))); } - tracing::info!("Loading {} from: {}", model_info.name, model_path.display()); + tracing::info!("Loading {} from: {}", model_id, model_path.display()); - // Device selection (from whisper.rs pattern) - let device = if let Ok(device) = Device::new_metal(0) { - tracing::info!("Using Metal device"); - device - } else { - tracing::info!("Using CPU device"); - Device::Cpu - }; + let backend = runtime.backend(); - // Load GGUF file - let mut file = std::fs::File::open(&model_path).map_err(|e| { - ProviderError::ExecutionError(format!("Failed to open model file: {}", e)) - })?; + let mut params = LlamaModelParams::default(); + if let Some(n_gpu_layers) = settings.n_gpu_layers { + params = params.with_n_gpu_layers(n_gpu_layers); + } + if settings.use_mlock { + params = params.with_use_mlock(true); + } + let model = LlamaModel::load_from_file(backend, &model_path, ¶ms) + .map_err(|e| ProviderError::ExecutionError(format!("Failed to load model: {}", e)))?; - // Read GGUF content - let content = candle_core::quantized::gguf_file::Content::read(&mut file).map_err(|e| { - ProviderError::ExecutionError(format!("Failed to read GGUF file: {}", e)) - })?; - - // Detect model architecture from ID - let model_id_lower = model_id.to_lowercase(); - let is_phi = model_id_lower.contains("phi"); - - // Load model weights based on architecture - // Try multiple architectures if name contains "phi" - let (model, eos_token_id) = if is_phi { - // Try Phi (Phi-2) first - match quantized_phi::ModelWeights::from_gguf(content, &mut file, &device) { - Ok(weights) => { - tracing::info!("Loaded with Phi architecture"); - (ModelWeights::Phi(weights), 50256) // Phi-2 EOS token - } - Err(e1) => { - tracing::info!("Phi architecture failed ({}), trying Phi-3", e1); - // Reopen file for second attempt - let mut file = std::fs::File::open(&model_path).map_err(|e| { - ProviderError::ExecutionError(format!("Failed to reopen model file: {}", e)) - })?; - let content = candle_core::quantized::gguf_file::Content::read(&mut file) - .map_err(|e| { - ProviderError::ExecutionError(format!( - "Failed to re-read GGUF file: {}", - e - )) - })?; - - match quantized_phi3::ModelWeights::from_gguf( - false, content, &mut file, &device, - ) { - Ok(weights) => { - tracing::info!("Loaded with Phi-3 architecture"); - (ModelWeights::Phi3(weights), 32000) // Phi-3 EOS token - } - Err(e2) => { - tracing::warn!( - "Phi-3 architecture failed ({}), falling back to Llama", - e2 - ); - // Try Llama as last resort - let mut file = std::fs::File::open(&model_path).map_err(|e| { - ProviderError::ExecutionError(format!( - "Failed to reopen model file: {}", - e - )) - })?; - let content = - candle_core::quantized::gguf_file::Content::read(&mut file) - .map_err(|e| { - ProviderError::ExecutionError(format!( - "Failed to re-read GGUF file: {}", - e - )) - })?; - - let weights = quantized_llama::ModelWeights::from_gguf( - content, &mut file, &device, - ) - .map_err(|e| { - ProviderError::ExecutionError(format!( - "Failed to load as Phi ({}), Phi-3 ({}), or Llama ({})", - e1, e2, e - )) - })?; - tracing::info!( - "Loaded Phi model with Llama architecture (may not work correctly)" - ); - (ModelWeights::Llama(weights), 50256) // Use Phi EOS token - } - } - } - } - } else { - tracing::info!("Using Llama architecture"); - let weights = quantized_llama::ModelWeights::from_gguf(content, &mut file, &device) - .map_err(|e| { + let template = match model.chat_template(None) { + Ok(t) => t, + Err(_) => { + tracing::warn!("Model has no embedded chat template, falling back to chatml"); + LlamaChatTemplate::new("chatml").map_err(|e| { ProviderError::ExecutionError(format!( - "Failed to load Llama model weights: {}", + "Failed to create fallback chat template: {}", e )) - })?; - (ModelWeights::Llama(weights), 128001) // Llama 3 EOS token - }; - - // Load tokenizer - let tokenizer = if tokenizer_path.exists() { - tracing::info!("Loading tokenizer from: {}", tokenizer_path.display()); - Tokenizer::from_file(&tokenizer_path).map_err(|e| { - ProviderError::ExecutionError(format!("Failed to load tokenizer: {}", e)) - })? - } else { - return Err(ProviderError::ExecutionError(format!( - "Tokenizer not found at {}. Please download the model again.", - tokenizer_path.display() - ))); - }; - - // Build list of EOS token IDs: start with architecture default, then add template-specific ones - let mut eos_token_ids = vec![eos_token_id]; - - // Look up token IDs for the chat template's EOS strings from tokenizer - for eos_str in model_info.chat_template.eos_strings() { - if let Some(token_id) = tokenizer.token_to_id(eos_str) { - if !eos_token_ids.contains(&token_id) { - eos_token_ids.push(token_id); - } - } else { - tracing::warn!("EOS string '{}' not found in tokenizer vocabulary", eos_str); + })? } - } + }; tracing::info!("Model loaded successfully"); - Ok(LoadedModel { - model, - tokenizer, - device, - eos_token_ids, - }) - } - - fn build_prompt( - &self, - system: &str, - messages: &[Message], - template: ChatTemplate, - tools: &[Tool], - ) -> String { - match template { - ChatTemplate::Llama3 => Self::format_llama3(system, messages, tools), - ChatTemplate::ChatML => Self::format_chatml(system, messages, tools), - ChatTemplate::Mistral => Self::format_mistral(system, messages, tools), - } - } - - /// Format message content for emulator, including text and tool responses - fn format_message_content_for_emulator(msg: &Message) -> String { - let mut parts = Vec::new(); - - for content in &msg.content { - match content { - MessageContent::Text(text) => { - parts.push(text.text.clone()); - } - MessageContent::ToolResponse(response) => { - // Include tool results in the prompt so model sees the output - match &response.tool_result { - Ok(result) => { - for content_item in &result.content { - if let Some(text_content) = content_item.as_text() { - parts.push(text_content.text.to_string()); - } - // Skip images and resources for now - } - } - Err(e) => { - parts.push(format!("Error: {}", e)); - } - } - } - _ => { - // Skip tool requests, images, etc. - } - } - } - - parts.join("\n") - } - - fn format_llama3(system: &str, messages: &[Message], tools: &[Tool]) -> String { - let mut prompt = String::from("<|begin_of_text|>"); - - // Add system message - if !system.is_empty() || !tools.is_empty() { - prompt.push_str("<|start_header_id|>system<|end_header_id|>\n\n"); - prompt.push_str(system); - - // Add tools if present - if !tools.is_empty() { - if !system.is_empty() { - prompt.push_str("\n\n"); - } - prompt.push_str("# Tools\n\nYou have access to the following tools:\n\n"); - for tool in tools { - let desc = tool - .description - .as_ref() - .map(|d| d.as_ref()) - .unwrap_or("No description"); - prompt.push_str(&format!("- {}: {}\n", tool.name, desc)); - } - } - - prompt.push_str("<|eot_id|>"); - } - - // Add conversation messages - for msg in messages { - let role = match msg.role { - Role::User => "user", - Role::Assistant => "assistant", - }; - - let content = Self::format_message_content_for_emulator(msg); - if !content.trim().is_empty() { - prompt.push_str(&format!("<|start_header_id|>{}<|end_header_id|>\n\n", role)); - prompt.push_str(&content); - prompt.push_str("<|eot_id|>"); - } - } - - // Add assistant prefix to prompt completion - prompt.push_str("<|start_header_id|>assistant<|end_header_id|>\n\n"); - prompt - } - - fn format_chatml(system: &str, messages: &[Message], _tools: &[Tool]) -> String { - let mut prompt = String::new(); - - // Add system message - if !system.is_empty() { - prompt.push_str("<|im_start|>system\n"); - prompt.push_str(system); - prompt.push_str("<|im_end|>\n"); - } - - // Add conversation messages - for msg in messages { - let role = match msg.role { - Role::User => "user", - Role::Assistant => "assistant", - }; - - let content = Self::format_message_content_for_emulator(msg); - if !content.trim().is_empty() { - prompt.push_str(&format!("<|im_start|>{}\n", role)); - prompt.push_str(&content); - prompt.push_str("<|im_end|>\n"); - } - } - - // Add assistant prefix - prompt.push_str("<|im_start|>assistant\n"); - prompt - } - - fn format_mistral(system: &str, messages: &[Message], _tools: &[Tool]) -> String { - let mut prompt = String::new(); - - // Mistral doesn't have a separate system role, prepend to first user message - let system_prefix = if !system.is_empty() { - format!("{}\n\n", system) - } else { - String::new() - }; - - // Add conversation messages - let mut first_user = true; - for msg in messages { - let content = Self::format_message_content_for_emulator(msg); - if content.trim().is_empty() { - continue; - } - - match msg.role { - Role::User => { - prompt.push_str("[INST] "); - if first_user { - prompt.push_str(&system_prefix); - first_user = false; - } - prompt.push_str(&content); - prompt.push_str(" [/INST]"); - } - Role::Assistant => { - prompt.push(' '); - prompt.push_str(&content); - prompt.push_str(""); - } - } - } - - // If no messages, still include system in first user turn - if first_user && !system.is_empty() { - prompt.push_str("[INST] "); - prompt.push_str(&system_prefix); - prompt.push_str("[/INST]"); - } - - prompt + Ok(LoadedModel { model, template }) } } @@ -687,19 +989,35 @@ impl ProviderDef for LocalInferenceProvider { where Self: Sized, { + use crate::providers::local_inference::local_model_registry::get_registry; + + let mut known_models: Vec<&str> = vec![ + "llama-3.2-1b", + "llama-3.2-3b", + "hermes-2-pro-7b", + "mistral-small-22b", + ]; + + // Add any registry models not already in the featured list + let mut dynamic_models = Vec::new(); + if let Ok(registry) = get_registry().lock() { + for entry in registry.list_models() { + if !known_models.contains(&entry.id.as_str()) { + dynamic_models.push(entry.id.clone()); + } + } + } + let dynamic_refs: Vec<&str> = dynamic_models.iter().map(|s| s.as_str()).collect(); + known_models.extend(dynamic_refs); + ProviderMetadata::new( PROVIDER_NAME, "Local Inference", - "Local inference using quantized GGUF models (Candle)", + "Local inference using quantized GGUF models (llama.cpp)", DEFAULT_MODEL, - vec![ - "llama-3.2-1b", - "llama-3.2-3b", - "hermes-2-pro-7b", - "mistral-small-22b", - ], - "https://github.com/huggingface/candle", - vec![], // No API keys required - models managed through UI + known_models, + "https://github.com/utilityai/llama-cpp-rs", + vec![], ) } @@ -726,17 +1044,25 @@ impl Provider for LocalInferenceProvider { _session_id: &str, _messages: &crate::conversation::Conversation, ) -> Result { - // Disable session naming for performance Ok("Local conversation".to_string()) } async fn fetch_supported_models(&self) -> Result, ProviderError> { - // Return all models - UI will show "(not downloaded)" for ones that aren't available - let all_models: Vec = available_local_models() + use crate::providers::local_inference::local_model_registry::get_registry; + + let mut all_models: Vec = available_local_models() .iter() .map(|m| m.id.to_string()) .collect(); + if let Ok(registry) = get_registry().lock() { + for entry in registry.list_models() { + if !all_models.contains(&entry.id) { + all_models.push(entry.id.clone()); + } + } + } + Ok(all_models) } @@ -750,7 +1076,6 @@ impl Provider for LocalInferenceProvider { ) -> Result<(Message, ProviderUsage), ProviderError> { use futures::StreamExt; - // Just call stream and accumulate results let mut stream = self .stream(session_id.unwrap_or(""), system, messages, tools) .await?; @@ -762,7 +1087,6 @@ impl Provider for LocalInferenceProvider { let (message_opt, usage_opt) = result?; if let Some(msg) = message_opt { - // Accumulate message content accumulated_message.id = msg.id.or(accumulated_message.id); for content in msg.content { accumulated_message.content.push(content); @@ -784,175 +1108,416 @@ impl Provider for LocalInferenceProvider { async fn stream( &self, _session_id: &str, - _system: &str, + system: &str, messages: &[Message], tools: &[Tool], ) -> Result { - // Get model metadata to determine chat template let model_config = &self.model_config; - let model_info = get_local_model(&model_config.model_name).ok_or_else(|| { - ProviderError::ExecutionError(format!("Model not found: {}", model_config.model_name)) - })?; - let template = model_info.chat_template; + let (_model_path, model_context_limit, model_settings) = + resolve_model_path(&model_config.model_name).ok_or_else(|| { + ProviderError::ExecutionError(format!( + "Model not found: {}", + model_config.model_name + )) + })?; - let tiny_prompt = load_tiny_model_prompt(); + // Ensure model is loaded — unload any other models first to free memory. + { + let mut model_lock = self.model.lock().await; + if model_lock.is_none() { + for slot in self.runtime.other_model_slots(&model_config.model_name) { + let mut other = slot.lock().await; + if other.is_some() { + tracing::info!("Unloading previous model to free memory"); + *other = None; + } + } - let prompt = self.build_prompt(&tiny_prompt, &messages, template, tools); - - let mut model_lock = self.model.lock().await; - if model_lock.is_none() { - *model_lock = Some(Self::load_model(&model_config.model_name).await?); + let model_id = model_config.model_name.clone(); + let settings_for_load = model_settings.clone(); + let runtime_for_load = self.runtime.clone(); + let loaded = tokio::task::spawn_blocking(move || { + Self::load_model_sync(&runtime_for_load, &model_id, &settings_for_load) + }) + .await + .map_err(|e| { + ProviderError::ExecutionError(format!("Model load task failed: {}", e)) + })??; + *model_lock = Some(loaded); + } } - // Clone Arc to move into the stream - let model_arc = self.model.clone(); - let model_name = model_config.model_name.clone(); + // Models that support native OpenAI-compatible tool-call JSON use the + // native path (template-based tool calling with JSON output). All other + // models use the emulator which parses `$ command` and ```execute blocks. + let use_emulator = !model_settings.native_tool_calling; + let system_prompt = if use_emulator { + load_tiny_model_prompt() + } else { + system.to_string() + }; - Ok(Box::pin(try_stream! { - // Generate a consistent message ID for all chunks - let message_id = Uuid::new_v4().to_string(); + // Build chat messages for the template + let mut chat_messages = + vec![ + LlamaChatMessage::new("system".to_string(), system_prompt.clone()).map_err( + |e| { + ProviderError::ExecutionError(format!( + "Failed to create system message: {}", + e + )) + }, + )?, + ]; - // Create streaming parser for emulator commands - let mut parser = StreamingEmulatorParser::new(); + // Check if Code Mode extension is available + let code_mode_enabled = tools.iter().any(|t| t.name == "code_execution__execute"); - // Get mutable access to model - let mut model_lock = model_arc.lock().await; - let loaded = model_lock.as_mut().ok_or_else(|| { - ProviderError::ExecutionError("Model not loaded".to_string()) - })?; + // Append tool descriptions to system prompt + if use_emulator && !tools.is_empty() { + let mut tool_desc = String::new(); - // Encode prompt - let prompt_tokens = loaded - .tokenizer - .encode(prompt.as_str(), false) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to encode prompt: {}", e)))? - .get_ids() - .to_vec(); + if code_mode_enabled { + // Build Code Mode instructions with available functions as + // Namespace.functionName() — matching how Code Mode exposes them. + tool_desc.push_str("\n\n# Running Code\n\n"); + tool_desc.push_str( + "You can call tools by writing code in a ```execute block. \ + The code runs immediately — do not explain it, just run it.\n\n", + ); + tool_desc.push_str("Example — counting files in /tmp:\n\n"); + tool_desc.push_str("```execute\nasync function run() {\n"); + tool_desc.push_str( + " const result = await Developer.shell({ command: \"ls -1 /tmp | wc -l\" });\n", + ); + tool_desc.push_str(" return result;\n}\n```\n\n"); + tool_desc.push_str("Rules:\n"); + tool_desc.push_str("- Code MUST define async function run() and return a result\n"); + tool_desc.push_str("- All function calls are async — use await\n"); + tool_desc + .push_str("- Use ```execute for tool calls, $ for simple shell one-liners\n\n"); + tool_desc.push_str("Available functions:\n\n"); - // PREFILL: Process entire prompt at once for speed - let input = Tensor::new(prompt_tokens.as_slice(), &loaded.device) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to create input tensor: {}", e)))? - .unsqueeze(0) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to unsqueeze input tensor: {}", e)))?; - - let logits = loaded.model.forward(&input, 0).map_err(|e| { - ProviderError::ExecutionError(format!("Prefill forward pass failed: {}", e)) - })?; - - // Quantized model returns [batch_size, vocab_size] directly for the last position - // Just squeeze to get [vocab_size] and sample - let logits = logits - .squeeze(0) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to squeeze batch dim: {}", e)))?; - - let mut next_token = logits - .argmax(0) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to sample token: {}", e)))? - .to_scalar::() - .map_err(|e| ProviderError::ExecutionError(format!("Failed to convert token: {}", e)))?; - - let decoded = loaded - .tokenizer - .decode(&[next_token], false) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to decode token: {}", e)))?; - - // Process first token through parser (skip if EOS) - let mut tool_call_emitted = false; - if !loaded.eos_token_ids.contains(&next_token) { - let parse_results = parser.process_chunk(&decoded); - for (text, command) in parse_results { - if let Some(text) = text { - let mut message = Message::assistant().with_text(&text); - message.id = Some(message_id.clone()); - yield (Some(message), None); + for tool in tools { + if tool.name.starts_with("code_execution__") { + continue; + } + let parts: Vec<&str> = tool.name.splitn(2, "__").collect(); + if parts.len() == 2 { + let namespace = { + let mut c = parts[0].chars(); + match c.next() { + None => String::new(), + Some(first) => first.to_uppercase().chain(c).collect::(), + } + }; + // Convert snake_case to camelCase + let camel_name: String = parts[1] + .split('_') + .enumerate() + .map(|(i, part)| { + if i == 0 { + part.to_string() + } else { + let mut c = part.chars(); + match c.next() { + None => String::new(), + Some(first) => first.to_uppercase().chain(c).collect(), + } + } + }) + .collect(); + let desc = tool.description.as_ref().map(|d| d.as_ref()).unwrap_or(""); + tool_desc.push_str(&format!("- {namespace}.{camel_name}(): {desc}\n")); } - if let Some(command) = command { - // Create tool request - let tool_id = Uuid::new_v4().to_string(); - let mut args = serde_json::Map::new(); - args.insert("command".to_string(), json!(command)); - - let tool_call = CallToolRequestParams { - meta: None, - task: None, - name: Cow::Borrowed("developer__shell"), - arguments: Some(args), - }; - - let mut message = Message::assistant(); - message.content.push(MessageContent::tool_request(tool_id, Ok(tool_call))); - message.id = Some(message_id.clone()); - yield (Some(message), None); - - // Stop after first tool call - tool_call_emitted = true; } + } else { + tool_desc.push_str("\n\n# Tools\n\nYou have access to the following tools:\n\n"); + for tool in tools { + let desc = tool + .description + .as_ref() + .map(|d| d.as_ref()) + .unwrap_or("No description"); + tool_desc.push_str(&format!("- {}: {}\n", tool.name, desc)); } } - // GENERATION LOOP: Generate remaining tokens (only if no tool call yet) - // Use model's context limit, cap output at 2K tokens to leave room for prompt - let max_output = model_info.context_limit.saturating_sub(prompt_tokens.len()).min(2048); - let mut output_token_count: i32 = 1; // Count the first token from prefill - if !tool_call_emitted { - for index in 0..max_output.saturating_sub(1) { - // Check for EOS tokens - if loaded.eos_token_ids.contains(&next_token) { - break; + chat_messages = vec![LlamaChatMessage::new( + "system".to_string(), + format!("{}{}", system_prompt, tool_desc), + ) + .map_err(|e| { + ProviderError::ExecutionError(format!("Failed to create system message: {}", e)) + })?]; + } + + for msg in messages { + let role = match msg.role { + Role::User => "user", + Role::Assistant => "assistant", + }; + let content = extract_text_content(msg); + if !content.trim().is_empty() { + chat_messages.push(LlamaChatMessage::new(role.to_string(), content).map_err( + |e| ProviderError::ExecutionError(format!("Failed to create message: {}", e)), + )?); + } + } + + let (full_tools_json, compact_tools) = if !use_emulator && !tools.is_empty() { + let full = format_tools(tools) + .ok() + .and_then(|spec| serde_json::to_string(&spec).ok()); + let compact = compact_tools_json(tools); + (full, compact) + } else { + (None, None) + }; + + let model_arc = self.model.clone(); + let runtime = self.runtime.clone(); + let model_name = model_config.model_name.clone(); + let context_limit = model_context_limit; + let settings = model_settings; + + // Channel for streaming tokens from blocking thread to async stream + let (tx, mut rx) = tokio::sync::mpsc::channel::< + Result<(Option, Option), ProviderError>, + >(32); + + tokio::task::spawn_blocking(move || { + let rt = tokio::runtime::Handle::current(); + + let model_guard = rt.block_on(model_arc.lock()); + let loaded = match model_guard.as_ref() { + Some(l) => l, + None => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError( + "Model not loaded".to_string(), + ))); + return; + } + }; + + let message_id = Uuid::new_v4().to_string(); + + if use_emulator { + // === Emulator path (Tiny/Small models) === + let prompt = + match loaded + .model + .apply_chat_template(&loaded.template, &chat_messages, true) + { + Ok(p) => p, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to apply chat template: {}", + e + )))); + return; + } + }; + + let tokens = match loaded.model.str_to_token(&prompt, AddBos::Never) { + Ok(t) => t, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to tokenize prompt: {}", + e + )))); + return; + } + }; + + let prompt_token_count = tokens.len(); + let n_ctx_train = loaded.model.n_ctx_train() as usize; + let memory_max_ctx = estimate_max_context_for_memory(&loaded.model, &runtime); + let effective_ctx = if let Some(ctx_size) = settings.context_size { + ctx_size as usize + } else { + effective_context_size( + prompt_token_count, + context_limit, + n_ctx_train, + memory_max_ctx, + ) + }; + if let Some(mem_max) = memory_max_ctx { + if prompt_token_count > mem_max { + let _ = + tx.blocking_send(Err(ProviderError::ContextLengthExceeded(format!( + "Prompt ({} tokens) exceeds estimated memory capacity ({} tokens). \ + Try a smaller model or reduce conversation length.", + prompt_token_count, mem_max, + )))); + return; + } + } + let ctx_params = build_context_params(effective_ctx as u32, &settings); + let mut ctx = match loaded.model.new_context(runtime.backend(), ctx_params) { + Ok(c) => c, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to create context: {}", + e + )))); + return; + } + }; + + let n_batch = ctx.n_batch() as usize; + for chunk in tokens.chunks(n_batch) { + let mut batch = match LlamaBatch::get_one(chunk) { + Ok(b) => b, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to create batch: {}", + e + )))); + return; + } + }; + if let Err(e) = ctx.decode(&mut batch) { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Prefill decode failed: {}", + e + )))); + return; + } } - // Single token input for generation - let input = Tensor::new(&[next_token], &loaded.device) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to create tensor: {}", e)))? - .unsqueeze(0) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to unsqueeze tensor: {}", e)))?; + let mut sampler = build_sampler(&settings); + let mut emulator_parser = StreamingEmulatorParser::new(code_mode_enabled); + let mut output_token_count: i32 = 0; + let mut tool_call_emitted = false; + let max_output = if let Some(max) = settings.max_output_tokens { + effective_ctx.saturating_sub(prompt_token_count).min(max) + } else { + effective_ctx.saturating_sub(prompt_token_count) + }; + let mut decoder = encoding_rs::UTF_8.new_decoder(); - // Position is prompt_len + already_generated_tokens - // We already generated 1 token from prefill, so add 1 - let pos = prompt_tokens.len() + index + 1; - let logits = loaded - .model - .forward(&input, pos) - .map_err(|e| ProviderError::ExecutionError(format!("Generation forward pass failed at pos {}: {}", pos, e)))?; + for _ in 0..max_output { + let token = sampler.sample(&ctx, -1); + sampler.accept(token); - let logits = logits.squeeze(0) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to squeeze logits: {}", e)))?; + if loaded.model.is_eog_token(token) { + break; + } - next_token = logits.argmax(0) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to sample token: {}", e)))? - .to_scalar::() - .map_err(|e| ProviderError::ExecutionError(format!("Failed to convert token: {}", e)))?; + output_token_count += 1; - // Check for EOS before decoding/yielding - if loaded.eos_token_ids.contains(&next_token) { - break; + let piece = match loaded.model.token_to_piece(token, &mut decoder, true, None) { + Ok(p) => p, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to decode token: {}", + e + )))); + return; + } + }; + + if !piece.is_empty() { + let actions = emulator_parser.process_chunk(&piece); + for action in actions { + match action { + EmulatorAction::Text(text) => { + let mut message = Message::assistant().with_text(&text); + message.id = Some(message_id.clone()); + if tx.blocking_send(Ok((Some(message), None))).is_err() { + return; + } + } + EmulatorAction::ShellCommand(command) => { + let tool_id = Uuid::new_v4().to_string(); + let mut args = serde_json::Map::new(); + args.insert("command".to_string(), json!(command)); + + let tool_call = CallToolRequestParams { + meta: None, + task: None, + name: Cow::Owned("developer__shell".to_string()), + arguments: Some(args), + }; + + let mut message = Message::assistant(); + message + .content + .push(MessageContent::tool_request(tool_id, Ok(tool_call))); + message.id = Some(message_id.clone()); + if tx.blocking_send(Ok((Some(message), None))).is_err() { + return; + } + tool_call_emitted = true; + } + EmulatorAction::ExecuteCode(code) => { + let tool_id = Uuid::new_v4().to_string(); + let wrapped = if code.contains("async function run()") { + code + } else { + format!("async function run() {{\n{}\n}}", code) + }; + let mut args = serde_json::Map::new(); + args.insert("code".to_string(), json!(wrapped)); + + let tool_call = CallToolRequestParams { + meta: None, + task: None, + name: Cow::Owned("code_execution__execute".to_string()), + arguments: Some(args), + }; + + let mut message = Message::assistant(); + message + .content + .push(MessageContent::tool_request(tool_id, Ok(tool_call))); + message.id = Some(message_id.clone()); + if tx.blocking_send(Ok((Some(message), None))).is_err() { + return; + } + tool_call_emitted = true; + } + } + } + } + + if tool_call_emitted { + break; + } + + let next_tokens = [token]; + let mut next_batch = match LlamaBatch::get_one(&next_tokens) { + Ok(b) => b, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to create batch: {}", + e + )))); + return; + } + }; + if let Err(e) = ctx.decode(&mut next_batch) { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Decode failed: {}", + e + )))); + return; + } } - // Count the generated token - output_token_count += 1; - - // Decode token - let mut decoded = loaded - .tokenizer - .decode(&[next_token], false) - .map_err(|e| ProviderError::ExecutionError(format!("Failed to decode token: {}", e)))?; - - // Strip EOS tokens from this chunk - for eos_str in template.eos_strings() { - decoded = decoded.replace(eos_str, ""); - } - - if !decoded.is_empty() { - // Process through parser - let parse_results = parser.process_chunk(&decoded); - for (text, command) in parse_results { - if let Some(text) = text { + let flush_results = emulator_parser.flush(); + for action in flush_results { + match action { + EmulatorAction::Text(text) => { let mut message = Message::assistant().with_text(&text); message.id = Some(message_id.clone()); - yield (Some(message), None); + if tx.blocking_send(Ok((Some(message), None))).is_err() { + return; + } } - if let Some(command) = command { - // Create tool request + EmulatorAction::ShellCommand(command) => { let tool_id = Uuid::new_v4().to_string(); let mut args = serde_json::Map::new(); args.insert("command".to_string(), json!(command)); @@ -960,62 +1525,343 @@ impl Provider for LocalInferenceProvider { let tool_call = CallToolRequestParams { meta: None, task: None, - name: Cow::Borrowed("developer__shell"), + name: Cow::Owned("developer__shell".to_string()), arguments: Some(args), }; let mut message = Message::assistant(); - message.content.push(MessageContent::tool_request(tool_id, Ok(tool_call))); + message + .content + .push(MessageContent::tool_request(tool_id, Ok(tool_call))); message.id = Some(message_id.clone()); - yield (Some(message), None); + if tx.blocking_send(Ok((Some(message), None))).is_err() { + return; + } + } + EmulatorAction::ExecuteCode(code) => { + let tool_id = Uuid::new_v4().to_string(); + let wrapped = if code.contains("async function run()") { + code + } else { + format!("async function run() {{\n{}\n}}", code) + }; + let mut args = serde_json::Map::new(); + args.insert("code".to_string(), json!(wrapped)); - // Stop generation after first tool call - tool_call_emitted = true; + let tool_call = CallToolRequestParams { + meta: None, + task: None, + name: Cow::Owned("code_execution__execute".to_string()), + arguments: Some(args), + }; + + let mut message = Message::assistant(); + message + .content + .push(MessageContent::tool_request(tool_id, Ok(tool_call))); + message.id = Some(message_id.clone()); + if tx.blocking_send(Ok((Some(message), None))).is_err() { + return; + } } } } - // Break out of generation loop after tool call - if tool_call_emitted { - break; - } - } - } + let input_tokens = prompt_token_count as i32; + let total_tokens = input_tokens + output_token_count; + let usage = Usage::new( + Some(input_tokens), + Some(output_token_count), + Some(total_tokens), + ); + let provider_usage = ProviderUsage::new(model_name, usage); + let _ = tx.blocking_send(Ok((None, Some(provider_usage)))); + } else { + // === Native tool-calling path (Medium/Large models) === + // Compute the context cap up front (before tokenizing) so the + // tool-budget check uses the same limit that inference will. + let min_generation_headroom = 512; + let n_ctx_train = loaded.model.n_ctx_train() as usize; + let memory_max_ctx = estimate_max_context_for_memory(&loaded.model, &runtime); + let context_cap = if let Some(ctx_size) = settings.context_size { + ctx_size as usize + } else { + let base = if context_limit > 0 { + context_limit + } else { + n_ctx_train + }; + match memory_max_ctx { + Some(mem_max) if mem_max < base => mem_max, + _ => base, + } + }; + let token_budget = context_cap.saturating_sub(min_generation_headroom); - // Flush any remaining parser buffer (handles incomplete commands at end of stream) - let flush_results = parser.flush(); - for (text, command) in flush_results { - if let Some(text) = text { - let mut message = Message::assistant().with_text(&text); - message.id = Some(message_id.clone()); - yield (Some(message), None); - } - if let Some(command) = command { - // Create tool request for the final command - let tool_id = Uuid::new_v4().to_string(); - let mut args = serde_json::Map::new(); - args.insert("command".to_string(), json!(command)); + // Try full tool schemas first; if the prompt overflows the + // context window, retry with compact definitions (name + + // description only, no parameter schemas). + let apply_template = |tools: Option<&str>| { + loaded.model.apply_chat_template_with_tools_oaicompat( + &loaded.template, + &chat_messages, + tools, + None, + true, + ) + }; - let tool_call = CallToolRequestParams { - meta: None, - task: None, - name: Cow::Borrowed("developer__shell"), - arguments: Some(args), + let template_result = match apply_template(full_tools_json.as_deref()) { + Ok(r) => { + let token_count = loaded + .model + .str_to_token(&r.prompt, AddBos::Never) + .map(|t| t.len()) + .unwrap_or(0); + if token_count > token_budget { + apply_template(compact_tools.as_deref()).unwrap_or(r) + } else { + r + } + } + Err(_) => match apply_template(compact_tools.as_deref()) { + Ok(r) => r, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to apply chat template: {}", + e + )))); + return; + } + }, + }; + + let tokens = match loaded + .model + .str_to_token(&template_result.prompt, AddBos::Never) + { + Ok(t) => t, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to tokenize prompt: {}", + e + )))); + return; + } + }; + + let prompt_token_count = tokens.len(); + + // effective_context_size may expand beyond the cap when the + // prompt itself already exceeds it, ensuring we at least + // attempt to fit the prompt plus minimal generation room. + let effective_ctx = if let Some(ctx_size) = settings.context_size { + ctx_size as usize + } else { + effective_context_size( + prompt_token_count, + context_limit, + n_ctx_train, + memory_max_ctx, + ) + }; + if let Some(mem_max) = memory_max_ctx { + if prompt_token_count > mem_max { + let _ = + tx.blocking_send(Err(ProviderError::ContextLengthExceeded(format!( + "Prompt ({} tokens) exceeds estimated memory capacity ({} tokens). \ + Try a smaller model or reduce conversation length.", + prompt_token_count, mem_max, + )))); + return; + } + } + let ctx_params = build_context_params(effective_ctx as u32, &settings); + let mut ctx = match loaded.model.new_context(runtime.backend(), ctx_params) { + Ok(c) => c, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to create context: {}", + e + )))); + return; + } + }; + + let n_batch = ctx.n_batch() as usize; + for chunk in tokens.chunks(n_batch) { + let mut batch = match LlamaBatch::get_one(chunk) { + Ok(b) => b, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to create batch: {}", + e + )))); + return; + } + }; + if let Err(e) = ctx.decode(&mut batch) { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Prefill decode failed: {}", + e + )))); + return; + } + } + + let mut sampler = build_sampler(&settings); + let mut output_token_count: i32 = 0; + let max_output = if let Some(max) = settings.max_output_tokens { + effective_ctx.saturating_sub(prompt_token_count).min(max) + } else { + effective_ctx.saturating_sub(prompt_token_count) + }; + let mut decoder = encoding_rs::UTF_8.new_decoder(); + let mut generated_text = String::new(); + + // For streaming, we buffer text and only send content that + // precedes any tool-call JSON. Once we see the start of a + // `{"tool_calls"` block we hold back everything from that + // point on so the raw JSON never reaches the user. + let mut streamed_len: usize = 0; + + for _ in 0..max_output { + let token = sampler.sample(&ctx, -1); + sampler.accept(token); + + if loaded.model.is_eog_token(token) { + break; + } + + output_token_count += 1; + + let piece = match loaded.model.token_to_piece(token, &mut decoder, true, None) { + Ok(p) => p, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to decode token: {}", + e + )))); + return; + } }; - let mut message = Message::assistant(); - message.content.push(MessageContent::tool_request(tool_id, Ok(tool_call))); - message.id = Some(message_id.clone()); - yield (Some(message), None); + if !piece.is_empty() { + generated_text.push_str(&piece); + + // Determine how far we can safely stream: + // - If a complete tool-call JSON block is found at + // the end, only stream the content before it. + // - Otherwise, hold back any unmatched `{` that could + // be the start of a tool-call block in progress. + let has_xml_tc = + split_content_and_xml_tool_calls(&generated_text).is_some(); + let (content, tc) = split_content_and_tool_calls(&generated_text); + let stream_up_to = if tc.is_some() { + content.len() + } else if has_xml_tc { + split_content_and_xml_tool_calls(&generated_text) + .map(|(c, _)| c.len()) + .unwrap_or(0) + } else { + safe_stream_end(&generated_text) + }; + if stream_up_to > streamed_len { + #[allow(clippy::string_slice)] + let new_text = &generated_text[streamed_len..stream_up_to]; + if !new_text.is_empty() { + let mut msg = Message::assistant().with_text(new_text); + msg.id = Some(message_id.clone()); + if tx.blocking_send(Ok((Some(msg), None))).is_err() { + return; + } + } + streamed_len = stream_up_to; + } + + let should_stop = template_result + .additional_stops + .iter() + .any(|stop| generated_text.ends_with(stop)); + if should_stop { + break; + } + } + + let next_tokens = [token]; + let mut next_batch = match LlamaBatch::get_one(&next_tokens) { + Ok(b) => b, + Err(e) => { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Failed to create batch: {}", + e + )))); + return; + } + }; + if let Err(e) = ctx.decode(&mut next_batch) { + let _ = tx.blocking_send(Err(ProviderError::ExecutionError(format!( + "Decode failed: {}", + e + )))); + return; + } } + + // Try XML tool call format first (e.g. qwen3-coder), then JSON + let (content, tool_call_msgs) = if let Some((xml_content, xml_calls)) = + split_content_and_xml_tool_calls(&generated_text) + { + let msgs = extract_xml_tool_call_messages(xml_calls, &message_id); + (xml_content, msgs) + } else { + let (json_content, tool_calls_json) = + split_content_and_tool_calls(&generated_text); + let msgs = tool_calls_json + .map(|tc| extract_tool_call_messages(&tc, &message_id)) + .unwrap_or_default(); + (json_content, msgs) + }; + + // Send any remaining content text we haven't streamed yet. + if content.len() > streamed_len { + #[allow(clippy::string_slice)] + let remaining = &content[streamed_len..]; + if !remaining.is_empty() { + let mut msg = Message::assistant().with_text(remaining); + msg.id = Some(message_id.clone()); + let _ = tx.blocking_send(Ok((Some(msg), None))); + } + } + + if !tool_call_msgs.is_empty() { + for msg in tool_call_msgs { + let _ = tx.blocking_send(Ok((Some(msg), None))); + } + } else if content.is_empty() && !generated_text.is_empty() { + let mut msg = Message::assistant().with_text(&generated_text); + msg.id = Some(message_id.clone()); + let _ = tx.blocking_send(Ok((Some(msg), None))); + } + + let input_tokens = prompt_token_count as i32; + let total_tokens = input_tokens + output_token_count; + let usage = Usage::new( + Some(input_tokens), + Some(output_token_count), + Some(total_tokens), + ); + let provider_usage = ProviderUsage::new(model_name, usage); + let _ = tx.blocking_send(Ok((None, Some(provider_usage)))); + } + }); + + Ok(Box::pin(try_stream! { + while let Some(result) = rx.recv().await { + let item = result?; + yield item; } - // Final yield with usage - let input_tokens = prompt_tokens.len() as i32; - let total_tokens = input_tokens + output_token_count; - let usage = Usage::new(Some(input_tokens), Some(output_token_count), Some(total_tokens)); - let provider_usage = ProviderUsage::new(model_name.clone(), usage); - yield (None, Some(provider_usage)); })) } @@ -1023,3 +1869,102 @@ impl Provider for LocalInferenceProvider { true } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_xml_tool_call_single() { + let text = "I'll search for that.\n\n\n\nlocal.*inference\n\n"; + let result = split_content_and_xml_tool_calls(text); + assert!(result.is_some()); + let (content, calls) = result.unwrap(); + assert_eq!(content, "I'll search for that."); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "search__files"); + assert_eq!(calls[0].1.get("pattern").unwrap(), "local.*inference"); + } + + #[test] + fn test_parse_xml_tool_call_multiple_params() { + let text = "\n\nls -la\n30\n\n"; + let result = split_content_and_xml_tool_calls(text); + assert!(result.is_some()); + let (content, calls) = result.unwrap(); + assert!(content.is_empty()); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "developer__shell"); + assert_eq!(calls[0].1.get("command").unwrap(), "ls -la"); + // 30 should be parsed as a number + assert_eq!(calls[0].1.get("timeout").unwrap(), &json!(30)); + } + + #[test] + fn test_parse_xml_tool_call_no_tool_call() { + let text = "Just some regular text with no tool calls."; + assert!(split_content_and_xml_tool_calls(text).is_none()); + } + + #[test] + fn test_parse_xml_tool_call_multiple_calls() { + let text = "Doing two things:\n\n\n1\n\n\n\n\nhello\n\n"; + let result = split_content_and_xml_tool_calls(text); + assert!(result.is_some()); + let (content, calls) = result.unwrap(); + assert_eq!(content, "Doing two things:"); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].0, "foo__bar"); + assert_eq!(calls[1].0, "baz__qux"); + } + + #[test] + fn test_parse_xml_tool_call_multiline_value() { + let text = "\n\ntest.py\ndef hello():\n print(\"world\")\n\n"; + let result = split_content_and_xml_tool_calls(text); + assert!(result.is_some()); + let (_content, calls) = result.unwrap(); + assert_eq!(calls[0].0, "developer__write_file"); + assert_eq!( + calls[0].1.get("content").unwrap(), + "def hello():\n print(\"world\")" + ); + } + + #[test] + fn test_safe_stream_end_holds_back_tool_call_tag() { + let text = "Some text before \n"; + let safe = safe_stream_end(text); + assert!(safe <= text.find("").unwrap()); + } + + #[test] + fn test_safe_stream_end_holds_back_partial_tag() { + let text = "Some text { + let call = req.tool_call.as_ref().unwrap(); + assert_eq!(&*call.name, "developer__shell"); + assert_eq!( + call.arguments.as_ref().unwrap().get("command").unwrap(), + "ls" + ); + } + _ => panic!("Expected ToolRequest"), + } + } +} diff --git a/crates/goose/src/providers/local_inference/hf_models.rs b/crates/goose/src/providers/local_inference/hf_models.rs new file mode 100644 index 0000000000..c61c435926 --- /dev/null +++ b/crates/goose/src/providers/local_inference/hf_models.rs @@ -0,0 +1,565 @@ +use anyhow::{bail, Result}; +use serde::{Deserialize, Serialize}; + +use utoipa::ToSchema; + +const HF_API_BASE: &str = "https://huggingface.co/api/models"; +const HF_DOWNLOAD_BASE: &str = "https://huggingface.co"; + +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +pub struct HfModelInfo { + pub repo_id: String, + pub author: String, + pub model_name: String, + pub downloads: u64, + pub gguf_files: Vec, +} + +/// A single downloadable GGUF file (used internally and for downloads). +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +pub struct HfGgufFile { + pub filename: String, + pub size_bytes: u64, + pub quantization: String, + pub download_url: String, +} + +/// A quantization variant — groups sharded files into one logical entry. +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +pub struct HfQuantVariant { + pub quantization: String, + pub size_bytes: u64, + pub filename: String, + pub download_url: String, + pub description: &'static str, + pub quality_rank: u8, +} + +#[derive(Debug, Deserialize)] +struct HfApiModel { + id: Option, + author: Option, + downloads: Option, + siblings: Option>, +} + +#[derive(Debug, Deserialize)] +struct HfApiSibling { + rfilename: String, + #[serde(default)] + size: Option, +} + +struct QuantInfo { + description: &'static str, + quality_rank: u8, +} + +fn quant_info(quant: &str) -> QuantInfo { + match quant { + "IQ1_S" => QuantInfo { + description: "Extremely small, very low quality", + quality_rank: 1, + }, + "IQ1_M" => QuantInfo { + description: "Extremely small, very low quality", + quality_rank: 2, + }, + "IQ2_XXS" => QuantInfo { + description: "Very small, low quality", + quality_rank: 3, + }, + "IQ2_XS" => QuantInfo { + description: "Very small, low quality", + quality_rank: 4, + }, + "IQ2_S" => QuantInfo { + description: "Very small, low quality", + quality_rank: 5, + }, + "IQ2_M" => QuantInfo { + description: "Very small, low quality", + quality_rank: 6, + }, + "Q2_K" | "Q2_K_S" => QuantInfo { + description: "Small, low quality", + quality_rank: 7, + }, + "IQ3_XXS" => QuantInfo { + description: "Very small, moderate quality loss", + quality_rank: 8, + }, + "IQ3_XS" | "IQ3_S" => QuantInfo { + description: "Small, moderate quality loss", + quality_rank: 9, + }, + "Q3_K_S" => QuantInfo { + description: "Small, moderate quality loss", + quality_rank: 10, + }, + "IQ3_M" => QuantInfo { + description: "Small, moderate quality loss", + quality_rank: 11, + }, + "Q3_K_M" => QuantInfo { + description: "Small, balanced quality/size", + quality_rank: 12, + }, + "Q3_K_L" => QuantInfo { + description: "Medium-small, decent quality", + quality_rank: 13, + }, + "IQ4_XS" => QuantInfo { + description: "Medium, good quality", + quality_rank: 14, + }, + "IQ4_NL" => QuantInfo { + description: "Medium, good quality", + quality_rank: 15, + }, + "Q4_0" => QuantInfo { + description: "Medium, good quality", + quality_rank: 16, + }, + "Q4_1" => QuantInfo { + description: "Medium, good quality", + quality_rank: 17, + }, + "Q4_K_S" => QuantInfo { + description: "Medium, good quality/size balance", + quality_rank: 18, + }, + "Q4_K_M" => QuantInfo { + description: "Medium, recommended balance of quality and size", + quality_rank: 19, + }, + "Q5_0" => QuantInfo { + description: "Medium-large, high quality", + quality_rank: 20, + }, + "Q5_1" => QuantInfo { + description: "Medium-large, high quality", + quality_rank: 21, + }, + "Q5_K_S" => QuantInfo { + description: "Medium-large, high quality", + quality_rank: 22, + }, + "Q5_K_M" => QuantInfo { + description: "Medium-large, very high quality", + quality_rank: 23, + }, + "Q6_K" => QuantInfo { + description: "Large, near-lossless quality", + quality_rank: 24, + }, + "Q8_0" => QuantInfo { + description: "Large, near-lossless quality", + quality_rank: 25, + }, + "F16" => QuantInfo { + description: "Full size, original quality (16-bit)", + quality_rank: 26, + }, + "BF16" => QuantInfo { + description: "Full size, original quality (bfloat16)", + quality_rank: 27, + }, + "F32" => QuantInfo { + description: "Full size, original quality (32-bit)", + quality_rank: 28, + }, + "MXFP4_MOE" => QuantInfo { + description: "Medium, mixed-precision 4-bit for MoE models", + quality_rank: 18, + }, + "TQ1_0" => QuantInfo { + description: "Tiny, ternary quantization", + quality_rank: 1, + }, + "Q2_K_XL" | "Q3_K_XL" | "Q4_K_XL" => QuantInfo { + description: "Extended-layer variant", + quality_rank: 15, + }, + "Q2_K_L" => QuantInfo { + description: "Small, low quality (large variant)", + quality_rank: 8, + }, + "Q4_K_L" => QuantInfo { + description: "Medium, good quality (large variant)", + quality_rank: 20, + }, + _ => QuantInfo { + description: "", + quality_rank: 15, + }, + } +} + +pub fn parse_quantization_from_filename(filename: &str) -> String { + parse_quantization(filename) +} + +fn parse_quantization(filename: &str) -> String { + // Strip directory prefix (e.g. "Q5_K_M/Model-Q5_K_M-00001-of-00002.gguf") + let basename = filename.rsplit('/').next().unwrap_or(filename); + let stem = basename.trim_end_matches(".gguf"); + + // Strip shard suffix like "-00001-of-00004" + let stem = if let Some(pos) = stem.rfind("-of-") { + stem.get(..pos) + .and_then(|s| s.rsplit_once('-').map(|(prefix, _)| prefix)) + .unwrap_or(stem) + } else { + stem + }; + + // The quantization tag is typically the last hyphen-separated component + // that looks like a quant identifier (starts with Q, IQ, F, BF, TQ, MXFP, etc.) + // e.g. "Qwen3-Coder-Next-Q4_K_M" -> "Q4_K_M" + // "Model-UD-IQ1_M" -> "IQ1_M" + if let Some((_, candidate)) = stem.rsplit_once('-') { + if looks_like_quant(candidate) { + return candidate.to_string(); + } + } + + // Fallback: try dot separator (e.g. "model.Q4_K_M") + if let Some((_, candidate)) = stem.rsplit_once('.') { + if looks_like_quant(candidate) { + return candidate.to_string(); + } + } + + "unknown".to_string() +} + +fn looks_like_quant(s: &str) -> bool { + let upper = s.to_uppercase(); + upper.starts_with("Q") + || upper.starts_with("IQ") + || upper.starts_with("TQ") + || upper.starts_with("MXFP") + || upper == "F16" + || upper == "F32" + || upper == "BF16" +} + +fn is_shard_file(filename: &str) -> bool { + // Matches patterns like "-00001-of-00003.gguf" + let basename = filename.rsplit('/').next().unwrap_or(filename); + let stem = basename.trim_end_matches(".gguf"); + if let Some(pos) = stem.rfind("-of-") { + stem.get(..pos) + .and_then(|before| before.rsplit('-').next()) + .map(|s| !s.is_empty() && s.chars().all(|c| c.is_ascii_digit())) + .unwrap_or(false) + } else { + false + } +} + +fn build_download_url(repo_id: &str, filename: &str) -> String { + format!("{}/{}/resolve/main/{}", HF_DOWNLOAD_BASE, repo_id, filename) +} + +/// Collect single-file GGUFs into quantization variants (sharded files are excluded). +fn group_into_variants(repo_id: &str, files: Vec) -> Vec { + let mut variants: Vec = files + .into_iter() + .filter(|s| { + s.rfilename.ends_with(".gguf") + && !is_shard_file(&s.rfilename) + && parse_quantization(&s.rfilename) != "unknown" + }) + .map(|s| { + let quant = parse_quantization(&s.rfilename); + let info = quant_info(&quant); + let download_url = build_download_url(repo_id, &s.rfilename); + HfQuantVariant { + quantization: quant, + size_bytes: s.size.unwrap_or(0), + filename: s.rfilename, + download_url, + description: info.description, + quality_rank: info.quality_rank, + } + }) + .collect(); + + variants.sort_by_key(|v| v.quality_rank); + variants +} + +pub async fn search_gguf_models(query: &str, limit: usize) -> Result> { + let client = reqwest::Client::new(); + let url = format!( + "{}?search={}&filter=gguf&sort=downloads&direction=-1&limit={}", + HF_API_BASE, query, limit + ); + + let response = client + .get(&url) + .header("User-Agent", "goose-ai-agent") + .send() + .await?; + + if !response.status().is_success() { + bail!("HuggingFace API returned status {}", response.status()); + } + + let models: Vec = response.json().await?; + + let results = models + .into_iter() + .filter_map(|m| { + let repo_id = m.id?; + let siblings = m.siblings.unwrap_or_default(); + + // The search endpoint may not include `siblings`; parse whatever + // is available. Files are fetched on-demand via `get_repo_gguf_variants`. + let gguf_files: Vec = siblings + .into_iter() + .filter(|s| s.rfilename.ends_with(".gguf")) + .map(|s| { + let quantization = parse_quantization(&s.rfilename); + let download_url = build_download_url(&repo_id, &s.rfilename); + HfGgufFile { + filename: s.rfilename, + size_bytes: s.size.unwrap_or(0), + quantization, + download_url, + } + }) + .collect(); + + let author = m + .author + .unwrap_or_else(|| repo_id.split('/').next().unwrap_or_default().to_string()); + let model_name = repo_id + .split('/') + .next_back() + .unwrap_or(&repo_id) + .to_string(); + + Some(HfModelInfo { + repo_id, + author, + model_name, + downloads: m.downloads.unwrap_or(0), + gguf_files, + }) + }) + .collect(); + + Ok(results) +} + +/// Fetch GGUF files for a repo and return them grouped by quantization. +pub async fn get_repo_gguf_variants(repo_id: &str) -> Result> { + let client = reqwest::Client::new(); + let url = format!("{}/{}?blobs=true", HF_API_BASE, repo_id); + + let response = client + .get(&url) + .header("User-Agent", "goose-ai-agent") + .send() + .await?; + + if !response.status().is_success() { + bail!( + "HuggingFace API returned status {} for repo {}", + response.status(), + repo_id + ); + } + + let model: HfApiModel = response.json().await?; + let siblings = model.siblings.unwrap_or_default(); + + Ok(group_into_variants(repo_id, siblings)) +} + +/// Fetch raw GGUF files (kept for resolve_model_spec). +pub async fn get_repo_gguf_files(repo_id: &str) -> Result> { + let client = reqwest::Client::new(); + let url = format!("{}/{}?blobs=true", HF_API_BASE, repo_id); + + let response = client + .get(&url) + .header("User-Agent", "goose-ai-agent") + .send() + .await?; + + if !response.status().is_success() { + bail!( + "HuggingFace API returned status {} for repo {}", + response.status(), + repo_id + ); + } + + let model: HfApiModel = response.json().await?; + let siblings = model.siblings.unwrap_or_default(); + + let files = siblings + .into_iter() + .filter(|s| s.rfilename.ends_with(".gguf")) + .filter(|s| !is_shard_file(&s.rfilename)) + .map(|s| { + let quantization = parse_quantization(&s.rfilename); + let download_url = build_download_url(repo_id, &s.rfilename); + HfGgufFile { + filename: s.rfilename, + size_bytes: s.size.unwrap_or(0), + quantization, + download_url, + } + }) + .collect(); + + Ok(files) +} + +/// Parse a model spec like "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M" into (repo_id, quantization). +pub fn parse_model_spec(spec: &str) -> Result<(String, String)> { + let (repo_id, quant) = spec.rsplit_once(':').ok_or_else(|| { + anyhow::anyhow!( + "Invalid model spec '{}': expected format 'user/repo:quantization'", + spec + ) + })?; + + if !repo_id.contains('/') { + bail!("Invalid repo_id '{}': expected format 'user/repo'", repo_id); + } + + Ok((repo_id.to_string(), quant.to_string())) +} + +/// Resolve a model spec to a specific GGUF file from the repo. +pub async fn resolve_model_spec(spec: &str) -> Result<(String, HfGgufFile)> { + let (repo_id, quant) = parse_model_spec(spec)?; + let files = get_repo_gguf_files(&repo_id).await?; + + let file = files + .into_iter() + .find(|f| f.quantization.eq_ignore_ascii_case(&quant)) + .ok_or_else(|| { + anyhow::anyhow!( + "No GGUF file with quantization '{}' found in {}", + quant, + repo_id + ) + })?; + + Ok((repo_id, file)) +} + +/// Recommend which quantization variant to use based on available memory. +pub fn recommend_variant( + variants: &[HfQuantVariant], + available_memory_bytes: u64, +) -> Option { + // We need ~10-20% overhead beyond model size for inference context. + // Pick the highest-quality variant that fits. + let usable = (available_memory_bytes as f64 * 0.85) as u64; + + let mut best: Option = None; + for (i, v) in variants.iter().enumerate() { + if v.size_bytes <= usable { + match best { + Some(bi) if variants[bi].quality_rank < v.quality_rank => best = Some(i), + None => best = Some(i), + _ => {} + } + } + } + best +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_quantization() { + assert_eq!(parse_quantization("Model-Q4_K_M.gguf"), "Q4_K_M"); + assert_eq!(parse_quantization("Model-Q8_0.gguf"), "Q8_0"); + assert_eq!(parse_quantization("Model-IQ4_NL.gguf"), "IQ4_NL"); + assert_eq!(parse_quantization("Model-F16.gguf"), "F16"); + assert_eq!(parse_quantization("random-name.gguf"), "unknown"); + } + + #[test] + fn test_parse_quantization_with_directory() { + assert_eq!( + parse_quantization("Q5_K_M/Model-Q5_K_M-00001-of-00002.gguf"), + "Q5_K_M" + ); + } + + #[test] + fn test_parse_quantization_extended_tags() { + assert_eq!(parse_quantization("Model-MXFP4_MOE.gguf"), "MXFP4_MOE"); + assert_eq!(parse_quantization("Model-UD-TQ1_0.gguf"), "TQ1_0"); + assert_eq!(parse_quantization("Model-Q2_K_L.gguf"), "Q2_K_L"); + assert_eq!(parse_quantization("Model-UD-Q4_K_XL.gguf"), "Q4_K_XL"); + assert_eq!(parse_quantization("Model-UD-IQ1_M.gguf"), "IQ1_M"); + } + + #[test] + fn test_is_shard_file() { + assert!(is_shard_file("Q5_K_M/Model-Q5_K_M-00001-of-00002.gguf")); + assert!(is_shard_file("Model-BF16-00003-of-00004.gguf")); + assert!(!is_shard_file("Model-Q4_K_M.gguf")); + } + + #[test] + fn test_parse_model_spec() { + let (repo, quant) = + parse_model_spec("bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M").unwrap(); + assert_eq!(repo, "bartowski/Llama-3.2-1B-Instruct-GGUF"); + assert_eq!(quant, "Q4_K_M"); + } + + #[test] + fn test_parse_model_spec_invalid() { + assert!(parse_model_spec("no-colon").is_err()); + assert!(parse_model_spec("noslash:Q4_K_M").is_err()); + } + + #[test] + fn test_recommend_variant() { + let variants = vec![ + HfQuantVariant { + quantization: "Q2_K".into(), + size_bytes: 2_000_000_000, + filename: "m-Q2_K.gguf".into(), + download_url: String::new(), + description: "Small", + quality_rank: 7, + }, + HfQuantVariant { + quantization: "Q4_K_M".into(), + size_bytes: 4_000_000_000, + filename: "m-Q4_K_M.gguf".into(), + download_url: String::new(), + description: "Medium", + quality_rank: 19, + }, + HfQuantVariant { + quantization: "Q8_0".into(), + size_bytes: 8_000_000_000, + filename: "m-Q8_0.gguf".into(), + download_url: String::new(), + description: "Large", + quality_rank: 25, + }, + ]; + + assert_eq!(recommend_variant(&variants, 5_000_000_000), Some(1)); + assert_eq!(recommend_variant(&variants, 10_000_000_000), Some(2)); + assert_eq!(recommend_variant(&variants, 1_000_000_000), None); + } +} diff --git a/crates/goose/src/providers/local_inference/local_model_registry.rs b/crates/goose/src/providers/local_inference/local_model_registry.rs new file mode 100644 index 0000000000..1092fed981 --- /dev/null +++ b/crates/goose/src/providers/local_inference/local_model_registry.rs @@ -0,0 +1,297 @@ +use crate::config::paths::Paths; +use anyhow::Result; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use std::sync::{Mutex, OnceLock}; +use utoipa::ToSchema; + +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +#[serde(tag = "type")] +pub enum SamplingConfig { + Greedy, + Temperature { + temperature: f32, + top_k: i32, + top_p: f32, + min_p: f32, + seed: Option, + }, + MirostatV2 { + tau: f32, + eta: f32, + seed: Option, + }, +} + +impl Default for SamplingConfig { + fn default() -> Self { + SamplingConfig::Temperature { + temperature: 0.8, + top_k: 40, + top_p: 0.95, + min_p: 0.05, + seed: None, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +pub struct ModelSettings { + pub context_size: Option, + pub max_output_tokens: Option, + #[serde(default)] + pub sampling: SamplingConfig, + #[serde(default = "default_repeat_penalty")] + pub repeat_penalty: f32, + #[serde(default = "default_repeat_last_n")] + pub repeat_last_n: i32, + #[serde(default)] + pub frequency_penalty: f32, + #[serde(default)] + pub presence_penalty: f32, + pub n_batch: Option, + pub n_gpu_layers: Option, + #[serde(default)] + pub use_mlock: bool, + pub flash_attention: Option, + pub n_threads: Option, + #[serde(default)] + pub native_tool_calling: bool, +} + +fn default_repeat_penalty() -> f32 { + 1.0 +} + +fn default_repeat_last_n() -> i32 { + 64 +} + +impl Default for ModelSettings { + fn default() -> Self { + Self { + context_size: None, + max_output_tokens: None, + sampling: SamplingConfig::default(), + repeat_penalty: 1.0, + repeat_last_n: 64, + frequency_penalty: 0.0, + presence_penalty: 0.0, + n_batch: None, + n_gpu_layers: None, + use_mlock: false, + flash_attention: None, + n_threads: None, + native_tool_calling: false, + } + } +} + +static REGISTRY: OnceLock> = OnceLock::new(); + +pub fn get_registry() -> &'static Mutex { + REGISTRY.get_or_init(|| { + let mut registry = LocalModelRegistry::load().unwrap_or_default(); + registry.migrate_legacy_models(); + registry.migrate_model_ids(); + Mutex::new(registry) + }) +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LocalModelEntry { + pub id: String, + pub display_name: String, + pub repo_id: String, + pub filename: String, + pub quantization: String, + pub local_path: PathBuf, + pub source_url: String, + #[serde(default)] + pub settings: ModelSettings, +} + +impl LocalModelEntry { + pub fn is_downloaded(&self) -> bool { + self.local_path.exists() + } + + pub fn file_size(&self) -> u64 { + std::fs::metadata(&self.local_path) + .map(|m| m.len()) + .unwrap_or(0) + } +} + +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct LocalModelRegistry { + pub models: Vec, +} + +fn registry_path() -> PathBuf { + Paths::in_data_dir("models/registry.json") +} + +/// The 4 legacy hardcoded model definitions for migration. +struct LegacyModel { + id: &'static str, + display_name: &'static str, + repo_id: &'static str, + filename: &'static str, + quantization: &'static str, + source_url: &'static str, +} + +const LEGACY_MODELS: &[LegacyModel] = &[ + LegacyModel { + id: "llama-3.2-1b", + display_name: "Llama 3.2 1B Instruct", + repo_id: "bartowski/Llama-3.2-1B-Instruct-GGUF", + filename: "Llama-3.2-1B-Instruct-Q4_K_M.gguf", + quantization: "Q4_K_M", + source_url: "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_K_M.gguf", + }, + LegacyModel { + id: "llama-3.2-3b", + display_name: "Llama 3.2 3B Instruct", + repo_id: "bartowski/Llama-3.2-3B-Instruct-GGUF", + filename: "Llama-3.2-3B-Instruct-Q4_K_M.gguf", + quantization: "Q4_K_M", + source_url: "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q4_K_M.gguf", + }, + LegacyModel { + id: "hermes-2-pro-7b", + display_name: "Hermes 2 Pro Llama-3 7B", + repo_id: "NousResearch/Hermes-2-Pro-Llama-3-8B-GGUF", + filename: "Hermes-2-Pro-Llama-3-8B-Q4_K_M.gguf", + quantization: "Q4_K_M", + source_url: "https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B-GGUF/resolve/main/Hermes-2-Pro-Llama-3-8B-Q4_K_M.gguf", + }, + LegacyModel { + id: "mistral-small-22b", + display_name: "Mistral Small 22B Instruct", + repo_id: "bartowski/Mistral-Small-Instruct-2409-GGUF", + filename: "Mistral-Small-Instruct-2409-Q4_K_M.gguf", + quantization: "Q4_K_M", + source_url: "https://huggingface.co/bartowski/Mistral-Small-Instruct-2409-GGUF/resolve/main/Mistral-Small-Instruct-2409-Q4_K_M.gguf", + }, +]; + +impl LocalModelRegistry { + pub fn load() -> Result { + let path = registry_path(); + if !path.exists() { + return Ok(Self::default()); + } + let data = std::fs::read_to_string(&path)?; + let registry: Self = serde_json::from_str(&data)?; + Ok(registry) + } + + pub fn save(&self) -> Result<()> { + let path = registry_path(); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let data = serde_json::to_string_pretty(self)?; + std::fs::write(&path, data)?; + Ok(()) + } + + /// Migrate model IDs from the old `author--repo--variant` format to the + /// HuggingFace-style `author/repo:variant` format. Only touches IDs that + /// contain `--` (the old separator); legacy featured models with short IDs + /// like `llama-3.2-1b` are left alone. + pub fn migrate_model_ids(&mut self) { + let mut changed = false; + for entry in &mut self.models { + if entry.id.contains("--") { + let new_id = model_id_from_repo(&entry.repo_id, &entry.quantization); + if entry.id != new_id { + entry.id = new_id; + changed = true; + } + } + } + if changed { + let _ = self.save(); + } + } + + /// Scan for legacy model files and add them to the registry if not already present. + pub fn migrate_legacy_models(&mut self) { + let mut changed = false; + for legacy in LEGACY_MODELS { + let legacy_path = Paths::in_data_dir("models").join(format!("{}.gguf", legacy.id)); + if legacy_path.exists() && !self.models.iter().any(|m| m.id == legacy.id) { + self.models.push(LocalModelEntry { + id: legacy.id.to_string(), + display_name: legacy.display_name.to_string(), + repo_id: legacy.repo_id.to_string(), + filename: legacy.filename.to_string(), + quantization: legacy.quantization.to_string(), + local_path: legacy_path, + source_url: legacy.source_url.to_string(), + settings: ModelSettings::default(), + }); + changed = true; + } + } + if changed { + let _ = self.save(); + } + } + + pub fn add_model(&mut self, entry: LocalModelEntry) -> Result<()> { + if let Some(existing) = self.models.iter_mut().find(|m| m.id == entry.id) { + *existing = entry; + } else { + self.models.push(entry); + } + self.save() + } + + pub fn remove_model(&mut self, id: &str) -> Result<()> { + self.models.retain(|m| m.id != id); + self.save() + } + + pub fn get_model(&self, id: &str) -> Option<&LocalModelEntry> { + self.models.iter().find(|m| m.id == id) + } + + pub fn get_model_settings(&self, id: &str) -> Option<&ModelSettings> { + self.models.iter().find(|m| m.id == id).map(|m| &m.settings) + } + + pub fn update_model_settings(&mut self, id: &str, settings: ModelSettings) -> Result<()> { + let entry = self + .models + .iter_mut() + .find(|m| m.id == id) + .ok_or_else(|| anyhow::anyhow!("Model not found: {}", id))?; + entry.settings = settings; + self.save() + } + + pub fn list_models(&self) -> &[LocalModelEntry] { + &self.models + } +} + +/// Generate a unique ID for a model from its repo_id and quantization. +/// Uses the HuggingFace convention: `author/model:variant`. +pub fn model_id_from_repo(repo_id: &str, quantization: &str) -> String { + format!("{}:{}", repo_id, quantization) +} + +/// Generate a display name from repo_id and quantization. +pub fn display_name_from_repo(repo_id: &str, quantization: &str) -> String { + let model_name = repo_id + .split('/') + .next_back() + .unwrap_or(repo_id) + .trim_end_matches("-GGUF") + .trim_end_matches("-gguf"); + format!("{} ({})", model_name, quantization) +} diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 3c13bdf96e..6843c9654f 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -1827,6 +1827,42 @@ } } }, + "/local-inference/download": { + "post": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "download_hf_model", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HfDownloadRequest" + } + } + }, + "required": true + }, + "responses": { + "202": { + "description": "Download started", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HfDownloadResponse" + } + } + } + }, + "400": { + "description": "Invalid request" + }, + "500": { + "description": "Internal server error" + } + } + } + }, "/local-inference/models": { "get": { "tags": [ @@ -1841,7 +1877,7 @@ "schema": { "type": "array", "items": { - "$ref": "#/components/schemas/LocalModelResponse" + "$ref": "#/components/schemas/ModelListItem" } } } @@ -1963,6 +1999,172 @@ } } }, + "/local-inference/models/{model_id}/settings": { + "get": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "get_model_settings", + "parameters": [ + { + "name": "model_id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Model settings", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelSettings" + } + } + } + }, + "404": { + "description": "Model not found" + } + } + }, + "put": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "update_model_settings", + "parameters": [ + { + "name": "model_id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelSettings" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Settings updated", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ModelSettings" + } + } + } + }, + "404": { + "description": "Model not found" + }, + "500": { + "description": "Failed to save settings" + } + } + } + }, + "/local-inference/repo/{author}/{repo}/files": { + "get": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "get_repo_files", + "parameters": [ + { + "name": "author", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "repo", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "GGUF quantization variants in repo", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RepoVariantsResponse" + } + } + } + }, + "500": { + "description": "Failed to fetch repo files" + } + } + } + }, + "/local-inference/search": { + "get": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "search_hf_models", + "parameters": [ + { + "name": "q", + "in": "query", + "description": "Search query", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "limit", + "in": "query", + "description": "Max results", + "required": false, + "schema": { + "type": "integer", + "nullable": true, + "minimum": 0 + } + } + ], + "responses": { + "200": { + "description": "Search results", + "content": { + "application/json": { + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/HfModelInfo" + } + } + } + } + }, + "500": { + "description": "Search failed" + } + } + } + }, "/mcp-ui-proxy": { "get": { "tags": [ @@ -4704,6 +4906,128 @@ } ] }, + "HfDownloadRequest": { + "type": "object", + "properties": { + "filename": { + "type": "string", + "nullable": true + }, + "repo_id": { + "type": "string", + "nullable": true + }, + "spec": { + "type": "string", + "nullable": true + } + } + }, + "HfDownloadResponse": { + "type": "object", + "required": [ + "model_id" + ], + "properties": { + "model_id": { + "type": "string" + } + } + }, + "HfGgufFile": { + "type": "object", + "description": "A single downloadable GGUF file (used internally and for downloads).", + "required": [ + "filename", + "size_bytes", + "quantization", + "download_url" + ], + "properties": { + "download_url": { + "type": "string" + }, + "filename": { + "type": "string" + }, + "quantization": { + "type": "string" + }, + "size_bytes": { + "type": "integer", + "format": "int64", + "minimum": 0 + } + } + }, + "HfModelInfo": { + "type": "object", + "required": [ + "repo_id", + "author", + "model_name", + "downloads", + "gguf_files" + ], + "properties": { + "author": { + "type": "string" + }, + "downloads": { + "type": "integer", + "format": "int64", + "minimum": 0 + }, + "gguf_files": { + "type": "array", + "items": { + "$ref": "#/components/schemas/HfGgufFile" + } + }, + "model_name": { + "type": "string" + }, + "repo_id": { + "type": "string" + } + } + }, + "HfQuantVariant": { + "type": "object", + "description": "A quantization variant — groups sharded files into one logical entry.", + "required": [ + "quantization", + "size_bytes", + "filename", + "download_url", + "description", + "quality_rank" + ], + "properties": { + "description": { + "type": "string" + }, + "download_url": { + "type": "string" + }, + "filename": { + "type": "string" + }, + "quality_rank": { + "type": "integer", + "format": "int32", + "minimum": 0 + }, + "quantization": { + "type": "string" + }, + "size_bytes": { + "type": "integer", + "format": "int64", + "minimum": 0 + } + } + }, "Icon": { "type": "object", "required": [ @@ -4897,44 +5221,33 @@ "size_mb", "context_limit", "url", - "tokenizer_url", "description", "tier" ], "properties": { "context_limit": { "type": "integer", - "description": "Maximum context window in tokens", "minimum": 0 }, "description": { - "type": "string", - "description": "Description and use case" + "type": "string" }, "id": { - "type": "string", - "description": "Model identifier (e.g., \"llama-3.2-1b\")" + "type": "string" }, "name": { - "type": "string", - "description": "Display name" + "type": "string" }, "size_mb": { "type": "integer", "format": "int32", - "description": "Model file size in MB", "minimum": 0 }, "tier": { "$ref": "#/components/schemas/ModelTier" }, - "tokenizer_url": { - "type": "string", - "description": "Download URL for the tokenizer JSON" - }, "url": { - "type": "string", - "description": "Download URL for the model GGUF file" + "type": "string" } } }, @@ -4948,44 +5261,33 @@ "size_mb", "context_limit", "url", - "tokenizer_url", "description", "tier" ], "properties": { "context_limit": { "type": "integer", - "description": "Maximum context window in tokens", "minimum": 0 }, "description": { - "type": "string", - "description": "Description and use case" + "type": "string" }, "id": { - "type": "string", - "description": "Model identifier (e.g., \"llama-3.2-1b\")" + "type": "string" }, "name": { - "type": "string", - "description": "Display name" + "type": "string" }, "size_mb": { "type": "integer", "format": "int32", - "description": "Model file size in MB", "minimum": 0 }, "tier": { "$ref": "#/components/schemas/ModelTier" }, - "tokenizer_url": { - "type": "string", - "description": "Download URL for the tokenizer JSON" - }, "url": { - "type": "string", - "description": "Download URL for the model GGUF file" + "type": "string" } } }, @@ -4993,12 +5295,16 @@ "type": "object", "required": [ "downloaded", - "recommended" + "recommended", + "featured" ], "properties": { "downloaded": { "type": "boolean" }, + "featured": { + "type": "boolean" + }, "recommended": { "type": "boolean" } @@ -5550,6 +5856,78 @@ } } }, + "ModelListItem": { + "oneOf": [ + { + "$ref": "#/components/schemas/LocalModelResponse" + }, + { + "$ref": "#/components/schemas/RegistryModelResponse" + } + ] + }, + "ModelSettings": { + "type": "object", + "properties": { + "context_size": { + "type": "integer", + "format": "int32", + "nullable": true, + "minimum": 0 + }, + "flash_attention": { + "type": "boolean", + "nullable": true + }, + "frequency_penalty": { + "type": "number", + "format": "float" + }, + "max_output_tokens": { + "type": "integer", + "nullable": true, + "minimum": 0 + }, + "n_batch": { + "type": "integer", + "format": "int32", + "nullable": true, + "minimum": 0 + }, + "n_gpu_layers": { + "type": "integer", + "format": "int32", + "nullable": true, + "minimum": 0 + }, + "n_threads": { + "type": "integer", + "format": "int32", + "nullable": true + }, + "native_tool_calling": { + "type": "boolean" + }, + "presence_penalty": { + "type": "number", + "format": "float" + }, + "repeat_last_n": { + "type": "integer", + "format": "int32" + }, + "repeat_penalty": { + "type": "number", + "format": "float" + }, + "sampling": { + "$ref": "#/components/schemas/SamplingConfig" + }, + "use_mlock": { + "type": "boolean" + } + } + }, "ModelTier": { "type": "string", "enum": [ @@ -6189,6 +6567,51 @@ } } }, + "RegistryModelResponse": { + "type": "object", + "required": [ + "id", + "display_name", + "repo_id", + "filename", + "quantization", + "size_bytes", + "downloaded", + "featured", + "settings" + ], + "properties": { + "display_name": { + "type": "string" + }, + "downloaded": { + "type": "boolean" + }, + "featured": { + "type": "boolean" + }, + "filename": { + "type": "string" + }, + "id": { + "type": "string" + }, + "quantization": { + "type": "string" + }, + "repo_id": { + "type": "string" + }, + "settings": { + "$ref": "#/components/schemas/ModelSettings" + }, + "size_bytes": { + "type": "integer", + "format": "int64", + "minimum": 0 + } + } + }, "RemoveExtensionRequest": { "type": "object", "required": [ @@ -6204,6 +6627,25 @@ } } }, + "RepoVariantsResponse": { + "type": "object", + "required": [ + "variants" + ], + "properties": { + "recommended_index": { + "type": "integer", + "nullable": true, + "minimum": 0 + }, + "variants": { + "type": "array", + "items": { + "$ref": "#/components/schemas/HfQuantVariant" + } + } + } + }, "ResourceContents": { "anyOf": [ { @@ -6395,6 +6837,97 @@ } } }, + "SamplingConfig": { + "oneOf": [ + { + "type": "object", + "required": [ + "type" + ], + "properties": { + "type": { + "type": "string", + "enum": [ + "Greedy" + ] + } + } + }, + { + "type": "object", + "required": [ + "temperature", + "top_k", + "top_p", + "min_p", + "type" + ], + "properties": { + "min_p": { + "type": "number", + "format": "float" + }, + "seed": { + "type": "integer", + "format": "int32", + "nullable": true, + "minimum": 0 + }, + "temperature": { + "type": "number", + "format": "float" + }, + "top_k": { + "type": "integer", + "format": "int32" + }, + "top_p": { + "type": "number", + "format": "float" + }, + "type": { + "type": "string", + "enum": [ + "Temperature" + ] + } + } + }, + { + "type": "object", + "required": [ + "tau", + "eta", + "type" + ], + "properties": { + "eta": { + "type": "number", + "format": "float" + }, + "seed": { + "type": "integer", + "format": "int32", + "nullable": true, + "minimum": 0 + }, + "tau": { + "type": "number", + "format": "float" + }, + "type": { + "type": "string", + "enum": [ + "MirostatV2" + ] + } + } + } + ], + "discriminator": { + "propertyName": "type" + } + }, "SavePromptRequest": { "type": "object", "required": [ diff --git a/ui/desktop/src/api/index.ts b/ui/desktop/src/api/index.ts index c72c7d389e..0a9369bf76 100644 --- a/ui/desktop/src/api/index.ts +++ b/ui/desktop/src/api/index.ts @@ -1,4 +1,4 @@ // This file is auto-generated by @hey-api/openapi-ts -export { addExtension, agentAddExtension, agentRemoveExtension, backupConfig, callTool, cancelDownload, cancelLocalModelDownload, checkProvider, configureProviderOauth, confirmToolAction, createCustomProvider, createRecipe, createSchedule, decodeRecipe, deleteLocalModel, deleteModel, deleteRecipe, deleteSchedule, deleteSession, detectProvider, diagnostics, downloadLocalModel, downloadModel, encodeRecipe, exportApp, exportSession, forkSession, getCustomProvider, getDictationConfig, getDownloadProgress, getExtensions, getLocalModelDownloadProgress, getPricing, getPrompt, getPrompts, getProviderModels, getSession, getSessionExtensions, getSessionInsights, getSlashCommands, getTools, getTunnelStatus, importApp, importSession, initConfig, inspectRunningJob, killRunningJob, listApps, listLocalModels, listModels, listRecipes, listSchedules, listSessions, mcpUiProxy, type Options, parseRecipe, pauseSchedule, providers, readAllConfig, readConfig, readResource, recipeToYaml, recoverConfig, removeConfig, removeCustomProvider, removeExtension, reply, resetPrompt, restartAgent, resumeAgent, runNowHandler, savePrompt, saveRecipe, scanRecipe, scheduleRecipe, searchSessions, sendTelemetryEvent, sessionsHandler, setConfigProvider, setRecipeSlashCommand, startAgent, startOpenrouterSetup, startTetrateSetup, startTunnel, status, stopAgent, stopTunnel, systemInfo, transcribeDictation, unpauseSchedule, updateAgentProvider, updateCustomProvider, updateFromSession, updateSchedule, updateSessionName, updateSessionUserRecipeValues, updateWorkingDir, upsertConfig, upsertPermissions, validateConfig } from './sdk.gen'; -export type { ActionRequired, ActionRequiredData, AddExtensionData, AddExtensionErrors, AddExtensionRequest, AddExtensionResponse, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponse, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponse, AgentRemoveExtensionResponses, Annotations, Author, AuthorRequest, BackupConfigData, BackupConfigErrors, BackupConfigResponse, BackupConfigResponses, CallToolData, CallToolErrors, CallToolRequest, CallToolResponse, CallToolResponse2, CallToolResponses, CancelDownloadData, CancelDownloadErrors, CancelDownloadResponses, CancelLocalModelDownloadData, CancelLocalModelDownloadErrors, CancelLocalModelDownloadResponses, ChatRequest, CheckProviderData, CheckProviderRequest, ClientOptions, CommandType, ConfigKey, ConfigKeyQuery, ConfigResponse, ConfigureProviderOauthData, ConfigureProviderOauthErrors, ConfigureProviderOauthResponses, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionRequest, ConfirmToolActionResponses, Content, Conversation, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponse, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeRequest, CreateRecipeResponse, CreateRecipeResponse2, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleRequest, CreateScheduleResponse, CreateScheduleResponses, CspMetadata, DeclarativeProviderConfig, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeRequest, DecodeRecipeResponse, DecodeRecipeResponse2, DecodeRecipeResponses, DeleteLocalModelData, DeleteLocalModelErrors, DeleteLocalModelResponses, DeleteModelData, DeleteModelErrors, DeleteModelResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeRequest, DeleteRecipeResponse, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponse, DeleteScheduleResponses, DeleteSessionData, DeleteSessionErrors, DeleteSessionResponses, DetectProviderData, DetectProviderErrors, DetectProviderRequest, DetectProviderResponse, DetectProviderResponse2, DetectProviderResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponse, DiagnosticsResponses, DictationProvider, DictationProviderStatus, DownloadLocalModelData, DownloadLocalModelErrors, DownloadLocalModelResponses, DownloadModelData, DownloadModelErrors, DownloadModelResponses, DownloadProgress, DownloadStatus, EmbeddedResource, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeRequest, EncodeRecipeResponse, EncodeRecipeResponse2, EncodeRecipeResponses, Envs, ErrorResponse, ExportAppData, ExportAppError, ExportAppErrors, ExportAppResponse, ExportAppResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponse, ExportSessionResponses, ExtensionConfig, ExtensionData, ExtensionEntry, ExtensionLoadResult, ExtensionQuery, ExtensionResponse, ForkRequest, ForkResponse, ForkSessionData, ForkSessionErrors, ForkSessionResponse, ForkSessionResponses, FrontendToolRequest, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponse, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponse, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponse, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponse, GetExtensionsResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponse, GetLocalModelDownloadProgressResponses, GetPricingData, GetPricingResponse, GetPricingResponses, GetPromptData, GetPromptErrors, GetPromptResponse, GetPromptResponses, GetPromptsData, GetPromptsResponse, GetPromptsResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponse, GetProviderModelsResponses, GetSessionData, GetSessionErrors, GetSessionExtensionsData, GetSessionExtensionsErrors, GetSessionExtensionsResponse, GetSessionExtensionsResponses, GetSessionInsightsData, GetSessionInsightsErrors, GetSessionInsightsResponse, GetSessionInsightsResponses, GetSessionResponse, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponse, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsQuery, GetToolsResponse, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponse, GetTunnelStatusResponses, GooseApp, Icon, ImageContent, ImportAppData, ImportAppError, ImportAppErrors, ImportAppRequest, ImportAppResponse, ImportAppResponse2, ImportAppResponses, ImportSessionData, ImportSessionErrors, ImportSessionRequest, ImportSessionResponse, ImportSessionResponses, InitConfigData, InitConfigErrors, InitConfigResponse, InitConfigResponses, InspectJobResponse, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponse, InspectRunningJobResponses, JsonObject, KillJobResponse, KillRunningJobData, KillRunningJobResponses, ListAppsData, ListAppsError, ListAppsErrors, ListAppsRequest, ListAppsResponse, ListAppsResponse2, ListAppsResponses, ListLocalModelsData, ListLocalModelsResponse, ListLocalModelsResponses, ListModelsData, ListModelsResponse, ListModelsResponses, ListRecipeResponse, ListRecipesData, ListRecipesErrors, ListRecipesResponse, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponse, ListSchedulesResponse2, ListSchedulesResponses, ListSessionsData, ListSessionsErrors, ListSessionsResponse, ListSessionsResponses, LoadedProvider, LocalLlmModel, LocalModelResponse, McpAppResource, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, Message, MessageContent, MessageEvent, MessageMetadata, ModelConfig, ModelInfo, ModelTier, ParseRecipeData, ParseRecipeError, ParseRecipeErrors, ParseRecipeRequest, ParseRecipeResponse, ParseRecipeResponse2, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponse, PauseScheduleResponses, Permission, PermissionLevel, PermissionsMetadata, PricingData, PricingQuery, PricingResponse, PrincipalType, PromptContentResponse, PromptsListResponse, ProviderDetails, ProviderEngine, ProviderMetadata, ProvidersData, ProvidersResponse, ProvidersResponse2, ProvidersResponses, ProviderType, RawAudioContent, RawEmbeddedResource, RawImageContent, RawResource, RawTextContent, ReadAllConfigData, ReadAllConfigResponse, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceRequest, ReadResourceResponse, ReadResourceResponse2, ReadResourceResponses, Recipe, RecipeManifest, RecipeParameter, RecipeParameterInputType, RecipeParameterRequirement, RecipeToYamlData, RecipeToYamlError, RecipeToYamlErrors, RecipeToYamlRequest, RecipeToYamlResponse, RecipeToYamlResponse2, RecipeToYamlResponses, RecoverConfigData, RecoverConfigErrors, RecoverConfigResponse, RecoverConfigResponses, RedactedThinkingContent, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponse, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponse, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionRequest, RemoveExtensionResponse, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponse, ReplyResponses, ResetPromptData, ResetPromptErrors, ResetPromptResponse, ResetPromptResponses, ResourceContents, ResourceMetadata, Response, RestartAgentData, RestartAgentErrors, RestartAgentRequest, RestartAgentResponse, RestartAgentResponse2, RestartAgentResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentRequest, ResumeAgentResponse, ResumeAgentResponse2, ResumeAgentResponses, RetryConfig, Role, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponse, RunNowHandlerResponses, RunNowResponse, SavePromptData, SavePromptErrors, SavePromptRequest, SavePromptResponse, SavePromptResponses, SaveRecipeData, SaveRecipeError, SaveRecipeErrors, SaveRecipeRequest, SaveRecipeResponse, SaveRecipeResponse2, SaveRecipeResponses, ScanRecipeData, ScanRecipeRequest, ScanRecipeResponse, ScanRecipeResponse2, ScanRecipeResponses, ScheduledJob, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeRequest, ScheduleRecipeResponses, SearchSessionsData, SearchSessionsErrors, SearchSessionsResponse, SearchSessionsResponses, SendTelemetryEventData, SendTelemetryEventResponses, Session, SessionDisplayInfo, SessionExtensionsResponse, SessionInsights, SessionListResponse, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponse, SessionsHandlerResponses, SessionsQuery, SessionType, SetConfigProviderData, SetProviderRequest, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, SetSlashCommandRequest, Settings, SetupResponse, SlashCommand, SlashCommandsResponse, StartAgentData, StartAgentError, StartAgentErrors, StartAgentRequest, StartAgentResponse, StartAgentResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponse, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponse, StartTetrateSetupResponses, StartTunnelData, StartTunnelError, StartTunnelErrors, StartTunnelResponse, StartTunnelResponses, StatusData, StatusResponse, StatusResponses, StopAgentData, StopAgentErrors, StopAgentRequest, StopAgentResponse, StopAgentResponses, StopTunnelData, StopTunnelError, StopTunnelErrors, StopTunnelResponses, SubRecipe, SuccessCheck, SystemInfo, SystemInfoData, SystemInfoResponse, SystemInfoResponses, SystemNotificationContent, SystemNotificationType, TelemetryEventRequest, Template, TextContent, ThinkingContent, TokenState, Tool, ToolAnnotations, ToolConfirmationRequest, ToolInfo, ToolPermission, ToolRequest, ToolResponse, TranscribeDictationData, TranscribeDictationErrors, TranscribeDictationResponse, TranscribeDictationResponses, TranscribeRequest, TranscribeResponse, TunnelInfo, TunnelState, UiMetadata, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponse, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderRequest, UpdateCustomProviderResponse, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionRequest, UpdateFromSessionResponses, UpdateProviderRequest, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleRequest, UpdateScheduleResponse, UpdateScheduleResponses, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameRequest, UpdateSessionNameResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesError, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesRequest, UpdateSessionUserRecipeValuesResponse, UpdateSessionUserRecipeValuesResponse2, UpdateSessionUserRecipeValuesResponses, UpdateWorkingDirData, UpdateWorkingDirErrors, UpdateWorkingDirRequest, UpdateWorkingDirResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigQuery, UpsertConfigResponse, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsQuery, UpsertPermissionsResponse, UpsertPermissionsResponses, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponse, ValidateConfigResponses, WhisperModelResponse, WindowProps } from './types.gen'; +export { addExtension, agentAddExtension, agentRemoveExtension, backupConfig, callTool, cancelDownload, cancelLocalModelDownload, checkProvider, configureProviderOauth, confirmToolAction, createCustomProvider, createRecipe, createSchedule, decodeRecipe, deleteLocalModel, deleteModel, deleteRecipe, deleteSchedule, deleteSession, detectProvider, diagnostics, downloadHfModel, downloadLocalModel, downloadModel, encodeRecipe, exportApp, exportSession, forkSession, getCustomProvider, getDictationConfig, getDownloadProgress, getExtensions, getLocalModelDownloadProgress, getModelSettings, getPricing, getPrompt, getPrompts, getProviderModels, getRepoFiles, getSession, getSessionExtensions, getSessionInsights, getSlashCommands, getTools, getTunnelStatus, importApp, importSession, initConfig, inspectRunningJob, killRunningJob, listApps, listLocalModels, listModels, listRecipes, listSchedules, listSessions, mcpUiProxy, type Options, parseRecipe, pauseSchedule, providers, readAllConfig, readConfig, readResource, recipeToYaml, recoverConfig, removeConfig, removeCustomProvider, removeExtension, reply, resetPrompt, restartAgent, resumeAgent, runNowHandler, savePrompt, saveRecipe, scanRecipe, scheduleRecipe, searchHfModels, searchSessions, sendTelemetryEvent, sessionsHandler, setConfigProvider, setRecipeSlashCommand, startAgent, startOpenrouterSetup, startTetrateSetup, startTunnel, status, stopAgent, stopTunnel, systemInfo, transcribeDictation, unpauseSchedule, updateAgentProvider, updateCustomProvider, updateFromSession, updateModelSettings, updateSchedule, updateSessionName, updateSessionUserRecipeValues, updateWorkingDir, upsertConfig, upsertPermissions, validateConfig } from './sdk.gen'; +export type { ActionRequired, ActionRequiredData, AddExtensionData, AddExtensionErrors, AddExtensionRequest, AddExtensionResponse, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponse, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponse, AgentRemoveExtensionResponses, Annotations, Author, AuthorRequest, BackupConfigData, BackupConfigErrors, BackupConfigResponse, BackupConfigResponses, CallToolData, CallToolErrors, CallToolRequest, CallToolResponse, CallToolResponse2, CallToolResponses, CancelDownloadData, CancelDownloadErrors, CancelDownloadResponses, CancelLocalModelDownloadData, CancelLocalModelDownloadErrors, CancelLocalModelDownloadResponses, ChatRequest, CheckProviderData, CheckProviderRequest, ClientOptions, CommandType, ConfigKey, ConfigKeyQuery, ConfigResponse, ConfigureProviderOauthData, ConfigureProviderOauthErrors, ConfigureProviderOauthResponses, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionRequest, ConfirmToolActionResponses, Content, Conversation, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponse, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeRequest, CreateRecipeResponse, CreateRecipeResponse2, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleRequest, CreateScheduleResponse, CreateScheduleResponses, CspMetadata, DeclarativeProviderConfig, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeRequest, DecodeRecipeResponse, DecodeRecipeResponse2, DecodeRecipeResponses, DeleteLocalModelData, DeleteLocalModelErrors, DeleteLocalModelResponses, DeleteModelData, DeleteModelErrors, DeleteModelResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeRequest, DeleteRecipeResponse, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponse, DeleteScheduleResponses, DeleteSessionData, DeleteSessionErrors, DeleteSessionResponses, DetectProviderData, DetectProviderErrors, DetectProviderRequest, DetectProviderResponse, DetectProviderResponse2, DetectProviderResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponse, DiagnosticsResponses, DictationProvider, DictationProviderStatus, DownloadHfModelData, DownloadHfModelErrors, DownloadHfModelResponse, DownloadHfModelResponses, DownloadLocalModelData, DownloadLocalModelErrors, DownloadLocalModelResponses, DownloadModelData, DownloadModelErrors, DownloadModelResponses, DownloadProgress, DownloadStatus, EmbeddedResource, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeRequest, EncodeRecipeResponse, EncodeRecipeResponse2, EncodeRecipeResponses, Envs, ErrorResponse, ExportAppData, ExportAppError, ExportAppErrors, ExportAppResponse, ExportAppResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponse, ExportSessionResponses, ExtensionConfig, ExtensionData, ExtensionEntry, ExtensionLoadResult, ExtensionQuery, ExtensionResponse, ForkRequest, ForkResponse, ForkSessionData, ForkSessionErrors, ForkSessionResponse, ForkSessionResponses, FrontendToolRequest, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponse, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponse, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponse, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponse, GetExtensionsResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponse, GetLocalModelDownloadProgressResponses, GetModelSettingsData, GetModelSettingsErrors, GetModelSettingsResponse, GetModelSettingsResponses, GetPricingData, GetPricingResponse, GetPricingResponses, GetPromptData, GetPromptErrors, GetPromptResponse, GetPromptResponses, GetPromptsData, GetPromptsResponse, GetPromptsResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponse, GetProviderModelsResponses, GetRepoFilesData, GetRepoFilesErrors, GetRepoFilesResponse, GetRepoFilesResponses, GetSessionData, GetSessionErrors, GetSessionExtensionsData, GetSessionExtensionsErrors, GetSessionExtensionsResponse, GetSessionExtensionsResponses, GetSessionInsightsData, GetSessionInsightsErrors, GetSessionInsightsResponse, GetSessionInsightsResponses, GetSessionResponse, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponse, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsQuery, GetToolsResponse, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponse, GetTunnelStatusResponses, GooseApp, HfDownloadRequest, HfDownloadResponse, HfGgufFile, HfModelInfo, HfQuantVariant, Icon, ImageContent, ImportAppData, ImportAppError, ImportAppErrors, ImportAppRequest, ImportAppResponse, ImportAppResponse2, ImportAppResponses, ImportSessionData, ImportSessionErrors, ImportSessionRequest, ImportSessionResponse, ImportSessionResponses, InitConfigData, InitConfigErrors, InitConfigResponse, InitConfigResponses, InspectJobResponse, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponse, InspectRunningJobResponses, JsonObject, KillJobResponse, KillRunningJobData, KillRunningJobResponses, ListAppsData, ListAppsError, ListAppsErrors, ListAppsRequest, ListAppsResponse, ListAppsResponse2, ListAppsResponses, ListLocalModelsData, ListLocalModelsResponse, ListLocalModelsResponses, ListModelsData, ListModelsResponse, ListModelsResponses, ListRecipeResponse, ListRecipesData, ListRecipesErrors, ListRecipesResponse, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponse, ListSchedulesResponse2, ListSchedulesResponses, ListSessionsData, ListSessionsErrors, ListSessionsResponse, ListSessionsResponses, LoadedProvider, LocalLlmModel, LocalModelResponse, McpAppResource, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, Message, MessageContent, MessageEvent, MessageMetadata, ModelConfig, ModelInfo, ModelListItem, ModelSettings, ModelTier, ParseRecipeData, ParseRecipeError, ParseRecipeErrors, ParseRecipeRequest, ParseRecipeResponse, ParseRecipeResponse2, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponse, PauseScheduleResponses, Permission, PermissionLevel, PermissionsMetadata, PricingData, PricingQuery, PricingResponse, PrincipalType, PromptContentResponse, PromptsListResponse, ProviderDetails, ProviderEngine, ProviderMetadata, ProvidersData, ProvidersResponse, ProvidersResponse2, ProvidersResponses, ProviderType, RawAudioContent, RawEmbeddedResource, RawImageContent, RawResource, RawTextContent, ReadAllConfigData, ReadAllConfigResponse, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceRequest, ReadResourceResponse, ReadResourceResponse2, ReadResourceResponses, Recipe, RecipeManifest, RecipeParameter, RecipeParameterInputType, RecipeParameterRequirement, RecipeToYamlData, RecipeToYamlError, RecipeToYamlErrors, RecipeToYamlRequest, RecipeToYamlResponse, RecipeToYamlResponse2, RecipeToYamlResponses, RecoverConfigData, RecoverConfigErrors, RecoverConfigResponse, RecoverConfigResponses, RedactedThinkingContent, RegistryModelResponse, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponse, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponse, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionRequest, RemoveExtensionResponse, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponse, ReplyResponses, RepoVariantsResponse, ResetPromptData, ResetPromptErrors, ResetPromptResponse, ResetPromptResponses, ResourceContents, ResourceMetadata, Response, RestartAgentData, RestartAgentErrors, RestartAgentRequest, RestartAgentResponse, RestartAgentResponse2, RestartAgentResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentRequest, ResumeAgentResponse, ResumeAgentResponse2, ResumeAgentResponses, RetryConfig, Role, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponse, RunNowHandlerResponses, RunNowResponse, SamplingConfig, SavePromptData, SavePromptErrors, SavePromptRequest, SavePromptResponse, SavePromptResponses, SaveRecipeData, SaveRecipeError, SaveRecipeErrors, SaveRecipeRequest, SaveRecipeResponse, SaveRecipeResponse2, SaveRecipeResponses, ScanRecipeData, ScanRecipeRequest, ScanRecipeResponse, ScanRecipeResponse2, ScanRecipeResponses, ScheduledJob, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeRequest, ScheduleRecipeResponses, SearchHfModelsData, SearchHfModelsErrors, SearchHfModelsResponse, SearchHfModelsResponses, SearchSessionsData, SearchSessionsErrors, SearchSessionsResponse, SearchSessionsResponses, SendTelemetryEventData, SendTelemetryEventResponses, Session, SessionDisplayInfo, SessionExtensionsResponse, SessionInsights, SessionListResponse, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponse, SessionsHandlerResponses, SessionsQuery, SessionType, SetConfigProviderData, SetProviderRequest, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, SetSlashCommandRequest, Settings, SetupResponse, SlashCommand, SlashCommandsResponse, StartAgentData, StartAgentError, StartAgentErrors, StartAgentRequest, StartAgentResponse, StartAgentResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponse, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponse, StartTetrateSetupResponses, StartTunnelData, StartTunnelError, StartTunnelErrors, StartTunnelResponse, StartTunnelResponses, StatusData, StatusResponse, StatusResponses, StopAgentData, StopAgentErrors, StopAgentRequest, StopAgentResponse, StopAgentResponses, StopTunnelData, StopTunnelError, StopTunnelErrors, StopTunnelResponses, SubRecipe, SuccessCheck, SystemInfo, SystemInfoData, SystemInfoResponse, SystemInfoResponses, SystemNotificationContent, SystemNotificationType, TelemetryEventRequest, Template, TextContent, ThinkingContent, TokenState, Tool, ToolAnnotations, ToolConfirmationRequest, ToolInfo, ToolPermission, ToolRequest, ToolResponse, TranscribeDictationData, TranscribeDictationErrors, TranscribeDictationResponse, TranscribeDictationResponses, TranscribeRequest, TranscribeResponse, TunnelInfo, TunnelState, UiMetadata, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponse, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderRequest, UpdateCustomProviderResponse, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionRequest, UpdateFromSessionResponses, UpdateModelSettingsData, UpdateModelSettingsErrors, UpdateModelSettingsResponse, UpdateModelSettingsResponses, UpdateProviderRequest, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleRequest, UpdateScheduleResponse, UpdateScheduleResponses, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameRequest, UpdateSessionNameResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesError, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesRequest, UpdateSessionUserRecipeValuesResponse, UpdateSessionUserRecipeValuesResponse2, UpdateSessionUserRecipeValuesResponses, UpdateWorkingDirData, UpdateWorkingDirErrors, UpdateWorkingDirRequest, UpdateWorkingDirResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigQuery, UpsertConfigResponse, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsQuery, UpsertPermissionsResponse, UpsertPermissionsResponses, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponse, ValidateConfigResponses, WhisperModelResponse, WindowProps } from './types.gen'; diff --git a/ui/desktop/src/api/sdk.gen.ts b/ui/desktop/src/api/sdk.gen.ts index 8760aac08c..a1e7d91f13 100644 --- a/ui/desktop/src/api/sdk.gen.ts +++ b/ui/desktop/src/api/sdk.gen.ts @@ -2,7 +2,7 @@ import type { Client, Options as Options2, TDataShape } from './client'; import { client } from './client.gen'; -import type { AddExtensionData, AddExtensionErrors, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponses, BackupConfigData, BackupConfigErrors, BackupConfigResponses, CallToolData, CallToolErrors, CallToolResponses, CancelDownloadData, CancelDownloadErrors, CancelDownloadResponses, CancelLocalModelDownloadData, CancelLocalModelDownloadErrors, CancelLocalModelDownloadResponses, CheckProviderData, ConfigureProviderOauthData, ConfigureProviderOauthErrors, ConfigureProviderOauthResponses, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionResponses, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleResponses, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeResponses, DeleteLocalModelData, DeleteLocalModelErrors, DeleteLocalModelResponses, DeleteModelData, DeleteModelErrors, DeleteModelResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponses, DeleteSessionData, DeleteSessionErrors, DeleteSessionResponses, DetectProviderData, DetectProviderErrors, DetectProviderResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponses, DownloadLocalModelData, DownloadLocalModelErrors, DownloadLocalModelResponses, DownloadModelData, DownloadModelErrors, DownloadModelResponses, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeResponses, ExportAppData, ExportAppErrors, ExportAppResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponses, ForkSessionData, ForkSessionErrors, ForkSessionResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponses, GetPricingData, GetPricingResponses, GetPromptData, GetPromptErrors, GetPromptResponses, GetPromptsData, GetPromptsResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponses, GetSessionData, GetSessionErrors, GetSessionExtensionsData, GetSessionExtensionsErrors, GetSessionExtensionsResponses, GetSessionInsightsData, GetSessionInsightsErrors, GetSessionInsightsResponses, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponses, ImportAppData, ImportAppErrors, ImportAppResponses, ImportSessionData, ImportSessionErrors, ImportSessionResponses, InitConfigData, InitConfigErrors, InitConfigResponses, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponses, KillRunningJobData, KillRunningJobResponses, ListAppsData, ListAppsErrors, ListAppsResponses, ListLocalModelsData, ListLocalModelsResponses, ListModelsData, ListModelsResponses, ListRecipesData, ListRecipesErrors, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponses, ListSessionsData, ListSessionsErrors, ListSessionsResponses, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, ParseRecipeData, ParseRecipeErrors, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponses, ProvidersData, ProvidersResponses, ReadAllConfigData, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceResponses, RecipeToYamlData, RecipeToYamlErrors, RecipeToYamlResponses, RecoverConfigData, RecoverConfigErrors, RecoverConfigResponses, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponses, ResetPromptData, ResetPromptErrors, ResetPromptResponses, RestartAgentData, RestartAgentErrors, RestartAgentResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentResponses, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponses, SavePromptData, SavePromptErrors, SavePromptResponses, SaveRecipeData, SaveRecipeErrors, SaveRecipeResponses, ScanRecipeData, ScanRecipeResponses, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeResponses, SearchSessionsData, SearchSessionsErrors, SearchSessionsResponses, SendTelemetryEventData, SendTelemetryEventResponses, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponses, SetConfigProviderData, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, StartAgentData, StartAgentErrors, StartAgentResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponses, StartTunnelData, StartTunnelErrors, StartTunnelResponses, StatusData, StatusResponses, StopAgentData, StopAgentErrors, StopAgentResponses, StopTunnelData, StopTunnelErrors, StopTunnelResponses, SystemInfoData, SystemInfoResponses, TranscribeDictationData, TranscribeDictationErrors, TranscribeDictationResponses, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionResponses, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleResponses, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesResponses, UpdateWorkingDirData, UpdateWorkingDirErrors, UpdateWorkingDirResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsResponses, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponses } from './types.gen'; +import type { AddExtensionData, AddExtensionErrors, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponses, BackupConfigData, BackupConfigErrors, BackupConfigResponses, CallToolData, CallToolErrors, CallToolResponses, CancelDownloadData, CancelDownloadErrors, CancelDownloadResponses, CancelLocalModelDownloadData, CancelLocalModelDownloadErrors, CancelLocalModelDownloadResponses, CheckProviderData, ConfigureProviderOauthData, ConfigureProviderOauthErrors, ConfigureProviderOauthResponses, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionResponses, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleResponses, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeResponses, DeleteLocalModelData, DeleteLocalModelErrors, DeleteLocalModelResponses, DeleteModelData, DeleteModelErrors, DeleteModelResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponses, DeleteSessionData, DeleteSessionErrors, DeleteSessionResponses, DetectProviderData, DetectProviderErrors, DetectProviderResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponses, DownloadHfModelData, DownloadHfModelErrors, DownloadHfModelResponses, DownloadLocalModelData, DownloadLocalModelErrors, DownloadLocalModelResponses, DownloadModelData, DownloadModelErrors, DownloadModelResponses, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeResponses, ExportAppData, ExportAppErrors, ExportAppResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponses, ForkSessionData, ForkSessionErrors, ForkSessionResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponses, GetModelSettingsData, GetModelSettingsErrors, GetModelSettingsResponses, GetPricingData, GetPricingResponses, GetPromptData, GetPromptErrors, GetPromptResponses, GetPromptsData, GetPromptsResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponses, GetRepoFilesData, GetRepoFilesErrors, GetRepoFilesResponses, GetSessionData, GetSessionErrors, GetSessionExtensionsData, GetSessionExtensionsErrors, GetSessionExtensionsResponses, GetSessionInsightsData, GetSessionInsightsErrors, GetSessionInsightsResponses, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponses, ImportAppData, ImportAppErrors, ImportAppResponses, ImportSessionData, ImportSessionErrors, ImportSessionResponses, InitConfigData, InitConfigErrors, InitConfigResponses, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponses, KillRunningJobData, KillRunningJobResponses, ListAppsData, ListAppsErrors, ListAppsResponses, ListLocalModelsData, ListLocalModelsResponses, ListModelsData, ListModelsResponses, ListRecipesData, ListRecipesErrors, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponses, ListSessionsData, ListSessionsErrors, ListSessionsResponses, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, ParseRecipeData, ParseRecipeErrors, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponses, ProvidersData, ProvidersResponses, ReadAllConfigData, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceResponses, RecipeToYamlData, RecipeToYamlErrors, RecipeToYamlResponses, RecoverConfigData, RecoverConfigErrors, RecoverConfigResponses, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponses, ResetPromptData, ResetPromptErrors, ResetPromptResponses, RestartAgentData, RestartAgentErrors, RestartAgentResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentResponses, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponses, SavePromptData, SavePromptErrors, SavePromptResponses, SaveRecipeData, SaveRecipeErrors, SaveRecipeResponses, ScanRecipeData, ScanRecipeResponses, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeResponses, SearchHfModelsData, SearchHfModelsErrors, SearchHfModelsResponses, SearchSessionsData, SearchSessionsErrors, SearchSessionsResponses, SendTelemetryEventData, SendTelemetryEventResponses, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponses, SetConfigProviderData, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, StartAgentData, StartAgentErrors, StartAgentResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponses, StartTunnelData, StartTunnelErrors, StartTunnelResponses, StatusData, StatusResponses, StopAgentData, StopAgentErrors, StopAgentResponses, StopTunnelData, StopTunnelErrors, StopTunnelResponses, SystemInfoData, SystemInfoResponses, TranscribeDictationData, TranscribeDictationErrors, TranscribeDictationResponses, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionResponses, UpdateModelSettingsData, UpdateModelSettingsErrors, UpdateModelSettingsResponses, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleResponses, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesResponses, UpdateWorkingDirData, UpdateWorkingDirErrors, UpdateWorkingDirResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsResponses, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponses } from './types.gen'; export type Options = Options2 & { /** @@ -308,6 +308,15 @@ export const startOpenrouterSetup = (optio export const startTetrateSetup = (options?: Options) => (options?.client ?? client).post({ url: '/handle_tetrate', ...options }); +export const downloadHfModel = (options: Options) => (options.client ?? client).post({ + url: '/local-inference/download', + ...options, + headers: { + 'Content-Type': 'application/json', + ...options.headers + } +}); + export const listLocalModels = (options?: Options) => (options?.client ?? client).get({ url: '/local-inference/models', ...options }); export const deleteLocalModel = (options: Options) => (options.client ?? client).delete({ url: '/local-inference/models/{model_id}', ...options }); @@ -318,6 +327,21 @@ export const getLocalModelDownloadProgress = (options: Options) => (options.client ?? client).post({ url: '/local-inference/models/{model_id}/download', ...options }); +export const getModelSettings = (options: Options) => (options.client ?? client).get({ url: '/local-inference/models/{model_id}/settings', ...options }); + +export const updateModelSettings = (options: Options) => (options.client ?? client).put({ + url: '/local-inference/models/{model_id}/settings', + ...options, + headers: { + 'Content-Type': 'application/json', + ...options.headers + } +}); + +export const getRepoFiles = (options: Options) => (options.client ?? client).get({ url: '/local-inference/repo/{author}/{repo}/files', ...options }); + +export const searchHfModels = (options: Options) => (options.client ?? client).get({ url: '/local-inference/search', ...options }); + export const mcpUiProxy = (options: Options) => (options.client ?? client).get({ url: '/mcp-ui-proxy', ...options }); export const createRecipe = (options: Options) => (options.client ?? client).post({ diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 189b410fb3..f87c78b722 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -433,6 +433,46 @@ export type GooseApp = McpAppResource & (WindowProps | null) & { prd?: string | null; }; +export type HfDownloadRequest = { + filename?: string | null; + repo_id?: string | null; + spec?: string | null; +}; + +export type HfDownloadResponse = { + model_id: string; +}; + +/** + * A single downloadable GGUF file (used internally and for downloads). + */ +export type HfGgufFile = { + download_url: string; + filename: string; + quantization: string; + size_bytes: number; +}; + +export type HfModelInfo = { + author: string; + downloads: number; + gguf_files: Array; + model_name: string; + repo_id: string; +}; + +/** + * A quantization variant — groups sharded files into one logical entry. + */ +export type HfQuantVariant = { + description: string; + download_url: string; + filename: string; + quality_rank: number; + quantization: string; + size_bytes: number; +}; + export type Icon = { mimeType?: string; sizes?: Array; @@ -499,69 +539,26 @@ export type LoadedProvider = { }; export type LocalLlmModel = { - /** - * Maximum context window in tokens - */ context_limit: number; - /** - * Description and use case - */ description: string; - /** - * Model identifier (e.g., "llama-3.2-1b") - */ id: string; - /** - * Display name - */ name: string; - /** - * Model file size in MB - */ size_mb: number; tier: ModelTier; - /** - * Download URL for the tokenizer JSON - */ - tokenizer_url: string; - /** - * Download URL for the model GGUF file - */ url: string; }; export type LocalModelResponse = { - /** - * Maximum context window in tokens - */ context_limit: number; - /** - * Description and use case - */ description: string; - /** - * Model identifier (e.g., "llama-3.2-1b") - */ id: string; - /** - * Display name - */ name: string; - /** - * Model file size in MB - */ size_mb: number; tier: ModelTier; - /** - * Download URL for the tokenizer JSON - */ - tokenizer_url: string; - /** - * Download URL for the model GGUF file - */ url: string; } & { downloaded: boolean; + featured: boolean; recommended: boolean; }; @@ -721,6 +718,24 @@ export type ModelInfo = { supports_cache_control?: boolean | null; }; +export type ModelListItem = LocalModelResponse | RegistryModelResponse; + +export type ModelSettings = { + context_size?: number | null; + flash_attention?: boolean | null; + frequency_penalty?: number; + max_output_tokens?: number | null; + n_batch?: number | null; + n_gpu_layers?: number | null; + n_threads?: number | null; + native_tool_calling?: boolean; + presence_penalty?: number; + repeat_last_n?: number; + repeat_penalty?: number; + sampling?: SamplingConfig; + use_mlock?: boolean; +}; + export type ModelTier = 'tiny' | 'small' | 'medium' | 'large'; export type ParseRecipeRequest = { @@ -952,11 +967,28 @@ export type RedactedThinkingContent = { data: string; }; +export type RegistryModelResponse = { + display_name: string; + downloaded: boolean; + featured: boolean; + filename: string; + id: string; + quantization: string; + repo_id: string; + settings: ModelSettings; + size_bytes: number; +}; + export type RemoveExtensionRequest = { name: string; session_id: string; }; +export type RepoVariantsResponse = { + recommended_index?: number | null; + variants: Array; +}; + export type ResourceContents = { _meta?: { [key: string]: unknown; @@ -1034,6 +1066,22 @@ export type RunNowResponse = { session_id: string; }; +export type SamplingConfig = { + type: 'Greedy'; +} | { + min_p: number; + seed?: number | null; + temperature: number; + top_k: number; + top_p: number; + type: 'Temperature'; +} | { + eta: number; + seed?: number | null; + tau: number; + type: 'MirostatV2'; +}; + export type SavePromptRequest = { content: string; }; @@ -2873,6 +2921,33 @@ export type StartTetrateSetupResponses = { export type StartTetrateSetupResponse = StartTetrateSetupResponses[keyof StartTetrateSetupResponses]; +export type DownloadHfModelData = { + body: HfDownloadRequest; + path?: never; + query?: never; + url: '/local-inference/download'; +}; + +export type DownloadHfModelErrors = { + /** + * Invalid request + */ + 400: unknown; + /** + * Internal server error + */ + 500: unknown; +}; + +export type DownloadHfModelResponses = { + /** + * Download started + */ + 202: HfDownloadResponse; +}; + +export type DownloadHfModelResponse = DownloadHfModelResponses[keyof DownloadHfModelResponses]; + export type ListLocalModelsData = { body?: never; path?: never; @@ -2884,7 +2959,7 @@ export type ListLocalModelsResponses = { /** * List of available local LLM models */ - 200: Array; + 200: Array; }; export type ListLocalModelsResponse = ListLocalModelsResponses[keyof ListLocalModelsResponses]; @@ -2991,6 +3066,118 @@ export type DownloadLocalModelResponses = { 202: unknown; }; +export type GetModelSettingsData = { + body?: never; + path: { + model_id: string; + }; + query?: never; + url: '/local-inference/models/{model_id}/settings'; +}; + +export type GetModelSettingsErrors = { + /** + * Model not found + */ + 404: unknown; +}; + +export type GetModelSettingsResponses = { + /** + * Model settings + */ + 200: ModelSettings; +}; + +export type GetModelSettingsResponse = GetModelSettingsResponses[keyof GetModelSettingsResponses]; + +export type UpdateModelSettingsData = { + body: ModelSettings; + path: { + model_id: string; + }; + query?: never; + url: '/local-inference/models/{model_id}/settings'; +}; + +export type UpdateModelSettingsErrors = { + /** + * Model not found + */ + 404: unknown; + /** + * Failed to save settings + */ + 500: unknown; +}; + +export type UpdateModelSettingsResponses = { + /** + * Settings updated + */ + 200: ModelSettings; +}; + +export type UpdateModelSettingsResponse = UpdateModelSettingsResponses[keyof UpdateModelSettingsResponses]; + +export type GetRepoFilesData = { + body?: never; + path: { + author: string; + repo: string; + }; + query?: never; + url: '/local-inference/repo/{author}/{repo}/files'; +}; + +export type GetRepoFilesErrors = { + /** + * Failed to fetch repo files + */ + 500: unknown; +}; + +export type GetRepoFilesResponses = { + /** + * GGUF quantization variants in repo + */ + 200: RepoVariantsResponse; +}; + +export type GetRepoFilesResponse = GetRepoFilesResponses[keyof GetRepoFilesResponses]; + +export type SearchHfModelsData = { + body?: never; + path?: never; + query: { + /** + * Search query + */ + q: string; + /** + * Max results + */ + limit?: number | null; + }; + url: '/local-inference/search'; +}; + +export type SearchHfModelsErrors = { + /** + * Search failed + */ + 500: unknown; +}; + +export type SearchHfModelsResponses = { + /** + * Search results + */ + 200: Array; +}; + +export type SearchHfModelsResponse = SearchHfModelsResponses[keyof SearchHfModelsResponses]; + export type McpUiProxyData = { body?: never; path?: never; diff --git a/ui/desktop/src/components/ToolCallWithResponse.tsx b/ui/desktop/src/components/ToolCallWithResponse.tsx index 5ad33862d7..a33f8b18db 100644 --- a/ui/desktop/src/components/ToolCallWithResponse.tsx +++ b/ui/desktop/src/components/ToolCallWithResponse.tsx @@ -527,7 +527,7 @@ function ToolCallView({ // Function to create a descriptive representation of what the tool is doing const getToolDescription = (): string | null => { - const args = toolCall.arguments as Record; + const args = (toolCall.arguments ?? {}) as Record; const toolName = getToolName(toolCall.name); const getStringValue = (value: ToolCallArgumentValue): string => { diff --git a/ui/desktop/src/components/settings/SettingsView.tsx b/ui/desktop/src/components/settings/SettingsView.tsx index 2e04ab651f..05a17d9ccb 100644 --- a/ui/desktop/src/components/settings/SettingsView.tsx +++ b/ui/desktop/src/components/settings/SettingsView.tsx @@ -9,10 +9,11 @@ import ConfigSettings from './config/ConfigSettings'; import PromptsSettingsSection from './PromptsSettingsSection'; import { ExtensionConfig } from '../../api'; import { MainPanelLayout } from '../Layout/MainPanelLayout'; -import { Bot, Share2, Monitor, MessageSquare, FileText, Keyboard } from 'lucide-react'; +import { Bot, Share2, Monitor, MessageSquare, FileText, Keyboard, HardDrive } from 'lucide-react'; import { useState, useEffect, useRef } from 'react'; import ChatSettingsSection from './chat/ChatSettingsSection'; import KeyboardShortcutsSection from './keyboard/KeyboardShortcutsSection'; +import LocalInferenceSection from './localInference/LocalInferenceSection'; import { CONFIGURATION_ENABLED } from '../../updates'; import { trackSettingsTabViewed } from '../../utils/analytics'; @@ -54,6 +55,7 @@ export default function SettingsView({ chat: 'chat', prompts: 'prompts', keyboard: 'keyboard', + 'local-inference': 'local-inference', }; const targetTab = sectionToTab[viewOptions.section]; @@ -112,6 +114,14 @@ export default function SettingsView({ Models + + + Local Inference + Chat @@ -155,6 +165,13 @@ export default function SettingsView({ + + + + { + if (bytes === 0) return 'unknown'; + if (bytes < 1024) return `${bytes}B`; + if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)}KB`; + if (bytes < 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(0)}MB`; + return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)}GB`; +}; + +const formatDownloads = (n: number): string => { + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; + if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`; + return `${n}`; +}; + +interface RepoData { + variants: HfQuantVariant[]; + recommendedIndex: number | null; +} + +interface Props { + onDownloadStarted: (modelId: string) => void; +} + +export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => { + const [query, setQuery] = useState(''); + const [results, setResults] = useState([]); + const [expandedRepo, setExpandedRepo] = useState(null); + const [repoData, setRepoData] = useState>({}); + const [searching, setSearching] = useState(false); + const [downloading, setDownloading] = useState>(new Set()); + const [loadingFiles, setLoadingFiles] = useState>(new Set()); + const [directSpec, setDirectSpec] = useState(''); + const [error, setError] = useState(null); + const debounceRef = useRef | null>(null); + + const doSearch = useCallback(async (q: string) => { + if (!q.trim()) { + setResults([]); + setError(null); + return; + } + setSearching(true); + setError(null); + try { + const response = await searchHfModels({ + query: { q, limit: 20 }, + }); + if (response.data) { + setResults(response.data); + if (response.data.length === 0) { + setError('No GGUF models found for this query.'); + } + } else { + console.error('Search response:', response); + const errMsg = response.error + ? `Search error: ${JSON.stringify(response.error)}` + : 'Search returned no data.'; + setError(errMsg); + } + } catch (e) { + console.error('Search failed:', e); + setError('Search failed. Please try again.'); + } finally { + setSearching(false); + } + }, []); + + const handleQueryChange = (value: string) => { + setQuery(value); + if (debounceRef.current) clearTimeout(debounceRef.current); + debounceRef.current = setTimeout(() => doSearch(value), 300); + }; + + const toggleRepo = async (repoId: string) => { + if (expandedRepo === repoId) { + setExpandedRepo(null); + return; + } + setExpandedRepo(repoId); + + if (!repoData[repoId]?.variants.length) { + setLoadingFiles((prev) => new Set(prev).add(repoId)); + try { + const [author, repo] = repoId.split('/'); + const response = await getRepoFiles({ + path: { author, repo }, + }); + if (response.data) { + const variants = response.data.variants; + setRepoData((prev) => ({ + ...prev, + [repoId]: { + variants, + recommendedIndex: response.data!.recommended_index ?? null, + }, + })); + if (variants.length === 0) { + setExpandedRepo(null); + setResults((prev) => prev.filter((m) => m.repo_id !== repoId)); + } + } + } catch (e) { + console.error('Failed to fetch repo files:', e); + } finally { + setLoadingFiles((prev) => { + const next = new Set(prev); + next.delete(repoId); + return next; + }); + } + } + }; + + const startDownload = async (repoId: string, filename: string) => { + const key = `${repoId}/${filename}`; + setDownloading((prev) => new Set(prev).add(key)); + try { + const response = await downloadHfModel({ + body: { repo_id: repoId, filename }, + }); + if (response.data) { + onDownloadStarted(response.data.model_id); + } else { + console.error('Download error:', response.error); + } + } catch (e) { + console.error('Download failed:', e); + } finally { + setDownloading((prev) => { + const next = new Set(prev); + next.delete(key); + return next; + }); + } + }; + + const startDirectDownload = async () => { + if (!directSpec.trim()) return; + const key = `direct:${directSpec}`; + setDownloading((prev) => new Set(prev).add(key)); + try { + const response = await downloadHfModel({ + body: { spec: directSpec.trim() }, + }); + if (response.data) { + onDownloadStarted(response.data.model_id); + setDirectSpec(''); + } + } catch (e) { + console.error('Direct download failed:', e); + } finally { + setDownloading((prev) => { + const next = new Set(prev); + next.delete(key); + return next; + }); + } + }; + + return ( +
+
+

Search HuggingFace

+
+ + handleQueryChange(e.target.value)} + placeholder="Search for GGUF models..." + className="w-full pl-9 pr-4 py-2 text-sm border border-border-subtle rounded-lg bg-background-default text-text-default placeholder:text-text-muted focus:outline-none focus:border-accent-primary" + /> + {searching && ( + + )} +
+
+ + {error && !searching && ( +

{error}

+ )} + + {results.length > 0 && ( +
+ {results.map((model) => { + const isExpanded = expandedRepo === model.repo_id; + const data = repoData[model.repo_id]; + const variants = data?.variants || []; + const recommendedIndex = data?.recommendedIndex ?? null; + + return ( +
+ + + {isExpanded && ( +
+ {loadingFiles.has(model.repo_id) && ( +
+ + Loading variants... +
+ )} + {variants.map((variant, idx) => { + const dlKey = `${model.repo_id}/${variant.filename}`; + const isStarting = downloading.has(dlKey); + const isRecommended = idx === recommendedIndex; + + return ( +
+
+
+ + {variant.quantization} + + + {formatBytes(variant.size_bytes)} + + {isRecommended && ( + + + Recommended + + )} +
+ {variant.description && ( + + {variant.description} + + )} +
+ +
+ ); + })} +
+ )} +
+ ); + })} +
+ )} + +
+

Direct Download

+

+ Specify a model directly: user/repo:quantization +

+
+ setDirectSpec(e.target.value)} + placeholder="bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M" + className="flex-1 px-3 py-2 text-sm border border-border-subtle rounded-lg bg-background-default text-text-default placeholder:text-text-muted focus:outline-none focus:border-accent-primary" + onKeyDown={(e) => { + if (e.key === 'Enter') startDirectDownload(); + }} + /> + +
+
+
+ ); +}; diff --git a/ui/desktop/src/components/settings/localInference/LocalInferenceSection.tsx b/ui/desktop/src/components/settings/localInference/LocalInferenceSection.tsx new file mode 100644 index 0000000000..3a68c21171 --- /dev/null +++ b/ui/desktop/src/components/settings/localInference/LocalInferenceSection.tsx @@ -0,0 +1,9 @@ +import { LocalInferenceSettings } from './LocalInferenceSettings'; + +export default function LocalInferenceSection() { + return ( +
+ +
+ ); +} diff --git a/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx b/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx index 739850d909..5902dec5bc 100644 --- a/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx +++ b/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx @@ -1,5 +1,5 @@ -import { useState, useEffect } from 'react'; -import { Download, Trash2, X, Check, ChevronDown, ChevronUp } from 'lucide-react'; +import { useState, useEffect, useCallback, useRef } from 'react'; +import { Download, Trash2, X, Check, ChevronDown, ChevronUp, Settings2 } from 'lucide-react'; import { Button } from '../../ui/button'; import { useConfig } from '../../ConfigContext'; import { @@ -8,9 +8,13 @@ import { getLocalModelDownloadProgress, cancelLocalModelDownload, deleteLocalModel, - type LocalModelResponse, type DownloadProgress, + type LocalModelResponse, + type RegistryModelResponse, + type ModelListItem, } from '../../../api'; +import { HuggingFaceModelSearch } from './HuggingFaceModelSearch'; +import { ModelSettingsPanel } from './ModelSettingsPanel'; const LOCAL_LLM_MODEL_CONFIG_KEY = 'LOCAL_LLM_MODEL'; @@ -21,12 +25,46 @@ const formatBytes = (bytes: number): string => { return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)}GB`; }; +function isFeaturedModel(item: ModelListItem): item is LocalModelResponse & { featured: boolean } { + return 'tier' in item; +} + +function isRegistryModel(item: ModelListItem): item is RegistryModelResponse { + return 'display_name' in item && !('tier' in item); +} + export const LocalInferenceSettings = () => { - const [models, setModels] = useState([]); + const [featuredModels, setFeaturedModels] = useState<(LocalModelResponse & { featured?: boolean })[]>([]); + const [registryModels, setRegistryModels] = useState([]); const [downloads, setDownloads] = useState>(new Map()); const [selectedModelId, setSelectedModelId] = useState(null); - const [showAllModels, setShowAllModels] = useState(false); + const [showAllFeatured, setShowAllFeatured] = useState(false); + const [settingsOpenFor, setSettingsOpenFor] = useState(null); const { read, upsert } = useConfig(); + const downloadSectionRef = useRef(null); + + const loadModels = useCallback(async () => { + try { + const response = await listLocalModels(); + if (response.data) { + const featured: (LocalModelResponse & { featured?: boolean })[] = []; + const registry: RegistryModelResponse[] = []; + + for (const item of response.data) { + if (isFeaturedModel(item)) { + featured.push(item); + } else if (isRegistryModel(item)) { + registry.push(item); + } + } + + setFeaturedModels(featured); + setRegistryModels(registry); + } + } catch (error) { + console.error('Failed to load models:', error); + } + }, []); useEffect(() => { loadModels(); @@ -55,26 +93,23 @@ export const LocalInferenceSettings = () => { setSelectedModelId(modelId); }; - const loadModels = async () => { - try { - const response = await listLocalModels(); - if (response.data) { - setModels(response.data); - } - } catch (error) { - console.error('Failed to load models:', error); - } - }; - - const startDownload = async (modelId: string) => { + const startFeaturedDownload = async (modelId: string) => { try { await downloadLocalModel({ path: { model_id: modelId } }); pollDownloadProgress(modelId); + scrollToDownloads(); } catch (error) { console.error('Failed to start download:', error); } }; + const scrollToDownloads = useCallback(() => { + // Wait a tick for the download section to render before scrolling. + requestAnimationFrame(() => { + downloadSectionRef.current?.scrollIntoView({ behavior: 'smooth', block: 'nearest' }); + }); + }, []); + const pollDownloadProgress = (modelId: string) => { const interval = setInterval(async () => { try { @@ -85,8 +120,7 @@ export const LocalInferenceSettings = () => { if (progress.status === 'completed') { clearInterval(interval); - await loadModels(); // Refresh model list - // Auto-select the model that was just downloaded + await loadModels(); await selectModel(modelId); } else if (progress.status === 'failed') { clearInterval(interval); @@ -115,9 +149,8 @@ export const LocalInferenceSettings = () => { } }; - const deleteModel = async (modelId: string) => { + const handleDeleteModel = async (modelId: string) => { if (!window.confirm('Delete this model? You can re-download it later.')) return; - try { await deleteLocalModel({ path: { model_id: modelId } }); if (selectedModelId === modelId) { @@ -130,172 +163,330 @@ export const LocalInferenceSettings = () => { } }; - const hasDownloadedNonRecommended = models.some( + const handleHfDownloadStarted = (modelId: string) => { + pollDownloadProgress(modelId); + scrollToDownloads(); + }; + + // Featured models display logic + const hasDownloadedNonRecommended = featuredModels.some( (model) => model.downloaded && !model.recommended ); - const displayedModels = showAllModels || hasDownloadedNonRecommended - ? models - : models.filter((m) => m.recommended); - const hasNonRecommendedModels = models.some((m) => !m.recommended); - const showToggleButton = hasNonRecommendedModels && !hasDownloadedNonRecommended; + const displayedFeatured = showAllFeatured || hasDownloadedNonRecommended + ? featuredModels + : featuredModels.filter((m) => m.recommended); + const hasNonRecommendedFeatured = featuredModels.some((m) => !m.recommended); + const showFeaturedToggle = hasNonRecommendedFeatured && !hasDownloadedNonRecommended; + + // Downloaded models from both featured and registry + const downloadedFeatured = featuredModels.filter((m) => m.downloaded); + const downloadedRegistry = registryModels.filter((m) => m.downloaded); + const hasDownloaded = downloadedFeatured.length > 0 || downloadedRegistry.length > 0; return ( -
+

Local Inference Models

- Download and manage local LLM models for inference without API keys. Supports GPU acceleration (Metal for Apple Silicon). + Download and manage local LLM models for inference without API keys. Search HuggingFace for any GGUF model or use the featured picks below.

-
- {displayedModels.map((model) => { - const progress = downloads.get(model.id); - const isDownloading = progress?.status === 'downloading'; - const isSelected = selectedModelId === model.id; - const canSelect = model.downloaded && !isDownloading; + {/* Active Downloads */} + {downloads.size > 0 && ( +
+

Downloading

+
+ {Array.from(downloads.entries()).map(([modelId, progress]) => { + if (progress.status === 'completed') return null; + return ( +
+
+ {modelId} + {progress.status === 'downloading' && ( + + )} +
+ {progress.status === 'downloading' && ( +
+
+
+
+
+ {formatBytes(progress.bytes_downloaded)} / {formatBytes(progress.total_bytes)} + {progress.progress_percent.toFixed(0)}% +
+
+ )} + {progress.status === 'failed' && ( +

{progress.error || 'Download failed'}

+ )} +
+ ); + })} +
+
+ )} - return ( -
-
-
-
- {canSelect && ( + {/* Downloaded Models */} + {hasDownloaded && ( +
+

Downloaded Models

+
+ {downloadedFeatured.map((model) => { + const isSelected = selectedModelId === model.id; + const showSettings = settingsOpenFor === model.id; + return ( +
+
+
selectModel(model.id)} className="cursor-pointer" /> - )} -

- {model.name} -

- - {model.size_mb}MB - - - {model.context_limit.toLocaleString()} tokens - - {model.recommended && ( - - Recommended - - )} - {isSelected && ( - - Active - - )} -
- -

- {model.description} -

- {model.recommended && ( -

- Recommended for your hardware -

- )} -
- -
- {model.downloaded ? ( - <> -
- - Downloaded -
+ {model.name} + {model.size_mb}MB + {model.recommended && ( + Recommended + )} +
+
+ - - ) : isDownloading ? ( - <> -
- {progress.progress_percent.toFixed(0)}% -
+
+
+ {showSettings && } +
+ ); + })} + + {downloadedRegistry.map((model) => { + const isSelected = selectedModelId === model.id; + const showSettings = settingsOpenFor === model.id; + return ( +
+
+
+ selectModel(model.id)} + className="cursor-pointer" + /> + {model.display_name} + {formatBytes(model.size_bytes)} +
+
- - ) : ( - - )} -
-
- - {isDownloading && progress && ( -
-
-
+ +
-
- - {formatBytes(progress.bytes_downloaded)} / {formatBytes(progress.total_bytes)} - - {progress.speed_bps && ( - {formatBytes(progress.speed_bps)}/s + {showSettings && } +
+ ); + })} +
+
+ )} + + {/* Featured Models */} +
+

Featured Models

+
+ {displayedFeatured.map((model) => { + const progress = downloads.get(model.id); + const isDownloading = progress?.status === 'downloading'; + + return ( +
+
+
+
+

{model.name}

+ {model.size_mb}MB + + {model.context_limit.toLocaleString()} tokens + + {model.recommended && ( + + Recommended + + )} +
+

{model.description}

+
+ +
+ {model.downloaded ? ( +
+ + Downloaded +
+ ) : isDownloading ? ( + <> +
+ {progress.progress_percent.toFixed(0)}% +
+ + + ) : ( + )}
- )} - {progress?.status === 'failed' && progress.error && ( -
{progress.error}
- )} + {isDownloading && progress && ( +
+
+
+
+
+ + {formatBytes(progress.bytes_downloaded)} / {formatBytes(progress.total_bytes)} + + {progress.speed_bps && {formatBytes(progress.speed_bps)}/s} +
+
+ )} + + {progress?.status === 'failed' && progress.error && ( +
{progress.error}
+ )} +
+ ); + })} +
+ + {showFeaturedToggle && ( + + )} +
+ + {/* Non-downloaded registry models being downloaded */} + {registryModels + .filter((m) => !m.downloaded && downloads.has(m.id)) + .map((model) => { + const progress = downloads.get(model.id); + if (!progress || progress.status !== 'downloading') return null; + return ( +
+
+
+ {model.display_name} + {progress.progress_percent.toFixed(0)}% +
+ +
+
+
+
+
+
); })} + + {/* HuggingFace Search */} +
+
- {showToggleButton && ( - - )} - - {models.length === 0 && ( -
- No models available -
+ {featuredModels.length === 0 && registryModels.length === 0 && ( +
No models available
)}
); diff --git a/ui/desktop/src/components/settings/localInference/ModelSettingsPanel.tsx b/ui/desktop/src/components/settings/localInference/ModelSettingsPanel.tsx new file mode 100644 index 0000000000..f177c95323 --- /dev/null +++ b/ui/desktop/src/components/settings/localInference/ModelSettingsPanel.tsx @@ -0,0 +1,417 @@ +import { useState, useEffect, useCallback } from 'react'; +import { RotateCcw } from 'lucide-react'; +import { Button } from '../../ui/button'; +import { Switch } from '../../ui/switch'; +import { + getModelSettings, + updateModelSettings, + type ModelSettings, + type SamplingConfig, +} from '../../../api'; + +const DEFAULT_SETTINGS: ModelSettings = { + context_size: null, + max_output_tokens: null, + sampling: { type: 'Temperature', temperature: 0.8, top_k: 40, top_p: 0.95, min_p: 0.05, seed: null }, + repeat_penalty: 1.0, + repeat_last_n: 64, + frequency_penalty: 0.0, + presence_penalty: 0.0, + n_batch: null, + n_gpu_layers: null, + use_mlock: false, + flash_attention: null, + n_threads: null, + native_tool_calling: false, +}; + +type SamplingType = SamplingConfig['type']; + +function NumberField({ + label, + description, + value, + onChange, + placeholder, + min, + max, + step, + allowNull, +}: { + label: string; + description?: string; + value: number | null | undefined; + onChange: (v: number | null) => void; + placeholder?: string; + min?: number; + max?: number; + step?: number; + allowNull?: boolean; +}) { + return ( +
+ + {description && {description}} + { + const raw = e.target.value; + if (raw === '' && allowNull) { + onChange(null); + } else { + const n = step && step < 1 ? parseFloat(raw) : parseInt(raw, 10); + if (!isNaN(n)) onChange(n); + } + }} + placeholder={placeholder ?? 'Auto'} + min={min} + max={max} + step={step} + /> +
+ ); +} + +function ToggleField({ + label, + description, + value, + onChange, +}: { + label: string; + description?: string; + value: boolean; + onChange: (v: boolean) => void; +}) { + return ( +
+
+
{label}
+ {description && {description}} +
+ +
+ ); +} + +function SelectField({ + label, + description, + value, + options, + onChange, +}: { + label: string; + description?: string; + value: T; + options: { value: T; label: string }[]; + onChange: (v: T) => void; +}) { + return ( +
+
+
{label}
+ {description && {description}} +
+ +
+ ); +} + +export const ModelSettingsPanel = ({ modelId }: { modelId: string }) => { + const [settings, setSettings] = useState(DEFAULT_SETTINGS); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + + const load = useCallback(async () => { + try { + const res = await getModelSettings({ path: { model_id: modelId } }); + if (res.data) setSettings(res.data); + } catch { + // use defaults + } finally { + setLoading(false); + } + }, [modelId]); + + useEffect(() => { + load(); + }, [load]); + + const save = async (updated: ModelSettings) => { + setSettings(updated); + setSaving(true); + try { + await updateModelSettings({ path: { model_id: modelId }, body: updated }); + } catch (e) { + console.error('Failed to save settings:', e); + } finally { + setSaving(false); + } + }; + + const resetDefaults = () => save(DEFAULT_SETTINGS); + + const updateField = (key: K, value: ModelSettings[K]) => { + save({ ...settings, [key]: value }); + }; + + const samplingType: SamplingType = settings.sampling?.type ?? 'Temperature'; + + const setSamplingType = (type: SamplingType) => { + let sampling: SamplingConfig; + if (type === 'Greedy') { + sampling = { type: 'Greedy' }; + } else if (type === 'MirostatV2') { + sampling = { type: 'MirostatV2', tau: 5.0, eta: 0.1, seed: null }; + } else { + sampling = { type: 'Temperature', temperature: 0.8, top_k: 40, top_p: 0.95, min_p: 0.05, seed: null }; + } + save({ ...settings, sampling }); + }; + + const updateSampling = (partial: Partial) => { + save({ ...settings, sampling: { ...settings.sampling!, ...partial } as SamplingConfig }); + }; + + if (loading) { + return
Loading settings...
; + } + + return ( +
+
+ + Model Settings {saving && '(saving...)'} + + +
+ + {/* Context & Generation */} +
+
Context & Generation
+
+ updateField('context_size', v)} + placeholder="Auto" + min={0} + allowNull + /> + updateField('max_output_tokens', v)} + placeholder="No limit" + min={1} + allowNull + /> +
+
+ + {/* Sampling */} +
+ setSamplingType(v)} + /> + + {samplingType === 'Temperature' && settings.sampling?.type === 'Temperature' && ( +
+ updateSampling({ temperature: v ?? 0.8 })} + min={0} + max={2} + step={0.05} + /> + updateSampling({ top_k: v ?? 40 })} + min={0} + /> + updateSampling({ top_p: v ?? 0.95 })} + min={0} + max={1} + step={0.01} + /> + updateSampling({ min_p: v ?? 0.05 })} + min={0} + max={1} + step={0.01} + /> + updateSampling({ seed: v })} + placeholder="Random" + min={0} + allowNull + /> +
+ )} + + {samplingType === 'MirostatV2' && settings.sampling?.type === 'MirostatV2' && ( +
+ updateSampling({ tau: v ?? 5.0 })} + min={0} + step={0.1} + /> + updateSampling({ eta: v ?? 0.1 })} + min={0} + max={1} + step={0.01} + /> + updateSampling({ seed: v })} + placeholder="Random" + min={0} + allowNull + /> +
+ )} +
+ + {/* Repetition Penalty */} +
+
Repetition Penalty
+
+ updateField('repeat_penalty', v ?? 1.0)} + min={0} + step={0.05} + /> + updateField('repeat_last_n', v ?? 64)} + min={0} + /> + updateField('frequency_penalty', v ?? 0.0)} + min={0} + max={2} + step={0.05} + /> + updateField('presence_penalty', v ?? 0.0)} + min={0} + max={2} + step={0.05} + /> +
+
+ + {/* Performance */} +
+
Performance
+
+ updateField('n_batch', v)} + placeholder="Auto" + min={1} + allowNull + /> + updateField('n_gpu_layers', v)} + placeholder="All" + min={0} + allowNull + /> + updateField('n_threads', v)} + placeholder="Auto" + min={1} + allowNull + /> +
+ updateField('use_mlock', v)} + /> + updateField('flash_attention', v === 'auto' ? null : v === 'on')} + /> +
+ {/* Tool Calling */} +
+
Tool Calling
+ updateField('native_tool_calling', v)} + /> +
+
+ ); +}; diff --git a/ui/desktop/src/components/settings/models/ModelsSection.tsx b/ui/desktop/src/components/settings/models/ModelsSection.tsx index 8cbbfba6e4..8e903e141b 100644 --- a/ui/desktop/src/components/settings/models/ModelsSection.tsx +++ b/ui/desktop/src/components/settings/models/ModelsSection.tsx @@ -11,7 +11,6 @@ import { toastError } from '../../../toasts'; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '../../ui/card'; import ResetProviderSection from '../reset_provider/ResetProviderSection'; -import { LocalInferenceSettings } from '../localInference/LocalInferenceSettings'; interface ModelsSectionProps { setView: (view: View) => void; @@ -103,11 +102,6 @@ export default function ModelsSection({ setView }: ModelsSectionProps) { - - - - - Reset Provider and Model