diff --git a/Cargo.lock b/Cargo.lock index 6e3ec2ba5a..c4252b9ea5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3307,6 +3307,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" @@ -3533,6 +3553,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" @@ -4224,6 +4253,7 @@ dependencies = [ "dashmap 6.1.0", "dirs 5.0.1", "dotenvy", + "encoding_rs", "env-lock", "etcetera 0.11.0", "fs2", @@ -4240,6 +4270,7 @@ dependencies = [ "jsonwebtoken", "keyring", "lazy_static", + "llama-cpp-2", "lru", "minijinja", "mockall", @@ -5702,6 +5733,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" diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index 4ac9eb2cab..26487713ae 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 { @@ -875,6 +882,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 @@ -964,6 +1003,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", Some(Command::ValidateExtensions { .. }) => "validate-extensions", None => "default_session", @@ -1400,6 +1440,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(), + size_bytes: file.size_bytes, + }; + + { + let mut registry = get_registry() + .lock() + .map_err(|_| anyhow::anyhow!("Failed to acquire registry lock"))?; + registry.add_model(entry)?; + } + + // Download + let manager = goose::download_manager::get_download_manager(); + manager + .download_model( + format!("{}-model", model_id), + file.download_url, + local_path, + None, + ) + .await?; + + // Poll progress + loop { + if let Some(progress) = manager.get_progress(&format!("{}-model", model_id)) { + match progress.status { + goose::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::download_manager::DownloadStatus::Completed => { + println!("\nDownloaded: {} (id: {})", display_name, model_id); + break; + } + goose::download_manager::DownloadStatus::Failed => { + let err = progress.error.unwrap_or_default(); + anyhow::bail!("Download failed: {}", err); + } + goose::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; @@ -1530,6 +1734,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, Some(Command::ValidateExtensions { file }) => { use goose::agents::validate_extensions::validate_bundled_extensions; match validate_bundled_extensions(&file) { diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index f33db83520..e128289bad 100644 --- a/crates/goose-server/src/openapi.rs +++ b/crates/goose-server/src/openapi.rs @@ -4,7 +4,7 @@ use goose::agents::ExtensionConfig; use goose::config::permission::PermissionLevel; use goose::config::ExtensionEntry; use goose::conversation::Conversation; -use goose::dictation::download_manager::{DownloadProgress, DownloadStatus}; +use goose::download_manager::{DownloadProgress, DownloadStatus}; use goose::model::ModelConfig; use goose::permission::permission_confirmation::{Permission, PrincipalType}; use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata, ProviderType}; @@ -424,6 +424,15 @@ derive_utoipa!(Icon as IconSchema); super::routes::dictation::get_download_progress, 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::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,6 +601,15 @@ derive_utoipa!(Icon as IconSchema); goose::dictation::providers::DictationProvider, super::routes::dictation::DictationProviderStatus, super::routes::dictation::WhisperModelResponse, + super::routes::local_inference::LocalModelResponse, + super::routes::local_inference::ModelDownloadStatus, + super::routes::local_inference::DownloadModelRequest, + 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/dictation.rs b/crates/goose-server/src/routes/dictation.rs index fb1c97aa71..8530d2b1c7 100644 --- a/crates/goose-server/src/routes/dictation.rs +++ b/crates/goose-server/src/routes/dictation.rs @@ -7,11 +7,11 @@ use axum::{ Json, Router, }; use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; -use goose::dictation::download_manager::{get_download_manager, DownloadProgress}; use goose::dictation::providers::{ is_configured, transcribe_local, transcribe_with_provider, DictationProvider, PROVIDERS, }; use goose::dictation::whisper; +use goose::download_manager::{get_download_manager, DownloadProgress}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; @@ -257,11 +257,16 @@ pub async fn download_model(Path(model_id): Path) -> Result, + }, + Downloaded, +} + +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +pub struct LocalModelResponse { + pub id: String, + pub display_name: String, + pub repo_id: String, + pub filename: String, + pub quantization: String, + pub size_bytes: u64, + pub status: ModelDownloadStatus, + pub recommended: bool, + pub settings: ModelSettings, +} + +async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> { + let mut entries_to_add = Vec::new(); + + for spec in FEATURED_MODELS { + let (repo_id, quantization) = match hf_models::parse_model_spec(spec) { + Ok(parts) => parts, + Err(_) => continue, + }; + + let model_id = model_id_from_repo(&repo_id, &quantization); + + { + let registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + if registry.has_model(&model_id) { + continue; + } + } + + let hf_file = match resolve_model_spec(spec).await { + Ok((_repo, file)) => file, + Err(_) => { + let filename = format!( + "{}-{}.gguf", + repo_id.split('/').next_back().unwrap_or("model"), + quantization + ); + HfGgufFile { + filename: filename.clone(), + size_bytes: 0, + quantization: quantization.to_string(), + download_url: format!( + "https://huggingface.co/{}/resolve/main/{}", + repo_id, filename + ), + } + } + }; + + let local_path = Paths::in_data_dir("models").join(&hf_file.filename); + + entries_to_add.push(LocalModelEntry { + id: model_id, + display_name: display_name_from_repo(&repo_id, &quantization), + repo_id, + filename: hf_file.filename, + quantization, + local_path, + source_url: hf_file.download_url, + settings: ModelSettings::default(), + size_bytes: hf_file.size_bytes, + }); + } + + if !entries_to_add.is_empty() { + let mut registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + registry.sync_with_featured(entries_to_add); + } + + Ok(()) +} + +#[utoipa::path( + get, + path = "/local-inference/models", + responses( + (status = 200, description = "List of available local LLM models", body = Vec) + ) +)] +pub async fn list_local_models( + axum::extract::State(state): axum::extract::State>, +) -> Result>, ErrorResponse> { + ensure_featured_models_in_registry().await?; + + let recommended_id = recommend_local_model(&state.inference_runtime); + + let registry = get_registry() + .lock() + .map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?; + + let mut models: Vec = Vec::new(); + + for entry in registry.list_models() { + let goose_status = entry.download_status(); + + let status = match goose_status { + RegistryDownloadStatus::NotDownloaded => ModelDownloadStatus::NotDownloaded, + RegistryDownloadStatus::Downloading { + progress_percent, + bytes_downloaded, + total_bytes, + speed_bps, + } => ModelDownloadStatus::Downloading { + progress_percent, + bytes_downloaded, + total_bytes, + speed_bps: Some(speed_bps), + }, + RegistryDownloadStatus::Downloaded => ModelDownloadStatus::Downloaded, + }; + + let size_bytes = entry.file_size(); + + models.push(LocalModelResponse { + 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, + status, + recommended: recommended_id == entry.id, + settings: entry.settings.clone(), + }); + } + + models.sort_by(|a, b| { + let a_downloaded = matches!(a.status, ModelDownloadStatus::Downloaded); + let b_downloaded = matches!(b.status, ModelDownloadStatus::Downloaded); + match (b_downloaded, a_downloaded) { + (true, false) => std::cmp::Ordering::Greater, + (false, true) => std::cmp::Ordering::Less, + _ => a.display_name.cmp(&b.display_name), + } + }); + + Ok(Json(models)) +} + +#[derive(Debug, Deserialize)] +pub struct SearchQuery { + pub q: String, + pub limit: Option, +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct RepoVariantsResponse { + pub variants: Vec, + pub recommended_index: 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)) +} + +#[utoipa::path( + get, + path = "/local-inference/repo/{author}/{repo}/files", + responses( + (status = 200, description = "GGUF files in the repo", body = RepoVariantsResponse) + ) +)] +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 DownloadModelRequest { + /// Model spec like "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M" + pub spec: String, +} + +#[utoipa::path( + post, + path = "/local-inference/download", + request_body = DownloadModelRequest, + responses( + (status = 202, description = "Download started", body = String), + (status = 400, description = "Invalid request") + ) +)] +pub async fn download_hf_model( + Json(req): Json, +) -> Result<(StatusCode, Json), ErrorResponse> { + let (repo_id, quantization) = hf_models::parse_model_spec(&req.spec) + .map_err(|e| ErrorResponse::bad_request(format!("Invalid spec format: {e}")))?; + + let (_repo, hf_file) = resolve_model_spec(&req.spec) + .await + .map_err(|e| ErrorResponse::bad_request(format!("Invalid spec: {}", e)))?; + + let model_id = model_id_from_repo(&repo_id, &quantization); + let local_path = Paths::in_data_dir("models").join(&hf_file.filename); + let download_url = hf_file.download_url.clone(); + + let entry = LocalModelEntry { + id: model_id.clone(), + display_name: display_name_from_repo(&repo_id, &quantization), + repo_id, + filename: hf_file.filename, + quantization, + local_path: local_path.clone(), + source_url: download_url.clone(), + settings: ModelSettings::default(), + size_bytes: hf_file.size_bytes, + }; + + { + 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 dm = get_download_manager(); + dm.download_model( + format!("{}-model", model_id), + download_url, + local_path, + None, + ) + .await + .map_err(|e| ErrorResponse::internal(format!("Download failed: {}", e)))?; + + Ok((StatusCode::ACCEPTED, Json(model_id))) +} + +#[utoipa::path( + get, + path = "/local-inference/models/{model_id}/download", + responses( + (status = 200, description = "Download progress", body = DownloadProgress), + (status = 404, description = "No active download") + ) +)] +pub async fn get_local_model_download_progress( + Path(model_id): Path, +) -> Result, ErrorResponse> { + let download_id = format!("{}-model", model_id); + debug!(model_id = %model_id, download_id = %download_id, "Getting download progress"); + + let manager = get_download_manager(); + + let model_progress = manager + .get_progress(&download_id) + .ok_or_else(|| ErrorResponse::not_found("No active download"))?; + + Ok(Json(model_progress)) +} + +#[utoipa::path( + delete, + path = "/local-inference/models/{model_id}/download", + responses( + (status = 200, description = "Download cancelled"), + (status = 404, description = "No active download") + ) +)] +pub async fn cancel_local_model_download( + Path(model_id): Path, +) -> Result { + let manager = get_download_manager(); + manager + .cancel_download(&format!("{}-model", model_id)) + .map_err(|e| ErrorResponse::internal(format!("{}", e)))?; + + Ok(StatusCode::OK) +} + +#[utoipa::path( + delete, + path = "/local-inference/models/{model_id}", + responses( + (status = 200, description = "Model deleted"), + (status = 404, description = "Model not found") + ) +)] +pub async fn delete_local_model(Path(model_id): Path) -> Result { + let local_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 local_path.exists() { + tokio::fs::remove_file(&local_path) + .await + .map_err(|e| ErrorResponse::internal(format!("Failed to delete: {}", e)))?; + } + + // Only remove non-featured models from registry (featured ones stay as placeholders) + if !is_featured_model(&model_id) { + 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::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", + get(get_local_model_download_progress), + ) + .route( + "/local-inference/models/{model_id}/download", + delete(cancel_local_model_download), + ) + .route( + "/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/routes/mod.rs b/crates/goose-server/src/routes/mod.rs index 9eb00449f1..cbb956849b 100644 --- a/crates/goose-server/src/routes/mod.rs +++ b/crates/goose-server/src/routes/mod.rs @@ -3,6 +3,7 @@ pub mod agent; pub mod config_management; pub mod dictation; pub mod errors; +pub mod local_inference; pub mod mcp_app_proxy; pub mod mcp_ui_proxy; pub mod prompts; @@ -30,6 +31,7 @@ pub fn configure(state: Arc, secret_key: String) -> Rout .merge(action_required::routes(state.clone())) .merge(agent::routes(state.clone())) .merge(dictation::routes(state.clone())) + .merge(local_inference::routes(state.clone())) .merge(config_management::routes(state.clone())) .merge(prompts::routes()) .merge(recipe::routes(state.clone())) diff --git a/crates/goose-server/src/routes/utils.rs b/crates/goose-server/src/routes/utils.rs index 713280b14a..1e97a0271d 100644 --- a/crates/goose-server/src/routes/utils.rs +++ b/crates/goose-server/src/routes/utils.rs @@ -94,6 +94,11 @@ pub fn inspect_keys( pub fn check_provider_configured(metadata: &ProviderMetadata, provider_type: ProviderType) -> bool { let config = Config::global(); + // Special override + if metadata.name == "local" { + return true; + } + if provider_type == ProviderType::Custom || provider_type == ProviderType::Declarative { if let Ok(loaded_provider) = load_provider(metadata.name.as_str()) { if !loaded_provider.config.requires_auth { diff --git a/crates/goose-server/src/state.rs b/crates/goose-server/src/state.rs index 57f104d47e..4cf836d42d 100644 --- a/crates/goose-server/src/state.rs +++ b/crates/goose-server/src/state.rs @@ -14,6 +14,7 @@ use tokio::task::JoinHandle; use crate::tunnel::TunnelManager; use goose::agents::ExtensionLoadResult; +use goose::providers::local_inference::InferenceRuntime; type ExtensionLoadingTasks = Arc>>>>>>>; @@ -26,6 +27,7 @@ pub struct AppState { recipe_session_tracker: Arc>>, pub tunnel_manager: Arc, pub extension_loading_tasks: ExtensionLoadingTasks, + pub inference_runtime: Arc, } fn spawn_developer(r: tokio::io::DuplexStream, w: tokio::io::DuplexStream) { @@ -57,6 +59,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 ded774e989..a0de55ef8f 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 @@ -126,14 +126,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 } @@ -156,6 +159,7 @@ path = "examples/agent.rs" name = "databricks_oauth" path = "examples/databricks_oauth.rs" + [[bin]] name = "build_canonical_models" path = "src/providers/canonical/build_canonical_models.rs" diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index d77668cf4d..51981780f8 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -236,7 +236,12 @@ async fn child_process_client( command.env("PATH", path); } - if let Some(dir) = working_dir { + // Use explicitly passed working_dir, falling back to GOOSE_WORKING_DIR env var + let effective_working_dir = working_dir + .map(|p| p.to_path_buf()) + .or_else(|| std::env::var("GOOSE_WORKING_DIR").ok().map(PathBuf::from)); + + if let Some(ref dir) = effective_working_dir { if dir.exists() && dir.is_dir() { tracing::info!("Setting MCP process working directory: {:?}", dir); command.current_dir(dir); @@ -710,6 +715,7 @@ impl ExtensionManager { })?; let mut context = self.context.clone(); context.extension_manager = Some(Arc::downgrade(self)); + (def.client_factory)(context) } ExtensionConfig::InlinePython { @@ -1584,6 +1590,16 @@ impl ExtensionManager { session_id: &str, working_dir: &std::path::Path, ) -> Option { + // Skip MOIM for models with small context windows to avoid consuming limited context + const MIN_CONTEXT_FOR_MOIM: usize = 32_000; + if let Ok(provider_guard) = self.provider.try_lock() { + if let Some(provider) = provider_guard.as_ref() { + if provider.get_model_config().context_limit() < MIN_CONTEXT_FOR_MOIM { + return None; + } + } + } + // Use minute-level granularity to prevent conversation changes every second let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:00").to_string(); let mut content = format!( diff --git a/crates/goose/src/dictation/mod.rs b/crates/goose/src/dictation/mod.rs index 9cef90aa8c..d14fb21642 100644 --- a/crates/goose/src/dictation/mod.rs +++ b/crates/goose/src/dictation/mod.rs @@ -1,3 +1,2 @@ -pub mod download_manager; pub mod providers; pub mod whisper; diff --git a/crates/goose/src/dictation/download_manager.rs b/crates/goose/src/download_manager.rs similarity index 82% rename from crates/goose/src/dictation/download_manager.rs rename to crates/goose/src/download_manager.rs index 8342ab5cc2..6fc4718cbc 100644 --- a/crates/goose/src/dictation/download_manager.rs +++ b/crates/goose/src/download_manager.rs @@ -1,12 +1,33 @@ -use crate::dictation::whisper::LOCAL_WHISPER_MODEL_CONFIG_KEY; 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 tracing::info; 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 @@ -78,8 +99,9 @@ impl DownloadManager { model_id: String, url: String, destination: PathBuf, + on_complete: Option>, ) -> Result<()> { - // Initialize progress + info!(model_id = %model_id, url = %url, destination = ?destination, "Starting model download"); { let mut downloads = self .downloads @@ -115,10 +137,13 @@ 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 { Ok(_) => { + info!(model_id = %model_id_clone, "Download completed successfully"); if let Ok(mut downloads) = downloads.lock() { if let Some(progress) = downloads.get_mut(&model_id_clone) { progress.status = DownloadStatus::Completed; @@ -126,10 +151,15 @@ impl DownloadManager { } } - let _ = crate::config::Config::global() - .set_param(LOCAL_WHISPER_MODEL_CONFIG_KEY, model_id_clone.clone()); + if let Some(callback) = on_complete { + callback(); + } } 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; @@ -166,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(); @@ -185,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(()); } @@ -226,6 +256,8 @@ impl DownloadManager { } file.flush().await?; + drop(file); + tokio::fs::rename(&partial_path, destination).await?; Ok(()) } diff --git a/crates/goose/src/lib.rs b/crates/goose/src/lib.rs index 579534ab69..322927d093 100644 --- a/crates/goose/src/lib.rs +++ b/crates/goose/src/lib.rs @@ -5,6 +5,7 @@ pub mod config; pub mod context_mgmt; pub mod conversation; pub mod dictation; +pub mod download_manager; pub mod execution; pub mod goose_apps; pub mod hints; diff --git a/crates/goose/src/prompt_template.rs b/crates/goose/src/prompt_template.rs index 1685799c65..d4ec7eb848 100644 --- a/crates/goose/src/prompt_template.rs +++ b/crates/goose/src/prompt_template.rs @@ -39,6 +39,14 @@ static TEMPLATE_REGISTRY: &[(&str, &str)] = &[ "plan.md", "Prompt used when goose creates step-by-step plans. CLI only", ), + ( + "tiny_model_system.md", + "System prompt for tiny local models using shell command emulation", + ), + ( + "session_name.md", + "System prompt for generating short session names from conversation history", + ), ]; /// Information about a template including its content and customization status diff --git a/crates/goose/src/prompts/session_name.md b/crates/goose/src/prompts/session_name.md new file mode 100644 index 0000000000..1a40184a4d --- /dev/null +++ b/crates/goose/src/prompts/session_name.md @@ -0,0 +1,6 @@ +Generate a short title (four words or less) that describes the topic of the user's messages. Reply with only the title, nothing else. + +Examples: +- "how do I reverse a list in python?" → Python list reversal +- "what's the weather in Tokyo?" → Tokyo weather +- "explain how transformers work in ML" → ML transformers explained \ No newline at end of file diff --git a/crates/goose/src/prompts/tiny_model_system.md b/crates/goose/src/prompts/tiny_model_system.md new file mode 100644 index 0000000000..08b1f116c0 --- /dev/null +++ b/crates/goose/src/prompts/tiny_model_system.md @@ -0,0 +1,22 @@ +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. + +The OS is {{os}}, the shell is {{shell}}, and the working 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. + +To run a shell command, start a new line with $: + +$ ls + +Keep your responses brief. State what you are doing, then do it. For example: + +User: how many files are in /tmp? +You: Let me check. +$ ls -1 /tmp | wc -l + +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/base.rs b/crates/goose/src/providers/base.rs index 99559dc11f..4a0c4c1bbc 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -17,10 +17,22 @@ use rmcp::model::Tool; use utoipa::ToSchema; use once_cell::sync::Lazy; +use regex::Regex; use std::ops::{Add, AddAssign}; use std::pin::Pin; +use std::sync::LazyLock; use std::sync::Mutex; +fn strip_xml_tags(text: &str) -> String { + static BLOCK_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?s)<([a-zA-Z][a-zA-Z0-9_]*)[^>]*>.*?").unwrap() + }); + static TAG_RE: LazyLock = + LazyLock::new(|| Regex::new(r"]*>").unwrap()); + let pass1 = BLOCK_RE.replace_all(text, ""); + TAG_RE.replace_all(&pass1, "").into_owned() +} + /// A global store for the current model being used, we use this as when a provider returns, it tells us the real model, not an alias pub static CURRENT_MODEL: Lazy>> = Lazy::new(|| Mutex::new(None)); @@ -594,20 +606,28 @@ pub trait Provider: Send + Sync { messages: &Conversation, ) -> Result { let context = self.get_initial_user_messages(messages); - let prompt = self.create_session_name_prompt(&context); - let message = Message::user().with_text(&prompt); + let system = crate::prompt_template::render_template( + "session_name.md", + &std::collections::HashMap::::new(), + ) + .map_err(|e| ProviderError::ContextLengthExceeded(e.to_string()))?; + + let user_text = format!( + "---BEGIN USER MESSAGES---\n{}\n---END USER MESSAGES---\n\nGenerate a short title for the above messages.", + context.join("\n") + ); + let message = Message::user().with_text(&user_text); let result = self - .complete_fast( - session_id, - "Reply with only a description in four words or less", - &[message], - &[], - ) + .complete_fast(session_id, &system, &[message], &[]) .await?; - let description = result + let raw: String = result .0 - .as_concat_text() + .content + .iter() + .filter_map(|c| c.as_text()) + .collect(); + let description = strip_xml_tags(&raw) .split_whitespace() .collect::>() .join(" "); @@ -615,21 +635,6 @@ pub trait Provider: Send + Sync { Ok(safe_truncate(&description, 100)) } - // Generate a prompt for a session name based on the conversation history - fn create_session_name_prompt(&self, context: &[String]) -> String { - // Create a prompt for a concise description - let mut prompt = "Based on the conversation so far, provide a concise description of this session in 4 words or less. This will be used for finding the session later in a UI with limited space - reply *ONLY* with the description".to_string(); - - if !context.is_empty() { - prompt = format!( - "Here are the first few user messages:\n{}\n\n{}", - context.join("\n"), - prompt - ); - } - prompt - } - /// Configure OAuth authentication for this provider /// /// This method is called when a provider has configuration keys marked with oauth_flow = true. @@ -719,6 +724,42 @@ mod tests { use test_case::test_case; use serde_json::json; + + #[test] + fn test_strip_xml_tags() { + assert_eq!(strip_xml_tags("reasoninganswer"), "answer"); + assert_eq!(strip_xml_tags("beforemidafter"), "beforeafter"); + assert_eq!(strip_xml_tags("xyz"), "z"); + assert_eq!(strip_xml_tags("no tags here"), "no tags here"); + assert_eq!(strip_xml_tags("a < b > c"), "a < b > c"); + assert_eq!(strip_xml_tags("überok"), "ok"); + assert_eq!(strip_xml_tags("日本語hello"), "hello"); + assert_eq!(strip_xml_tags(""), ""); + assert_eq!(strip_xml_tags("<>stuff"), "<>stuff"); + // attributes + assert_eq!( + strip_xml_tags(r#"reasoninganswer"#), + "answer" + ); + // self-closing tags + assert_eq!(strip_xml_tags("
self closing"), "self closing"); + // orphan closing tags + assert_eq!(strip_xml_tags("orphan tag"), "orphan tag"); + // multiline content + assert_eq!( + strip_xml_tags("\nline1\nline2\nresult"), + "result" + ); + } + + #[test] + fn test_usage_creation() { + let usage = Usage::new(Some(10), Some(20), Some(30)); + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.output_tokens, Some(20)); + assert_eq!(usage.total_tokens, Some(30)); + } + fn content_from_str(s: String) -> MessageContent { if let Some(img_data) = s.strip_prefix("*img:") { MessageContent::image(format!("http://example.com/{}", img_data), "image/png") @@ -798,11 +839,10 @@ mod tests { #[tokio::test] async fn test_collect_stream_defaults_usage() { - // Should not error when usage is missing let stream = create_test_stream(vec!["Hello".to_string()]); let (msg, usage) = collect_stream(Box::pin(stream)).await.unwrap(); assert_eq!(content_to_strings(&msg), vec!["Hello"]); - assert_eq!(usage.model, "unknown"); // Default usage + assert_eq!(usage.model, "unknown"); } #[test] diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index c1f69e61d6..8300f5a6df 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -16,6 +16,7 @@ use super::{ google::GoogleProvider, lead_worker::LeadWorkerProvider, litellm::LiteLLMProvider, + local_inference::LocalInferenceProvider, ollama::OllamaProvider, openai::OpenAiProvider, openrouter::OpenRouterProvider, @@ -47,6 +48,7 @@ async fn init_registry() -> RwLock { registry.register::(true); registry.register::(false); registry.register::(false); + registry.register::(false); registry.register::(true); registry.register::(true); registry.register::(true); diff --git a/crates/goose/src/providers/local_inference.rs b/crates/goose/src/providers/local_inference.rs new file mode 100644 index 0000000000..32503e9b3b --- /dev/null +++ b/crates/goose/src/providers/local_inference.rs @@ -0,0 +1,647 @@ +pub mod hf_models; +mod inference_emulated_tools; +mod inference_engine; +mod inference_native_tools; +pub mod local_model_registry; +mod tool_parsing; + +use inference_emulated_tools::{ + build_emulator_tool_description, generate_with_emulated_tools, load_tiny_model_prompt, +}; +use inference_engine::GenerationContext; +use inference_engine::LoadedModel; +use inference_native_tools::generate_with_native_tools; +use tool_parsing::compact_tools_json; + +use crate::config::ExtensionConfig; +use crate::conversation::message::{Message, MessageContent}; +use crate::model::ModelConfig; +use crate::providers::base::{ + MessageStream, Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage, +}; +use crate::providers::errors::ProviderError; +use crate::providers::formats::openai::format_tools; +use crate::providers::utils::RequestLog; +use anyhow::Result; +use async_stream::try_stream; +use async_trait::async_trait; +use futures::future::BoxFuture; +use llama_cpp_2::llama_backend::LlamaBackend; +use llama_cpp_2::model::params::LlamaModelParams; +use llama_cpp_2::model::{LlamaChatMessage, LlamaChatTemplate, LlamaModel}; +use llama_cpp_2::{list_llama_ggml_backend_devices, LlamaBackendDeviceType}; +use rmcp::model::{Role, Tool}; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::{Arc, Mutex as StdMutex, Weak}; +use tokio::sync::Mutex; +use uuid::Uuid; + +const SHELL_TOOL: &str = "developer__shell"; +const CODE_EXECUTION_TOOL: &str = "code_execution__execute"; + +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; + } + // Safety invariant: the Weak::upgrade() check and LlamaBackend::init() + // both execute inside this same mutex guard, so there is no window where + // another thread could drop the Arc and re-enter concurrently. + // BackendAlreadyInitialized therefore means LlamaBackend::drop() did not + // reset the C library's init flag — a llama-cpp-rs bug, not a race. + let backend = match LlamaBackend::init() { + Ok(b) => b, + Err(llama_cpp_2::LlamaCppError::BackendAlreadyInitialized) => { + unreachable!( + "LlamaBackend already initialized but Weak was dead; \ + the mutex guard prevents concurrent re-init" + ) + } + 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 = "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M"; + +pub const LOCAL_LLM_MODEL_CONFIG_KEY: &str = "LOCAL_LLM_MODEL"; + +/// Resolve model path, context limit, and settings for a model ID from the registry. +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 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())); + } + } + + None +} + +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 { + devices + .iter() + .filter(|d| d.device_type == LlamaBackendDeviceType::Cpu) + .map(|d| d.memory_free as u64) + .max() + .unwrap_or(0) + } +} + +pub fn recommend_local_model(runtime: &InferenceRuntime) -> String { + use local_model_registry::{get_registry, is_featured_model, FEATURED_MODELS}; + + let available_memory = available_inference_memory_bytes(runtime); + + if let Ok(registry) = get_registry().lock() { + let mut models: Vec<_> = registry + .list_models() + .iter() + .filter(|m| is_featured_model(&m.id) && m.size_bytes > 0) + .collect(); + models.sort_by(|a, b| b.size_bytes.cmp(&a.size_bytes)); + + // Return largest that fits in available memory + for model in &models { + if available_memory >= model.size_bytes { + return model.id.clone(); + } + } + + // If nothing fits, return smallest + if let Some(smallest) = models.last() { + return smallest.id.clone(); + } + } + + // Fallback to first featured model + FEATURED_MODELS[0].to_string() +} + +fn build_openai_messages_json(system: &str, messages: &[Message]) -> String { + use crate::providers::formats::openai::format_messages; + use crate::providers::utils::ImageFormat; + + let mut arr: Vec = vec![json!({"role": "system", "content": system})]; + arr.extend(format_messages(messages, &ImageFormat::OpenAi)); + serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_string()) +} + +/// Convert a message into plain text for the emulator path's chat history. +/// +/// This is the emulator-path counterpart of [`format_messages`] used by the native +/// path. It reconstructs the text-based tool syntax that the emulator prompt teaches +/// the model: +/// +/// - `ToolRequest` with a `"command"` argument → `$ command` +/// - `ToolRequest` with a `"code"` argument → `` ```execute\n…\n``` `` +/// - `ToolResponse` → `Command output:\n…` +/// +/// Only `developer__shell` and `code_execution__execute` style tool calls are +/// recognized (by argument shape, not tool name). Tool calls from other extensions +/// (e.g. custom MCP tools made by a native-tool-calling model earlier in the +/// conversation) are silently dropped, since the emulator path has no syntax to +/// represent them. +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 `ProviderUsage` and write the request log entry. +fn finalize_usage( + log: &mut RequestLog, + model_name: String, + path_label: &str, + prompt_token_count: usize, + output_token_count: i32, + extra_log_fields: Option<(&str, &str)>, +) -> ProviderUsage { + 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 mut log_json = serde_json::json!({ + "path": path_label, + "prompt_tokens": input_tokens, + "output_tokens": output_token_count, + }); + if let Some((key, value)) = extra_log_fields { + log_json[key] = serde_json::json!(value); + } + let _ = log.write(&log_json, Some(&usage)); + ProviderUsage::new(model_name, usage) +} + +type StreamSender = + tokio::sync::mpsc::Sender, Option), ProviderError>>; + +pub struct LocalInferenceProvider { + runtime: Arc, + model: ModelSlot, + model_config: ModelConfig, + name: String, +} + +impl LocalInferenceProvider { + pub async fn from_env(model: ModelConfig, _extensions: Vec) -> Result { + let runtime = InferenceRuntime::get_or_init(); + let model_slot = runtime.get_or_create_model_slot(&model.model_name); + Ok(Self { + runtime, + model: model_slot, + model_config: model, + name: PROVIDER_NAME.to_string(), + }) + } + + 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)))?; + + if !model_path.exists() { + return Err(ProviderError::ExecutionError(format!( + "Model not downloaded: {}. Please download it from Settings > Local Inference.", + model_id + ))); + } + + tracing::info!("Loading {} from: {}", model_id, model_path.display()); + + let backend = runtime.backend(); + + 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(e.to_string()))?; + + 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 create fallback chat template: {}", + e + )) + })? + } + }; + + tracing::info!("Model loaded successfully"); + + Ok(LoadedModel { model, template }) + } +} + +impl ProviderDef for LocalInferenceProvider { + type Provider = Self; + + fn metadata() -> ProviderMetadata + where + Self: Sized, + { + use crate::providers::local_inference::local_model_registry::{ + get_registry, FEATURED_MODELS, + }; + + let mut known_models: Vec<&str> = FEATURED_MODELS.to_vec(); + + // 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 (llama.cpp)", + DEFAULT_MODEL, + known_models, + "https://github.com/utilityai/llama-cpp-rs", + vec![], + ) + } + + fn from_env( + model: ModelConfig, + extensions: Vec, + ) -> BoxFuture<'static, Result> + where + Self: Sized, + { + Box::pin(Self::from_env(model, extensions)) + } +} + +#[async_trait] +impl Provider for LocalInferenceProvider { + fn get_name(&self) -> &str { + &self.name + } + + fn get_model_config(&self) -> ModelConfig { + self.model_config.clone() + } + + async fn fetch_supported_models(&self) -> Result, ProviderError> { + use crate::providers::local_inference::local_model_registry::get_registry; + + let mut all_models: Vec = Vec::new(); + + if let Ok(registry) = get_registry().lock() { + for entry in registry.list_models() { + all_models.push(entry.id.clone()); + } + } + + Ok(all_models) + } + + async fn stream( + &self, + model_config: &ModelConfig, + _session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + 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 + )) + })?; + + // 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 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(e.to_string()))??; + *model_lock = Some(loaded); + } + } + + // 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. + // Only use emulator when there are actually tools to emulate - utility calls + // like compaction and session naming pass empty tools and should preserve + // their system prompts. + let use_emulator = !model_settings.native_tool_calling && !tools.is_empty(); + let system_prompt = if use_emulator { + load_tiny_model_prompt() + } else { + system.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 + )) + }, + )?, + ]; + + let code_mode_enabled = tools.iter().any(|t| t.name == CODE_EXECUTION_TOOL); + + if use_emulator && !tools.is_empty() { + let tool_desc = build_emulator_tool_description(tools, code_mode_enabled); + 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 oai_messages_json = if model_settings.use_jinja { + Some(build_openai_messages_json(&system_prompt, messages)) + } else { + 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; + + let log_payload = serde_json::json!({ + "system": &system_prompt, + "messages": messages.iter().map(|m| { + serde_json::json!({ + "role": match m.role { Role::User => "user", Role::Assistant => "assistant" }, + "content": extract_text_content(m), + }) + }).collect::>(), + "tools": tools.iter().map(|t| &t.name).collect::>(), + "settings": { + "use_jinja": settings.use_jinja, + "native_tool_calling": settings.native_tool_calling, + "context_size": settings.context_size, + "sampling": settings.sampling, + }, + }); + + let mut log = RequestLog::start(&self.model_config, &log_payload) + .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; + + let (tx, mut rx) = tokio::sync::mpsc::channel::< + Result<(Option, Option), ProviderError>, + >(32); + + tokio::task::spawn_blocking(move || { + // Macro to log errors before sending them through the channel + macro_rules! send_err { + ($err:expr) => {{ + let err = $err; + let msg = match &err { + ProviderError::ExecutionError(s) => s.as_str(), + ProviderError::ContextLengthExceeded(s) => s.as_str(), + _ => "unknown error", + }; + let _ = log.error(msg); + let _ = tx.blocking_send(Err(err)); + return; + }}; + } + + let model_guard = model_arc.blocking_lock(); + let loaded = match model_guard.as_ref() { + Some(l) => l, + None => { + send_err!(ProviderError::ExecutionError( + "Model not loaded".to_string() + )); + } + }; + + let message_id = Uuid::new_v4().to_string(); + + let mut gen_ctx = GenerationContext { + loaded, + runtime: &runtime, + chat_messages: &chat_messages, + settings: &settings, + context_limit, + model_name, + message_id: &message_id, + tx: &tx, + log: &mut log, + }; + + let result = if use_emulator { + generate_with_emulated_tools(&mut gen_ctx, code_mode_enabled) + } else { + generate_with_native_tools( + &mut gen_ctx, + &oai_messages_json, + full_tools_json.as_deref(), + compact_tools.as_deref(), + ) + }; + + if let Err(err) = result { + let msg = match &err { + ProviderError::ExecutionError(s) => s.as_str(), + ProviderError::ContextLengthExceeded(s) => s.as_str(), + _ => "unknown error", + }; + let _ = log.error(msg); + let _ = tx.blocking_send(Err(err)); + } + }); + + Ok(Box::pin(try_stream! { + while let Some(result) = rx.recv().await { + let item = result?; + yield item; + } + + })) + } +} 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..9729fcd832 --- /dev/null +++ b/crates/goose/src/providers/local_inference/hf_models.rs @@ -0,0 +1,486 @@ +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, +} + +const QUANT_TABLE: &[(&str, &str, u8)] = &[ + ("IQ1_S", "Extremely small, very low quality", 1), + ("IQ1_M", "Extremely small, very low quality", 2), + ("IQ2_XXS", "Very small, low quality", 3), + ("IQ2_XS", "Very small, low quality", 4), + ("IQ2_S", "Very small, low quality", 5), + ("IQ2_M", "Very small, low quality", 6), + ("Q2_K", "Small, low quality", 7), + ("Q2_K_S", "Small, low quality", 7), + ("IQ3_XXS", "Very small, moderate quality loss", 8), + ("IQ3_XS", "Small, moderate quality loss", 9), + ("IQ3_S", "Small, moderate quality loss", 9), + ("Q3_K_S", "Small, moderate quality loss", 10), + ("IQ3_M", "Small, moderate quality loss", 11), + ("Q3_K_M", "Small, balanced quality/size", 12), + ("Q3_K_L", "Medium-small, decent quality", 13), + ("IQ4_XS", "Medium, good quality", 14), + ("IQ4_NL", "Medium, good quality", 15), + ("Q4_0", "Medium, good quality", 16), + ("Q4_1", "Medium, good quality", 17), + ("Q4_K_S", "Medium, good quality/size balance", 18), + ( + "Q4_K_M", + "Medium, recommended balance of quality and size", + 19, + ), + ("Q5_0", "Medium-large, high quality", 20), + ("Q5_1", "Medium-large, high quality", 21), + ("Q5_K_S", "Medium-large, high quality", 22), + ("Q5_K_M", "Medium-large, very high quality", 23), + ("Q6_K", "Large, near-lossless quality", 24), + ("Q8_0", "Large, near-lossless quality", 25), + ("F16", "Full size, original quality (16-bit)", 26), + ("BF16", "Full size, original quality (bfloat16)", 27), + ("F32", "Full size, original quality (32-bit)", 28), + ( + "MXFP4_MOE", + "Medium, mixed-precision 4-bit for MoE models", + 18, + ), + ("TQ1_0", "Tiny, ternary quantization", 1), + ("Q2_K_XL", "Extended-layer variant", 15), + ("Q3_K_XL", "Extended-layer variant", 15), + ("Q4_K_XL", "Extended-layer variant", 15), + ("Q2_K_L", "Small, low quality (large variant)", 8), + ("Q4_K_L", "Medium, good quality (large variant)", 20), +]; + +fn quant_info(quant: &str) -> QuantInfo { + QUANT_TABLE + .iter() + .find(|(name, _, _)| *name == quant) + .map(|(_, description, quality_rank)| QuantInfo { + description, + quality_rank: *quality_rank, + }) + .unwrap_or(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/inference_emulated_tools.rs b/crates/goose/src/providers/local_inference/inference_emulated_tools.rs new file mode 100644 index 0000000000..4f9fb0f870 --- /dev/null +++ b/crates/goose/src/providers/local_inference/inference_emulated_tools.rs @@ -0,0 +1,685 @@ +//! Tool call emulation for models without native tool-calling support. +//! +//! The model is prompted to emit shell commands as `$ command` on a new line and +//! code blocks as `` ```execute `` fenced blocks. A streaming parser detects these +//! patterns and converts them into tool-call messages. +//! +//! # Known false-positive scenarios +//! +//! Because detection is purely text-based, the parser can misinterpret model output: +//! +//! - **`$` at line start in explanatory text.** If the model writes a line starting +//! with `$` as an example (e.g. "$ is the jQuery selector"), it will be treated as +//! a shell command. Mid-sentence `$` (e.g. "costs $50") is safe — only `\n$` or +//! `$` at the very start of output triggers command detection. +//! +//! - **`` ```execute `` in explanatory code fences.** If the model uses this exact +//! fence tag in prose, the content will be executed. Standard `` ```js `` or +//! `` ```python `` fences are not affected. +//! +//! These are inherent to text-based tool emulation. Models with native tool-calling +//! support should use the `inference_native_tools` path instead. + +use crate::conversation::message::{Message, MessageContent}; +use crate::providers::errors::ProviderError; +use llama_cpp_2::model::AddBos; +use rmcp::model::{CallToolRequestParams, Tool}; +use serde_json::json; +use std::borrow::Cow; +use uuid::Uuid; + +use super::inference_engine::{ + create_and_prefill_context, generation_loop, validate_and_compute_context, GenerationContext, + TokenAction, +}; +use super::{finalize_usage, StreamSender, CODE_EXECUTION_TOOL, SHELL_TOOL}; + +const HOLD_BACK_CODE_MODE: usize = " ```execute\n".len(); +const HOLD_BACK_SHELL_ONLY: usize = "\n$".len(); + +pub(super) fn load_tiny_model_prompt() -> String { + use std::env; + + let os = if cfg!(target_os = "macos") { + "macos" + } else if cfg!(target_os = "linux") { + "linux" + } else if cfg!(target_os = "windows") { + "windows" + } else { + "unknown" + }; + + let working_directory = env::current_dir() + .map(|p| p.display().to_string()) + .unwrap_or_else(|_| "unknown".to_string()); + + let shell = env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string()); + + let context = json!({ + "os": os, + "working_directory": working_directory, + "shell": shell, + }); + + crate::prompt_template::render_template("tiny_model_system.md", &context).unwrap_or_else(|e| { + tracing::warn!("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() + }) +} + +pub(super) fn build_emulator_tool_description(tools: &[Tool], code_mode_enabled: bool) -> String { + let mut tool_desc = String::new(); + + if code_mode_enabled { + 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"); + + 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::(), + } + }; + 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")); + } + } + } 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)); + } + } + + tool_desc +} + +enum EmulatorAction { + Text(String), + ShellCommand(String), + ExecuteCode(String), +} + +enum ParserState { + Normal, + InCommand, + InExecuteBlock, +} + +struct StreamingEmulatorParser { + buffer: String, + state: ParserState, + code_mode_enabled: bool, +} + +impl StreamingEmulatorParser { + fn new(code_mode_enabled: bool) -> Self { + Self { + buffer: String::new(), + state: ParserState::Normal, + code_mode_enabled, + } + } + + fn process_chunk(&mut self, chunk: &str) -> Vec { + self.buffer.push_str(chunk); + let mut results = Vec::new(); + + loop { + 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; + } + } + 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)); + } + } 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 { + let hold_back = if self.code_mode_enabled { + HOLD_BACK_CODE_MODE + } else { + HOLD_BACK_SHELL_ONLY + }; + 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; + } + } + } + } + + results + } + + fn flush(&mut self) -> Vec { + let mut results = Vec::new(); + + if !self.buffer.is_empty() { + 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())); + } + } + ParserState::InExecuteBlock => { + 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.state = ParserState::Normal; + } + + results + } +} + +fn send_emulator_action( + action: &EmulatorAction, + message_id: &str, + tx: &StreamSender, +) -> Result { + match action { + EmulatorAction::Text(text) => { + let mut message = Message::assistant().with_text(text); + message.id = Some(message_id.to_string()); + tx.blocking_send(Ok((Some(message), None))) + .map_err(|_| ())?; + Ok(false) + } + 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::Borrowed(SHELL_TOOL), + arguments: Some(args), + }; + let mut message = Message::assistant(); + message + .content + .push(MessageContent::tool_request(tool_id, Ok(tool_call))); + message.id = Some(message_id.to_string()); + tx.blocking_send(Ok((Some(message), None))) + .map_err(|_| ())?; + Ok(true) + } + EmulatorAction::ExecuteCode(code) => { + let tool_id = Uuid::new_v4().to_string(); + let wrapped = if code.contains("async function run()") { + code.clone() + } 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::Borrowed(CODE_EXECUTION_TOOL), + arguments: Some(args), + }; + let mut message = Message::assistant(); + message + .content + .push(MessageContent::tool_request(tool_id, Ok(tool_call))); + message.id = Some(message_id.to_string()); + tx.blocking_send(Ok((Some(message), None))) + .map_err(|_| ())?; + Ok(true) + } + } +} + +pub(super) fn generate_with_emulated_tools( + ctx: &mut GenerationContext<'_>, + code_mode_enabled: bool, +) -> Result<(), ProviderError> { + let prompt = ctx + .loaded + .model + .apply_chat_template(&ctx.loaded.template, ctx.chat_messages, true) + .map_err(|e| { + ProviderError::ExecutionError(format!("Failed to apply chat template: {}", e)) + })?; + + let tokens = ctx + .loaded + .model + .str_to_token(&prompt, AddBos::Never) + .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; + + let (prompt_token_count, effective_ctx) = validate_and_compute_context( + ctx.loaded, + ctx.runtime, + tokens.len(), + ctx.context_limit, + ctx.settings, + )?; + let mut llama_ctx = create_and_prefill_context( + ctx.loaded, + ctx.runtime, + &tokens, + effective_ctx, + ctx.settings, + )?; + + let message_id = ctx.message_id; + let tx = ctx.tx; + let mut emulator_parser = StreamingEmulatorParser::new(code_mode_enabled); + let mut tool_call_emitted = false; + let mut send_failed = false; + + let output_token_count = generation_loop( + &ctx.loaded.model, + &mut llama_ctx, + ctx.settings, + prompt_token_count, + effective_ctx, + |piece| { + let actions = emulator_parser.process_chunk(piece); + for action in actions { + match send_emulator_action(&action, message_id, tx) { + Ok(is_tool) => { + if is_tool { + tool_call_emitted = true; + } + } + Err(_) => { + send_failed = true; + return Ok(TokenAction::Stop); + } + } + } + if tool_call_emitted { + Ok(TokenAction::Stop) + } else { + Ok(TokenAction::Continue) + } + }, + )?; + + if !send_failed { + for action in emulator_parser.flush() { + if send_emulator_action(&action, message_id, tx).is_err() { + break; + } + } + } + + let provider_usage = finalize_usage( + ctx.log, + std::mem::take(&mut ctx.model_name), + "emulator", + prompt_token_count, + output_token_count, + None, + ); + let _ = ctx.tx.blocking_send(Ok((None, Some(provider_usage)))); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Collect all actions from feeding chunks through the parser, then flushing. + fn parse_chunks(chunks: &[&str], code_mode: bool) -> Vec { + let mut parser = StreamingEmulatorParser::new(code_mode); + let mut actions = Vec::new(); + for chunk in chunks { + actions.extend(parser.process_chunk(chunk)); + } + actions.extend(parser.flush()); + actions + } + + fn parse_all(input: &str, code_mode: bool) -> Vec { + parse_chunks(&[input], code_mode) + } + + fn assert_text(action: &EmulatorAction, expected: &str) { + match action { + EmulatorAction::Text(t) => assert_eq!(t.trim(), expected.trim(), "text mismatch"), + other => panic!("expected Text, got {:?}", action_label(other)), + } + } + + fn assert_shell(action: &EmulatorAction, expected: &str) { + match action { + EmulatorAction::ShellCommand(cmd) => { + assert_eq!(cmd, expected, "shell command mismatch") + } + other => panic!("expected ShellCommand, got {:?}", action_label(other)), + } + } + + fn assert_execute(action: &EmulatorAction, expected: &str) { + match action { + EmulatorAction::ExecuteCode(code) => { + assert_eq!(code.trim(), expected.trim(), "execute code mismatch") + } + other => panic!("expected ExecuteCode, got {:?}", action_label(other)), + } + } + + fn action_label(a: &EmulatorAction) -> &'static str { + match a { + EmulatorAction::Text(_) => "Text", + EmulatorAction::ShellCommand(_) => "ShellCommand", + EmulatorAction::ExecuteCode(_) => "ExecuteCode", + } + } + + #[test] + fn plain_text_no_tools() { + let actions = parse_all("Hello, world!", false); + // Hold-back may split text across actions; concatenate all text + let all_text: String = actions + .iter() + .map(|a| match a { + EmulatorAction::Text(t) => t.as_str(), + _ => panic!("expected only Text actions"), + }) + .collect(); + assert_eq!(all_text.trim(), "Hello, world!"); + } + + #[test] + fn single_shell_command() { + let actions = parse_all("$ ls -la\n", false); + assert_eq!(actions.len(), 1); + assert_shell(&actions[0], "ls -la"); + } + + #[test] + fn text_then_shell_command() { + let actions = parse_all("Let me check:\n$ ls -la\n", false); + assert!(actions.len() >= 2); + assert_text(&actions[0], "Let me check:"); + assert_shell(&actions[actions.len() - 1], "ls -la"); + } + + #[test] + fn shell_command_at_start_of_output() { + let actions = parse_all("$ whoami\n", false); + assert_eq!(actions.len(), 1); + assert_shell(&actions[0], "whoami"); + } + + #[test] + fn shell_command_without_trailing_newline() { + // Flush should handle unterminated command + let actions = parse_all("$ whoami", false); + assert_eq!(actions.len(), 1); + assert_shell(&actions[0], "whoami"); + } + + #[test] + fn dollar_sign_mid_sentence_is_not_command() { + let actions = parse_all("It costs $50 per month", false); + for action in &actions { + assert!( + matches!(action, EmulatorAction::Text(_)), + "mid-sentence $ should not trigger a shell command" + ); + } + let all_text: String = actions + .iter() + .filter_map(|a| match a { + EmulatorAction::Text(t) => Some(t.as_str()), + _ => None, + }) + .collect(); + assert_eq!(all_text.trim(), "It costs $50 per month"); + } + + #[test] + fn execute_block() { + let input = "Here's the code:\n```execute\nconsole.log('hi');\n```\n"; + let actions = parse_all(input, true); + assert!(actions.len() >= 2); + assert_text(&actions[0], "Here's the code:"); + assert_execute(&actions[actions.len() - 1], "console.log('hi');"); + } + + #[test] + fn execute_block_not_detected_without_code_mode() { + let input = "```execute\nconsole.log('hi');\n```\n"; + let actions = parse_all(input, false); + // Should be treated as plain text + for action in &actions { + assert!(matches!(action, EmulatorAction::Text(_))); + } + } + + #[test] + fn dollar_split_across_chunks() { + // The \n and $ arrive in separate chunks + let actions = parse_chunks(&["Let me check\n", "$ ls -la\n"], false); + let shells: Vec<_> = actions + .iter() + .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) + .collect(); + assert_eq!(shells.len(), 1); + assert_shell(shells[0], "ls -la"); + } + + #[test] + fn execute_fence_split_across_chunks() { + let actions = parse_chunks(&["Here:\n```ex", "ecute\nlet x = 1;\n", "```\n"], true); + let executes: Vec<_> = actions + .iter() + .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) + .collect(); + assert_eq!(executes.len(), 1); + assert_execute(executes[0], "let x = 1;"); + } + + #[test] + fn multiple_commands_on_separate_lines() { + // In practice, generation stops after the first tool call. But the + // parser should detect commands separated by \n$ when fed as chunks. + let actions = parse_chunks(&["Here:\n$ cd /tmp\n", "Done.\n$ ls\n"], false); + let shells: Vec<_> = actions + .iter() + .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) + .collect(); + assert_eq!(shells.len(), 2); + assert_shell(shells[0], "cd /tmp"); + assert_shell(shells[1], "ls"); + } + + #[test] + fn regular_code_fence_not_treated_as_execute() { + let input = "```python\nprint('hi')\n```\n"; + let actions = parse_all(input, true); + for action in &actions { + assert!( + matches!(action, EmulatorAction::Text(_)), + "regular code fence should be text" + ); + } + } + + #[test] + fn empty_command_ignored() { + let actions = parse_all("$\n", false); + // Empty command after $ should not produce a ShellCommand + let shells: Vec<_> = actions + .iter() + .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) + .collect(); + assert_eq!(shells.len(), 0); + } + + #[test] + fn token_by_token_streaming() { + // Simulate LLM generating one token at a time + let input = "$ echo hello\n"; + let chars: Vec = input.chars().map(|c| c.to_string()).collect(); + let chunks: Vec<&str> = chars.iter().map(|s| s.as_str()).collect(); + let actions = parse_chunks(&chunks, false); + let shells: Vec<_> = actions + .iter() + .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) + .collect(); + assert_eq!(shells.len(), 1); + assert_shell(shells[0], "echo hello"); + } + + #[test] + fn execute_block_with_multiline_code() { + let input = "```execute\nasync function run() {\n const r = await Developer.shell({ command: \"ls\" });\n return r;\n}\n```\n"; + let actions = parse_all(input, true); + let executes: Vec<_> = actions + .iter() + .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) + .collect(); + assert_eq!(executes.len(), 1); + match executes[0] { + EmulatorAction::ExecuteCode(code) => { + assert!(code.contains("async function run()")); + assert!(code.contains("Developer.shell")); + } + _ => unreachable!(), + } + } + + #[test] + fn unclosed_execute_block_flushed() { + // Model stops generating mid-block + let input = "```execute\nlet x = 1;"; + let actions = parse_all(input, true); + let executes: Vec<_> = actions + .iter() + .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) + .collect(); + assert_eq!(executes.len(), 1); + assert_execute(executes[0], "let x = 1;"); + } +} diff --git a/crates/goose/src/providers/local_inference/inference_engine.rs b/crates/goose/src/providers/local_inference/inference_engine.rs new file mode 100644 index 0000000000..06256fca90 --- /dev/null +++ b/crates/goose/src/providers/local_inference/inference_engine.rs @@ -0,0 +1,387 @@ +use crate::providers::errors::ProviderError; +use crate::providers::local_inference::local_model_registry::ModelSettings; +use crate::providers::utils::RequestLog; +use llama_cpp_2::context::params::LlamaContextParams; +use llama_cpp_2::llama_batch::LlamaBatch; +use llama_cpp_2::model::{LlamaChatMessage, LlamaChatTemplate, LlamaModel}; +use llama_cpp_2::sampling::LlamaSampler; +use std::num::NonZeroU32; + +use super::{InferenceRuntime, StreamSender}; + +pub(super) struct GenerationContext<'a> { + pub loaded: &'a LoadedModel, + pub runtime: &'a InferenceRuntime, + pub chat_messages: &'a [LlamaChatMessage], + pub settings: &'a ModelSettings, + pub context_limit: usize, + pub model_name: String, + pub message_id: &'a str, + pub tx: &'a StreamSender, + pub log: &'a mut RequestLog, +} + +pub(super) struct LoadedModel { + pub model: LlamaModel, + pub template: LlamaChatTemplate, +} + +/// 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. +pub(super) fn estimate_max_context_for_memory( + model: &LlamaModel, + runtime: &InferenceRuntime, +) -> Option { + let available = super::available_inference_memory_bytes(runtime); + if available == 0 { + return None; + } + + // Reserve memory for computation scratch buffers (attention, etc.) and other overhead. + // The compute buffer can be 40-50% of the KV cache size for large models, so we + // conservatively use only half the available memory for the KV cache. + let usable = (available as f64 * 0.5) 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; + } + + // For MLA (Multi-head Latent Attention) models like DeepSeek/GLM, the actual KV cache + // dimensions differ from n_head_kv * head_dim. Read the true dimensions from GGUF metadata. + let arch = model + .meta_val_str("general.architecture") + .unwrap_or_default(); + let head_dim = n_embd / n_head; + let k_per_head = model + .meta_val_str(&format!("{arch}.attention.key_length")) + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(head_dim); + let v_per_head = model + .meta_val_str(&format!("{arch}.attention.value_length")) + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(head_dim); + + // Total KV dimensions across all KV heads, times n_layer, times 2 bytes (f16) per element + let bytes_per_token = (k_per_head + v_per_head) * n_head_kv * n_layer * 2; + + if bytes_per_token == 0 { + return None; + } + + Some((usable / bytes_per_token) as usize) +} + +pub(super) fn context_cap( + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + context_limit: usize, + n_ctx_train: usize, + memory_max_ctx: Option, +) -> usize { + if let Some(ctx_size) = settings.context_size { + return ctx_size as usize; + } + + let limit = if context_limit > 0 { + context_limit + } else { + n_ctx_train + }; + + 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, + } +} + +pub(super) fn effective_context_size( + prompt_token_count: usize, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + context_limit: usize, + n_ctx_train: usize, + memory_max_ctx: Option, +) -> usize { + let limit = context_cap(settings, context_limit, n_ctx_train, memory_max_ctx); + let min_generation_headroom = 512; + let needed = prompt_token_count + min_generation_headroom; + if needed > limit { + tracing::warn!( + "Prompt ({} tokens) + headroom exceeds context limit ({}), capping to limit", + prompt_token_count, + limit, + ); + } + needed.min(limit) +} + +pub(super) 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 { + let policy = if flash_attn { 1 } else { 0 }; + params = params.with_flash_attention_policy(policy); + } + + params +} + +pub(super) 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) + } +} + +/// Validate prompt tokens against memory limits and compute the effective +/// context size. Returns `(prompt_token_count, effective_ctx)`. +pub(super) fn validate_and_compute_context( + loaded: &LoadedModel, + runtime: &InferenceRuntime, + prompt_token_count: usize, + context_limit: usize, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, +) -> Result<(usize, usize), ProviderError> { + 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 = effective_context_size( + prompt_token_count, + settings, + context_limit, + n_ctx_train, + memory_max_ctx, + ); + if let Some(mem_max) = memory_max_ctx { + if prompt_token_count > mem_max { + return Err(ProviderError::ContextLengthExceeded(format!( + "Prompt ({} tokens) exceeds estimated memory capacity ({} tokens). \ + Try a smaller model or reduce conversation length.", + prompt_token_count, mem_max, + ))); + } + } + Ok((prompt_token_count, effective_ctx)) +} + +/// Create a llama context and prefill (decode) all prompt tokens. +pub(super) fn create_and_prefill_context<'model>( + loaded: &'model LoadedModel, + runtime: &InferenceRuntime, + tokens: &[llama_cpp_2::token::LlamaToken], + effective_ctx: usize, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, +) -> Result, ProviderError> { + let ctx_params = build_context_params(effective_ctx as u32, settings); + let mut ctx = loaded + .model + .new_context(runtime.backend(), ctx_params) + .map_err(|e| ProviderError::ExecutionError(format!("Failed to create context: {}", e)))?; + + let n_batch = ctx.n_batch() as usize; + for chunk in tokens.chunks(n_batch) { + let mut batch = LlamaBatch::get_one(chunk) + .map_err(|e| ProviderError::ExecutionError(format!("Failed to create batch: {}", e)))?; + ctx.decode(&mut batch) + .map_err(|e| ProviderError::ExecutionError(format!("Prefill decode failed: {}", e)))?; + } + + Ok(ctx) +} + +/// Action to take after processing a generated token piece. +pub(super) enum TokenAction { + Continue, + Stop, +} + +/// Run the autoregressive generation loop. Calls `on_piece` for each non-empty +/// token piece. The callback returns `TokenAction::Stop` to break early. +/// Returns the total number of generated tokens. +pub(super) fn generation_loop( + model: &LlamaModel, + ctx: &mut llama_cpp_2::context::LlamaContext<'_>, + settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + prompt_token_count: usize, + effective_ctx: usize, + mut on_piece: impl FnMut(&str) -> Result, +) -> Result { + let mut sampler = build_sampler(settings); + 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 output_token_count: i32 = 0; + + for _ in 0..max_output { + let token = sampler.sample(ctx, -1); + sampler.accept(token); + + if model.is_eog_token(token) { + break; + } + + output_token_count += 1; + + let piece = model + .token_to_piece(token, &mut decoder, true, None) + .map_err(|e| ProviderError::ExecutionError(format!("Failed to decode token: {}", e)))?; + + if !piece.is_empty() && matches!(on_piece(&piece)?, TokenAction::Stop) { + break; + } + + let next_tokens = [token]; + let mut next_batch = LlamaBatch::get_one(&next_tokens) + .map_err(|e| ProviderError::ExecutionError(format!("Failed to create batch: {}", e)))?; + ctx.decode(&mut next_batch) + .map_err(|e| ProviderError::ExecutionError(format!("Decode failed: {}", e)))?; + } + + Ok(output_token_count) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::providers::local_inference::local_model_registry::ModelSettings; + + fn default_settings() -> ModelSettings { + ModelSettings::default() + } + + #[test] + fn test_effective_context_size_basic() { + assert_eq!( + effective_context_size(100, &default_settings(), 4096, 4096, None), + 612 + ); + } + + #[test] + fn test_effective_context_size_capped_by_limit() { + assert_eq!( + effective_context_size(100, &default_settings(), 1024, 8192, None), + 612 + ); + } + + #[test] + fn test_effective_context_size_capped_by_memory() { + assert_eq!( + effective_context_size(100, &default_settings(), 4096, 4096, Some(800)), + 612 + ); + } + + #[test] + fn test_effective_context_size_memory_smaller_than_needed() { + assert_eq!( + effective_context_size(600, &default_settings(), 4096, 4096, Some(700)), + 700 + ); + } + + #[test] + fn test_effective_context_size_zero_limit_uses_train() { + assert_eq!( + effective_context_size(100, &default_settings(), 0, 2048, None), + 612 + ); + } + + #[test] + fn test_effective_context_size_prompt_exceeds_all_limits() { + assert_eq!( + effective_context_size(5000, &default_settings(), 4096, 4096, None), + 4096 + ); + } + + #[test] + fn test_context_cap_with_settings_override() { + let mut settings = default_settings(); + settings.context_size = Some(2048); + assert_eq!(context_cap(&settings, 4096, 8192, Some(1024)), 2048); + } + + #[test] + fn test_context_cap_without_override() { + assert_eq!(context_cap(&default_settings(), 4096, 8192, None), 4096); + } + + #[test] + fn test_context_cap_memory_limited() { + assert_eq!( + context_cap(&default_settings(), 4096, 8192, Some(2048)), + 2048 + ); + } +} diff --git a/crates/goose/src/providers/local_inference/inference_native_tools.rs b/crates/goose/src/providers/local_inference/inference_native_tools.rs new file mode 100644 index 0000000000..656a6e08cf --- /dev/null +++ b/crates/goose/src/providers/local_inference/inference_native_tools.rs @@ -0,0 +1,197 @@ +use crate::conversation::message::Message; +use crate::providers::errors::ProviderError; +use llama_cpp_2::model::AddBos; +use llama_cpp_2::openai::OpenAIChatTemplateParams; + +use super::finalize_usage; +use super::inference_engine::{ + context_cap, create_and_prefill_context, estimate_max_context_for_memory, generation_loop, + validate_and_compute_context, GenerationContext, TokenAction, +}; +use super::tool_parsing::{ + extract_tool_call_messages, extract_xml_tool_call_messages, safe_stream_end, + split_content_and_tool_calls, split_content_and_xml_tool_calls, +}; + +pub(super) fn generate_with_native_tools( + ctx: &mut GenerationContext<'_>, + oai_messages_json: &Option, + full_tools_json: Option<&str>, + compact_tools: Option<&str>, +) -> Result<(), ProviderError> { + let min_generation_headroom = 512; + let n_ctx_train = ctx.loaded.model.n_ctx_train() as usize; + let memory_max_ctx = estimate_max_context_for_memory(&ctx.loaded.model, ctx.runtime); + let cap = context_cap(ctx.settings, ctx.context_limit, n_ctx_train, memory_max_ctx); + let token_budget = cap.saturating_sub(min_generation_headroom); + + let apply_template = |tools: Option<&str>| { + if let Some(ref messages_json) = oai_messages_json { + let params = OpenAIChatTemplateParams { + messages_json: messages_json.as_str(), + tools_json: tools, + tool_choice: None, + json_schema: None, + grammar: None, + reasoning_format: None, + chat_template_kwargs: None, + add_generation_prompt: true, + use_jinja: true, + parallel_tool_calls: false, + enable_thinking: false, + add_bos: false, + add_eos: false, + parse_tool_calls: true, + }; + ctx.loaded + .model + .apply_chat_template_oaicompat(&ctx.loaded.template, ¶ms) + } else { + ctx.loaded.model.apply_chat_template_with_tools_oaicompat( + &ctx.loaded.template, + ctx.chat_messages, + tools, + None, + true, + ) + } + }; + + let template_result = match apply_template(full_tools_json) { + Ok(r) => { + let token_count = ctx + .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).unwrap_or(r) + } else { + r + } + } + Err(_) => apply_template(compact_tools).map_err(|e| { + ProviderError::ExecutionError(format!("Failed to apply chat template: {}", e)) + })?, + }; + + let _ = ctx.log.write( + &serde_json::json!({"applied_prompt": &template_result.prompt}), + None, + ); + + let tokens = ctx + .loaded + .model + .str_to_token(&template_result.prompt, AddBos::Never) + .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; + + let (prompt_token_count, effective_ctx) = validate_and_compute_context( + ctx.loaded, + ctx.runtime, + tokens.len(), + ctx.context_limit, + ctx.settings, + )?; + let mut llama_ctx = create_and_prefill_context( + ctx.loaded, + ctx.runtime, + &tokens, + effective_ctx, + ctx.settings, + )?; + + let message_id = ctx.message_id; + let tx = ctx.tx; + let mut generated_text = String::new(); + let mut streamed_len: usize = 0; + + let output_token_count = generation_loop( + &ctx.loaded.model, + &mut llama_ctx, + ctx.settings, + prompt_token_count, + effective_ctx, + |piece| { + generated_text.push_str(piece); + + 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.to_string()); + if tx.blocking_send(Ok((Some(msg), None))).is_err() { + return Ok(TokenAction::Stop); + } + } + streamed_len = stream_up_to; + } + + let should_stop = template_result + .additional_stops + .iter() + .any(|stop| generated_text.ends_with(stop)); + if should_stop { + Ok(TokenAction::Stop) + } else { + Ok(TokenAction::Continue) + } + }, + )?; + + 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) + }; + + 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.to_string()); + 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.to_string()); + let _ = tx.blocking_send(Ok((Some(msg), None))); + } + + let provider_usage = finalize_usage( + ctx.log, + std::mem::take(&mut ctx.model_name), + "native", + prompt_token_count, + output_token_count, + Some(("generated_text", &generated_text)), + ); + let _ = ctx.tx.blocking_send(Ok((None, Some(provider_usage)))); + Ok(()) +} 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..898a138912 --- /dev/null +++ b/crates/goose/src/providers/local_inference/local_model_registry.rs @@ -0,0 +1,319 @@ +use crate::config::paths::Paths; +use crate::download_manager::{get_download_manager, DownloadStatus}; +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, + #[serde(default)] + pub use_jinja: 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, + use_jinja: false, + } + } +} + +/// Featured models — HuggingFace specs in "author/repo-GGUF:quantization" format. +pub const FEATURED_MODELS: &[&str] = &[ + "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M", + "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", + "bartowski/Hermes-2-Pro-Mistral-7B-GGUF:Q4_K_M", + "bartowski/Mistral-Small-24B-Instruct-2501-GGUF:Q4_K_M", +]; + +/// Check if a model ID corresponds to a featured model. +pub fn is_featured_model(model_id: &str) -> bool { + use super::hf_models::parse_model_spec; + FEATURED_MODELS.iter().any(|spec| { + if let Ok((repo_id, quant)) = parse_model_spec(spec) { + model_id_from_repo(&repo_id, &quant) == model_id + } else { + false + } + }) +} + +static REGISTRY: OnceLock> = OnceLock::new(); + +pub fn get_registry() -> &'static Mutex { + REGISTRY.get_or_init(|| { + let registry = LocalModelRegistry::load().unwrap_or_default(); + 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, + #[serde(default)] + pub size_bytes: u64, +} + +impl LocalModelEntry { + pub fn is_downloaded(&self) -> bool { + self.local_path.exists() + } + + pub fn is_downloading(&self) -> bool { + let download_id = format!("{}-model", self.id); + let manager = get_download_manager(); + manager.get_progress(&download_id).is_some() + } + + pub fn download_status(&self) -> ModelDownloadStatus { + if self.local_path.exists() { + return ModelDownloadStatus::Downloaded; + } + + let download_id = format!("{}-model", self.id); + let manager = get_download_manager(); + if let Some(progress) = manager.get_progress(&download_id) { + return match progress.status { + DownloadStatus::Downloading => ModelDownloadStatus::Downloading { + progress_percent: progress.progress_percent, + bytes_downloaded: progress.bytes_downloaded, + total_bytes: progress.total_bytes, + speed_bps: progress.speed_bps.unwrap_or(0), + }, + DownloadStatus::Completed => ModelDownloadStatus::Downloaded, + DownloadStatus::Failed | DownloadStatus::Cancelled => { + ModelDownloadStatus::NotDownloaded + } + }; + } + + ModelDownloadStatus::NotDownloaded + } + + pub fn file_size(&self) -> u64 { + if self.size_bytes > 0 { + return self.size_bytes; + } + std::fs::metadata(&self.local_path) + .map(|m| m.len()) + .unwrap_or(0) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum ModelDownloadStatus { + NotDownloaded, + Downloading { + progress_percent: f32, + bytes_downloaded: u64, + total_bytes: u64, + speed_bps: u64, + }, + Downloaded, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct LocalModelRegistry { + pub models: Vec, +} + +impl LocalModelRegistry { + fn registry_path() -> PathBuf { + Paths::in_data_dir("models/registry.json") + } + + pub fn load() -> Result { + let path = Self::registry_path(); + if path.exists() { + let lock_path = path.with_extension("json.lock"); + let lock_file = std::fs::File::create(&lock_path)?; + fs2::FileExt::lock_shared(&lock_file)?; + let contents = std::fs::read_to_string(&path)?; + fs2::FileExt::unlock(&lock_file)?; + let registry: LocalModelRegistry = serde_json::from_str(&contents)?; + Ok(registry) + } else { + Ok(Self::default()) + } + } + + pub fn save(&self) -> Result<()> { + let path = Self::registry_path(); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + let lock_path = path.with_extension("json.lock"); + let lock_file = std::fs::File::create(&lock_path)?; + fs2::FileExt::lock_exclusive(&lock_file)?; + + let mut tmp = tempfile::NamedTempFile::new_in(path.parent().unwrap())?; + let contents = serde_json::to_string_pretty(self)?; + std::io::Write::write_all(&mut tmp, contents.as_bytes())?; + tmp.persist(&path)?; + + fs2::FileExt::unlock(&lock_file)?; + Ok(()) + } + + /// Sync registry with featured models: + /// add any featured models that are missing, remove non-downloaded non-featured models. + pub fn sync_with_featured(&mut self, featured_entries: Vec) { + let mut changed = false; + + for entry in featured_entries { + if !self.models.iter().any(|m| m.id == entry.id) { + self.models.push(entry); + changed = true; + } + } + + let before_len = self.models.len(); + self.models + .retain(|m| m.is_downloaded() || m.is_downloading() || is_featured_model(&m.id)); + if self.models.len() != before_len { + 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 has_model(&self, id: &str) -> bool { + self.models.iter().any(|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. +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/crates/goose/src/providers/local_inference/tool_parsing.rs b/crates/goose/src/providers/local_inference/tool_parsing.rs new file mode 100644 index 0000000000..cf6124a6e1 --- /dev/null +++ b/crates/goose/src/providers/local_inference/tool_parsing.rs @@ -0,0 +1,601 @@ +use crate::conversation::message::{Message, MessageContent}; +use rmcp::model::{CallToolRequestParams, Tool}; +use serde_json::{json, Value}; +use std::borrow::Cow; +use uuid::Uuid; + +pub(super) 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() +} + +/// 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)] +pub(super) 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. +pub(super) 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). +pub(super) 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)] +pub(super) 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)> { + // Try V... format first + if let Some(result) = parse_xml_function_format(block) { + return Some(result); + } + // Try GLM-style: TOOL_NAMEKV... + parse_xml_arg_key_value_format(block) +} + +fn parse_xml_function_format(block: &str) -> Option<(String, serde_json::Map)> { + 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; + + 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)) +} + +/// Parse GLM-style tool calls: `NAMEKV...` +/// Also handles zero-argument calls like just `NAME`. +fn parse_xml_arg_key_value_format(block: &str) -> Option<(String, serde_json::Map)> { + let func_name_end = block.find("").unwrap_or(block.len()); + // Safe: find returns a byte offset at the start of an ASCII '<' character, + // and block.len() is always a valid boundary. + #[allow(clippy::string_slice)] + let func_name = block[..func_name_end].trim().to_string(); + if func_name.is_empty() { + return None; + } + + let mut args = serde_json::Map::new(); + #[allow(clippy::string_slice)] + let mut rest = &block[func_name_end..]; + + while let Some((_, after_key_open)) = rest.split_once("") { + let Some((key, after_key_close)) = after_key_open.split_once("") else { + break; + }; + let key = key.trim().to_string(); + + let Some((_, after_val_open)) = after_key_close.split_once("") else { + break; + }; + let (value, after_val_close) = after_val_open + .split_once("") + .unwrap_or((after_val_open, "")); + let value = value.trim(); + + let json_value = + serde_json::from_str(value).unwrap_or_else(|_| Value::String(value.to_string())); + args.insert(key, json_value); + + rest = after_val_close; + } + + Some((func_name, args)) +} + +pub(super) 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() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + const SHELL_TOOL: &str = "developer__shell"; + + #[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, SHELL_TOOL); + 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 result = split_content_and_xml_tool_calls(text); + assert!(result.is_some()); + let (content, calls) = result.unwrap(); + assert_eq!(content, "Let me check."); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "execute"); + assert_eq!( + calls[0].1.get("code").unwrap(), + "async function run() { return 1; }" + ); + // tool_graph should be parsed as JSON array + assert!(calls[0].1.get("tool_graph").unwrap().is_array()); + } + + #[test] + fn test_extract_xml_tool_call_messages() { + let calls = vec![( + SHELL_TOOL.to_string(), + serde_json::Map::from_iter(vec![("command".to_string(), json!("ls"))]), + )]; + let msgs = extract_xml_tool_call_messages(calls, "test-id"); + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0].id, Some("test-id".to_string())); + match &msgs[0].content[0] { + MessageContent::ToolRequest(req) => { + let call = req.tool_call.as_ref().unwrap(); + assert_eq!(&*call.name, SHELL_TOOL); + assert_eq!( + call.arguments.as_ref().unwrap().get("command").unwrap(), + "ls" + ); + } + _ => panic!("Expected ToolRequest"), + } + } + + #[test] + fn test_split_content_and_tool_calls_with_tool() { + let text = "Here is the result.\n{\"tool_calls\": [{\"function\": {\"name\": \"shell\", \"arguments\": \"{}\"}, \"id\": \"abc\"}]}"; + let (content, tc) = split_content_and_tool_calls(text); + assert_eq!(content, "Here is the result."); + assert!(tc.is_some()); + let parsed: serde_json::Value = serde_json::from_str(&tc.unwrap()).unwrap(); + assert_eq!(parsed["tool_calls"].as_array().unwrap().len(), 1); + } + + #[test] + fn test_split_content_and_tool_calls_no_tool() { + let text = "Just regular text, no JSON."; + let (content, tc) = split_content_and_tool_calls(text); + assert_eq!(content, text); + assert!(tc.is_none()); + } + + #[test] + fn test_split_content_and_tool_calls_json_without_tool_calls_key() { + let text = "{\"key\": \"value\"}"; + let (content, tc) = split_content_and_tool_calls(text); + assert_eq!(content, text); + assert!(tc.is_none()); + } + + #[test] + fn test_extract_tool_call_messages_openai_format() { + let json_str = r#"{"tool_calls": [{"function": {"name": "developer__shell", "arguments": "{\"command\": \"ls\"}"}, "id": "call-1"}]}"#; + let msgs = extract_tool_call_messages(json_str, "msg-1"); + assert_eq!(msgs.len(), 1); + match &msgs[0].content[0] { + MessageContent::ToolRequest(req) => { + let call = req.tool_call.as_ref().unwrap(); + assert_eq!(&*call.name, SHELL_TOOL); + assert_eq!( + call.arguments.as_ref().unwrap().get("command").unwrap(), + "ls" + ); + } + _ => panic!("Expected ToolRequest"), + } + } + + #[test] + fn test_extract_tool_call_messages_native_format() { + let json_str = r#"{"tool_calls": [{"name": "developer__shell", "arguments": {"command": "ls"}, "id": "call-2"}]}"#; + let msgs = extract_tool_call_messages(json_str, "msg-2"); + assert_eq!(msgs.len(), 1); + match &msgs[0].content[0] { + MessageContent::ToolRequest(req) => { + let call = req.tool_call.as_ref().unwrap(); + assert_eq!(&*call.name, SHELL_TOOL); + } + _ => panic!("Expected ToolRequest"), + } + } + + #[test] + fn test_extract_tool_call_messages_invalid_json() { + assert!(extract_tool_call_messages("not json", "msg-3").is_empty()); + } + + #[test] + fn test_extract_tool_call_messages_empty_name_skipped() { + let json_str = r#"{"tool_calls": [{"name": "", "arguments": {}, "id": "x"}]}"#; + assert!(extract_tool_call_messages(json_str, "msg-4").is_empty()); + } + + #[test] + fn test_compact_tools_json_produces_minimal_output() { + use rmcp::model::Tool; + use rmcp::object; + + let tools = vec![Tool::new( + "developer__shell".to_string(), + "Run shell commands".to_string(), + object!({"type": "object", "properties": {"command": {"type": "string"}}}), + )]; + let result = compact_tools_json(&tools); + assert!(result.is_some()); + let parsed: Vec = serde_json::from_str(&result.unwrap()).unwrap(); + assert_eq!(parsed.len(), 1); + let func = &parsed[0]["function"]; + assert_eq!(func["name"], "developer__shell"); + assert_eq!(func["description"], "Run shell commands"); + // Should not contain full parameter schemas + assert!(func.get("parameters").is_none()); + } + + #[test] + fn test_compact_tools_json_empty() { + let result = compact_tools_json(&[]); + assert!(result.is_some()); + let parsed: Vec = serde_json::from_str(&result.unwrap()).unwrap(); + assert!(parsed.is_empty()); + } + + #[test] + fn test_safe_stream_end_balanced_braces() { + let text = "Result: {\"key\": \"value\"} done"; + assert_eq!(safe_stream_end(text), text.len()); + } + + #[test] + fn test_safe_stream_end_unbalanced_open_brace() { + let text = "Some text {\"tool_calls\": ["; + assert_eq!(safe_stream_end(text), "Some text ".len()); + } + + #[test] + fn test_safe_stream_end_empty() { + assert_eq!(safe_stream_end(""), 0); + } + + #[test] + fn test_safe_stream_end_no_braces() { + let text = "plain text here"; + assert_eq!(safe_stream_end(text), text.len()); + } +} diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 6e4862f0df..581fc2747e 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -23,6 +23,7 @@ pub mod google; mod init; pub mod lead_worker; pub mod litellm; +pub mod local_inference; pub mod oauth; pub mod ollama; pub mod openai; diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index 30a53b5c54..7b6104d658 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -7,16 +7,13 @@ use super::utils::{ImageFormat, RequestLog}; use crate::config::declarative_providers::DeclarativeProviderConfig; use crate::config::GooseMode; use crate::conversation::message::Message; -use crate::conversation::Conversation; use crate::model::ModelConfig; use crate::providers::formats::ollama::{create_request, response_to_streaming_message_ollama}; -use crate::utils::safe_truncate; use anyhow::{Error, Result}; use async_stream::try_stream; use async_trait::async_trait; use futures::future::BoxFuture; use futures::TryStreamExt; -use regex::Regex; use reqwest::Response; use rmcp::model::Tool; use serde_json::{json, Value}; @@ -215,30 +212,6 @@ impl Provider for OllamaProvider { self.model.clone() } - async fn generate_session_name( - &self, - session_id: &str, - messages: &Conversation, - ) -> Result { - let context = self.get_initial_user_messages(messages); - let message = Message::user().with_text(self.create_session_name_prompt(&context)); - let model_config = self.get_model_config(); - let result = self - .complete( - &model_config, - session_id, - "You are a title generator. Output only the requested title of 4 words or less, with no additional text, reasoning, or explanations.", - &[message], - &[], - ) - .await?; - - let mut description = result.0.as_concat_text(); - description = Self::filter_reasoning_tokens(&description); - - Ok(safe_truncate(&description, 100)) - } - async fn stream( &self, model_config: &ModelConfig, @@ -318,46 +291,6 @@ impl Provider for OllamaProvider { } } -impl OllamaProvider { - fn filter_reasoning_tokens(text: &str) -> String { - let mut filtered = text.to_string(); - - let reasoning_patterns = [ - r".*?", - r".*?", - r"Let me think.*?\n", - r"I need to.*?\n", - r"First, I.*?\n", - r"Okay, .*?\n", - r"So, .*?\n", - r"Well, .*?\n", - r"Hmm, .*?\n", - r"Actually, .*?\n", - r"Based on.*?I think", - r"Looking at.*?I would say", - ]; - - for pattern in reasoning_patterns { - if let Ok(re) = Regex::new(pattern) { - filtered = re.replace_all(&filtered, "").to_string(); - } - } - filtered = filtered - .replace("", "") - .replace("", "") - .replace("", "") - .replace("", ""); - filtered = filtered - .lines() - .map(|line| line.trim()) - .filter(|line| !line.is_empty()) - .collect::>() - .join(" "); - - filtered - } -} - /// Ollama-specific streaming handler with XML tool call fallback. /// Uses the Ollama format module which buffers text when XML tool calls are detected, /// preventing duplicate content from being emitted to the UI. diff --git a/crates/goose/tests/local_inference_integration.rs b/crates/goose/tests/local_inference_integration.rs new file mode 100644 index 0000000000..35002904a5 --- /dev/null +++ b/crates/goose/tests/local_inference_integration.rs @@ -0,0 +1,128 @@ +//! Integration tests for LocalInferenceProvider. +//! +//! These tests require a downloaded GGUF model and are ignored by default. +//! Run with: cargo test -p goose --test local_inference_integration -- --ignored + +use futures::StreamExt; +use goose::conversation::message::Message; +use goose::model::ModelConfig; +use goose::providers::create; +use std::time::Instant; + +const TEST_MODEL: &str = "llama-3.2-1b"; + +#[tokio::test] +#[ignore] +async fn test_local_inference_stream_produces_output() { + let model_config = ModelConfig::new(TEST_MODEL).expect("valid model config"); + let provider = create("local", model_config.clone(), Vec::new()) + .await + .expect("provider creation should succeed"); + + let system = "You are a helpful assistant. Be brief."; + let messages = vec![Message::user().with_text("Say hello.")]; + + let mut stream = provider + .stream(&model_config, "test-session", system, &messages, &[]) + .await + .expect("stream should start"); + + let mut got_text = false; + let mut got_usage = false; + + while let Some(result) = stream.next().await { + let (msg, usage) = result.expect("stream item should be Ok"); + if msg.is_some() { + got_text = true; + } + if let Some(u) = usage { + got_usage = true; + let usage_inner = u.usage; + assert!( + usage_inner.input_tokens.unwrap_or(0) > 0, + "should have input tokens" + ); + assert!( + usage_inner.output_tokens.unwrap_or(0) > 0, + "should have output tokens" + ); + } + } + + assert!(got_text, "stream should produce at least one text message"); + assert!(got_usage, "stream should produce usage info"); +} + +#[tokio::test] +#[ignore] +async fn test_local_inference_cold_and_warm_performance() { + let model_config = ModelConfig::new(TEST_MODEL).expect("valid model config"); + let provider = create("local", model_config.clone(), Vec::new()) + .await + .expect("provider creation should succeed"); + + // Cold start (includes model loading) + let messages = vec![Message::user().with_text("what is the capital of Moldova?")]; + let start = Instant::now(); + let (response, _usage) = provider + .complete(&model_config, "test-session", "", &messages, &[]) + .await + .expect("cold completion should succeed"); + let cold_elapsed = start.elapsed(); + + let text = response.as_concat_text(); + assert!(!text.is_empty(), "cold start should produce a response"); + println!( + "Cold start: {cold_elapsed:.2?}, response length: {}", + text.len() + ); + + // Warm run (model already loaded) + let messages2 = vec![Message::user().with_text("what is the capital of France?")]; + let start2 = Instant::now(); + let (response2, _usage2) = provider + .complete(&model_config, "test-session", "", &messages2, &[]) + .await + .expect("warm completion should succeed"); + let warm_elapsed = start2.elapsed(); + + let text2 = response2.as_concat_text(); + assert!(!text2.is_empty(), "warm run should produce a response"); + println!( + "Warm run: {warm_elapsed:.2?}, response length: {}", + text2.len() + ); + assert!( + warm_elapsed < cold_elapsed, + "warm run ({warm_elapsed:.2?}) should be faster than cold start ({cold_elapsed:.2?})" + ); +} + +#[tokio::test] +#[ignore] +async fn test_local_inference_large_prompt() { + let model_config = ModelConfig::new(TEST_MODEL).expect("valid model config"); + let provider = create("local", model_config.clone(), Vec::new()) + .await + .expect("provider creation should succeed"); + + // Build a large prompt (~3500 tokens) to exercise prefill performance + let padding = "You are Goose, a highly capable AI assistant.\n".repeat(80); + let prompt = format!("{padding}\nNow answer this: what is the capital of Moldova?"); + let messages = vec![Message::user().with_text(&prompt)]; + + let start = Instant::now(); + let (response, _usage) = provider + .complete(&model_config, "test-session", "", &messages, &[]) + .await + .expect("large prompt completion should succeed"); + let elapsed = start.elapsed(); + + let text = response.as_concat_text(); + assert!(!text.is_empty(), "large prompt should produce a response"); + println!( + "Large prompt: {elapsed:.2?}, prompt ~{} chars, response length: {}", + prompt.len(), + text.len() + ); +} diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index fb9712d723..2fee277007 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -1827,6 +1827,308 @@ } } }, + "/local-inference/download": { + "post": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "download_hf_model", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DownloadModelRequest" + } + } + }, + "required": true + }, + "responses": { + "202": { + "description": "Download started", + "content": { + "text/plain": { + "schema": { + "type": "string" + } + } + } + }, + "400": { + "description": "Invalid request" + } + } + } + }, + "/local-inference/models": { + "get": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "list_local_models", + "responses": { + "200": { + "description": "List of available local LLM models", + "content": { + "application/json": { + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/LocalModelResponse" + } + } + } + } + } + } + } + }, + "/local-inference/models/{model_id}": { + "delete": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "delete_local_model", + "parameters": [ + { + "name": "model_id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Model deleted" + }, + "404": { + "description": "Model not found" + } + } + } + }, + "/local-inference/models/{model_id}/download": { + "get": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "get_local_model_download_progress", + "parameters": [ + { + "name": "model_id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Download progress", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DownloadProgress" + } + } + } + }, + "404": { + "description": "No active download" + } + } + }, + "delete": { + "tags": [ + "super::routes::local_inference" + ], + "operationId": "cancel_local_model_download", + "parameters": [ + { + "name": "model_id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Download cancelled" + }, + "404": { + "description": "No active download" + } + } + } + }, + "/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 files in the repo", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RepoVariantsResponse" + } + } + } + } + } + } + }, + "/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": [ @@ -3954,6 +4256,18 @@ } } }, + "DownloadModelRequest": { + "type": "object", + "required": [ + "spec" + ], + "properties": { + "spec": { + "type": "string", + "description": "Model spec like \"bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M\"" + } + } + }, "DownloadProgress": { "type": "object", "required": [ @@ -4588,6 +4902,100 @@ } ] }, + "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": [ @@ -4773,6 +5181,51 @@ } } }, + "LocalModelResponse": { + "type": "object", + "required": [ + "id", + "display_name", + "repo_id", + "filename", + "quantization", + "size_bytes", + "status", + "recommended", + "settings" + ], + "properties": { + "display_name": { + "type": "string" + }, + "filename": { + "type": "string" + }, + "id": { + "type": "string" + }, + "quantization": { + "type": "string" + }, + "recommended": { + "type": "boolean" + }, + "repo_id": { + "type": "string" + }, + "settings": { + "$ref": "#/components/schemas/ModelSettings" + }, + "size_bytes": { + "type": "integer", + "format": "int64", + "minimum": 0 + }, + "status": { + "$ref": "#/components/schemas/ModelDownloadStatus" + } + } + }, "McpAppResource": { "type": "object", "description": "MCP App Resource\nRepresents a UI resource that can be rendered in an MCP App", @@ -5293,6 +5746,78 @@ } } }, + "ModelDownloadStatus": { + "oneOf": [ + { + "type": "object", + "required": [ + "state" + ], + "properties": { + "state": { + "type": "string", + "enum": [ + "NotDownloaded" + ] + } + } + }, + { + "type": "object", + "required": [ + "progress_percent", + "bytes_downloaded", + "total_bytes", + "state" + ], + "properties": { + "bytes_downloaded": { + "type": "integer", + "format": "int64", + "minimum": 0 + }, + "progress_percent": { + "type": "number", + "format": "float" + }, + "speed_bps": { + "type": "integer", + "format": "int64", + "nullable": true, + "minimum": 0 + }, + "state": { + "type": "string", + "enum": [ + "Downloading" + ] + }, + "total_bytes": { + "type": "integer", + "format": "int64", + "minimum": 0 + } + } + }, + { + "type": "object", + "required": [ + "state" + ], + "properties": { + "state": { + "type": "string", + "enum": [ + "Downloaded" + ] + } + } + } + ], + "discriminator": { + "propertyName": "state" + } + }, "ModelInfo": { "type": "object", "description": "Information about a model's capabilities", @@ -5417,6 +5942,71 @@ } } }, + "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_jinja": { + "type": "boolean" + }, + "use_mlock": { + "type": "boolean" + } + } + }, "ParseRecipeRequest": { "type": "object", "required": [ @@ -6005,6 +6595,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": [ { @@ -6196,6 +6805,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 9a24bb7ac0..2e65e54472 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, checkProvider, configureProviderOauth, confirmToolAction, createCustomProvider, createRecipe, createSchedule, decodeRecipe, deleteModel, deleteRecipe, deleteSchedule, deleteSession, detectProvider, diagnostics, downloadModel, encodeRecipe, exportApp, exportSession, forkSession, getCanonicalModelInfo, getCustomProvider, getDictationConfig, getDownloadProgress, getExtensions, getPrompt, getPrompts, getProviderModels, getSession, getSessionExtensions, getSessionInsights, getSlashCommands, getTools, getTunnelStatus, importApp, importSession, initConfig, inspectRunningJob, killRunningJob, listApps, 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, 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, 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, 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, GetCanonicalModelInfoData, GetCanonicalModelInfoResponse, GetCanonicalModelInfoResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponse, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponse, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponse, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponse, GetExtensionsResponses, 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, ListModelsData, ListModelsResponse, ListModelsResponses, ListRecipeResponse, ListRecipesData, ListRecipesErrors, ListRecipesResponse, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponse, ListSchedulesResponse2, ListSchedulesResponses, ListSessionsData, ListSessionsErrors, ListSessionsResponse, ListSessionsResponses, LoadedProvider, McpAppResource, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, Message, MessageContent, MessageEvent, MessageMetadata, ModelConfig, ModelInfo, ModelInfoData, ModelInfoQuery, ModelInfoResponse, ParseRecipeData, ParseRecipeError, ParseRecipeErrors, ParseRecipeRequest, ParseRecipeResponse, ParseRecipeResponse2, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponse, PauseScheduleResponses, Permission, PermissionLevel, PermissionsMetadata, 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, ReasoningContent, 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, TaskSupport, TelemetryEventRequest, Template, TextContent, ThinkingContent, TokenState, Tool, ToolAnnotations, ToolConfirmationRequest, ToolExecution, 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, downloadModel, encodeRecipe, exportApp, exportSession, forkSession, getCanonicalModelInfo, getCustomProvider, getDictationConfig, getDownloadProgress, getExtensions, getLocalModelDownloadProgress, getModelSettings, 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, DownloadModelData, DownloadModelErrors, DownloadModelRequest, 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, GetCanonicalModelInfoData, GetCanonicalModelInfoResponse, GetCanonicalModelInfoResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponse, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponse, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponse, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponse, GetExtensionsResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponse, GetLocalModelDownloadProgressResponses, GetModelSettingsData, GetModelSettingsErrors, GetModelSettingsResponse, GetModelSettingsResponses, GetPromptData, GetPromptErrors, GetPromptResponse, GetPromptResponses, GetPromptsData, GetPromptsResponse, GetPromptsResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponse, GetProviderModelsResponses, GetRepoFilesData, 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, 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, LocalModelResponse, McpAppResource, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, Message, MessageContent, MessageEvent, MessageMetadata, ModelConfig, ModelDownloadStatus, ModelInfo, ModelInfoData, ModelInfoQuery, ModelInfoResponse, ModelSettings, ParseRecipeData, ParseRecipeError, ParseRecipeErrors, ParseRecipeRequest, ParseRecipeResponse, ParseRecipeResponse2, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponse, PauseScheduleResponses, Permission, PermissionLevel, PermissionsMetadata, 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, ReasoningContent, 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, 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, TaskSupport, TelemetryEventRequest, Template, TextContent, ThinkingContent, TokenState, Tool, ToolAnnotations, ToolConfirmationRequest, ToolExecution, 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 012124b3b4..da2dc61d4b 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, CheckProviderData, ConfigureProviderOauthData, ConfigureProviderOauthErrors, ConfigureProviderOauthResponses, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionResponses, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleResponses, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeResponses, DeleteModelData, DeleteModelErrors, DeleteModelResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponses, DeleteSessionData, DeleteSessionErrors, DeleteSessionResponses, DetectProviderData, DetectProviderErrors, DetectProviderResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponses, DownloadModelData, DownloadModelErrors, DownloadModelResponses, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeResponses, ExportAppData, ExportAppErrors, ExportAppResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponses, ForkSessionData, ForkSessionErrors, ForkSessionResponses, GetCanonicalModelInfoData, GetCanonicalModelInfoResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponses, 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, 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, DownloadModelData, DownloadModelErrors, DownloadModelResponses, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeResponses, ExportAppData, ExportAppErrors, ExportAppResponses, ExportSessionData, ExportSessionErrors, ExportSessionResponses, ForkSessionData, ForkSessionErrors, ForkSessionResponses, GetCanonicalModelInfoData, GetCanonicalModelInfoResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponses, GetModelSettingsData, GetModelSettingsErrors, GetModelSettingsResponses, GetPromptData, GetPromptErrors, GetPromptResponses, GetPromptsData, GetPromptsResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponses, GetRepoFilesData, 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,38 @@ 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 }); + +export const cancelLocalModelDownload = (options: Options) => (options.client ?? client).delete({ url: '/local-inference/models/{model_id}/download', ...options }); + +export const getLocalModelDownloadProgress = (options: Options) => (options.client ?? client).get({ 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 22524f9403..a4bc570102 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -232,6 +232,13 @@ export type DictationProviderStatus = { uses_provider_config: boolean; }; +export type DownloadModelRequest = { + /** + * Model spec like "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M" + */ + spec: string; +}; + export type DownloadProgress = { /** * Bytes downloaded so far @@ -446,6 +453,36 @@ export type GooseApp = McpAppResource & (WindowProps | null) & { prd?: string | null; }; +/** + * 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; @@ -511,6 +548,18 @@ export type LoadedProvider = { is_editable: boolean; }; +export type LocalModelResponse = { + display_name: string; + filename: string; + id: string; + quantization: string; + recommended: boolean; + repo_id: string; + settings: ModelSettings; + size_bytes: number; + status: ModelDownloadStatus; +}; + /** * MCP App Resource * Represents a UI resource that can be rendered in an MCP App @@ -638,6 +687,18 @@ export type ModelConfig = { toolshim_model?: string | null; }; +export type ModelDownloadStatus = { + state: 'NotDownloaded'; +} | { + bytes_downloaded: number; + progress_percent: number; + speed_bps?: number | null; + state: 'Downloading'; + total_bytes: number; +} | { + state: 'Downloaded'; +}; + /** * Information about a model's capabilities */ @@ -690,6 +751,23 @@ export type ModelInfoResponse = { source: string; }; +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_jinja?: boolean; + use_mlock?: boolean; +}; + export type ParseRecipeRequest = { content: string; }; @@ -909,6 +987,11 @@ export type RemoveExtensionRequest = { session_id: string; }; +export type RepoVariantsResponse = { + recommended_index?: number | null; + variants: Array; +}; + export type ResourceContents = { _meta?: { [key: string]: unknown; @@ -986,6 +1069,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; }; @@ -2836,6 +2935,221 @@ export type StartTetrateSetupResponses = { export type StartTetrateSetupResponse = StartTetrateSetupResponses[keyof StartTetrateSetupResponses]; +export type DownloadHfModelData = { + body: DownloadModelRequest; + path?: never; + query?: never; + url: '/local-inference/download'; +}; + +export type DownloadHfModelErrors = { + /** + * Invalid request + */ + 400: unknown; +}; + +export type DownloadHfModelResponses = { + /** + * Download started + */ + 202: string; +}; + +export type DownloadHfModelResponse = DownloadHfModelResponses[keyof DownloadHfModelResponses]; + +export type ListLocalModelsData = { + body?: never; + path?: never; + query?: never; + url: '/local-inference/models'; +}; + +export type ListLocalModelsResponses = { + /** + * List of available local LLM models + */ + 200: Array; +}; + +export type ListLocalModelsResponse = ListLocalModelsResponses[keyof ListLocalModelsResponses]; + +export type DeleteLocalModelData = { + body?: never; + path: { + model_id: string; + }; + query?: never; + url: '/local-inference/models/{model_id}'; +}; + +export type DeleteLocalModelErrors = { + /** + * Model not found + */ + 404: unknown; +}; + +export type DeleteLocalModelResponses = { + /** + * Model deleted + */ + 200: unknown; +}; + +export type CancelLocalModelDownloadData = { + body?: never; + path: { + model_id: string; + }; + query?: never; + url: '/local-inference/models/{model_id}/download'; +}; + +export type CancelLocalModelDownloadErrors = { + /** + * No active download + */ + 404: unknown; +}; + +export type CancelLocalModelDownloadResponses = { + /** + * Download cancelled + */ + 200: unknown; +}; + +export type GetLocalModelDownloadProgressData = { + body?: never; + path: { + model_id: string; + }; + query?: never; + url: '/local-inference/models/{model_id}/download'; +}; + +export type GetLocalModelDownloadProgressErrors = { + /** + * No active download + */ + 404: unknown; +}; + +export type GetLocalModelDownloadProgressResponses = { + /** + * Download progress + */ + 200: DownloadProgress; +}; + +export type GetLocalModelDownloadProgressResponse = GetLocalModelDownloadProgressResponses[keyof GetLocalModelDownloadProgressResponses]; + +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 GetRepoFilesResponses = { + /** + * GGUF files in the 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/LocalModelSetup.tsx b/ui/desktop/src/components/LocalModelSetup.tsx new file mode 100644 index 0000000000..4fe469d3a8 --- /dev/null +++ b/ui/desktop/src/components/LocalModelSetup.tsx @@ -0,0 +1,397 @@ +import { useState, useEffect, useCallback, useRef } from 'react'; +import { useConfig } from './ConfigContext'; +import { + listLocalModels, + downloadHfModel, + getLocalModelDownloadProgress, + cancelLocalModelDownload, + type DownloadProgress, + type LocalModelResponse, +} from '../api'; +import { toastService } from '../toasts'; +import { trackOnboardingSetupFailed } from '../utils/analytics'; +import { Goose } from './icons'; + +interface LocalModelSetupProps { + onSuccess: () => void; + onCancel: () => void; +} + +const formatBytes = (bytes: number): string => { + 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 formatSize = (bytes: number): string => { + const mb = bytes / (1024 * 1024); + return mb >= 1024 ? `${(mb / 1024).toFixed(1)}GB` : `${mb.toFixed(0)}MB`; +}; + +type SetupPhase = 'loading' | 'select' | 'downloading' | 'error'; + +export function LocalModelSetup({ onSuccess, onCancel }: LocalModelSetupProps) { + const { upsert } = useConfig(); + const [phase, setPhase] = useState('loading'); + const [models, setModels] = useState([]); + const [selectedModelId, setSelectedModelId] = useState(null); + const [downloadProgress, setDownloadProgress] = useState(null); + const [errorMessage, setErrorMessage] = useState(null); + const [showAllModels, setShowAllModels] = useState(false); + const pollRef = useRef | null>(null); + + const cleanup = useCallback(() => { + if (pollRef.current) { + clearInterval(pollRef.current); + pollRef.current = null; + } + }, []); + + useEffect(() => cleanup, [cleanup]); + + useEffect(() => { + const load = async () => { + try { + const response = await listLocalModels(); + if (response.data) { + setModels(response.data); + + const alreadyDownloaded = response.data.find((m) => m.status.state === 'Downloaded'); + if (alreadyDownloaded) { + setSelectedModelId(alreadyDownloaded.id); + } else { + const recommended = response.data.find((m: LocalModelResponse) => m.recommended); + if (recommended) setSelectedModelId(recommended.id); + } + } + } catch (error) { + console.error('Failed to load local models:', error); + setErrorMessage('Failed to load available models. Please try again.'); + setPhase('error'); + return; + } + setPhase('select'); + }; + load(); + }, []); + + const finishSetup = async (modelId: string) => { + await upsert('GOOSE_PROVIDER', 'local', false); + await upsert('GOOSE_MODEL', modelId, false); + toastService.success({ + title: 'Local Model Ready', + msg: `Running entirely on your machine with ${modelId}.`, + }); + onSuccess(); + }; + + const startDownload = async (modelId: string) => { + setPhase('downloading'); + setDownloadProgress(null); + setErrorMessage(null); + + const model = models.find((m) => m.id === modelId); + if (!model) { + setErrorMessage('Model not found'); + setPhase('error'); + return; + } + + try { + await downloadHfModel({ body: { spec: model.id } }); + } catch (error) { + console.error('Failed to start download:', error); + setErrorMessage('Failed to start download. Please try again.'); + trackOnboardingSetupFailed('local', 'download_start_failed'); + setPhase('error'); + return; + } + + pollRef.current = setInterval(async () => { + try { + const response = await getLocalModelDownloadProgress({ path: { model_id: modelId } }); + if (response.data) { + setDownloadProgress(response.data); + if (response.data.status === 'completed') { + cleanup(); + await finishSetup(modelId); + } else if (response.data.status === 'failed') { + cleanup(); + setErrorMessage(response.data.error || 'Download failed.'); + trackOnboardingSetupFailed('local', response.data.error || 'download_failed'); + setPhase('error'); + } else if (response.data.status === 'cancelled') { + cleanup(); + setPhase('select'); + } + } + } catch { + cleanup(); + setErrorMessage('Lost connection to download. Please try again.'); + trackOnboardingSetupFailed('local', 'progress_poll_failed'); + setPhase('error'); + } + }, 500); + }; + + const handleCancel = async () => { + if (phase === 'downloading' && selectedModelId) { + cleanup(); + try { + await cancelLocalModelDownload({ path: { model_id: selectedModelId } }); + } catch { + // best-effort + } + setDownloadProgress(null); + setPhase('select'); + } else { + onCancel(); + } + }; + + const handlePrimaryAction = async () => { + if (!selectedModelId) return; + const model = models.find((m) => m.id === selectedModelId); + if (!model) return; + if (model.status.state === 'Downloaded') { + await finishSetup(model.id); + } else { + await startDownload(model.id); + } + }; + + const recommended = models.find((m) => m.recommended); + const otherModels = models.filter((m) => m.id !== recommended?.id); + const selectedModel = models.find((m) => m.id === selectedModelId); + + if (phase === 'loading') { + return ( +
+
+

Checking available models...

+
+ ); + } + + return ( +
+ {/* Header */} +
+
+ +
+

Run Locally

+

+ Download a model to run Goose entirely on your machine — no API keys, no accounts, completely free and private. +

+
+ + {/* Error state */} + {phase === 'error' && ( +
+
+

{errorMessage}

+
+ + +
+ )} + + {/* Model selection */} + {phase === 'select' && ( +
+ {/* Recommended model card */} + {recommended && ( +
setSelectedModelId(recommended.id)} + className={`relative w-full p-4 sm:p-6 border rounded-xl cursor-pointer transition-all duration-200 group ${ + selectedModelId === recommended.id + ? 'border-blue-500 bg-blue-500/5' + : 'border-border-subtle hover:border-border-default' + }`} + > +
+ + Best for your machine + +
+
+ setSelectedModelId(recommended.id)} + className="cursor-pointer flex-shrink-0 mt-1" + /> +
+
+ + {recommended.display_name} + + {recommended.status.state === 'Downloaded' && ( + + Ready + + )} +
+

+ {formatSize(recommended.size_bytes)} +

+
+
+
+ )} + + {/* Expandable other models */} + {otherModels.length > 0 && ( +
+ + + {showAllModels && ( +
+ {otherModels.map((model) => ( +
setSelectedModelId(model.id)} + className={`w-full p-4 border rounded-xl cursor-pointer transition-all duration-200 ${ + selectedModelId === model.id + ? 'border-blue-500 bg-blue-500/5' + : 'border-border-subtle hover:border-border-default' + }`} + > +
+ setSelectedModelId(model.id)} + className="cursor-pointer flex-shrink-0 mt-0.5" + /> +
+
+ {model.display_name} + {formatSize(model.size_bytes)} + {model.status.state === 'Downloaded' && ( + + Ready + + )} +
+
+
+
+ ))} +
+ )} +
+ )} + + {/* Primary action */} + + + +
+ )} + + {/* Downloading state */} + {phase === 'downloading' && selectedModel && ( +
+
+

+ Downloading {selectedModel.display_name} +

+ + {downloadProgress ? ( +
+ {/* Progress bar */} +
+
+
+ + {/* Stats row */} +
+ + {formatBytes(downloadProgress.bytes_downloaded)} of{' '} + {formatBytes(downloadProgress.total_bytes)} + + {downloadProgress.progress_percent.toFixed(0)}% +
+ +
+ {downloadProgress.speed_bps ? ( + {formatBytes(downloadProgress.speed_bps)}/s + ) : ( + + )} + {downloadProgress.eta_seconds != null && downloadProgress.eta_seconds > 0 && ( + + ~{downloadProgress.eta_seconds < 60 + ? `${Math.round(downloadProgress.eta_seconds)}s` + : `${Math.round(downloadProgress.eta_seconds / 60)}m`}{' '} + remaining + + )} +
+
+ ) : ( +
+
+ Starting download... +
+ )} +
+ + +
+ )} +
+ ); +} diff --git a/ui/desktop/src/components/McpApps/McpAppRenderer.tsx b/ui/desktop/src/components/McpApps/McpAppRenderer.tsx index cd1646fcbb..4733727359 100644 --- a/ui/desktop/src/components/McpApps/McpAppRenderer.tsx +++ b/ui/desktop/src/components/McpApps/McpAppRenderer.tsx @@ -15,7 +15,7 @@ * - "standalone" — Goose-specific mode for dedicated Electron windows */ -import { AppRenderer, type RequestHandlerExtra } from '@mcp-ui/client'; +import { AppRenderer } from '@mcp-ui/client'; import type { McpUiDisplayMode, McpUiHostContext, @@ -23,7 +23,7 @@ import type { McpUiResourcePermissions, McpUiSizeChangedNotification, } from '@modelcontextprotocol/ext-apps/app-bridge'; -import type { CallToolResult, JSONRPCRequest } from '@modelcontextprotocol/sdk/types.js'; +import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; import { useCallback, useEffect, useMemo, useReducer, useRef, useState } from 'react'; import { callTool, readResource } from '../../api'; import { AppEvents } from '../../constants/events'; @@ -40,8 +40,6 @@ import { McpAppToolInputPartial, McpAppToolResult, DimensionLayout, - SamplingCreateMessageParams, - SamplingCreateMessageResponse, } from './types'; const DEFAULT_IFRAME_HEIGHT = 200; @@ -267,13 +265,7 @@ export default function McpAppRenderer({ const containerRef = useRef(null); const [containerWidth, setContainerWidth] = useState(0); const [containerHeight, setContainerHeight] = useState(0); - const [apiHost, setApiHost] = useState(null); - const [secretKey, setSecretKey] = useState(null); - useEffect(() => { - window.electron.getGoosedHostPort().then(setApiHost); - window.electron.getSecretKey().then(setSecretKey); - }, []); // Fetch the resource from the extension to get HTML and metadata (CSP, permissions, etc.). // If cachedHtml is provided we show it immediately; the fetch updates metadata and @@ -535,42 +527,6 @@ export default function McpAppRenderer({ return () => observer.disconnect(); }, []); - const handleFallbackRequest = useCallback( - async (request: JSONRPCRequest, _extra: RequestHandlerExtra) => { - if (request.method === 'sampling/createMessage') { - if (!sessionId || !apiHost || !secretKey) { - throw new Error('Session not initialized for sampling request'); - } - const { messages, systemPrompt, maxTokens } = - request.params as unknown as SamplingCreateMessageParams; - const response = await fetch(`${apiHost}/sessions/${sessionId}/sampling/message`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'X-Secret-Key': secretKey, - }, - body: JSON.stringify({ - messages: messages.map((m) => ({ - role: m.role, - content: m.content, - })), - systemPrompt, - maxTokens, - }), - }); - if (!response.ok) { - throw new Error(`Sampling request failed: ${response.statusText}`); - } - return (await response.json()) as SamplingCreateMessageResponse; - } - return { - status: 'error' as const, - message: `Unhandled JSON-RPC method: ${request.method ?? ''}`, - }; - }, - [sessionId, apiHost, secretKey] - ); - const handleError = useCallback((err: Error) => { console.error('[MCP App Error]:', err); dispatch({ type: 'ERROR', message: errorMessage(err) }); @@ -687,7 +643,6 @@ export default function McpAppRenderer({ onReadResource={handleReadResource} onLoggingMessage={handleLoggingMessage} onSizeChanged={handleSizeChanged} - onFallbackRequest={handleFallbackRequest} onError={handleError} /> ); diff --git a/ui/desktop/src/components/ProviderGuard.tsx b/ui/desktop/src/components/ProviderGuard.tsx index e1f6dd8421..21094510df 100644 --- a/ui/desktop/src/components/ProviderGuard.tsx +++ b/ui/desktop/src/components/ProviderGuard.tsx @@ -8,6 +8,7 @@ import { startChatGptCodexSetup } from '../utils/chatgptCodexSetup'; import WelcomeGooseLogo from './WelcomeGooseLogo'; import { toastService } from '../toasts'; import { OllamaSetup } from './OllamaSetup'; +import { LocalModelSetup } from './LocalModelSetup'; import ApiKeyTester from './ApiKeyTester'; import { SwitchModelModal } from './settings/models/subcomponents/SwitchModelModal'; import { createNavigationHandler } from '../utils/navigationUtils'; @@ -34,6 +35,7 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG const [hasProvider, setHasProvider] = useState(false); const [showFirstTimeSetup, setShowFirstTimeSetup] = useState(false); const [showOllamaSetup, setShowOllamaSetup] = useState(false); + const [showLocalModelSetup, setShowLocalModelSetup] = useState(false); const [userInActiveSetup, setUserInActiveSetup] = useState(false); const [showSwitchModelModal, setShowSwitchModelModal] = useState(false); const [switchModelProvider, setSwitchModelProvider] = useState(null); @@ -200,6 +202,19 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG setShowOllamaSetup(false); }; + const handleLocalModelComplete = () => { + trackOnboardingCompleted('local'); + setShowLocalModelSetup(false); + setShowFirstTimeSetup(false); + setHasProvider(true); + navigate('/', { replace: true }); + }; + + const handleLocalModelCancel = () => { + trackOnboardingAbandoned('local_model_setup'); + setShowLocalModelSetup(false); + }; + const handleRetrySetup = (setupType: 'openrouter' | 'tetrate' | 'chatgpt_codex') => { if (setupType === 'openrouter') { setOpenRouterSetupState(null); @@ -285,6 +300,23 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG return ; } + if (showLocalModelSetup) { + return ( +
+
+
+
+ +
+
+
+
+ ); + } + if (!hasProvider && showFirstTimeSetup) { return (
@@ -316,6 +348,48 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG }} /> + {/* Run Locally Card */} +
+
+ + Free & Private + +
+
{ + trackOnboardingProviderSelected('local'); + setShowLocalModelSetup(true); + }} + className="w-full p-4 sm:p-6 bg-transparent border rounded-xl transition-all duration-200 cursor-pointer group" + > +
+
+ + Run Locally + +
+
+ + + +
+
+

