diff --git a/Cargo.lock b/Cargo.lock index c22c0c45e..0b6214379 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4894,7 +4894,6 @@ dependencies = [ "dirs", "dotenvy", "dtor 1.0.5", - "encoding_rs", "env-lock", "etcetera 0.11.0", "fs-err 3.3.0", @@ -4902,11 +4901,11 @@ dependencies = [ "futures", "gethostname", "goose-acp-macros", + "goose-download-manager", "goose-mcp", "goose-providers", "goose-sdk-types", "goose-test-support", - "hf-hub 1.0.0-rc.1", "http 1.4.2", "icu_calendar", "icu_locale", @@ -4920,13 +4919,8 @@ dependencies = [ "jsonwebtoken", "keyring", "libc", - "llama-cpp-2", - "llama-cpp-sys-2", "lru 0.18.0", "minijinja", - "mlx-lm", - "mlx-lm-utils", - "mlx-rs", "mockall", "nanoid", "nostr", @@ -5069,6 +5063,55 @@ dependencies = [ "zip 8.6.0", ] +[[package]] +name = "goose-download-manager" +version = "1.40.0" +dependencies = [ + "anyhow", + "once_cell", + "reqwest 0.13.4", + "serde", + "tokio", + "tracing", + "utoipa 4.2.3", +] + +[[package]] +name = "goose-local-inference" +version = "1.40.0" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "base64 0.22.1", + "chrono", + "encoding_rs", + "etcetera 0.11.0", + "fs2", + "futures", + "goose-download-manager", + "goose-provider-types", + "goose-sdk-types", + "hf-hub 1.0.0-rc.1", + "include_dir", + "llama-cpp-2", + "llama-cpp-sys-2", + "minijinja", + "mlx-lm", + "mlx-lm-utils", + "mlx-rs", + "regex", + "reqwest 0.13.4", + "rmcp", + "serde", + "serde_json", + "tempfile", + "tokio", + "tracing", + "utoipa 4.2.3", + "uuid", +] + [[package]] name = "goose-mcp" version = "1.40.0" @@ -5100,7 +5143,7 @@ dependencies = [ ] [[package]] -name = "goose-providers" +name = "goose-provider-types" version = "1.40.0" dependencies = [ "anyhow", @@ -5111,14 +5154,10 @@ dependencies = [ "env-lock", "futures", "once_cell", - "pem", - "pkcs1", - "pkcs8", "rand 0.10.1", "regex", "reqwest 0.13.4", "rmcp", - "sec1", "serde", "serde_json", "strum 0.28.0", @@ -5127,12 +5166,40 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-stream", - "tokio-util", "tracing", "unicode-normalization", - "url", "utoipa 4.2.3", "uuid", +] + +[[package]] +name = "goose-providers" +version = "1.40.0" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "chrono", + "env-lock", + "futures", + "goose-local-inference", + "goose-provider-types", + "pem", + "pkcs1", + "pkcs8", + "reqwest 0.13.4", + "rmcp", + "sec1", + "serde", + "serde_json", + "tempfile", + "test-case", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "url", + "utoipa 4.2.3", "wiremock", ] diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index 20eab6303..fe1efccbc 100644 --- a/crates/goose-cli/src/cli.rs +++ b/crates/goose-cli/src/cli.rs @@ -2038,6 +2038,8 @@ async fn handle_local_models_command(command: LocalModelsCommand) -> Result<()> use goose::providers::local_inference::hf_models; use goose::providers::local_inference::local_model_registry::get_registry; + goose::providers::local_inference::configure_huggingface_auth(); + match command { LocalModelsCommand::Search { query, limit } => { println!("Searching HuggingFace for '{}'...", query); diff --git a/crates/goose-download-manager/Cargo.toml b/crates/goose-download-manager/Cargo.toml new file mode 100644 index 000000000..1d99b49ed --- /dev/null +++ b/crates/goose-download-manager/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "goose-download-manager" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description.workspace = true + +[lints] +workspace = true + +[dependencies] +anyhow = { workspace = true } +once_cell = { workspace = true } +reqwest = { workspace = true, features = ["stream"] } +serde = { workspace = true } +tokio = { workspace = true, features = ["fs", "io-util", "time"] } +tracing = { workspace = true } +utoipa = { workspace = true } diff --git a/crates/goose-download-manager/src/lib.rs b/crates/goose-download-manager/src/lib.rs new file mode 100644 index 000000000..89942e527 --- /dev/null +++ b/crates/goose-download-manager/src/lib.rs @@ -0,0 +1,691 @@ +use anyhow::Result; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +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 orphaned `.part` files in the given directory (and one level of subdirectories). +/// Preserves `.part` files whose final destination is in `registered_paths` so that +/// in-progress shard downloads can resume after a restart. +pub fn cleanup_partial_downloads( + dir: &Path, + registered_paths: &std::collections::HashSet, +) { + let should_keep = |part_path: &Path| -> bool { + // Derive the final path by stripping the trailing ".part" extension + let final_path = part_path.with_extension(""); + registered_paths.contains(&final_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") && !should_keep(&path) { + let _ = std::fs::remove_file(&path); + } + if path.is_dir() { + if let Ok(sub_entries) = std::fs::read_dir(&path) { + for sub in sub_entries.flatten() { + let sub_path = sub.path(); + if sub_path.extension().is_some_and(|e| e == "part") + && !should_keep(&sub_path) + { + let _ = std::fs::remove_file(&sub_path); + } + } + } + } + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +pub struct DownloadProgress { + /// Model ID being downloaded + pub model_id: String, + /// Download status + pub status: DownloadStatus, + /// Bytes downloaded so far + pub bytes_downloaded: u64, + /// Total bytes to download + pub total_bytes: u64, + /// Download progress percentage (0-100) + pub progress_percent: f32, + /// Download speed in bytes per second + pub speed_bps: Option, + /// Estimated time remaining in seconds + pub eta_seconds: Option, + /// Error message if failed + pub error: Option, + /// Whether the background download task has exited + #[serde(skip)] + pub task_exited: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum DownloadStatus { + Downloading, + Completed, + Failed, + Cancelled, +} + +type DownloadMap = Arc>>; + +pub struct DownloadManager { + downloads: DownloadMap, +} + +impl Default for DownloadManager { + fn default() -> Self { + Self::new() + } +} + +impl DownloadManager { + pub fn new() -> Self { + Self { + downloads: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub fn get_progress(&self, model_id: &str) -> Option { + self.downloads.lock().ok()?.get(model_id).cloned() + } + + pub fn is_downloading(&self, model_id: &str) -> bool { + self.get_progress(model_id) + .is_some_and(|progress| progress.status == DownloadStatus::Downloading) + } + + pub fn list_progress(&self) -> Vec { + self.downloads + .lock() + .map(|downloads| downloads.values().cloned().collect()) + .unwrap_or_default() + } + + pub fn set_progress(&self, progress: DownloadProgress) { + if let Ok(mut downloads) = self.downloads.lock() { + downloads.insert(progress.model_id.clone(), progress); + } + } + + pub fn reserve_download(&self, progress: DownloadProgress) -> Result { + let mut downloads = self + .downloads + .lock() + .map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?; + + if let Some(existing) = downloads.get(&progress.model_id) { + if existing.status == DownloadStatus::Downloading + || (existing.status == DownloadStatus::Cancelled && !existing.task_exited) + { + return Ok(false); + } + } + + downloads.insert(progress.model_id.clone(), progress); + Ok(true) + } + + pub fn update_progress(&self, model_id: &str, update: impl FnOnce(&mut DownloadProgress)) { + if let Ok(mut downloads) = self.downloads.lock() { + if let Some(progress) = downloads.get_mut(model_id) { + update(progress); + } + } + } + + pub fn cancel_download(&self, model_id: &str) -> Result<()> { + let mut downloads = self + .downloads + .lock() + .map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?; + + if let Some(progress) = downloads.get_mut(model_id) { + progress.status = DownloadStatus::Cancelled; + Ok(()) + } else { + anyhow::bail!("Download not found") + } + } + + pub async fn download_model( + &self, + model_id: String, + url: String, + destination: PathBuf, + on_complete: Option>, + ) -> Result<()> { + self.download_model_sharded(model_id, vec![(url, destination)], 0, on_complete) + .await + } + + pub async fn download_model_with_bearer_token( + &self, + model_id: String, + url: String, + destination: PathBuf, + bearer_token: Option, + on_complete: Option>, + ) -> Result<()> { + self.download_model_sharded_with_bearer_token( + model_id, + vec![(url, destination)], + 0, + bearer_token, + on_complete, + ) + .await + } + + pub async fn download_model_sharded( + &self, + model_id: String, + files: Vec<(String, PathBuf)>, + total_size_hint: u64, + on_complete: Option>, + ) -> Result<()> { + self.download_model_sharded_with_bearer_token( + model_id, + files, + total_size_hint, + None, + on_complete, + ) + .await + } + + pub async fn download_model_sharded_with_bearer_token( + &self, + model_id: String, + files: Vec<(String, PathBuf)>, + total_size_hint: u64, + bearer_token: Option, + on_complete: Option>, + ) -> Result<()> { + info!(model_id = %model_id, file_count = files.len(), "Starting model download"); + { + let mut downloads = self + .downloads + .lock() + .map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?; + + if let Some(existing) = downloads.get(&model_id) { + if existing.status == DownloadStatus::Downloading { + anyhow::bail!("Download already in progress"); + } + if existing.status == DownloadStatus::Cancelled && !existing.task_exited { + anyhow::bail!( + "Download is being cancelled; wait for it to finish before restarting" + ); + } + } + + downloads.insert( + model_id.clone(), + DownloadProgress { + model_id: model_id.clone(), + status: DownloadStatus::Downloading, + bytes_downloaded: 0, + total_bytes: total_size_hint, + progress_percent: 0.0, + speed_bps: None, + eta_seconds: None, + error: None, + task_exited: false, + }, + ); + } + + // Create parent directories for all files + for (_, dest) in &files { + if let Some(parent) = dest.parent() { + tokio::fs::create_dir_all(parent) + .await + .map_err(|e| anyhow::anyhow!("Failed to create directory: {}", e))?; + } + } + + let downloads = self.downloads.clone(); + let model_id_clone = model_id.clone(); + let files_for_cleanup: Vec = files.iter().map(|(_, d)| d.clone()).collect(); + + tokio::spawn(async move { + let result = Self::download_files_sequentially( + &files, + &downloads, + &model_id_clone, + bearer_token.as_deref(), + ) + .await; + + match result { + 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; + progress.progress_percent = 100.0; + progress.task_exited = true; + } + } + + if let Some(callback) = on_complete { + callback(); + } + } + Err(e) => { + for dest in &files_for_cleanup { + let partial = partial_path_for(dest); + let _ = tokio::fs::remove_file(&partial).await; + } + + if let Ok(mut downloads) = downloads.lock() { + if let Some(progress) = downloads.get_mut(&model_id_clone) { + if progress.status != DownloadStatus::Cancelled { + progress.status = DownloadStatus::Failed; + } + progress.error = Some(e.to_string()); + progress.task_exited = true; + } + } + } + } + }); + + Ok(()) + } + + const MAX_RETRIES: u32 = 10; + const RETRY_BASE_DELAY: std::time::Duration = std::time::Duration::from_secs(2); + const RETRY_MAX_DELAY: std::time::Duration = std::time::Duration::from_secs(60); + + async fn cancellable_sleep( + delay: std::time::Duration, + downloads: &DownloadMap, + model_id: &str, + ) -> Result<(), anyhow::Error> { + let check_interval = std::time::Duration::from_millis(500); + let start = std::time::Instant::now(); + while start.elapsed() < delay { + if Self::is_cancelled(downloads, model_id) { + anyhow::bail!("Download cancelled"); + } + let remaining = delay.saturating_sub(start.elapsed()); + tokio::time::sleep(std::cmp::min(check_interval, remaining)).await; + } + Ok(()) + } + + fn is_cancelled(downloads: &DownloadMap, model_id: &str) -> bool { + if let Ok(downloads) = downloads.lock() { + if let Some(progress) = downloads.get(model_id) { + return progress.status == DownloadStatus::Cancelled; + } + } + false + } + + #[allow(clippy::too_many_arguments)] + /// Download multiple files sequentially, tracking cumulative progress under one model_id. + async fn download_files_sequentially( + files: &[(String, PathBuf)], + downloads: &DownloadMap, + model_id: &str, + bearer_token: Option<&str>, + ) -> Result<(), anyhow::Error> { + let client = reqwest::Client::builder() + .connect_timeout(std::time::Duration::from_secs(30)) + .read_timeout(std::time::Duration::from_secs(120)) + .build()?; + + // HEAD each file to get accurate total size. Only replace the hint if + // every file returned a size; partial results would underestimate. + let mut total: u64 = 0; + let mut all_resolved = true; + for (url, _) in files { + let size = Self::apply_bearer_token(client.head(url), bearer_token) + .send() + .await + .ok() + .and_then(|r| r.content_length()) + .unwrap_or(0); + if size == 0 { + all_resolved = false; + } + total += size; + } + if all_resolved && total > 0 { + if let Ok(mut dl) = downloads.lock() { + if let Some(progress) = dl.get_mut(model_id) { + progress.total_bytes = total; + } + } + } + + let start_time = std::time::Instant::now(); + let mut cumulative_bytes: u64 = 0; + // Account for already-downloaded shards + for (_, dest) in files { + let partial = partial_path_for(dest); + if dest.exists() { + if let Ok(meta) = tokio::fs::metadata(dest).await { + cumulative_bytes += meta.len(); + } + } else if partial.exists() { + if let Ok(meta) = tokio::fs::metadata(&partial).await { + cumulative_bytes += meta.len(); + } + } + } + let bytes_at_start = cumulative_bytes; + + for (url, destination) in files { + if Self::is_cancelled(downloads, model_id) { + anyhow::bail!("Download cancelled"); + } + + // Skip already-completed shards + if destination.exists() { + continue; + } + + Self::download_one_file( + &client, + url, + destination, + downloads, + model_id, + &mut cumulative_bytes, + start_time, + bytes_at_start, + bearer_token, + ) + .await?; + } + + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + async fn download_one_file( + client: &reqwest::Client, + url: &str, + destination: &Path, + downloads: &DownloadMap, + model_id: &str, + cumulative_bytes: &mut u64, + start_time: std::time::Instant, + bytes_at_start: u64, + bearer_token: Option<&str>, + ) -> Result<(), anyhow::Error> { + let partial_path = partial_path_for(destination); + let mut retries = 0u32; + + let mut file_bytes: u64 = if partial_path.exists() { + tokio::fs::metadata(&partial_path).await?.len() + } else { + 0 + }; + + // Get this file's total size + let mut file_total: u64 = Self::apply_bearer_token(client.head(url), bearer_token) + .send() + .await + .ok() + .and_then(|r| r.content_length()) + .unwrap_or(0); + + // If partial matches expected size exactly, promote it + if file_total > 0 && file_bytes == file_total { + tokio::fs::rename(&partial_path, destination).await?; + // cumulative_bytes already accounts for this file from the pre-scan + return Ok(()); + } + + // If partial is oversized or remote changed, discard and re-download + if file_total > 0 && file_bytes > file_total { + info!(model_id = %model_id, file_bytes, file_total, "Partial file oversized, re-downloading"); + *cumulative_bytes = cumulative_bytes.saturating_sub(file_bytes); + file_bytes = 0; + let _ = tokio::fs::remove_file(&partial_path).await; + } + + loop { + if Self::is_cancelled(downloads, model_id) { + let _ = tokio::fs::remove_file(&partial_path).await; + anyhow::bail!("Download cancelled"); + } + + let mut request = Self::apply_bearer_token(client.get(url), bearer_token); + if file_bytes > 0 { + request = request.header("Range", format!("bytes={}-", file_bytes)); + } + + let response = match request.send().await { + Ok(r) => r, + Err(e) => { + if retries >= Self::MAX_RETRIES { + anyhow::bail!("Download failed after {} retries: {}", retries, e); + } + retries += 1; + let delay = std::cmp::min( + Self::RETRY_BASE_DELAY * 2u32.saturating_pow(retries - 1), + Self::RETRY_MAX_DELAY, + ); + info!(model_id = %model_id, retry = retries, delay_secs = ?delay.as_secs(), error = %e, "Retrying download after connection error"); + Self::cancellable_sleep(delay, downloads, model_id).await?; + continue; + } + }; + + let status = response.status(); + if status == reqwest::StatusCode::RANGE_NOT_SATISFIABLE { + if file_total > 0 && file_bytes == file_total { + break; + } + *cumulative_bytes = cumulative_bytes.saturating_sub(file_bytes); + file_bytes = 0; + let _ = tokio::fs::remove_file(&partial_path).await; + continue; + } + + if !status.is_success() && status != reqwest::StatusCode::PARTIAL_CONTENT { + let is_transient = status.is_server_error() + || status == reqwest::StatusCode::REQUEST_TIMEOUT + || status == reqwest::StatusCode::TOO_MANY_REQUESTS; + + if !is_transient || retries >= Self::MAX_RETRIES { + anyhow::bail!("Failed to download: HTTP {}", status); + } + retries += 1; + let delay = std::cmp::min( + Self::RETRY_BASE_DELAY * 2u32.saturating_pow(retries - 1), + Self::RETRY_MAX_DELAY, + ); + info!(model_id = %model_id, retry = retries, http_status = %status, "Retrying download after transient HTTP error"); + Self::cancellable_sleep(delay, downloads, model_id).await?; + continue; + } + + if file_bytes > 0 && status == reqwest::StatusCode::OK { + info!(model_id = %model_id, "Server ignored Range header, restarting file from scratch"); + // Subtract already-counted partial bytes from cumulative + *cumulative_bytes = cumulative_bytes.saturating_sub(file_bytes); + file_bytes = 0; + let _ = tokio::fs::remove_file(&partial_path).await; + } + + // If HEAD didn't return this file's size, learn it from the GET response. + // This block only fires once per file (file_total stays non-zero after), + // so retries don't double-count. Since download_files_sequentially's HEAD + // pass contributed 0 for this file, we add the discovered size to the + // shared total so progress/ETA are accurate. + if file_total == 0 { + let new_file_total = if file_bytes > 0 { + response + .headers() + .get("content-range") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.rsplit('/').next()) + .and_then(|s| s.parse::().ok()) + } else { + response.content_length() + }; + if let Some(t) = new_file_total { + file_total = t; + if let Ok(mut dl) = downloads.lock() { + if let Some(progress) = dl.get_mut(model_id) { + progress.total_bytes = progress.total_bytes.saturating_add(t); + } + } + } + } + + let mut file = tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&partial_path) + .await?; + + let file_len = tokio::fs::metadata(&partial_path).await?.len(); + if file_len != file_bytes { + file.set_len(file_bytes).await?; + } + + let mut stream_error = false; + let mut resp = response; + + loop { + let chunk_result = resp.chunk().await; + match chunk_result { + Ok(Some(chunk)) => { + if Self::is_cancelled(downloads, model_id) { + let _ = tokio::fs::remove_file(&partial_path).await; + anyhow::bail!("Download cancelled"); + } + + file.write_all(&chunk).await?; + let chunk_len = chunk.len() as u64; + file_bytes += chunk_len; + *cumulative_bytes += chunk_len; + + let elapsed = start_time.elapsed().as_secs_f64(); + let bytes_this_session = cumulative_bytes.saturating_sub(bytes_at_start); + let speed_bps = if elapsed > 0.0 { + Some((bytes_this_session as f64 / elapsed) as u64) + } else { + None + }; + + let current_total = if let Ok(dl) = downloads.lock() { + dl.get(model_id).map(|p| p.total_bytes).unwrap_or(0) + } else { + 0 + }; + + let eta_seconds = if let Some(speed) = speed_bps { + if speed > 0 && current_total > 0 { + Some(current_total.saturating_sub(*cumulative_bytes) / speed) + } else { + None + } + } else { + None + }; + + if let Ok(mut dl) = downloads.lock() { + if let Some(progress) = dl.get_mut(model_id) { + progress.bytes_downloaded = *cumulative_bytes; + progress.progress_percent = if current_total > 0 { + (*cumulative_bytes as f64 / current_total as f64 * 100.0) as f32 + } else { + 0.0 + }; + progress.speed_bps = speed_bps; + progress.eta_seconds = eta_seconds; + } + } + } + Ok(None) => break, + Err(e) => { + info!(model_id = %model_id, bytes = *cumulative_bytes, error = %e, "Download stream interrupted, will retry"); + stream_error = true; + break; + } + } + } + + file.flush().await?; + drop(file); + + if stream_error { + if retries >= Self::MAX_RETRIES { + anyhow::bail!( + "Download failed after {} retries due to stream interruption", + retries + ); + } + retries += 1; + let delay = std::cmp::min( + Self::RETRY_BASE_DELAY * 2u32.saturating_pow(retries - 1), + Self::RETRY_MAX_DELAY, + ); + info!(model_id = %model_id, retry = retries, delay_secs = ?delay.as_secs(), "Retrying download with resume"); + Self::cancellable_sleep(delay, downloads, model_id).await?; + continue; + } + + break; + } + + tokio::fs::rename(&partial_path, destination).await?; + Ok(()) + } + + pub fn clear_completed(&self, model_id: &str) { + if let Ok(mut downloads) = self.downloads.lock() { + if let Some(progress) = downloads.get(model_id) { + let is_terminal = progress.status == DownloadStatus::Completed + || progress.status == DownloadStatus::Failed + || progress.status == DownloadStatus::Cancelled; + if is_terminal && progress.task_exited { + downloads.remove(model_id); + } + } + } + } + + fn apply_bearer_token( + request: reqwest::RequestBuilder, + bearer_token: Option<&str>, + ) -> reqwest::RequestBuilder { + if let Some(token) = bearer_token.filter(|token| !token.is_empty()) { + request.header("Authorization", format!("Bearer {}", token)) + } else { + request + } + } +} + +static DOWNLOAD_MANAGER: once_cell::sync::Lazy = + once_cell::sync::Lazy::new(DownloadManager::new); + +pub fn get_download_manager() -> &'static DownloadManager { + &DOWNLOAD_MANAGER +} diff --git a/crates/goose-local-inference/Cargo.toml b/crates/goose-local-inference/Cargo.toml new file mode 100644 index 000000000..f69db30c2 --- /dev/null +++ b/crates/goose-local-inference/Cargo.toml @@ -0,0 +1,54 @@ +[package] +name = "goose-local-inference" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description.workspace = true + +[lints] +workspace = true + +[features] +default = [] +cuda = ["llama-cpp-2/cuda"] +vulkan = ["llama-cpp-2/vulkan"] +mlx = ["dep:mlx-rs", "dep:mlx-lm", "dep:mlx-lm-utils"] + +[dependencies] +anyhow = { workspace = true } +async-stream = { workspace = true } +async-trait = { workspace = true } +base64 = { workspace = true } +chrono = { workspace = true } +etcetera = { workspace = true } +encoding_rs = { version = "0.8.35", default-features = false } +fs2 = { workspace = true } +futures = { workspace = true } +goose-download-manager = { path = "../goose-download-manager" } +goose-provider-types = { path = "../goose-provider-types", default-features = false } +goose-sdk-types = { path = "../goose-sdk-types", default-features = false } +hf-hub = { version = "1.0.0-rc.1", default-features = false } +include_dir = { workspace = true } +llama-cpp-2 = { workspace = true } +llama-cpp-sys-2 = { workspace = true } +minijinja = { version = "2.18", default-features = false, features = ["loader", "multi_template", "serde"] } +reqwest = { workspace = true, features = ["json", "stream"] } +regex = { workspace = true } +rmcp = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["fs", "rt-multi-thread", "sync", "time"] } +tracing = { workspace = true } +utoipa = { workspace = true, features = ["chrono"] } +uuid = { workspace = true, features = ["v4", "std"] } +tempfile = { workspace = true } + +mlx-rs = { version = "0.25.3", default-features = false, features = ["accelerate", "metal", "safetensors"], optional = true } +mlx-lm = { git = "https://github.com/jh-block/mlx-lm", optional = true } +mlx-lm-utils = { git = "https://github.com/jh-block/mlx-lm", optional = true } + +[target.'cfg(target_os = "macos")'.dependencies] +llama-cpp-2 = { workspace = true, features = ["sampler", "metal", "mtmd"] } diff --git a/crates/goose/src/providers/local_inference/backend.rs b/crates/goose-local-inference/src/backend.rs similarity index 85% rename from crates/goose/src/providers/local_inference/backend.rs rename to crates/goose-local-inference/src/backend.rs index 78a2cc3a9..2187945a7 100644 --- a/crates/goose/src/providers/local_inference/backend.rs +++ b/crates/goose-local-inference/src/backend.rs @@ -1,10 +1,10 @@ use rmcp::model::Tool; use std::any::Any; -use crate::conversation::message::Message; -use crate::providers::local_inference::local_model_registry::ModelSettings; -use goose_providers::errors::ProviderError; -use goose_providers::request_log::RequestLogHandle; +use crate::local_model_registry::ModelSettings; +use goose_provider_types::conversation::message::Message; +use goose_provider_types::errors::ProviderError; +use goose_provider_types::request_log::RequestLogHandle; use super::{ResolvedModelPaths, StreamSender}; diff --git a/crates/goose-local-inference/src/config_resolver.rs b/crates/goose-local-inference/src/config_resolver.rs new file mode 100644 index 000000000..3c5a61022 --- /dev/null +++ b/crates/goose-local-inference/src/config_resolver.rs @@ -0,0 +1,30 @@ +use anyhow::Result; +use std::sync::OnceLock; + +pub type StringParamResolver = fn(&'static str) -> Result>; +pub type BoolParamResolver = fn(&'static str) -> Result>; + +static STRING_PARAM_RESOLVER: OnceLock = OnceLock::new(); +static BOOL_PARAM_RESOLVER: OnceLock = OnceLock::new(); + +pub fn set_string_param_resolver(resolve_param: StringParamResolver) { + let _ = STRING_PARAM_RESOLVER.set(resolve_param); +} + +pub fn set_bool_param_resolver(resolve_param: BoolParamResolver) { + let _ = BOOL_PARAM_RESOLVER.set(resolve_param); +} + +pub fn string_param(key: &'static str) -> Result> { + match STRING_PARAM_RESOLVER.get() { + Some(resolve_param) => resolve_param(key), + None => Ok(None), + } +} + +pub fn bool_param(key: &'static str) -> Result> { + match BOOL_PARAM_RESOLVER.get() { + Some(resolve_param) => resolve_param(key), + None => Ok(None), + } +} diff --git a/crates/goose/src/providers/local_inference/hf_models.rs b/crates/goose-local-inference/src/hf_models.rs similarity index 99% rename from crates/goose/src/providers/local_inference/hf_models.rs rename to crates/goose-local-inference/src/hf_models.rs index bb31fb8de..8f458b323 100644 --- a/crates/goose/src/providers/local_inference/hf_models.rs +++ b/crates/goose-local-inference/src/hf_models.rs @@ -8,7 +8,7 @@ use super::local_model_registry::{get_registry, model_id_from_repo, LocalModelSt use serde::{Deserialize, Serialize}; use std::sync::{Arc, Mutex}; -use crate::providers::huggingface_auth; +use crate::huggingface_auth; use utoipa::ToSchema; diff --git a/crates/goose-local-inference/src/huggingface_auth.rs b/crates/goose-local-inference/src/huggingface_auth.rs new file mode 100644 index 000000000..c1510e062 --- /dev/null +++ b/crates/goose-local-inference/src/huggingface_auth.rs @@ -0,0 +1,64 @@ +use crate::{config_resolver, paths::Paths}; +use anyhow::Result; +use chrono::{DateTime, Utc}; +use futures::future::BoxFuture; +use serde::{Deserialize, Serialize}; +use std::path::{Path, PathBuf}; +use std::sync::OnceLock; + +pub const HUGGINGFACE_TOKEN_SECRET_KEY: &str = "HF_TOKEN"; +pub const HUGGINGFACE_OAUTH_CACHE_PATH: &str = "huggingface/oauth/tokens.json"; + +pub type TokenResolver = fn() -> BoxFuture<'static, Result>>; + +static TOKEN_RESOLVER: OnceLock = OnceLock::new(); + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HuggingFaceTokenData { + pub access_token: String, + #[serde(default)] + pub refresh_token: Option, + #[serde(default)] + pub expires_at: Option>, +} + +impl HuggingFaceTokenData { + pub fn is_expired(&self) -> bool { + self.expires_at + .is_some_and(|expires_at| expires_at <= Utc::now()) + } +} + +pub fn oauth_cache_path() -> PathBuf { + Paths::in_config_dir(HUGGINGFACE_OAUTH_CACHE_PATH) +} + +fn load_oauth_token_from_path(path: &Path) -> Option { + let contents = std::fs::read_to_string(path).ok()?; + serde_json::from_str(&contents).ok() +} + +pub fn usable_oauth_token() -> Option { + let token = load_oauth_token_from_path(&oauth_cache_path())?; + (!token.is_expired()).then_some(token.access_token) +} + +pub fn hf_token_secret() -> Result> { + Ok(config_resolver::string_param(HUGGINGFACE_TOKEN_SECRET_KEY)? + .filter(|token| !token.trim().is_empty())) +} + +pub fn set_token_resolver(resolve_token: TokenResolver) { + let _ = TOKEN_RESOLVER.set(resolve_token); +} + +pub async fn resolve_token_async() -> Result> { + if let Some(resolve_token) = TOKEN_RESOLVER.get() { + return resolve_token().await; + } + + if let Some(token) = usable_oauth_token() { + return Ok(Some(token)); + } + hf_token_secret() +} diff --git a/crates/goose-local-inference/src/lib.rs b/crates/goose-local-inference/src/lib.rs new file mode 100644 index 000000000..5c04a2330 --- /dev/null +++ b/crates/goose-local-inference/src/lib.rs @@ -0,0 +1,761 @@ +pub mod config_resolver; +pub use goose_download_manager as download_manager; +pub mod huggingface_auth; +pub mod paths; +pub mod prompt_template; +pub mod provider_utils; + +mod backend; +pub mod hf_models; +mod llamacpp; +pub mod local_model_registry; +pub mod management; +mod mlx; +pub(crate) mod multimodal; +#[cfg(feature = "mlx")] +mod native_tool_parsing; +#[cfg(feature = "mlx")] +mod tool_emulation; +mod tool_parsing; + +use anyhow::Result; +use async_stream::try_stream; +use async_trait::async_trait; +use backend::{BackendLoadedModel, LocalInferenceBackend}; +use goose_provider_types::base::{MessageStream, Provider, ProviderDescriptor, ProviderMetadata}; +use goose_provider_types::conversation::message::{Message, MessageContent}; +use goose_provider_types::conversation::token_usage::{ProviderUsage, Usage}; +use goose_provider_types::errors::ProviderError; +use goose_provider_types::images::ImageFormat; +use goose_provider_types::model::ModelConfig; +use goose_provider_types::request_log::{start_log, LoggerHandleExt, RequestLogHandle}; +use llamacpp::{LlamaCppBackend, LLAMACPP_BACKEND_ID}; +use local_model_registry::ChatTemplate; +use mlx::{MlxBackend, MLX_BACKEND_ID}; +use rmcp::model::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; + +type ModelSlot = Arc>>>; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct ModelCacheKey { + backend_id: &'static str, + model_id: String, + chat_template: ChatTemplate, +} + +impl ModelCacheKey { + fn new( + backend_id: &'static str, + model_id: impl Into, + chat_template: ChatTemplate, + ) -> Self { + Self { + backend_id, + model_id: model_id.into(), + chat_template, + } + } +} + +pub struct InferenceRuntime { + models: StdMutex>, + backends: HashMap<&'static str, Arc>, +} + +pub fn builtin_chat_template_names() -> Vec { + llamacpp::builtin_chat_template_names() +} + +/// Global weak reference used to share a single `InferenceRuntime` across +/// all providers and management APIs. Only a `Weak` is stored here — strong +/// `Arc`s live in providers and the local-inference management layer. 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() -> Result> { + let mut guard = RUNTIME.lock().expect("runtime lock poisoned"); + if let Some(runtime) = guard.upgrade() { + return Ok(runtime); + } + let llamacpp_backend: Arc = Arc::new(LlamaCppBackend::new()?); + let mlx_backend: Arc = Arc::new(MlxBackend::new()); + let mut backends = HashMap::new(); + backends.insert(LLAMACPP_BACKEND_ID, llamacpp_backend); + backends.insert(MLX_BACKEND_ID, mlx_backend); + let runtime = Arc::new(Self { + models: StdMutex::new(HashMap::new()), + backends, + }); + *guard = Arc::downgrade(&runtime); + Ok(runtime) + } + + fn default_backend(&self) -> &dyn LocalInferenceBackend { + self.backends + .get(LLAMACPP_BACKEND_ID) + .expect("default local inference backend registered") + .as_ref() + } + + fn backend_for_model( + &self, + resolved: &ResolvedModelPaths, + ) -> Result, ProviderError> { + let backend_id = resolved + .backend_id + .as_deref() + .unwrap_or(LLAMACPP_BACKEND_ID); + self.backends.get(backend_id).cloned().ok_or_else(|| { + ProviderError::ExecutionError(format!( + "Local inference backend '{}' unavailable", + backend_id + )) + }) + } + + fn get_or_create_model_slot(&self, key: ModelCacheKey) -> ModelSlot { + let mut map = self.models.lock().expect("model cache lock poisoned"); + map.entry(key) + .or_insert_with(|| Arc::new(Mutex::new(None))) + .clone() + } + + fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec { + let map = self.models.lock().expect("model cache lock poisoned"); + map.iter() + .filter(|(key, _)| *key != keep_key) + .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"; + +#[derive(Clone)] +pub(crate) struct ResolvedModelPaths { + pub model_path: PathBuf, + pub context_limit: usize, + pub settings: crate::local_model_registry::ModelSettings, + pub mmproj_path: Option, + pub backend_id: Option, + pub draft_model_path: Option, +} + +fn resolve_model_local_path(model_id: &str) -> Option { + use crate::local_model_registry::get_registry; + + get_registry() + .lock() + .ok()? + .get_model(model_id) + .map(|entry| entry.local_path.clone()) +} + +/// Resolve model path, context limit, settings, and mmproj path for a model ID from the registry. +fn resolve_model_path(model_id: &str) -> Option { + use crate::local_model_registry::{default_settings_for_model, 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; + let mut settings = entry.settings.clone(); + let defaults = default_settings_for_model(model_id); + settings.vision_capable = defaults.vision_capable; + settings.mmproj_size_bytes = entry.mmproj_size_bytes; + let mmproj_path = entry.mmproj_path.as_ref().filter(|p| p.exists()).cloned(); + let backend_id = entry + .backend_id + .clone() + .or_else(|| settings.backend_id.clone()); + let draft_model = settings + .draft_model + .clone() + .or_else(|| { + config_resolver::string_param("GOOSE_LOCAL_DRAFT_MODEL") + .ok() + .flatten() + }) + .filter(|draft_model| draft_model != model_id); + let draft_model_path = draft_model.as_deref().and_then(resolve_model_local_path); + return Some(ResolvedModelPaths { + model_path: entry.local_path.clone(), + context_limit: ctx, + settings, + mmproj_path, + backend_id, + draft_model_path, + }); + } + } + + None +} + +pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 { + runtime.default_backend().available_memory_bytes() +} + +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].spec.to_string() +} + +fn build_openai_messages_json( + system: &str, + messages: &[Message], + media_marker: Option<&str>, +) -> String { + use goose_provider_types::formats::openai::format_messages; + + let mut arr: Vec = vec![json!({"role": "system", "content": system})]; + arr.extend(format_messages(messages, &ImageFormat::OpenAi)); + strip_image_parts_from_messages(&mut arr); + if let Some(marker) = media_marker { + convert_text_media_markers(&mut arr, marker); + } + serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_string()) +} + +fn build_openai_text_messages_json( + system: &str, + messages: &[Message], + media_marker: Option<&str>, +) -> String { + let mut arr: Vec = vec![json!({"role": "system", "content": system})]; + arr.extend(messages.iter().filter_map(|m| { + let content = extract_text_content(m); + if content.trim().is_empty() { + return None; + } + let role = match m.role { + rmcp::model::Role::User => "user", + rmcp::model::Role::Assistant => "assistant", + }; + Some(json!({"role": role, "content": content})) + })); + if let Some(marker) = media_marker { + convert_text_media_markers(&mut arr, marker); + } + serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_string()) +} + +fn convert_text_media_markers(messages: &mut [Value], marker: &str) { + if marker.is_empty() { + return; + } + + for msg in messages { + let Some(content) = msg.get_mut("content") else { + continue; + }; + + if let Some(text) = content.as_str() { + if let Some(parts) = split_media_marker_text(text, marker) { + *content = json!(parts); + } + continue; + } + + let Some(content_parts) = content.as_array_mut() else { + continue; + }; + let mut updated = Vec::new(); + let mut changed = false; + for part in content_parts.iter() { + if part.get("type").and_then(|v| v.as_str()) == Some("text") { + if let Some(text) = part.get("text").and_then(|v| v.as_str()) { + if let Some(parts) = split_media_marker_text(text, marker) { + updated.extend(parts); + changed = true; + continue; + } + } + } + updated.push(part.clone()); + } + if changed { + *content_parts = updated; + } + } +} + +fn split_media_marker_text(text: &str, marker: &str) -> Option> { + let mut parts = Vec::new(); + let mut rest = text; + let mut found_marker = false; + while let Some((before, after)) = rest.split_once(marker) { + found_marker = true; + let before = before.strip_suffix('\n').unwrap_or(before); + if !before.is_empty() { + parts.push(json!({"type": "text", "text": before})); + } + parts.push(json!({"type": "media_marker", "text": marker})); + rest = after; + rest = rest.strip_prefix('\n').unwrap_or(rest); + } + if !found_marker { + return None; + } + if !rest.is_empty() { + parts.push(json!({"type": "text", "text": rest})); + } + Some(parts) +} + +/// Remove `image_url` content parts from OpenAI-format messages JSON, replacing +/// each with a text note. This prevents an FFI crash in llama.cpp which does not +/// accept `image_url` content-part types. +fn strip_image_parts_from_messages(messages: &mut [Value]) { + let mut stripped = false; + for msg in messages.iter_mut() { + if let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) { + for part in content.iter_mut() { + if part.get("type").and_then(|t| t.as_str()) == Some("image_url") { + *part = json!({ + "type": "text", + "text": "[Image attached — image input is not supported with the currently selected model]" + }); + stripped = true; + } + } + } + } + if stripped { + tracing::warn!("Stripped image content parts from messages — vision encoder not available for this model"); + } +} + +/// 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_typescript\n…\n``` `` +/// - `ToolResponse` → `Command output:\n…` +/// +/// Only `developer__shell` and `code_execution__execute_typescript` 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) => { + let text = strip_info_messages(&text.text); + if !text.trim().is_empty() { + parts.push(text); + } + } + 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_typescript\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)); + } + }, + MessageContent::Image(_) => { + parts.push( + "[Image attached — image input is not supported with the currently selected model]" + .to_string(), + ); + } + _ => {} + } + } + + parts.join("\n") +} + +fn strip_info_messages(text: &str) -> String { + let mut remaining = text; + let mut output = String::new(); + + while let Some((before, after_start)) = remaining.split_once("") { + output.push_str(before); + if let Some((_, after_end)) = after_start.split_once("") { + remaining = after_end; + } else { + remaining = ""; + break; + } + } + + output.push_str(remaining); + output.trim().to_string() +} + +/// Build a `ProviderUsage` and write the request log entry. +fn finalize_usage( + log: &mut Option>, + 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, + name: String, +} + +impl LocalInferenceProvider { + pub async fn from_env() -> Result { + let runtime = InferenceRuntime::get_or_init()?; + Ok(Self { + runtime, + name: PROVIDER_NAME.to_string(), + }) + } +} + +impl ProviderDescriptor for LocalInferenceProvider { + fn metadata() -> ProviderMetadata + where + Self: Sized, + { + use crate::local_model_registry::{get_registry, FEATURED_MODELS}; + + let mut known_models: Vec<&str> = FEATURED_MODELS.iter().map(|m| m.spec).collect(); + + // 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![], + ) + } +} + +#[async_trait] +impl Provider for LocalInferenceProvider { + fn get_name(&self) -> &str { + &self.name + } + + async fn fetch_supported_models(&self) -> Result, ProviderError> { + use crate::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, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + let resolved = resolve_model_path(&model_config.model_name).ok_or_else(|| { + ProviderError::ExecutionError(format!("Model not found: {}", model_config.model_name)) + })?; + let backend = self.runtime.backend_for_model(&resolved)?; + let model_context_limit = resolved.context_limit; + let model_settings = resolved.settings.clone(); + let cache_key = ModelCacheKey::new( + backend.id(), + model_config.model_name.clone(), + model_settings.chat_template.clone(), + ); + let model_slot = self.runtime.get_or_create_model_slot(cache_key.clone()); + + // Ensure model is loaded — unload any other models first to free memory. + { + let mut model_lock = model_slot.lock().await; + if model_lock.is_none() { + for slot in self.runtime.other_model_slots(&cache_key) { + 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 resolved_for_load = resolved.clone(); + let settings_for_load = model_settings.clone(); + let backend_for_load = backend.clone(); + let loaded = tokio::task::spawn_blocking(move || { + backend_for_load.load_model(&model_id, &resolved_for_load, &settings_for_load) + }) + .await + .map_err(|e| ProviderError::ExecutionError(e.to_string()))??; + *model_lock = Some(loaded); + } + } + + // Allow request_params to override thinking + let mut model_settings = model_settings; + if let Some(false) = model_config + .request_param::("enable_thinking") + .or_else(|| { + config_resolver::bool_param("GOOSE_LOCAL_ENABLE_THINKING") + .ok() + .flatten() + }) + { + model_settings.enable_thinking = false; + } + + let model_arc = model_slot.clone(); + let backend = backend.clone(); + let model_name = model_config.model_name.clone(); + let temperature = model_config.temperature; + let max_tokens = model_config.max_tokens; + let context_limit = model_context_limit; + let settings = model_settings; + let resolved_model = resolved.clone(); + let system = system.to_string(); + let messages = messages.to_vec(); + let tools = tools.to_vec(); + let log_payload = serde_json::json!({ + "system": &system, + "messages": messages.iter().map(|m| { + serde_json::json!({ + "role": match m.role { rmcp::model::Role::User => "user", rmcp::model::Role::Assistant => "assistant" }, + "content": extract_text_content(m), + }) + }).collect::>(), + "tools": tools.iter().map(|t| &t.name).collect::>(), + "settings": { + "tool_calling": settings.tool_calling, + "chat_template": settings.chat_template, + "context_size": settings.context_size, + "sampling": settings.sampling, + }, + }); + + let mut log = start_log(model_config, &log_payload)?; + + 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 mut model_guard = model_arc.blocking_lock(); + let loaded = match model_guard.as_mut() { + Some(l) => l, + None => { + send_err!(ProviderError::ExecutionError( + "Model not loaded".to_string() + )); + } + }; + + let message_id = Uuid::new_v4().to_string(); + + let request = backend::LocalGenerationRequest { + model_name, + system: &system, + messages: &messages, + tools: &tools, + settings: &settings, + temperature, + max_tokens, + context_limit, + resolved_model: &resolved_model, + draft_model_path: resolved_model.draft_model_path.clone(), + message_id: &message_id, + tx: &tx, + log: &mut log, + }; + + let result = backend.generate(loaded.as_mut(), request); + + 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; + } + + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn converts_marker_in_string_content_to_media_marker_part() { + let mut messages = vec![json!({ + "role": "user", + "content": "look\n<__media__>\nclosely", + })]; + + convert_text_media_markers(&mut messages, "<__media__>"); + + assert_eq!( + messages[0]["content"], + json!([ + {"type": "text", "text": "look"}, + {"type": "media_marker", "text": "<__media__>"}, + {"type": "text", "text": "closely"}, + ]) + ); + } + + #[test] + fn converts_marker_inside_text_content_parts() { + let mut messages = vec![json!({ + "role": "user", + "content": [ + {"type": "text", "text": "<__media__>describe"}, + {"type": "text", "text": "next"}, + {"type": "media_marker", "text": "<__media__>"}, + ], + })]; + + convert_text_media_markers(&mut messages, "<__media__>"); + + assert_eq!( + messages[0]["content"], + json!([ + {"type": "media_marker", "text": "<__media__>"}, + {"type": "text", "text": "describe"}, + {"type": "text", "text": "next"}, + {"type": "media_marker", "text": "<__media__>"}, + ]) + ); + } +} diff --git a/crates/goose/src/providers/local_inference/llamacpp/inference_emulated_tools.rs b/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs similarity index 99% rename from crates/goose/src/providers/local_inference/llamacpp/inference_emulated_tools.rs rename to crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs index 703a51a19..ba0b2e7d8 100644 --- a/crates/goose/src/providers/local_inference/llamacpp/inference_emulated_tools.rs +++ b/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs @@ -20,8 +20,8 @@ //! 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 goose_providers::errors::ProviderError; +use goose_provider_types::conversation::message::{Message, MessageContent}; +use goose_provider_types::errors::ProviderError; use rmcp::model::{CallToolRequestParams, Tool}; use serde_json::json; use std::borrow::Cow; diff --git a/crates/goose/src/providers/local_inference/llamacpp/inference_engine.rs b/crates/goose-local-inference/src/llamacpp/inference_engine.rs similarity index 95% rename from crates/goose/src/providers/local_inference/llamacpp/inference_engine.rs rename to crates/goose-local-inference/src/llamacpp/inference_engine.rs index 166f6476c..95c66c3ac 100644 --- a/crates/goose/src/providers/local_inference/llamacpp/inference_engine.rs +++ b/crates/goose-local-inference/src/llamacpp/inference_engine.rs @@ -1,9 +1,9 @@ -use crate::providers::local_inference::backend::LocalInferenceBackend; -use crate::providers::local_inference::local_model_registry::ModelSettings; -use crate::providers::local_inference::multimodal::ExtractedImage; -use goose_providers::errors::ProviderError; -use goose_providers::request_log::{LoggerHandleExt, RequestLogHandle}; -use goose_providers::thinking::{FilterOut, ThinkFilter}; +use crate::backend::LocalInferenceBackend; +use crate::local_model_registry::ModelSettings; +use crate::multimodal::ExtractedImage; +use goose_provider_types::errors::ProviderError; +use goose_provider_types::request_log::{LoggerHandleExt, RequestLogHandle}; +use goose_provider_types::thinking::{FilterOut, ThinkFilter}; use llama_cpp_2::context::params::LlamaContextParams; use llama_cpp_2::llama_batch::LlamaBatch; use llama_cpp_2::model::{AddBos, ChatTemplateResult, LlamaChatTemplate, LlamaModel}; @@ -246,7 +246,7 @@ pub(super) fn estimate_max_context_for_memory( } pub(super) fn context_cap( - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, context_limit: usize, n_ctx_train: usize, memory_max_ctx: Option, @@ -276,7 +276,7 @@ pub(super) fn context_cap( pub(super) fn effective_context_size( prompt_token_count: usize, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, context_limit: usize, n_ctx_train: usize, memory_max_ctx: Option, @@ -296,7 +296,7 @@ pub(super) fn effective_context_size( pub(super) fn build_context_params( ctx_size: u32, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, ) -> LlamaContextParams { let mut params = LlamaContextParams::default().with_n_ctx(NonZeroU32::new(ctx_size)); @@ -315,10 +315,8 @@ pub(super) fn build_context_params( 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; +pub(super) fn build_sampler(settings: &crate::local_model_registry::ModelSettings) -> LlamaSampler { + use crate::local_model_registry::SamplingConfig; let has_penalties = settings.repeat_penalty != 1.0 || settings.frequency_penalty != 0.0 @@ -371,7 +369,7 @@ pub(super) fn validate_and_compute_context( backend: &LlamaCppBackend, prompt_token_count: usize, context_limit: usize, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, ) -> Result<(usize, usize), ProviderError> { let n_ctx_train = loaded.model.n_ctx_train() as usize; let mmproj_overhead = if loaded.mtmd_ctx.is_some() { @@ -412,7 +410,7 @@ pub(super) fn create_and_prefill_context<'model>( backend: &LlamaCppBackend, tokens: &[llama_cpp_2::token::LlamaToken], effective_ctx: usize, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, ) -> Result, ProviderError> { let ctx_params = build_context_params(effective_ctx as u32, settings); let mut ctx = loaded @@ -633,7 +631,7 @@ pub(super) enum TokenAction { pub(super) fn generation_loop( model: &LlamaModel, ctx: &mut llama_cpp_2::context::LlamaContext<'_>, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, prompt_token_count: usize, effective_ctx: usize, mut on_piece: impl FnMut(&str) -> Result, @@ -695,7 +693,7 @@ pub(super) fn generation_loop( #[cfg(test)] mod tests { use super::*; - use crate::providers::local_inference::local_model_registry::ModelSettings; + use crate::local_model_registry::ModelSettings; fn default_settings() -> ModelSettings { ModelSettings::default() diff --git a/crates/goose/src/providers/local_inference/llamacpp/inference_native_tools.rs b/crates/goose-local-inference/src/llamacpp/inference_native_tools.rs similarity index 99% rename from crates/goose/src/providers/local_inference/llamacpp/inference_native_tools.rs rename to crates/goose-local-inference/src/llamacpp/inference_native_tools.rs index 63d1e721a..5ac76d059 100644 --- a/crates/goose/src/providers/local_inference/llamacpp/inference_native_tools.rs +++ b/crates/goose-local-inference/src/llamacpp/inference_native_tools.rs @@ -1,5 +1,5 @@ -use crate::conversation::message::{Message, MessageContent}; -use goose_providers::errors::ProviderError; +use goose_provider_types::conversation::message::{Message, MessageContent}; +use goose_provider_types::errors::ProviderError; use rmcp::model::CallToolRequestParams; use serde_json::Value; use std::borrow::Cow; diff --git a/crates/goose/src/providers/local_inference/llamacpp/mod.rs b/crates/goose-local-inference/src/llamacpp/mod.rs similarity index 96% rename from crates/goose/src/providers/local_inference/llamacpp/mod.rs rename to crates/goose-local-inference/src/llamacpp/mod.rs index 9a77d181a..bdfd8cf41 100644 --- a/crates/goose/src/providers/local_inference/llamacpp/mod.rs +++ b/crates/goose-local-inference/src/llamacpp/mod.rs @@ -18,19 +18,13 @@ use self::inference_emulated_tools::{ }; use self::inference_engine::{GenerationContext, LoadedChatTemplates, LoadedModel}; use self::inference_native_tools::generate_with_native_tools; -use crate::providers::local_inference::backend::{ - BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend, -}; -use crate::providers::local_inference::local_model_registry::{ - ChatTemplate, ModelSettings, ToolCallingMode, -}; -use crate::providers::local_inference::multimodal::ExtractedImage; -use crate::providers::local_inference::tool_parsing::compact_tools_json; -use crate::providers::local_inference::{ - build_openai_messages_json, build_openai_text_messages_json, ResolvedModelPaths, -}; -use goose_providers::errors::ProviderError; -use goose_providers::formats::openai::format_tools; +use crate::backend::{BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend}; +use crate::local_model_registry::{ChatTemplate, ModelSettings, ToolCallingMode}; +use crate::multimodal::ExtractedImage; +use crate::tool_parsing::compact_tools_json; +use crate::{build_openai_messages_json, build_openai_text_messages_json, ResolvedModelPaths}; +use goose_provider_types::errors::ProviderError; +use goose_provider_types::formats::openai::format_tools; pub(super) const LLAMACPP_BACKEND_ID: &str = "llamacpp"; @@ -342,7 +336,7 @@ impl LlamaCppBackend { fn init_mtmd_context( model: &LlamaModel, mmproj_path: &Option, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, ) -> Option { use llama_cpp_2::mtmd::{MtmdContext, MtmdContextParams}; @@ -383,7 +377,7 @@ impl LocalInferenceBackend for LlamaCppBackend { &self, model_id: &str, resolved: &ResolvedModelPaths, - settings: &crate::providers::local_inference::local_model_registry::ModelSettings, + settings: &crate::local_model_registry::ModelSettings, ) -> Result, ProviderError> { let model_path = &resolved.model_path; diff --git a/crates/goose/src/providers/local_inference/local_model_registry.rs b/crates/goose-local-inference/src/local_model_registry.rs similarity index 99% rename from crates/goose/src/providers/local_inference/local_model_registry.rs rename to crates/goose-local-inference/src/local_model_registry.rs index ed5a7a008..893d4e369 100644 --- a/crates/goose/src/providers/local_inference/local_model_registry.rs +++ b/crates/goose-local-inference/src/local_model_registry.rs @@ -1,5 +1,5 @@ -use crate::config::paths::Paths; use crate::download_manager::{get_download_manager, DownloadStatus}; +use crate::paths::Paths; use anyhow::Result; use serde::{Deserialize, Serialize}; use std::path::PathBuf; diff --git a/crates/goose/src/providers/local_inference/management.rs b/crates/goose-local-inference/src/management.rs similarity index 99% rename from crates/goose/src/providers/local_inference/management.rs rename to crates/goose-local-inference/src/management.rs index 3b602921d..4517124d6 100644 --- a/crates/goose/src/providers/local_inference/management.rs +++ b/crates/goose-local-inference/src/management.rs @@ -11,9 +11,9 @@ use super::{ available_inference_memory_bytes, builtin_chat_template_names, recommend_local_model, InferenceRuntime, }; -use crate::config::paths::Paths; use crate::download_manager::{get_download_manager, DownloadProgress, DownloadStatus}; -use crate::providers::huggingface_auth; +use crate::huggingface_auth; +use crate::paths::Paths; use anyhow::{anyhow, Result}; use futures::future::join_all; use goose_sdk_types::custom_requests::{ diff --git a/crates/goose/src/providers/local_inference/mlx.rs b/crates/goose-local-inference/src/mlx.rs similarity index 90% rename from crates/goose/src/providers/local_inference/mlx.rs rename to crates/goose-local-inference/src/mlx.rs index 0051e5e91..2deb4b3f2 100644 --- a/crates/goose/src/providers/local_inference/mlx.rs +++ b/crates/goose-local-inference/src/mlx.rs @@ -10,29 +10,30 @@ mod imp { use mlx_rs::transforms::eval; use serde_json::json; - use crate::conversation::message::Message; - use crate::providers::base::{DraftStats, ProviderStats, ProviderUsage, Usage}; - use crate::providers::local_inference::backend::{ - BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend, - }; - use crate::providers::local_inference::local_model_registry::{ModelSettings, ToolCallingMode}; - use crate::providers::local_inference::native_tool_parsing::message_from_native_tool_text; - use crate::providers::local_inference::tool_emulation::{ + use crate::backend::{BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend}; + use crate::local_model_registry::{ModelSettings, ToolCallingMode}; + use crate::native_tool_parsing::message_from_native_tool_text; + use crate::provider_utils::filter_extensions_from_system_prompt; + use crate::tool_emulation::{ build_emulator_tool_description, load_tiny_model_prompt, message_for_emulator_action, StreamingEmulatorParser, CODE_EXECUTION_TOOL, }; - use crate::providers::local_inference::{extract_text_content, ResolvedModelPaths}; - use crate::providers::utils::filter_extensions_from_system_prompt; - use goose_providers::errors::ProviderError; - use goose_providers::formats::openai; - use goose_providers::images::ImageFormat; + use crate::{extract_text_content, ResolvedModelPaths}; + use goose_provider_types::conversation::message::Message; + use goose_provider_types::conversation::token_usage::{ + DraftStats, ProviderStats, ProviderUsage, Usage, + }; + use goose_provider_types::errors::ProviderError; + use goose_provider_types::formats::openai; + use goose_provider_types::images::ImageFormat; + use goose_provider_types::request_log::LoggerHandleExt; - pub(in crate::providers::local_inference) const MLX_BACKEND_ID: &str = "mlx"; + pub(crate) const MLX_BACKEND_ID: &str = "mlx"; - pub(in crate::providers::local_inference) struct MlxBackend; + pub(crate) struct MlxBackend; impl MlxBackend { - pub(in crate::providers::local_inference) fn new() -> Self { + pub(crate) fn new() -> Self { Self } } @@ -511,14 +512,11 @@ mod imp { fn temperature(settings: &ModelSettings) -> f32 { match &settings.sampling { - crate::providers::local_inference::local_model_registry::SamplingConfig::Greedy => 0.0, - crate::providers::local_inference::local_model_registry::SamplingConfig::Temperature { - temperature, - .. - } => *temperature, - crate::providers::local_inference::local_model_registry::SamplingConfig::MirostatV2 { - .. - } => 0.0, + crate::local_model_registry::SamplingConfig::Greedy => 0.0, + crate::local_model_registry::SamplingConfig::Temperature { temperature, .. } => { + *temperature + } + crate::local_model_registry::SamplingConfig::MirostatV2 { .. } => 0.0, } } @@ -553,19 +551,17 @@ mod imp { #[cfg(not(feature = "mlx"))] mod imp { - use crate::providers::local_inference::backend::{ - BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend, - }; - use crate::providers::local_inference::local_model_registry::ModelSettings; - use crate::providers::local_inference::ResolvedModelPaths; - use goose_providers::errors::ProviderError; + use crate::backend::{BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend}; + use crate::local_model_registry::ModelSettings; + use crate::ResolvedModelPaths; + use goose_provider_types::errors::ProviderError; - pub(in crate::providers::local_inference) const MLX_BACKEND_ID: &str = "mlx"; + pub(crate) const MLX_BACKEND_ID: &str = "mlx"; - pub(in crate::providers::local_inference) struct MlxBackend; + pub(crate) struct MlxBackend; impl MlxBackend { - pub(in crate::providers::local_inference) fn new() -> Self { + pub(crate) fn new() -> Self { Self } } @@ -604,4 +600,4 @@ mod imp { } } -pub(in crate::providers::local_inference) use imp::{MlxBackend, MLX_BACKEND_ID}; +pub(crate) use imp::{MlxBackend, MLX_BACKEND_ID}; diff --git a/crates/goose/src/providers/local_inference/multimodal.rs b/crates/goose-local-inference/src/multimodal.rs similarity index 98% rename from crates/goose/src/providers/local_inference/multimodal.rs rename to crates/goose-local-inference/src/multimodal.rs index 1a464cd02..2389ace53 100644 --- a/crates/goose/src/providers/local_inference/multimodal.rs +++ b/crates/goose-local-inference/src/multimodal.rs @@ -1,8 +1,8 @@ use base64::prelude::*; use serde_json::Value; -use crate::conversation::message::{Message, MessageContent}; -use goose_providers::errors::ProviderError; +use goose_provider_types::conversation::message::{Message, MessageContent}; +use goose_provider_types::errors::ProviderError; #[derive(Debug)] pub struct ExtractedImage { diff --git a/crates/goose/src/providers/local_inference/native_tool_parsing.rs b/crates/goose-local-inference/src/native_tool_parsing.rs similarity index 97% rename from crates/goose/src/providers/local_inference/native_tool_parsing.rs rename to crates/goose-local-inference/src/native_tool_parsing.rs index ffbe7f71d..3773c6b95 100644 --- a/crates/goose/src/providers/local_inference/native_tool_parsing.rs +++ b/crates/goose-local-inference/src/native_tool_parsing.rs @@ -1,7 +1,7 @@ -use crate::conversation::message::{Message, MessageContent}; -use goose_providers::errors::ProviderError; -use goose_providers::formats::openai::is_valid_function_name; -use goose_providers::json::safely_parse_json; +use goose_provider_types::conversation::message::{Message, MessageContent}; +use goose_provider_types::errors::ProviderError; +use goose_provider_types::formats::openai::is_valid_function_name; +use goose_provider_types::json::safely_parse_json; use rmcp::model::{object, CallToolRequestParams, ErrorCode, ErrorData}; use serde_json::{json, Value}; use std::borrow::Cow; diff --git a/crates/goose-local-inference/src/paths.rs b/crates/goose-local-inference/src/paths.rs new file mode 100644 index 000000000..4a79cbf30 --- /dev/null +++ b/crates/goose-local-inference/src/paths.rs @@ -0,0 +1,88 @@ +use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs}; +use std::path::PathBuf; + +pub struct Paths; + +impl Paths { + fn get_dir(dir_type: DirType) -> PathBuf { + if let Ok(test_root) = std::env::var("GOOSE_PATH_ROOT") { + let base = PathBuf::from(test_root); + match dir_type { + DirType::Config => base.join("config"), + DirType::Data => base.join("data"), + DirType::State => base.join("state"), + DirType::Plugins => base.join(".agents").join("plugins"), + DirType::Agents => base.join(".agents").join("agents"), + DirType::AgentsHome => base.join(".agents"), + } + } else { + // NOTE: "Block" is kept here for backwards compatibility with existing + // user config/data directories (e.g. ~/Library/Application Support/Block/goose/). + // Changing this would orphan existing installations. + let strategy = choose_app_strategy(AppStrategyArgs { + top_level_domain: "Block".to_string(), + author: "Block".to_string(), + app_name: "goose".to_string(), + }) + .expect("goose requires a home dir"); + + match dir_type { + DirType::Config => strategy.config_dir(), + DirType::Data => strategy.data_dir(), + DirType::State => strategy.state_dir().unwrap_or(strategy.data_dir()), + DirType::Plugins => strategy.home_dir().join(".agents").join("plugins"), + DirType::Agents => strategy.home_dir().join(".agents").join("agents"), + DirType::AgentsHome => strategy.home_dir().join(".agents"), + } + } + } + + pub fn config_dir() -> PathBuf { + Self::get_dir(DirType::Config) + } + + pub fn data_dir() -> PathBuf { + Self::get_dir(DirType::Data) + } + + pub fn state_dir() -> PathBuf { + Self::get_dir(DirType::State) + } + + pub fn plugins_dir() -> PathBuf { + Self::get_dir(DirType::Plugins) + } + + pub fn agents_dir() -> PathBuf { + Self::get_dir(DirType::Agents) + } + + pub fn agents_home_dir() -> PathBuf { + Self::get_dir(DirType::AgentsHome) + } + + pub fn in_agents_home_dir(subpath: &str) -> PathBuf { + Self::agents_home_dir().join(subpath) + } + + pub fn in_state_dir(subpath: &str) -> PathBuf { + Self::state_dir().join(subpath) + } + + pub fn in_config_dir(subpath: &str) -> PathBuf { + Self::config_dir().join(subpath) + } + + pub fn in_data_dir(subpath: &str) -> PathBuf { + Self::data_dir().join(subpath) + } +} + +enum DirType { + Config, + Data, + State, + Plugins, + Agents, + AgentsHome, +} diff --git a/crates/goose-local-inference/src/prompt_template.rs b/crates/goose-local-inference/src/prompt_template.rs new file mode 100644 index 000000000..046fc2402 --- /dev/null +++ b/crates/goose-local-inference/src/prompt_template.rs @@ -0,0 +1,43 @@ +use include_dir::{include_dir, Dir}; +use minijinja::{Environment, Error as MiniJinjaError, Value as MJValue}; +use serde::Serialize; + +use crate::paths::Paths; + +static CORE_PROMPTS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/prompts"); + +pub fn render_string( + template_str: &str, + context: &T, +) -> Result { + let mut env = Environment::new(); + env.set_trim_blocks(true); + env.set_lstrip_blocks(true); + env.add_template("template", template_str)?; + let tmpl = env.get_template("template")?; + let ctx = MJValue::from_serialize(context); + let rendered = tmpl.render(ctx)?; + Ok(rendered.trim().to_string()) +} + +pub fn render_template(name: &str, context: &T) -> Result { + let user_path = Paths::config_dir().join("prompts").join(name); + let template_str = if user_path.exists() { + std::fs::read_to_string(&user_path).map_err(|e| { + MiniJinjaError::new( + minijinja::ErrorKind::InvalidOperation, + format!("Failed to read user template: {}", e), + ) + })? + } else { + let file = CORE_PROMPTS_DIR.get_file(name).ok_or_else(|| { + MiniJinjaError::new( + minijinja::ErrorKind::TemplateNotFound, + format!("Built-in template '{}' not found", name), + ) + })?; + String::from_utf8_lossy(file.contents()).to_string() + }; + + render_string(&template_str, context) +} diff --git a/crates/goose-local-inference/src/prompts/tiny_model_system.md b/crates/goose-local-inference/src/prompts/tiny_model_system.md new file mode 100644 index 000000000..2a05d8416 --- /dev/null +++ b/crates/goose-local-inference/src/prompts/tiny_model_system.md @@ -0,0 +1,22 @@ +You are goose, an autonomous AI agent created by AAIF (Agentic AI Foundation). 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-local-inference/src/provider_utils.rs b/crates/goose-local-inference/src/provider_utils.rs new file mode 100644 index 000000000..1300c1f66 --- /dev/null +++ b/crates/goose-local-inference/src/provider_utils.rs @@ -0,0 +1,24 @@ +pub fn filter_extensions_from_system_prompt(system: &str) -> String { + let Some(extensions_start) = system.find("# Extensions") else { + return system.to_string(); + }; + + let Some(after_extensions) = system.get(extensions_start + 1..) else { + return system.to_string(); + }; + + if let Some(next_section_pos) = after_extensions.find("\n# ") { + let Some(before) = system.get(..extensions_start) else { + return system.to_string(); + }; + let Some(after) = system.get(extensions_start + next_section_pos + 1..) else { + return system.to_string(); + }; + format!("{}{}", before.trim_end(), after) + } else { + system + .get(..extensions_start) + .map(|s| s.trim_end().to_string()) + .unwrap_or_else(|| system.to_string()) + } +} diff --git a/crates/goose/src/providers/local_inference/tool_emulation.rs b/crates/goose-local-inference/src/tool_emulation.rs similarity index 99% rename from crates/goose/src/providers/local_inference/tool_emulation.rs rename to crates/goose-local-inference/src/tool_emulation.rs index 9bb645d6c..ac528be23 100644 --- a/crates/goose/src/providers/local_inference/tool_emulation.rs +++ b/crates/goose-local-inference/src/tool_emulation.rs @@ -4,7 +4,7 @@ //! as `$ command` on a new line and code blocks as ```execute_typescript fenced blocks. //! The parser converts those patterns into Goose tool-call messages. -use crate::conversation::message::{Message, MessageContent}; +use goose_provider_types::conversation::message::{Message, MessageContent}; use rmcp::model::{CallToolRequestParams, Tool}; use serde_json::json; use std::borrow::Cow; diff --git a/crates/goose/src/providers/local_inference/tool_parsing.rs b/crates/goose-local-inference/src/tool_parsing.rs similarity index 100% rename from crates/goose/src/providers/local_inference/tool_parsing.rs rename to crates/goose-local-inference/src/tool_parsing.rs diff --git a/crates/goose-provider-types/Cargo.toml b/crates/goose-provider-types/Cargo.toml new file mode 100644 index 000000000..ce0ead3fe --- /dev/null +++ b/crates/goose-provider-types/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "goose-provider-types" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description.workspace = true + +[lints] +workspace = true + +[dependencies] +anyhow = { workspace = true } +async-stream = { workspace = true } +async-trait = { workspace = true } +base64 = { workspace = true } +chrono = { workspace = true } +futures = { workspace = true } +once_cell = { workspace = true } +rand = { workspace = true } +regex = { workspace = true, features = ["unicode"] } +reqwest = { workspace = true } +rmcp = { workspace = true, features = ["server", "macros"] } +serde = { workspace = true } +serde_json = { workspace = true } +strum = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["time"] } +tracing = { workspace = true } +unicode-normalization = { version = "0.1.22", default-features = false, features = ["std"] } +utoipa = { workspace = true, features = ["chrono"] } +uuid = { workspace = true, features = ["v4", "std"] } + +[dev-dependencies] +env-lock = { workspace = true } +tempfile = { workspace = true } +test-case = { workspace = true } +tokio = { workspace = true, features = ["rt-multi-thread"] } +tokio-stream = { workspace = true } diff --git a/crates/goose-providers/src/base.rs b/crates/goose-provider-types/src/base.rs similarity index 100% rename from crates/goose-providers/src/base.rs rename to crates/goose-provider-types/src/base.rs diff --git a/crates/goose-providers/src/canonical.rs b/crates/goose-provider-types/src/canonical.rs similarity index 100% rename from crates/goose-providers/src/canonical.rs rename to crates/goose-provider-types/src/canonical.rs diff --git a/crates/goose-providers/src/canonical/README.md b/crates/goose-provider-types/src/canonical/README.md similarity index 100% rename from crates/goose-providers/src/canonical/README.md rename to crates/goose-provider-types/src/canonical/README.md diff --git a/crates/goose-providers/src/canonical/catalog.rs b/crates/goose-provider-types/src/canonical/catalog.rs similarity index 100% rename from crates/goose-providers/src/canonical/catalog.rs rename to crates/goose-provider-types/src/canonical/catalog.rs diff --git a/crates/goose-providers/src/canonical/data/canonical_mapping_report.json b/crates/goose-provider-types/src/canonical/data/canonical_mapping_report.json similarity index 100% rename from crates/goose-providers/src/canonical/data/canonical_mapping_report.json rename to crates/goose-provider-types/src/canonical/data/canonical_mapping_report.json diff --git a/crates/goose-providers/src/canonical/data/canonical_models.json b/crates/goose-provider-types/src/canonical/data/canonical_models.json similarity index 100% rename from crates/goose-providers/src/canonical/data/canonical_models.json rename to crates/goose-provider-types/src/canonical/data/canonical_models.json diff --git a/crates/goose-providers/src/canonical/data/provider_metadata.json b/crates/goose-provider-types/src/canonical/data/provider_metadata.json similarity index 100% rename from crates/goose-providers/src/canonical/data/provider_metadata.json rename to crates/goose-provider-types/src/canonical/data/provider_metadata.json diff --git a/crates/goose-providers/src/canonical/model.rs b/crates/goose-provider-types/src/canonical/model.rs similarity index 100% rename from crates/goose-providers/src/canonical/model.rs rename to crates/goose-provider-types/src/canonical/model.rs diff --git a/crates/goose-providers/src/canonical/name_builder.rs b/crates/goose-provider-types/src/canonical/name_builder.rs similarity index 100% rename from crates/goose-providers/src/canonical/name_builder.rs rename to crates/goose-provider-types/src/canonical/name_builder.rs diff --git a/crates/goose-providers/src/canonical/registry.rs b/crates/goose-provider-types/src/canonical/registry.rs similarity index 100% rename from crates/goose-providers/src/canonical/registry.rs rename to crates/goose-provider-types/src/canonical/registry.rs diff --git a/crates/goose-providers/src/conversation.rs b/crates/goose-provider-types/src/conversation.rs similarity index 100% rename from crates/goose-providers/src/conversation.rs rename to crates/goose-provider-types/src/conversation.rs diff --git a/crates/goose-providers/src/conversation/message.rs b/crates/goose-provider-types/src/conversation/message.rs similarity index 100% rename from crates/goose-providers/src/conversation/message.rs rename to crates/goose-provider-types/src/conversation/message.rs diff --git a/crates/goose-providers/src/conversation/token_usage.rs b/crates/goose-provider-types/src/conversation/token_usage.rs similarity index 100% rename from crates/goose-providers/src/conversation/token_usage.rs rename to crates/goose-provider-types/src/conversation/token_usage.rs diff --git a/crates/goose-providers/src/conversation/tool_result_serde.rs b/crates/goose-provider-types/src/conversation/tool_result_serde.rs similarity index 100% rename from crates/goose-providers/src/conversation/tool_result_serde.rs rename to crates/goose-provider-types/src/conversation/tool_result_serde.rs diff --git a/crates/goose-providers/src/errors.rs b/crates/goose-provider-types/src/errors.rs similarity index 100% rename from crates/goose-providers/src/errors.rs rename to crates/goose-provider-types/src/errors.rs diff --git a/crates/goose-providers/src/formats.rs b/crates/goose-provider-types/src/formats.rs similarity index 100% rename from crates/goose-providers/src/formats.rs rename to crates/goose-provider-types/src/formats.rs diff --git a/crates/goose-providers/src/formats/anthropic.rs b/crates/goose-provider-types/src/formats/anthropic.rs similarity index 100% rename from crates/goose-providers/src/formats/anthropic.rs rename to crates/goose-provider-types/src/formats/anthropic.rs diff --git a/crates/goose-providers/src/formats/ollama.rs b/crates/goose-provider-types/src/formats/ollama.rs similarity index 100% rename from crates/goose-providers/src/formats/ollama.rs rename to crates/goose-provider-types/src/formats/ollama.rs diff --git a/crates/goose-providers/src/formats/openai.rs b/crates/goose-provider-types/src/formats/openai.rs similarity index 100% rename from crates/goose-providers/src/formats/openai.rs rename to crates/goose-provider-types/src/formats/openai.rs diff --git a/crates/goose-providers/src/formats/openai_responses.rs b/crates/goose-provider-types/src/formats/openai_responses.rs similarity index 100% rename from crates/goose-providers/src/formats/openai_responses.rs rename to crates/goose-provider-types/src/formats/openai_responses.rs diff --git a/crates/goose-providers/src/goose_mode.rs b/crates/goose-provider-types/src/goose_mode.rs similarity index 100% rename from crates/goose-providers/src/goose_mode.rs rename to crates/goose-provider-types/src/goose_mode.rs diff --git a/crates/goose-providers/src/images.rs b/crates/goose-provider-types/src/images.rs similarity index 100% rename from crates/goose-providers/src/images.rs rename to crates/goose-provider-types/src/images.rs diff --git a/crates/goose-providers/src/json.rs b/crates/goose-provider-types/src/json.rs similarity index 100% rename from crates/goose-providers/src/json.rs rename to crates/goose-provider-types/src/json.rs diff --git a/crates/goose-provider-types/src/lib.rs b/crates/goose-provider-types/src/lib.rs new file mode 100644 index 000000000..3ab49ef87 --- /dev/null +++ b/crates/goose-provider-types/src/lib.rs @@ -0,0 +1,15 @@ +pub mod base; +pub mod canonical; +pub mod conversation; +pub mod errors; +pub mod formats; +pub mod goose_mode; +pub mod images; +pub mod json; +pub(crate) mod mcp_utils; +pub mod model; +pub mod permission; +pub mod request_log; +pub mod retry; +pub mod thinking; +pub mod utils; diff --git a/crates/goose-providers/src/mcp_utils.rs b/crates/goose-provider-types/src/mcp_utils.rs similarity index 100% rename from crates/goose-providers/src/mcp_utils.rs rename to crates/goose-provider-types/src/mcp_utils.rs diff --git a/crates/goose-providers/src/model.rs b/crates/goose-provider-types/src/model.rs similarity index 100% rename from crates/goose-providers/src/model.rs rename to crates/goose-provider-types/src/model.rs diff --git a/crates/goose-providers/src/permission.rs b/crates/goose-provider-types/src/permission.rs similarity index 100% rename from crates/goose-providers/src/permission.rs rename to crates/goose-provider-types/src/permission.rs diff --git a/crates/goose-providers/src/request_log.rs b/crates/goose-provider-types/src/request_log.rs similarity index 100% rename from crates/goose-providers/src/request_log.rs rename to crates/goose-provider-types/src/request_log.rs diff --git a/crates/goose-providers/src/retry.rs b/crates/goose-provider-types/src/retry.rs similarity index 100% rename from crates/goose-providers/src/retry.rs rename to crates/goose-provider-types/src/retry.rs diff --git a/crates/goose-providers/src/thinking.rs b/crates/goose-provider-types/src/thinking.rs similarity index 100% rename from crates/goose-providers/src/thinking.rs rename to crates/goose-provider-types/src/thinking.rs diff --git a/crates/goose-providers/src/utils.rs b/crates/goose-provider-types/src/utils.rs similarity index 100% rename from crates/goose-providers/src/utils.rs rename to crates/goose-provider-types/src/utils.rs diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index 0ef45601d..f1e2b6f21 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -13,6 +13,10 @@ workspace = true [features] default = [] +local-inference = ["dep:goose-local-inference"] +cuda = ["local-inference", "goose-local-inference/cuda"] +vulkan = ["local-inference", "goose-local-inference/vulkan"] +mlx = ["local-inference", "goose-local-inference/mlx"] rustls-tls = [ "reqwest/rustls", "rmcp/reqwest", @@ -29,26 +33,20 @@ native-tls = [ [dependencies] anyhow = { workspace = true } async-stream = { workspace = true } -base64 = { workspace = true } chrono = { workspace = true } futures = { workspace = true } -once_cell = { workspace = true } -regex = { workspace = true, features = ["unicode"] } +goose-provider-types = { path = "../goose-provider-types", default-features = false } +goose-local-inference = { path = "../goose-local-inference", default-features = false, optional = true } reqwest = { workspace = true } rmcp = { workspace = true, features = ["server", "macros"] } serde = { workspace = true } serde_json = { workspace = true } -thiserror = { workspace = true } tracing = { workspace = true } -unicode-normalization = { version = "0.1.22", default-features = false, features = ["std"] } utoipa = { workspace = true, features = ["chrono"] } -uuid = { workspace = true, features = ["v4", "std"] } async-trait = { workspace = true } -strum = { workspace = true } tokio = { workspace = true } tokio-stream = { workspace = true, features = ["io-util"] } tokio-util = { workspace = true, features = ["compat"] } -rand = { workspace = true } url = { workspace = true } pem = { version = "3.0.2", default-features = false, features = ["std"], optional = true } pkcs1 = { version = "0.7.5", default-features = false, features = ["pkcs8", "std"], optional = true } diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index f84486bdc..50fb511e4 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -1,22 +1,13 @@ pub mod anthropic; pub mod api_client; -pub mod base; -pub mod canonical; -pub mod conversation; +pub use goose_provider_types::{ + base, canonical, conversation, errors, formats, goose_mode, images, json, model, permission, + request_log, retry, thinking, utils, +}; pub mod declarative; -pub mod errors; -pub mod formats; -pub mod goose_mode; pub mod http_status; -pub mod images; -pub mod json; -pub(crate) mod mcp_utils; -pub mod model; +#[cfg(feature = "local-inference")] +pub mod local_inference; pub mod ollama; pub mod openai; pub mod openai_compatible; -pub mod permission; -pub mod request_log; -pub mod retry; -pub mod thinking; -pub mod utils; diff --git a/crates/goose-providers/src/local_inference.rs b/crates/goose-providers/src/local_inference.rs new file mode 100644 index 000000000..954840b85 --- /dev/null +++ b/crates/goose-providers/src/local_inference.rs @@ -0,0 +1 @@ +pub use goose_local_inference::*; diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 0f0fdd25f..6e7511c97 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -24,9 +24,7 @@ local-inference = [ "dep:candle-core", "dep:candle-nn", "dep:candle-transformers", - "dep:hf-hub", - "dep:llama-cpp-2", - "dep:llama-cpp-sys-2", + "goose-providers/local-inference", "dep:tokenizers", "dep:symphonia", "dep:rubato", @@ -39,9 +37,9 @@ aws-providers = [ "dep:aws-sdk-sagemakerruntime", "dep:smithy-transport-reqwest", ] -cuda = ["local-inference", "candle-core/cuda", "candle-nn/cuda", "llama-cpp-2/cuda"] -vulkan = ["local-inference", "llama-cpp-2/vulkan"] -mlx = ["local-inference", "dep:mlx-rs", "dep:mlx-lm", "dep:mlx-lm-utils"] +cuda = ["local-inference", "candle-core/cuda", "candle-nn/cuda", "goose-providers/cuda"] +vulkan = ["local-inference", "goose-providers/vulkan"] +mlx = ["local-inference", "goose-providers/mlx"] rustls-tls = [ "dep:rustls", "dep:rcgen", @@ -56,7 +54,6 @@ rustls-tls = [ "jsonwebtoken/aws_lc_rs", "oauth2/reqwest", "oauth2/rustls-tls", - "hf-hub/rustls-tls", "goose-providers/rustls-tls", ] native-tls = [ @@ -135,6 +132,7 @@ once_cell = { workspace = true } etcetera = { workspace = true } fs-err = { version = "3.1", default-features = false } goose-providers = { path = "../goose-providers", default-features = false } +goose-download-manager = { path = "../goose-download-manager" } goose-sdk-types = { path = "../goose-sdk-types" } rand = { workspace = true } utoipa = { workspace = true, features = ["chrono"] } @@ -182,8 +180,6 @@ symphonia = { version = "0.5", default-features = false, features = ["aac", "adp rubato = { version = "0.16", default-features = false, optional = true } sys-info = { version = "0.9", default-features = false } -llama-cpp-2 = { workspace = true, optional = true } - schemars = { workspace = true, features = [ "derive", ] } @@ -203,7 +199,6 @@ tree-sitter-swift = { workspace = true } tree-sitter-typescript = { workspace = true } which = { workspace = true } pulldown-cmark = { version = "0.13", default-features = false } -encoding_rs = { version = "0.8.35", default-features = false } pastey = { version = "0.2", default-features = false } shell-words = { workspace = true } goose-acp-macros = { path = "../goose-acp-macros", default-features = false } @@ -218,10 +213,6 @@ axum-server = { version = "0.8", default-features = false, optional = true } aws-lc-rs = { version = "1.17", default-features = false, optional = true } openssl = { version = "0.10.66", default-features = false, optional = true } pem = { version = "3.0.2", default-features = false, features = ["std"], optional = true } -hf-hub = { version = "1.0.0-rc.1", default-features = false, optional = true } -mlx-rs = { version = "0.25.3", default-features = false, features = ["accelerate", "metal", "safetensors"], optional = true } -mlx-lm = { git = "https://github.com/jh-block/mlx-lm", optional = true } -mlx-lm-utils = { git = "https://github.com/jh-block/mlx-lm", optional = true } pctx_code_mode = { version = "0.3", default-features = false, optional = true } @@ -230,7 +221,6 @@ pctx_code_mode = { version = "0.3", default-features = false, optional = true } # They are just here to pin the version, and can be removed if PCTX updates temporal_rs icu_calendar = { version = "=2.1.1", default-features = false } icu_locale = { version = "=2.1.1", default-features = false } -llama-cpp-sys-2 = { workspace = true, optional = true } image = { version = "0.24.9", default-features = false, features = ["png", "jpeg", "gif", "webp"] } subtle = { version = "2.5", default-features = false, features = ["std"] } gethostname = "1.1.0" @@ -243,7 +233,6 @@ keyring = { workspace = true, features = ["windows-native"], optional = true } [target.'cfg(target_os = "macos")'.dependencies] candle-core = { workspace = true, features = ["metal"], optional = true } candle-nn = { workspace = true, features = ["metal"], optional = true } -llama-cpp-2 = { workspace = true, features = ["sampler", "metal", "mtmd"], optional = true } keyring = { workspace = true, features = ["apple-native"], optional = true } [target.'cfg(target_os = "linux")'.dependencies] diff --git a/crates/goose/src/acp/server/local_inference.rs b/crates/goose/src/acp/server/local_inference.rs index 06cb8b5d7..9796c200b 100644 --- a/crates/goose/src/acp/server/local_inference.rs +++ b/crates/goose/src/acp/server/local_inference.rs @@ -12,6 +12,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::list_models() .await .internal_err() @@ -27,6 +28,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::download_model(req) .await .invalid_params_err() @@ -45,6 +47,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::download_progress(&req.model_id) .map(|progress| LocalInferenceModelDownloadProgressResponse { progress }) .internal_err() @@ -63,6 +66,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::cancel_download(&req.model_id) .internal_err()?; Ok(EmptyResponse {}) @@ -81,6 +85,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::delete_model(&req.model_id) .invalid_params_err()?; Ok(EmptyResponse {}) @@ -99,6 +104,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::get_model_settings(&req.model_id) .invalid_params_err() } @@ -116,6 +122,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::update_model_settings( &req.model_id, req.settings, @@ -136,6 +143,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::search_huggingface_models( req.query, req.limit, ) @@ -156,6 +164,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); crate::providers::local_inference::management::huggingface_repo_variants(req.repo_id) .await .internal_err() @@ -174,6 +183,7 @@ impl GooseAcpAgent { ) -> Result { #[cfg(feature = "local-inference")] { + crate::providers::local_inference::configure_huggingface_auth(); Ok(crate::providers::local_inference::management::list_builtin_chat_templates()) } diff --git a/crates/goose/src/download_manager.rs b/crates/goose/src/download_manager.rs index 89942e527..14cd9ddb7 100644 --- a/crates/goose/src/download_manager.rs +++ b/crates/goose/src/download_manager.rs @@ -1,691 +1,14 @@ -use anyhow::Result; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::path::{Path, PathBuf}; -use std::sync::{Arc, Mutex}; -use tokio::io::AsyncWriteExt; -use tracing::info; -use utoipa::ToSchema; +pub use goose_download_manager::*; -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()), - ) -} +#[cfg(test)] +mod tests { + #[cfg(feature = "local-inference")] + #[test] + fn local_inference_uses_same_download_manager() { + let goose_manager = crate::download_manager::get_download_manager() as *const _; + let local_manager = + crate::providers::local_inference::download_manager::get_download_manager() as *const _; -/// Remove orphaned `.part` files in the given directory (and one level of subdirectories). -/// Preserves `.part` files whose final destination is in `registered_paths` so that -/// in-progress shard downloads can resume after a restart. -pub fn cleanup_partial_downloads( - dir: &Path, - registered_paths: &std::collections::HashSet, -) { - let should_keep = |part_path: &Path| -> bool { - // Derive the final path by stripping the trailing ".part" extension - let final_path = part_path.with_extension(""); - registered_paths.contains(&final_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") && !should_keep(&path) { - let _ = std::fs::remove_file(&path); - } - if path.is_dir() { - if let Ok(sub_entries) = std::fs::read_dir(&path) { - for sub in sub_entries.flatten() { - let sub_path = sub.path(); - if sub_path.extension().is_some_and(|e| e == "part") - && !should_keep(&sub_path) - { - let _ = std::fs::remove_file(&sub_path); - } - } - } - } - } + assert_eq!(goose_manager, local_manager); } } - -#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] -pub struct DownloadProgress { - /// Model ID being downloaded - pub model_id: String, - /// Download status - pub status: DownloadStatus, - /// Bytes downloaded so far - pub bytes_downloaded: u64, - /// Total bytes to download - pub total_bytes: u64, - /// Download progress percentage (0-100) - pub progress_percent: f32, - /// Download speed in bytes per second - pub speed_bps: Option, - /// Estimated time remaining in seconds - pub eta_seconds: Option, - /// Error message if failed - pub error: Option, - /// Whether the background download task has exited - #[serde(skip)] - pub task_exited: bool, -} - -#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)] -#[serde(rename_all = "lowercase")] -pub enum DownloadStatus { - Downloading, - Completed, - Failed, - Cancelled, -} - -type DownloadMap = Arc>>; - -pub struct DownloadManager { - downloads: DownloadMap, -} - -impl Default for DownloadManager { - fn default() -> Self { - Self::new() - } -} - -impl DownloadManager { - pub fn new() -> Self { - Self { - downloads: Arc::new(Mutex::new(HashMap::new())), - } - } - - pub fn get_progress(&self, model_id: &str) -> Option { - self.downloads.lock().ok()?.get(model_id).cloned() - } - - pub fn is_downloading(&self, model_id: &str) -> bool { - self.get_progress(model_id) - .is_some_and(|progress| progress.status == DownloadStatus::Downloading) - } - - pub fn list_progress(&self) -> Vec { - self.downloads - .lock() - .map(|downloads| downloads.values().cloned().collect()) - .unwrap_or_default() - } - - pub fn set_progress(&self, progress: DownloadProgress) { - if let Ok(mut downloads) = self.downloads.lock() { - downloads.insert(progress.model_id.clone(), progress); - } - } - - pub fn reserve_download(&self, progress: DownloadProgress) -> Result { - let mut downloads = self - .downloads - .lock() - .map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?; - - if let Some(existing) = downloads.get(&progress.model_id) { - if existing.status == DownloadStatus::Downloading - || (existing.status == DownloadStatus::Cancelled && !existing.task_exited) - { - return Ok(false); - } - } - - downloads.insert(progress.model_id.clone(), progress); - Ok(true) - } - - pub fn update_progress(&self, model_id: &str, update: impl FnOnce(&mut DownloadProgress)) { - if let Ok(mut downloads) = self.downloads.lock() { - if let Some(progress) = downloads.get_mut(model_id) { - update(progress); - } - } - } - - pub fn cancel_download(&self, model_id: &str) -> Result<()> { - let mut downloads = self - .downloads - .lock() - .map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?; - - if let Some(progress) = downloads.get_mut(model_id) { - progress.status = DownloadStatus::Cancelled; - Ok(()) - } else { - anyhow::bail!("Download not found") - } - } - - pub async fn download_model( - &self, - model_id: String, - url: String, - destination: PathBuf, - on_complete: Option>, - ) -> Result<()> { - self.download_model_sharded(model_id, vec![(url, destination)], 0, on_complete) - .await - } - - pub async fn download_model_with_bearer_token( - &self, - model_id: String, - url: String, - destination: PathBuf, - bearer_token: Option, - on_complete: Option>, - ) -> Result<()> { - self.download_model_sharded_with_bearer_token( - model_id, - vec![(url, destination)], - 0, - bearer_token, - on_complete, - ) - .await - } - - pub async fn download_model_sharded( - &self, - model_id: String, - files: Vec<(String, PathBuf)>, - total_size_hint: u64, - on_complete: Option>, - ) -> Result<()> { - self.download_model_sharded_with_bearer_token( - model_id, - files, - total_size_hint, - None, - on_complete, - ) - .await - } - - pub async fn download_model_sharded_with_bearer_token( - &self, - model_id: String, - files: Vec<(String, PathBuf)>, - total_size_hint: u64, - bearer_token: Option, - on_complete: Option>, - ) -> Result<()> { - info!(model_id = %model_id, file_count = files.len(), "Starting model download"); - { - let mut downloads = self - .downloads - .lock() - .map_err(|_| anyhow::anyhow!("Failed to acquire lock"))?; - - if let Some(existing) = downloads.get(&model_id) { - if existing.status == DownloadStatus::Downloading { - anyhow::bail!("Download already in progress"); - } - if existing.status == DownloadStatus::Cancelled && !existing.task_exited { - anyhow::bail!( - "Download is being cancelled; wait for it to finish before restarting" - ); - } - } - - downloads.insert( - model_id.clone(), - DownloadProgress { - model_id: model_id.clone(), - status: DownloadStatus::Downloading, - bytes_downloaded: 0, - total_bytes: total_size_hint, - progress_percent: 0.0, - speed_bps: None, - eta_seconds: None, - error: None, - task_exited: false, - }, - ); - } - - // Create parent directories for all files - for (_, dest) in &files { - if let Some(parent) = dest.parent() { - tokio::fs::create_dir_all(parent) - .await - .map_err(|e| anyhow::anyhow!("Failed to create directory: {}", e))?; - } - } - - let downloads = self.downloads.clone(); - let model_id_clone = model_id.clone(); - let files_for_cleanup: Vec = files.iter().map(|(_, d)| d.clone()).collect(); - - tokio::spawn(async move { - let result = Self::download_files_sequentially( - &files, - &downloads, - &model_id_clone, - bearer_token.as_deref(), - ) - .await; - - match result { - 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; - progress.progress_percent = 100.0; - progress.task_exited = true; - } - } - - if let Some(callback) = on_complete { - callback(); - } - } - Err(e) => { - for dest in &files_for_cleanup { - let partial = partial_path_for(dest); - let _ = tokio::fs::remove_file(&partial).await; - } - - if let Ok(mut downloads) = downloads.lock() { - if let Some(progress) = downloads.get_mut(&model_id_clone) { - if progress.status != DownloadStatus::Cancelled { - progress.status = DownloadStatus::Failed; - } - progress.error = Some(e.to_string()); - progress.task_exited = true; - } - } - } - } - }); - - Ok(()) - } - - const MAX_RETRIES: u32 = 10; - const RETRY_BASE_DELAY: std::time::Duration = std::time::Duration::from_secs(2); - const RETRY_MAX_DELAY: std::time::Duration = std::time::Duration::from_secs(60); - - async fn cancellable_sleep( - delay: std::time::Duration, - downloads: &DownloadMap, - model_id: &str, - ) -> Result<(), anyhow::Error> { - let check_interval = std::time::Duration::from_millis(500); - let start = std::time::Instant::now(); - while start.elapsed() < delay { - if Self::is_cancelled(downloads, model_id) { - anyhow::bail!("Download cancelled"); - } - let remaining = delay.saturating_sub(start.elapsed()); - tokio::time::sleep(std::cmp::min(check_interval, remaining)).await; - } - Ok(()) - } - - fn is_cancelled(downloads: &DownloadMap, model_id: &str) -> bool { - if let Ok(downloads) = downloads.lock() { - if let Some(progress) = downloads.get(model_id) { - return progress.status == DownloadStatus::Cancelled; - } - } - false - } - - #[allow(clippy::too_many_arguments)] - /// Download multiple files sequentially, tracking cumulative progress under one model_id. - async fn download_files_sequentially( - files: &[(String, PathBuf)], - downloads: &DownloadMap, - model_id: &str, - bearer_token: Option<&str>, - ) -> Result<(), anyhow::Error> { - let client = reqwest::Client::builder() - .connect_timeout(std::time::Duration::from_secs(30)) - .read_timeout(std::time::Duration::from_secs(120)) - .build()?; - - // HEAD each file to get accurate total size. Only replace the hint if - // every file returned a size; partial results would underestimate. - let mut total: u64 = 0; - let mut all_resolved = true; - for (url, _) in files { - let size = Self::apply_bearer_token(client.head(url), bearer_token) - .send() - .await - .ok() - .and_then(|r| r.content_length()) - .unwrap_or(0); - if size == 0 { - all_resolved = false; - } - total += size; - } - if all_resolved && total > 0 { - if let Ok(mut dl) = downloads.lock() { - if let Some(progress) = dl.get_mut(model_id) { - progress.total_bytes = total; - } - } - } - - let start_time = std::time::Instant::now(); - let mut cumulative_bytes: u64 = 0; - // Account for already-downloaded shards - for (_, dest) in files { - let partial = partial_path_for(dest); - if dest.exists() { - if let Ok(meta) = tokio::fs::metadata(dest).await { - cumulative_bytes += meta.len(); - } - } else if partial.exists() { - if let Ok(meta) = tokio::fs::metadata(&partial).await { - cumulative_bytes += meta.len(); - } - } - } - let bytes_at_start = cumulative_bytes; - - for (url, destination) in files { - if Self::is_cancelled(downloads, model_id) { - anyhow::bail!("Download cancelled"); - } - - // Skip already-completed shards - if destination.exists() { - continue; - } - - Self::download_one_file( - &client, - url, - destination, - downloads, - model_id, - &mut cumulative_bytes, - start_time, - bytes_at_start, - bearer_token, - ) - .await?; - } - - Ok(()) - } - - #[allow(clippy::too_many_arguments)] - async fn download_one_file( - client: &reqwest::Client, - url: &str, - destination: &Path, - downloads: &DownloadMap, - model_id: &str, - cumulative_bytes: &mut u64, - start_time: std::time::Instant, - bytes_at_start: u64, - bearer_token: Option<&str>, - ) -> Result<(), anyhow::Error> { - let partial_path = partial_path_for(destination); - let mut retries = 0u32; - - let mut file_bytes: u64 = if partial_path.exists() { - tokio::fs::metadata(&partial_path).await?.len() - } else { - 0 - }; - - // Get this file's total size - let mut file_total: u64 = Self::apply_bearer_token(client.head(url), bearer_token) - .send() - .await - .ok() - .and_then(|r| r.content_length()) - .unwrap_or(0); - - // If partial matches expected size exactly, promote it - if file_total > 0 && file_bytes == file_total { - tokio::fs::rename(&partial_path, destination).await?; - // cumulative_bytes already accounts for this file from the pre-scan - return Ok(()); - } - - // If partial is oversized or remote changed, discard and re-download - if file_total > 0 && file_bytes > file_total { - info!(model_id = %model_id, file_bytes, file_total, "Partial file oversized, re-downloading"); - *cumulative_bytes = cumulative_bytes.saturating_sub(file_bytes); - file_bytes = 0; - let _ = tokio::fs::remove_file(&partial_path).await; - } - - loop { - if Self::is_cancelled(downloads, model_id) { - let _ = tokio::fs::remove_file(&partial_path).await; - anyhow::bail!("Download cancelled"); - } - - let mut request = Self::apply_bearer_token(client.get(url), bearer_token); - if file_bytes > 0 { - request = request.header("Range", format!("bytes={}-", file_bytes)); - } - - let response = match request.send().await { - Ok(r) => r, - Err(e) => { - if retries >= Self::MAX_RETRIES { - anyhow::bail!("Download failed after {} retries: {}", retries, e); - } - retries += 1; - let delay = std::cmp::min( - Self::RETRY_BASE_DELAY * 2u32.saturating_pow(retries - 1), - Self::RETRY_MAX_DELAY, - ); - info!(model_id = %model_id, retry = retries, delay_secs = ?delay.as_secs(), error = %e, "Retrying download after connection error"); - Self::cancellable_sleep(delay, downloads, model_id).await?; - continue; - } - }; - - let status = response.status(); - if status == reqwest::StatusCode::RANGE_NOT_SATISFIABLE { - if file_total > 0 && file_bytes == file_total { - break; - } - *cumulative_bytes = cumulative_bytes.saturating_sub(file_bytes); - file_bytes = 0; - let _ = tokio::fs::remove_file(&partial_path).await; - continue; - } - - if !status.is_success() && status != reqwest::StatusCode::PARTIAL_CONTENT { - let is_transient = status.is_server_error() - || status == reqwest::StatusCode::REQUEST_TIMEOUT - || status == reqwest::StatusCode::TOO_MANY_REQUESTS; - - if !is_transient || retries >= Self::MAX_RETRIES { - anyhow::bail!("Failed to download: HTTP {}", status); - } - retries += 1; - let delay = std::cmp::min( - Self::RETRY_BASE_DELAY * 2u32.saturating_pow(retries - 1), - Self::RETRY_MAX_DELAY, - ); - info!(model_id = %model_id, retry = retries, http_status = %status, "Retrying download after transient HTTP error"); - Self::cancellable_sleep(delay, downloads, model_id).await?; - continue; - } - - if file_bytes > 0 && status == reqwest::StatusCode::OK { - info!(model_id = %model_id, "Server ignored Range header, restarting file from scratch"); - // Subtract already-counted partial bytes from cumulative - *cumulative_bytes = cumulative_bytes.saturating_sub(file_bytes); - file_bytes = 0; - let _ = tokio::fs::remove_file(&partial_path).await; - } - - // If HEAD didn't return this file's size, learn it from the GET response. - // This block only fires once per file (file_total stays non-zero after), - // so retries don't double-count. Since download_files_sequentially's HEAD - // pass contributed 0 for this file, we add the discovered size to the - // shared total so progress/ETA are accurate. - if file_total == 0 { - let new_file_total = if file_bytes > 0 { - response - .headers() - .get("content-range") - .and_then(|v| v.to_str().ok()) - .and_then(|s| s.rsplit('/').next()) - .and_then(|s| s.parse::().ok()) - } else { - response.content_length() - }; - if let Some(t) = new_file_total { - file_total = t; - if let Ok(mut dl) = downloads.lock() { - if let Some(progress) = dl.get_mut(model_id) { - progress.total_bytes = progress.total_bytes.saturating_add(t); - } - } - } - } - - let mut file = tokio::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&partial_path) - .await?; - - let file_len = tokio::fs::metadata(&partial_path).await?.len(); - if file_len != file_bytes { - file.set_len(file_bytes).await?; - } - - let mut stream_error = false; - let mut resp = response; - - loop { - let chunk_result = resp.chunk().await; - match chunk_result { - Ok(Some(chunk)) => { - if Self::is_cancelled(downloads, model_id) { - let _ = tokio::fs::remove_file(&partial_path).await; - anyhow::bail!("Download cancelled"); - } - - file.write_all(&chunk).await?; - let chunk_len = chunk.len() as u64; - file_bytes += chunk_len; - *cumulative_bytes += chunk_len; - - let elapsed = start_time.elapsed().as_secs_f64(); - let bytes_this_session = cumulative_bytes.saturating_sub(bytes_at_start); - let speed_bps = if elapsed > 0.0 { - Some((bytes_this_session as f64 / elapsed) as u64) - } else { - None - }; - - let current_total = if let Ok(dl) = downloads.lock() { - dl.get(model_id).map(|p| p.total_bytes).unwrap_or(0) - } else { - 0 - }; - - let eta_seconds = if let Some(speed) = speed_bps { - if speed > 0 && current_total > 0 { - Some(current_total.saturating_sub(*cumulative_bytes) / speed) - } else { - None - } - } else { - None - }; - - if let Ok(mut dl) = downloads.lock() { - if let Some(progress) = dl.get_mut(model_id) { - progress.bytes_downloaded = *cumulative_bytes; - progress.progress_percent = if current_total > 0 { - (*cumulative_bytes as f64 / current_total as f64 * 100.0) as f32 - } else { - 0.0 - }; - progress.speed_bps = speed_bps; - progress.eta_seconds = eta_seconds; - } - } - } - Ok(None) => break, - Err(e) => { - info!(model_id = %model_id, bytes = *cumulative_bytes, error = %e, "Download stream interrupted, will retry"); - stream_error = true; - break; - } - } - } - - file.flush().await?; - drop(file); - - if stream_error { - if retries >= Self::MAX_RETRIES { - anyhow::bail!( - "Download failed after {} retries due to stream interruption", - retries - ); - } - retries += 1; - let delay = std::cmp::min( - Self::RETRY_BASE_DELAY * 2u32.saturating_pow(retries - 1), - Self::RETRY_MAX_DELAY, - ); - info!(model_id = %model_id, retry = retries, delay_secs = ?delay.as_secs(), "Retrying download with resume"); - Self::cancellable_sleep(delay, downloads, model_id).await?; - continue; - } - - break; - } - - tokio::fs::rename(&partial_path, destination).await?; - Ok(()) - } - - pub fn clear_completed(&self, model_id: &str) { - if let Ok(mut downloads) = self.downloads.lock() { - if let Some(progress) = downloads.get(model_id) { - let is_terminal = progress.status == DownloadStatus::Completed - || progress.status == DownloadStatus::Failed - || progress.status == DownloadStatus::Cancelled; - if is_terminal && progress.task_exited { - downloads.remove(model_id); - } - } - } - } - - fn apply_bearer_token( - request: reqwest::RequestBuilder, - bearer_token: Option<&str>, - ) -> reqwest::RequestBuilder { - if let Some(token) = bearer_token.filter(|token| !token.is_empty()) { - request.header("Authorization", format!("Bearer {}", token)) - } else { - request - } - } -} - -static DOWNLOAD_MANAGER: once_cell::sync::Lazy = - once_cell::sync::Lazy::new(DownloadManager::new); - -pub fn get_download_manager() -> &'static DownloadManager { - &DOWNLOAD_MANAGER -} diff --git a/crates/goose/src/providers/local_inference.rs b/crates/goose/src/providers/local_inference.rs index 8f84811ee..c309ebf81 100644 --- a/crates/goose/src/providers/local_inference.rs +++ b/crates/goose/src/providers/local_inference.rs @@ -1,770 +1,46 @@ -mod backend; -pub mod hf_models; -mod llamacpp; -pub mod local_model_registry; -pub mod management; -mod mlx; -pub(crate) mod multimodal; -#[cfg(feature = "mlx")] -mod native_tool_parsing; -#[cfg(feature = "mlx")] -mod tool_emulation; -mod tool_parsing; +pub use goose_providers::local_inference::*; use crate::config::ExtensionConfig; -use crate::conversation::message::{Message, MessageContent}; -use crate::providers::base::{MessageStream, Provider, ProviderDef, ProviderMetadata}; +use crate::providers::api_client::TlsConfig; +use crate::providers::base::ProviderDef; use anyhow::Result; -use async_stream::try_stream; -use async_trait::async_trait; -use backend::{BackendLoadedModel, LocalInferenceBackend}; use futures::future::BoxFuture; -use goose_providers::conversation::token_usage::{ProviderUsage, Usage}; -use goose_providers::errors::ProviderError; -use goose_providers::images::ImageFormat; -use goose_providers::model::ModelConfig; -use goose_providers::request_log::{start_log, LoggerHandleExt, RequestLogHandle}; -use llamacpp::{LlamaCppBackend, LLAMACPP_BACKEND_ID}; -use local_model_registry::ChatTemplate; -use mlx::{MlxBackend, MLX_BACKEND_ID}; -use rmcp::model::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; -type ModelSlot = Arc>>>; - -#[derive(Clone, Debug, Eq, Hash, PartialEq)] -struct ModelCacheKey { - backend_id: &'static str, - model_id: String, - chat_template: ChatTemplate, +fn resolve_huggingface_token() -> BoxFuture<'static, Result>> { + Box::pin(crate::providers::huggingface_auth::resolve_token_async()) } -impl ModelCacheKey { - fn new( - backend_id: &'static str, - model_id: impl Into, - chat_template: ChatTemplate, - ) -> Self { - Self { - backend_id, - model_id: model_id.into(), - chat_template, - } - } +fn resolve_string_param(key: &'static str) -> Result> { + Ok(crate::config::Config::global() + .get_param::(key) + .ok()) } -pub struct InferenceRuntime { - models: StdMutex>, - backends: HashMap<&'static str, Arc>, +fn resolve_bool_param(key: &'static str) -> Result> { + Ok(crate::config::Config::global().get_param::(key).ok()) } -pub fn builtin_chat_template_names() -> Vec { - llamacpp::builtin_chat_template_names() +pub fn configure_local_inference() { + huggingface_auth::set_token_resolver(resolve_huggingface_token); + config_resolver::set_string_param_resolver(resolve_string_param); + config_resolver::set_bool_param_resolver(resolve_bool_param); } -/// Global weak reference used to share a single `InferenceRuntime` across -/// all providers and management APIs. Only a `Weak` is stored here — strong -/// `Arc`s live in providers and the local-inference management layer. 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() -> Result> { - let mut guard = RUNTIME.lock().expect("runtime lock poisoned"); - if let Some(runtime) = guard.upgrade() { - return Ok(runtime); - } - let llamacpp_backend: Arc = Arc::new(LlamaCppBackend::new()?); - let mlx_backend: Arc = Arc::new(MlxBackend::new()); - let mut backends = HashMap::new(); - backends.insert(LLAMACPP_BACKEND_ID, llamacpp_backend); - backends.insert(MLX_BACKEND_ID, mlx_backend); - let runtime = Arc::new(Self { - models: StdMutex::new(HashMap::new()), - backends, - }); - *guard = Arc::downgrade(&runtime); - Ok(runtime) - } - - fn default_backend(&self) -> &dyn LocalInferenceBackend { - self.backends - .get(LLAMACPP_BACKEND_ID) - .expect("default local inference backend registered") - .as_ref() - } - - fn backend_for_model( - &self, - resolved: &ResolvedModelPaths, - ) -> Result, ProviderError> { - let backend_id = resolved - .backend_id - .as_deref() - .unwrap_or(LLAMACPP_BACKEND_ID); - self.backends.get(backend_id).cloned().ok_or_else(|| { - ProviderError::ExecutionError(format!( - "Local inference backend '{}' unavailable", - backend_id - )) - }) - } - - fn get_or_create_model_slot(&self, key: ModelCacheKey) -> ModelSlot { - let mut map = self.models.lock().expect("model cache lock poisoned"); - map.entry(key) - .or_insert_with(|| Arc::new(Mutex::new(None))) - .clone() - } - - fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec { - let map = self.models.lock().expect("model cache lock poisoned"); - map.iter() - .filter(|(key, _)| *key != keep_key) - .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"; - -#[derive(Clone)] -pub(super) struct ResolvedModelPaths { - pub model_path: PathBuf, - pub context_limit: usize, - pub settings: crate::providers::local_inference::local_model_registry::ModelSettings, - pub mmproj_path: Option, - pub backend_id: Option, - pub draft_model_path: Option, -} - -fn resolve_model_local_path(model_id: &str) -> Option { - use crate::providers::local_inference::local_model_registry::get_registry; - - get_registry() - .lock() - .ok()? - .get_model(model_id) - .map(|entry| entry.local_path.clone()) -} - -/// Resolve model path, context limit, settings, and mmproj path for a model ID from the registry. -fn resolve_model_path(model_id: &str) -> Option { - use crate::providers::local_inference::local_model_registry::{ - default_settings_for_model, 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; - let mut settings = entry.settings.clone(); - let defaults = default_settings_for_model(model_id); - settings.vision_capable = defaults.vision_capable; - settings.mmproj_size_bytes = entry.mmproj_size_bytes; - let mmproj_path = entry.mmproj_path.as_ref().filter(|p| p.exists()).cloned(); - let backend_id = entry - .backend_id - .clone() - .or_else(|| settings.backend_id.clone()); - let draft_model = settings - .draft_model - .clone() - .or_else(|| std::env::var("GOOSE_LOCAL_DRAFT_MODEL").ok()) - .filter(|draft_model| draft_model != model_id); - let draft_model_path = draft_model.as_deref().and_then(resolve_model_local_path); - return Some(ResolvedModelPaths { - model_path: entry.local_path.clone(), - context_limit: ctx, - settings, - mmproj_path, - backend_id, - draft_model_path, - }); - } - } - - None -} - -pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 { - runtime.default_backend().available_memory_bytes() -} - -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].spec.to_string() -} - -fn build_openai_messages_json( - system: &str, - messages: &[Message], - media_marker: Option<&str>, -) -> String { - use goose_providers::formats::openai::format_messages; - - let mut arr: Vec = vec![json!({"role": "system", "content": system})]; - arr.extend(format_messages(messages, &ImageFormat::OpenAi)); - strip_image_parts_from_messages(&mut arr); - if let Some(marker) = media_marker { - convert_text_media_markers(&mut arr, marker); - } - serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_string()) -} - -fn build_openai_text_messages_json( - system: &str, - messages: &[Message], - media_marker: Option<&str>, -) -> String { - let mut arr: Vec = vec![json!({"role": "system", "content": system})]; - arr.extend(messages.iter().filter_map(|m| { - let content = extract_text_content(m); - if content.trim().is_empty() { - return None; - } - let role = match m.role { - rmcp::model::Role::User => "user", - rmcp::model::Role::Assistant => "assistant", - }; - Some(json!({"role": role, "content": content})) - })); - if let Some(marker) = media_marker { - convert_text_media_markers(&mut arr, marker); - } - serde_json::to_string(&arr).unwrap_or_else(|_| "[]".to_string()) -} - -fn convert_text_media_markers(messages: &mut [Value], marker: &str) { - if marker.is_empty() { - return; - } - - for msg in messages { - let Some(content) = msg.get_mut("content") else { - continue; - }; - - if let Some(text) = content.as_str() { - if let Some(parts) = split_media_marker_text(text, marker) { - *content = json!(parts); - } - continue; - } - - let Some(content_parts) = content.as_array_mut() else { - continue; - }; - let mut updated = Vec::new(); - let mut changed = false; - for part in content_parts.iter() { - if part.get("type").and_then(|v| v.as_str()) == Some("text") { - if let Some(text) = part.get("text").and_then(|v| v.as_str()) { - if let Some(parts) = split_media_marker_text(text, marker) { - updated.extend(parts); - changed = true; - continue; - } - } - } - updated.push(part.clone()); - } - if changed { - *content_parts = updated; - } - } -} - -fn split_media_marker_text(text: &str, marker: &str) -> Option> { - let mut parts = Vec::new(); - let mut rest = text; - let mut found_marker = false; - while let Some((before, after)) = rest.split_once(marker) { - found_marker = true; - let before = before.strip_suffix('\n').unwrap_or(before); - if !before.is_empty() { - parts.push(json!({"type": "text", "text": before})); - } - parts.push(json!({"type": "media_marker", "text": marker})); - rest = after; - rest = rest.strip_prefix('\n').unwrap_or(rest); - } - if !found_marker { - return None; - } - if !rest.is_empty() { - parts.push(json!({"type": "text", "text": rest})); - } - Some(parts) -} - -/// Remove `image_url` content parts from OpenAI-format messages JSON, replacing -/// each with a text note. This prevents an FFI crash in llama.cpp which does not -/// accept `image_url` content-part types. -fn strip_image_parts_from_messages(messages: &mut [Value]) { - let mut stripped = false; - for msg in messages.iter_mut() { - if let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) { - for part in content.iter_mut() { - if part.get("type").and_then(|t| t.as_str()) == Some("image_url") { - *part = json!({ - "type": "text", - "text": "[Image attached — image input is not supported with the currently selected model]" - }); - stripped = true; - } - } - } - } - if stripped { - tracing::warn!("Stripped image content parts from messages — vision encoder not available for this model"); - } -} - -/// 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_typescript\n…\n``` `` -/// - `ToolResponse` → `Command output:\n…` -/// -/// Only `developer__shell` and `code_execution__execute_typescript` 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) => { - let text = strip_info_messages(&text.text); - if !text.trim().is_empty() { - parts.push(text); - } - } - 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_typescript\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)); - } - }, - MessageContent::Image(_) => { - parts.push( - "[Image attached — image input is not supported with the currently selected model]" - .to_string(), - ); - } - _ => {} - } - } - - parts.join("\n") -} - -fn strip_info_messages(text: &str) -> String { - let mut remaining = text; - let mut output = String::new(); - - while let Some((before, after_start)) = remaining.split_once("") { - output.push_str(before); - if let Some((_, after_end)) = after_start.split_once("") { - remaining = after_end; - } else { - remaining = ""; - break; - } - } - - output.push_str(remaining); - output.trim().to_string() -} - -/// Build a `ProviderUsage` and write the request log entry. -fn finalize_usage( - log: &mut Option>, - 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, - name: String, -} - -impl LocalInferenceProvider { - pub async fn from_env(_extensions: Vec) -> Result { - let runtime = InferenceRuntime::get_or_init()?; - Ok(Self { - runtime, - name: PROVIDER_NAME.to_string(), - }) - } -} - -impl goose_providers::base::ProviderDescriptor for LocalInferenceProvider { - 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.iter().map(|m| m.spec).collect(); - - // 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![], - ) - } +pub fn configure_huggingface_auth() { + configure_local_inference(); } impl ProviderDef for LocalInferenceProvider { type Provider = Self; fn from_env( - extensions: Vec, - _tls_config: Option, + _extensions: Vec, + _tls_config: Option, ) -> BoxFuture<'static, Result> where Self: Sized, { - Box::pin(Self::from_env(extensions)) - } -} - -#[async_trait] -impl Provider for LocalInferenceProvider { - fn get_name(&self) -> &str { - &self.name - } - - 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, - system: &str, - messages: &[Message], - tools: &[Tool], - ) -> Result { - let resolved = resolve_model_path(&model_config.model_name).ok_or_else(|| { - ProviderError::ExecutionError(format!("Model not found: {}", model_config.model_name)) - })?; - let backend = self.runtime.backend_for_model(&resolved)?; - let model_context_limit = resolved.context_limit; - let model_settings = resolved.settings.clone(); - let cache_key = ModelCacheKey::new( - backend.id(), - model_config.model_name.clone(), - model_settings.chat_template.clone(), - ); - let model_slot = self.runtime.get_or_create_model_slot(cache_key.clone()); - - // Ensure model is loaded — unload any other models first to free memory. - { - let mut model_lock = model_slot.lock().await; - if model_lock.is_none() { - for slot in self.runtime.other_model_slots(&cache_key) { - 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 resolved_for_load = resolved.clone(); - let settings_for_load = model_settings.clone(); - let backend_for_load = backend.clone(); - let loaded = tokio::task::spawn_blocking(move || { - backend_for_load.load_model(&model_id, &resolved_for_load, &settings_for_load) - }) - .await - .map_err(|e| ProviderError::ExecutionError(e.to_string()))??; - *model_lock = Some(loaded); - } - } - - // Allow request_params to override thinking - let mut model_settings = model_settings; - if let Some(false) = model_config - .request_param::("enable_thinking") - .or_else(|| { - crate::config::Config::global() - .get_param("GOOSE_LOCAL_ENABLE_THINKING") - .ok() - }) - { - model_settings.enable_thinking = false; - } - - let model_arc = model_slot.clone(); - let backend = backend.clone(); - let model_name = model_config.model_name.clone(); - let temperature = model_config.temperature; - let max_tokens = model_config.max_tokens; - let context_limit = model_context_limit; - let settings = model_settings; - let resolved_model = resolved.clone(); - let system = system.to_string(); - let messages = messages.to_vec(); - let tools = tools.to_vec(); - let log_payload = serde_json::json!({ - "system": &system, - "messages": messages.iter().map(|m| { - serde_json::json!({ - "role": match m.role { rmcp::model::Role::User => "user", rmcp::model::Role::Assistant => "assistant" }, - "content": extract_text_content(m), - }) - }).collect::>(), - "tools": tools.iter().map(|t| &t.name).collect::>(), - "settings": { - "tool_calling": settings.tool_calling, - "chat_template": settings.chat_template, - "context_size": settings.context_size, - "sampling": settings.sampling, - }, - }); - - let mut log = start_log(model_config, &log_payload)?; - - 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 mut model_guard = model_arc.blocking_lock(); - let loaded = match model_guard.as_mut() { - Some(l) => l, - None => { - send_err!(ProviderError::ExecutionError( - "Model not loaded".to_string() - )); - } - }; - - let message_id = Uuid::new_v4().to_string(); - - let request = backend::LocalGenerationRequest { - model_name, - system: &system, - messages: &messages, - tools: &tools, - settings: &settings, - temperature, - max_tokens, - context_limit, - resolved_model: &resolved_model, - draft_model_path: resolved_model.draft_model_path.clone(), - message_id: &message_id, - tx: &tx, - log: &mut log, - }; - - let result = backend.generate(loaded.as_mut(), request); - - 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; - } - - })) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn converts_marker_in_string_content_to_media_marker_part() { - let mut messages = vec![json!({ - "role": "user", - "content": "look\n<__media__>\nclosely", - })]; - - convert_text_media_markers(&mut messages, "<__media__>"); - - assert_eq!( - messages[0]["content"], - json!([ - {"type": "text", "text": "look"}, - {"type": "media_marker", "text": "<__media__>"}, - {"type": "text", "text": "closely"}, - ]) - ); - } - - #[test] - fn converts_marker_inside_text_content_parts() { - let mut messages = vec![json!({ - "role": "user", - "content": [ - {"type": "text", "text": "<__media__>describe"}, - {"type": "text", "text": "next"}, - {"type": "media_marker", "text": "<__media__>"}, - ], - })]; - - convert_text_media_markers(&mut messages, "<__media__>"); - - assert_eq!( - messages[0]["content"], - json!([ - {"type": "media_marker", "text": "<__media__>"}, - {"type": "text", "text": "describe"}, - {"type": "text", "text": "next"}, - {"type": "media_marker", "text": "<__media__>"}, - ]) - ); + configure_local_inference(); + Box::pin(Self::from_env()) } } diff --git a/deny.toml b/deny.toml index a9e0ba19a..52370e3e9 100644 --- a/deny.toml +++ b/deny.toml @@ -6,8 +6,13 @@ yanked = "deny" unmaintained = "none" unsound = "none" -# Ignore proc-macro-error unmaintained warning -# See: https://github.com/aaif-goose/goose/issues/7008 ignore = [ - "RUSTSEC-2023-0071", # rsa: Marvin Attack timing sidechannel (no safe upgrade available, via jsonwebtoken) + # rsa: Marvin Attack timing sidechannel. No safe upgrade is available, and + # the reachable path is through jsonwebtoken. + "RUSTSEC-2023-0071", + + # quick-xml: quadratic duplicate-attribute checks. Current paths are through + # docx-rs, umya-spreadsheet, and bat/plist, and their latest releases do not + # yet expose a quick-xml >= 0.41.0 upgrade path. + "RUSTSEC-2026-0194", ]