+ Download a model and run entirely on your machine. No API keys, no accounts. +

+
+
+ {/* ChatGPT Subscription Card - Full Width */}
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({ + + + + { // eslint-disable-next-line react-hooks/exhaustive-deps }, []); - // Determine if we should show all models by default (if non-recommended models are downloaded) - useEffect(() => { - if (models.length === 0) return; - - const hasDownloadedNonRecommended = models.some( - (model) => model.downloaded && !model.recommended - ); - - if (hasDownloadedNonRecommended && !showAllModels) { - setShowAllModels(true); - } - }, [models, showAllModels]); - const loadSelectedModel = async () => { try { const value = await read(LOCAL_WHISPER_MODEL_CONFIG_KEY, false); @@ -145,8 +132,14 @@ export const LocalModelManager = () => { } }; - const displayedModels = showAllModels ? models : models.filter((m) => m.recommended); + const hasDownloadedNonRecommended = models.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; return (
@@ -267,7 +260,7 @@ export const LocalModelManager = () => { })}
- {hasNonRecommendedModels && ( + {showToggleButton && ( diff --git a/ui/desktop/src/components/settings/localInference/HuggingFaceModelSearch.tsx b/ui/desktop/src/components/settings/localInference/HuggingFaceModelSearch.tsx new file mode 100644 index 0000000000..241921555d --- /dev/null +++ b/ui/desktop/src/components/settings/localInference/HuggingFaceModelSearch.tsx @@ -0,0 +1,358 @@ +import { useState, useCallback, useRef } from 'react'; +import { Search, Download, ChevronDown, ChevronUp, Loader2, Star } from 'lucide-react'; +import { Button } from '../../ui/button'; +import { + searchHfModels, + getRepoFiles, + downloadHfModel, + type HfModelInfo, + type HfQuantVariant, +} from '../../../api'; + +const formatBytes = (bytes: number): string => { + 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) { + // Pre-fetch variants for all results and filter out repos with no suitable quantizations + const modelsWithVariants = await Promise.all( + response.data.map(async (model) => { + try { + const [author, repo] = model.repo_id.split('/'); + const filesResponse = await getRepoFiles({ path: { author, repo } }); + if (filesResponse.data && filesResponse.data.variants.length > 0) { + return { model, data: filesResponse.data }; + } + } catch { + // Skip repos we can't fetch + } + return null; + }) + ); + + const validResults = modelsWithVariants.filter(Boolean) as { + model: HfModelInfo; + data: { variants: HfQuantVariant[]; recommended_index?: number | null }; + }[]; + + setResults(validResults.map((r) => r.model)); + setRepoData((prev) => { + const next = { ...prev }; + for (const r of validResults) { + next[r.model.repo_id] = { + variants: r.data.variants, + recommendedIndex: r.data.recommended_index ?? null, + }; + } + return next; + }); + + if (validResults.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, + }, + })); + } + } 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, quantization: string) => { + const spec = `${repoId}:${quantization}`; + setDownloading((prev) => new Set(prev).add(spec)); + try { + const response = await downloadHfModel({ + body: { spec }, + }); + if (response.data) { + onDownloadStarted(response.data); + } + } catch (e) { + console.error('Download failed:', e); + } finally { + setDownloading((prev) => { + const next = new Set(prev); + next.delete(spec); + return next; + }); + } + }; + + const startDirectDownload = async () => { + const spec = directSpec.trim(); + if (!spec) return; + const key = `direct:${spec}`; + setDownloading((prev) => new Set(prev).add(key)); + try { + const response = await downloadHfModel({ + body: { spec }, + }); + if (response.data) { + onDownloadStarted(response.data); + 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.quantization}`; + 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 new file mode 100644 index 0000000000..dc97643585 --- /dev/null +++ b/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx @@ -0,0 +1,410 @@ +import { useState, useEffect, useCallback, useRef } from 'react'; +import { Download, Trash2, X, ChevronDown, ChevronUp, Settings2 } from 'lucide-react'; +import { Button } from '../../ui/button'; +import { useModelAndProvider } from '../../ModelAndProviderContext'; +import { + listLocalModels, + downloadHfModel, + getLocalModelDownloadProgress, + cancelLocalModelDownload, + deleteLocalModel, + setConfigProvider, + type DownloadProgress, + type LocalModelResponse, +} from '../../../api'; +import { HuggingFaceModelSearch } from './HuggingFaceModelSearch'; +import { ModelSettingsPanel } from './ModelSettingsPanel'; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, +} from '../../ui/dialog'; + +const formatBytes = (bytes: number): string => { + 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`; +}; + +export const LocalInferenceSettings = () => { + const [models, setModels] = useState([]); + const [downloads, setDownloads] = useState>(new Map()); + const [showAllFeatured, setShowAllFeatured] = useState(false); + const [settingsOpenFor, setSettingsOpenFor] = useState(null); + const { currentModel, currentProvider, setProviderAndModel } = useModelAndProvider(); + const downloadSectionRef = useRef(null); + const selectedModelId = currentProvider === 'local' ? currentModel : null; + + const getDisplayName = useCallback( + (modelId: string): string => { + const model = models.find((m) => m.id === modelId); + return model?.display_name || modelId; + }, + [models] + ); + + const loadModels = useCallback(async () => { + try { + const response = await listLocalModels(); + if (response.data) { + setModels(response.data); + } + } catch (error) { + console.error('Failed to load models:', error); + } + }, []); + + // Check for any in-progress downloads when models list changes + const detectActiveDownloads = useCallback(async () => { + for (const model of models) { + if (downloads.has(model.id)) continue; + // Check models that the API reports as downloading + if (model.status.state === 'Downloading') { + pollDownloadProgress(model.id); + } + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [models, downloads]); + + useEffect(() => { + loadModels(); + }, [loadModels]); + + useEffect(() => { + if (models.length > 0) { + detectActiveDownloads(); + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [models]); + + const selectModel = async (modelId: string) => { + setProviderAndModel('local', modelId); + try { + await setConfigProvider({ + body: { provider: 'local', model: modelId }, + throwOnError: true, + }); + } catch (error) { + console.error('Failed to select model:', error); + } + }; + + const startFeaturedDownload = async (modelId: string) => { + const model = models.find((m) => m.id === modelId); + if (!model) return; + try { + await downloadHfModel({ body: { spec: model.id } }); + pollDownloadProgress(modelId); + scrollToDownloads(); + } catch (error) { + console.error('Failed to start download:', error); + } + }; + + const scrollToDownloads = useCallback(() => { + requestAnimationFrame(() => { + downloadSectionRef.current?.scrollIntoView({ behavior: 'smooth', block: 'nearest' }); + }); + }, []); + + const pollDownloadProgress = (modelId: string) => { + const interval = setInterval(async () => { + try { + const response = await getLocalModelDownloadProgress({ path: { model_id: modelId } }); + if (response.data) { + const progress = response.data; + setDownloads((prev) => new Map(prev).set(modelId, progress)); + + if (progress.status === 'completed') { + clearInterval(interval); + setDownloads((prev) => { + const next = new Map(prev); + next.delete(modelId); + return next; + }); + await loadModels(); + await selectModel(modelId); + } else if (progress.status === 'failed') { + clearInterval(interval); + await loadModels(); + } + } else { + clearInterval(interval); + } + } catch { + clearInterval(interval); + } + }, 1000); + }; + + const cancelDownload = async (modelId: string) => { + try { + await cancelLocalModelDownload({ path: { model_id: modelId } }); + setDownloads((prev) => { + const next = new Map(prev); + next.delete(modelId); + return next; + }); + } catch (error) { + console.error('Failed to cancel download:', error); + } + }; + + 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 } }); + await loadModels(); + } catch (error) { + console.error('Failed to delete model:', error); + } + }; + + const handleHfDownloadStarted = (modelId: string) => { + pollDownloadProgress(modelId); + loadModels(); + scrollToDownloads(); + }; + + const isDownloaded = (model: LocalModelResponse) => model.status.state === 'Downloaded'; + const isNotDownloaded = (model: LocalModelResponse) => + model.status.state === 'NotDownloaded' && !downloads.has(model.id); + + const downloadedModels = models.filter(isDownloaded); + const notDownloadedModels = models.filter(isNotDownloaded); + const recommendedModels = notDownloadedModels.filter((m) => m.recommended); + const displayedFeatured = showAllFeatured ? notDownloadedModels : recommendedModels; + const showFeaturedToggle = notDownloadedModels.length > recommendedModels.length; + + return ( +
+
+

Local Inference Models

+

+ Download and manage local LLM models for inference without API keys. Search HuggingFace + for any GGUF model or use the featured picks below. +

+
+ + {/* Active Downloads */} + {downloads.size > 0 && ( +
+

Downloading

+
+ {Array.from(downloads.entries()).map(([modelId, progress]) => { + if (progress.status === 'completed') return null; + const displayName = getDisplayName(modelId); + return ( +
+
+ + {displayName} + + {progress.status === 'downloading' && ( + + )} +
+ {progress.status === 'downloading' && ( +
+
+
+
+
+ + {formatBytes(progress.bytes_downloaded)} /{' '} + {formatBytes(progress.total_bytes)} ( + {progress.progress_percent.toFixed(0)}%) + + + {progress.eta_seconds != null && progress.eta_seconds > 0 && ( + + {progress.eta_seconds < 60 + ? `${Math.round(progress.eta_seconds)}s` + : `${Math.round(progress.eta_seconds / 60)}m`}{' '} + remaining + + )} + {progress.speed_bps != null && progress.speed_bps > 0 && ( + {formatBytes(progress.speed_bps)}/s + )} + +
+
+ )} + {progress.status === 'failed' && ( +

+ {progress.error || 'Download failed'} +

+ )} +
+ ); + })} +
+
+ )} + + {/* Downloaded Models */} + {downloadedModels.length > 0 && ( +
+

Downloaded Models

+
+ {downloadedModels.map((model) => { + const isSelected = selectedModelId === model.id; + return ( +
+
+
+ selectModel(model.id)} + className="cursor-pointer" + /> + + {model.display_name} + + + {formatBytes(model.size_bytes)} + + {model.recommended && ( + + Recommended + + )} +
+
+ + +
+
+
+ ); + })} +
+
+ )} + + {/* Featured Models (not yet downloaded) */} + {displayedFeatured.length > 0 && ( +
+

Featured Models

+
+ {displayedFeatured.map((model) => ( +
+
+
+
+

+ {model.display_name} +

+ + {formatBytes(model.size_bytes)} + + {model.recommended && ( + + Recommended + + )} +
+
+ +
+
+ ))} +
+ + {showFeaturedToggle && ( + + )} +
+ )} + + {/* HuggingFace Search */} +
+ +
+ + {models.length === 0 && ( +
No models available
+ )} + + { + if (!open) setSettingsOpenFor(null); + }} + > + + + Model Settings +

{getDisplayName(settingsOpenFor || '')}

+ + {settingsOpenFor && } + + +
+ ); +}; 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..6ab726cc3e --- /dev/null +++ b/ui/desktop/src/components/settings/localInference/ModelSettingsPanel.tsx @@ -0,0 +1,415 @@ +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 ( +
+
+ {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/modelInterface.ts b/ui/desktop/src/components/settings/models/modelInterface.ts index f5fb73758a..0be12e2244 100644 --- a/ui/desktop/src/components/settings/models/modelInterface.ts +++ b/ui/desktop/src/components/settings/models/modelInterface.ts @@ -1,4 +1,4 @@ -import { ProviderDetails, getProviderModels } from '../../../api'; +import { ProviderDetails, getProviderModels, listLocalModels } from '../../../api'; import { errorMessage as getErrorMessage } from '../../../utils/conversionUtils'; export default interface Model { @@ -54,6 +54,16 @@ export async function fetchModelsForProviders( ): Promise { const modelPromises = activeProviders.map(async (p) => { try { + // For local provider, use listLocalModels and filter to only downloaded models + if (p.name === 'local') { + const response = await listLocalModels(); + const allModels = response.data || []; + const downloadedModels = allModels + .filter((m) => m.status.state === 'Downloaded') + .map((m) => m.id); + return { provider: p, models: downloadedModels, error: null }; + } + const response = await getProviderModels({ path: { name: p.name }, throwOnError: true, diff --git a/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx b/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx index 6d294a18b9..4783e77c75 100644 --- a/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx +++ b/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx @@ -486,7 +486,36 @@ export const SwitchModelModal = ({ {provider && ( <> - {providerErrors[provider] ? ( + {provider === 'local' && + !loadingModels && + filteredModelOptions.flatMap((g) => g.options).filter((o) => o.value !== 'custom') + .length === 0 ? ( + /* Show special UI for local provider when no models are downloaded */ +
+
+
+

+ Local models need to be downloaded first +

+
+ To use local inference, you need to download a model to your computer + first. Go to Settings → Models to manage local models. +
+
+ +
+
+ ) : providerErrors[provider] ? ( /* Show error message when provider failed to connect */
diff --git a/ui/desktop/src/utils/analytics.ts b/ui/desktop/src/utils/analytics.ts index 73ee208f4c..6bf157e499 100644 --- a/ui/desktop/src/utils/analytics.ts +++ b/ui/desktop/src/utils/analytics.ts @@ -70,7 +70,7 @@ export type AnalyticsEvent = | { name: 'onboarding_provider_selected'; properties: { - method: 'api_key' | 'openrouter' | 'tetrate' | 'chatgpt_codex' | 'ollama' | 'other'; + method: 'api_key' | 'openrouter' | 'tetrate' | 'chatgpt_codex' | 'ollama' | 'local' | 'other'; }; } | { @@ -80,7 +80,7 @@ export type AnalyticsEvent = | { name: 'onboarding_abandoned'; properties: { step: string; duration_seconds?: number } } | { name: 'onboarding_setup_failed'; - properties: { provider: 'openrouter' | 'tetrate' | 'chatgpt_codex'; error_message?: string }; + properties: { provider: 'openrouter' | 'tetrate' | 'chatgpt_codex' | 'local'; error_message?: string }; } | { name: 'error_occurred'; @@ -284,7 +284,7 @@ export function trackOnboardingStarted(): void { } export function trackOnboardingProviderSelected( - method: 'api_key' | 'openrouter' | 'tetrate' | 'chatgpt_codex' | 'ollama' | 'other' + method: 'api_key' | 'openrouter' | 'tetrate' | 'chatgpt_codex' | 'ollama' | 'local' | 'other' ): void { trackEvent({ name: 'onboarding_provider_selected', @@ -317,7 +317,7 @@ export function trackOnboardingAbandoned(step: string): void { } export function trackOnboardingSetupFailed( - provider: 'openrouter' | 'tetrate' | 'chatgpt_codex', + provider: 'openrouter' | 'tetrate' | 'chatgpt_codex' | 'local', errorMessage?: string ): void { trackEvent({