vrr/goose
2
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-07-09 16:09:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 22

Switch the local inference provider MLX backend to use the safemlx crate (#10304)

Browse source
This commit is contained in:
Jasper 2026-07-07 14:01:15 -07:00 • committed by GitHub
parent f96f62d985
commit 1bced6616b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
58 changed files with 1835 additions and 646 deletions
Download patch file Download diff file Expand all files Collapse all files

318
Cargo.lock generated
View file

@ -1201,7 +1201,7 @@ dependencies = [
"bincode", "bincode",
"bytesize", "bytesize",
"clircle", "clircle",
"console 0.16.3", "console",
"content_inspector", "content_inspector",
"encoding_rs", "encoding_rs",
"flate2", "flate2",
@ -1255,26 +1255,6 @@ dependencies = [
"serde", "serde",
] ]
[[package]]
name = "bindgen"
version = "0.70.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f49d8fed880d473ea71efb9bf597651e77201bdd4893efe54c9e5d65ae04ce6f"
dependencies = [
"bitflags 2.11.1",
"cexpr",
"clang-sys",
"itertools 0.13.0",
"log",
"prettyplease",
"proc-macro2",
"quote",
"regex",
"rustc-hash 1.1.0",
"shlex 1.3.0",
"syn 2.0.117",
]
[[package]] [[package]]
name = "bindgen" name = "bindgen"
version = "0.72.1" version = "0.72.1"
@ -2302,8 +2282,8 @@ version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4529f45438fc25ca048b242d5c48e2d3ce9a521e2a5a9123d9737d8520b030dd" checksum = "4529f45438fc25ca048b242d5c48e2d3ce9a521e2a5a9123d9737d8520b030dd"
dependencies = [ dependencies = [
"console 0.16.3", "console",
"indicatif 0.18.4", "indicatif",
"once_cell", "once_cell",
"strsim", "strsim",
"textwrap", "textwrap",
@ -2478,19 +2458,6 @@ dependencies = [
"winnow 1.0.3", "winnow 1.0.3",
] ]
[[package]]
name = "console"
version = "0.15.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8"
dependencies = [
"encode_unicode",
"libc",
"once_cell",
"unicode-width 0.2.2",
"windows-sys 0.59.0",
]
[[package]] [[package]]
name = "console" name = "console"
version = "0.16.3" version = "0.16.3"
@ -5024,7 +4991,7 @@ dependencies = [
"clap_mangen", "clap_mangen",
"cliclack", "cliclack",
"comfy-table", "comfy-table",
"console 0.16.3", "console",
"dotenvy", "dotenvy",
"env-lock", "env-lock",
"etcetera 0.11.0", "etcetera 0.11.0",
@ -5032,7 +4999,7 @@ dependencies = [
"goose", "goose",
"goose-mcp", "goose-mcp",
"goose-providers", "goose-providers",
"indicatif 0.18.4", "indicatif",
"open", "open",
"rand 0.10.1", "rand 0.10.1",
"regex", "regex",
@ -5089,17 +5056,17 @@ dependencies = [
"goose-download-manager", "goose-download-manager",
"goose-provider-types", "goose-provider-types",
"goose-sdk-types", "goose-sdk-types",
"hf-hub 1.0.0-rc.1", "hf-hub",
"include_dir", "include_dir",
"llama-cpp-2", "llama-cpp-2",
"llama-cpp-sys-2", "llama-cpp-sys-2",
"minijinja", "minijinja",
"mlx-lm",
"mlx-lm-utils",
"mlx-rs",
"regex", "regex",
"reqwest 0.13.4", "reqwest 0.13.4",
"rmcp", "rmcp",
"safemlx",
"safemlx-lm",
"safemlx-lm-utils",
"serde", "serde",
"serde_json", "serde_json",
"tempfile", "tempfile",
@ -5447,25 +5414,6 @@ dependencies = [
"arrayvec", "arrayvec",
] ]
[[package]]
name = "hf-hub"
version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97"
dependencies = [
"dirs",
"http 1.4.2",
"indicatif 0.17.11",
"libc",
"log",
"rand 0.9.4",
"serde",
"serde_json",
"thiserror 2.0.18",
"ureq",
"windows-sys 0.60.2",
]
[[package]] [[package]]
name = "hf-hub" name = "hf-hub"
version = "1.0.0-rc.1" version = "1.0.0-rc.1"
@ -6070,26 +6018,13 @@ dependencies = [
"serde_core", "serde_core",
] ]
[[package]]
name = "indicatif"
version = "0.17.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
dependencies = [
"console 0.15.11",
"number_prefix",
"portable-atomic",
"unicode-width 0.2.2",
"web-time",
]
[[package]] [[package]]
name = "indicatif" name = "indicatif"
version = "0.18.4" version = "0.18.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb"
dependencies = [ dependencies = [
"console 0.16.3", "console",
"portable-atomic", "portable-atomic",
"unicode-width 0.2.2", "unicode-width 0.2.2",
"unit-prefix", "unit-prefix",
@ -6611,7 +6546,7 @@ version = "0.1.146"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b291e4bc2d10c43cd8dec16d49b6104cb3cb125f596ec380a753a5db1d965dd" checksum = "9b291e4bc2d10c43cd8dec16d49b6104cb3cb125f596ec380a753a5db1d965dd"
dependencies = [ dependencies = [
"bindgen 0.72.1", "bindgen",
"cc", "cc",
"cmake", "cmake",
"find_cuda_helper", "find_cuda_helper",
@ -6862,99 +6797,6 @@ dependencies = [
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
[[package]]
name = "mlx-internal-macros"
version = "0.25.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7a7c4444d624bf6b93db5cc22ebff4fdfa13593fd56154fe33b1f302a557c2c6"
dependencies = [
"darling 0.21.3",
"itertools 0.14.0",
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "mlx-lm"
version = "0.0.1"
source = "git+https://github.com/jh-block/mlx-lm#4cb8572045f55f170eb6e41005a3f8a730673c42"
dependencies = [
"anyhow",
"clap",
"idna_adapter",
"minijinja",
"mlx-lm-utils",
"mlx-macros",
"mlx-rs",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokenizers 0.22.2",
]
[[package]]
name = "mlx-lm-utils"
version = "0.0.1"
source = "git+https://github.com/jh-block/mlx-lm#4cb8572045f55f170eb6e41005a3f8a730673c42"
dependencies = [
"minijinja",
"minijinja-contrib",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokenizers 0.22.2",
]
[[package]]
name = "mlx-macros"
version = "0.25.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a819ee8b4434690572b6feb9c3ef0b6e90137e4190b340cf00150703b410aaf9"
dependencies = [
"darling 0.21.3",
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "mlx-rs"
version = "0.25.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f3a0f592c5839b0237b3072530b1d3c503923579a2009ee0761df3edd1b1f27b"
dependencies = [
"bytemuck",
"dyn-clone",
"half",
"itertools 0.14.0",
"libc",
"mach-sys",
"mlx-internal-macros",
"mlx-macros",
"mlx-sys",
"num-complex",
"num-traits",
"num_enum",
"parking_lot",
"paste",
"safetensors 0.6.2",
"smallvec",
"strum 0.27.2",
"thiserror 2.0.18",
]
[[package]]
name = "mlx-sys"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2e3bc3880111918b2d5018f845d48fd995f9901f16efc81d1fcfd2f4210b8219"
dependencies = [
"bindgen 0.70.1",
"cc",
"cmake",
]
[[package]] [[package]]
name = "mockall" name = "mockall"
version = "0.15.0" version = "0.15.0"
@ -7398,19 +7240,13 @@ dependencies = [
"syn 2.0.117", "syn 2.0.117",
] ]
[[package]]
name = "number_prefix"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
[[package]] [[package]]
name = "oauth2" name = "oauth2"
version = "5.0.0" version = "5.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d"
dependencies = [ dependencies = [
"base64 0.21.7", "base64 0.22.1",
"chrono", "chrono",
"getrandom 0.2.17", "getrandom 0.2.17",
"http 1.4.2", "http 1.4.2",
@ -8564,7 +8400,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"itertools 0.13.0", "itertools 0.14.0",
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.117", "syn 2.0.117",
@ -9347,7 +9183,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b"
dependencies = [ dependencies = [
"aws-lc-rs", "aws-lc-rs",
"log",
"once_cell", "once_cell",
"ring", "ring",
"rustls-pki-types", "rustls-pki-types",
@ -9462,6 +9297,99 @@ version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3944826ff8fa8093089aba3acb4ef44b9446a99a16f3bf4e74af3f77d340ab7d" checksum = "3944826ff8fa8093089aba3acb4ef44b9446a99a16f3bf4e74af3f77d340ab7d"
[[package]]
name = "safemlx"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f501c4f9748dfe9ccc905e7e13ae88873d96a00735268c69258e158da07a777f"
dependencies = [
"bytemuck",
"dyn-clone",
"half",
"itertools 0.14.0",
"libc",
"mach-sys",
"num-complex",
"num-traits",
"num_enum",
"parking_lot",
"paste",
"safemlx-internal-macros",
"safemlx-macros",
"safemlx-sys",
"safetensors 0.6.2",
"smallvec",
"strum 0.27.2",
"thiserror 2.0.18",
]
[[package]]
name = "safemlx-internal-macros"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9708b1312176963a4ea6412859e6655e56a83fe6ab936142715aa0e9fd662fe5"
dependencies = [
"darling 0.21.3",
"itertools 0.14.0",
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "safemlx-lm"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78f450654386009003b745adfdae7af2c7f44ae749d75621a2a0b07607945e9f"
dependencies = [
"anyhow",
"clap",
"idna_adapter",
"minijinja",
"safemlx",
"safemlx-lm-utils",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokenizers 0.22.2",
]
[[package]]
name = "safemlx-lm-utils"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "583d39124dae0b4a93ac7528d16d543f5e717c21365e6b034a8a74d8e1b4eb3a"
dependencies = [
"minijinja",
"minijinja-contrib",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokenizers 0.22.2",
]
[[package]]
name = "safemlx-macros"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d549c5fdbe69b904ef84c5cf352add203de1093b3c636e8133d94c24cfc33ce0"
dependencies = [
"darling 0.21.3",
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "safemlx-sys"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "552285b007dab687890f5c49382e52e69a3bb1359bb2c8f2c2743743509af071"
dependencies = [
"cc",
"cmake",
]
[[package]] [[package]]
name = "safetensors" name = "safetensors"
version = "0.4.5" version = "0.4.5"
@ -10317,17 +10245,6 @@ dependencies = [
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
[[package]]
name = "socks"
version = "0.3.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b"
dependencies = [
"byteorder",
"libc",
"winapi",
]
[[package]] [[package]]
name = "sourcemap" name = "sourcemap"
version = "9.3.2" version = "9.3.2"
@ -11777,8 +11694,6 @@ dependencies = [
"derive_builder", "derive_builder",
"esaxx-rs", "esaxx-rs",
"getrandom 0.3.4", "getrandom 0.3.4",
"hf-hub 0.4.3",
"indicatif 0.18.4",
"itertools 0.14.0", "itertools 0.14.0",
"log", "log",
"macro_rules_attribute", "macro_rules_attribute",
@ -12810,25 +12725,6 @@ version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
[[package]]
name = "ureq"
version = "2.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d"
dependencies = [
"base64 0.22.1",
"flate2",
"log",
"once_cell",
"rustls",
"rustls-pki-types",
"serde",
"serde_json",
"socks",
"url",
"webpki-roots 0.26.11",
]
[[package]] [[package]]
name = "url" name = "url"
version = "2.5.8" version = "2.5.8"
@ -12948,7 +12844,7 @@ version = "145.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6f84468f393984251db025e944544590a8fb429d5088b78102bd72f09232f0da" checksum = "6f84468f393984251db025e944544590a8fb429d5088b78102bd72f09232f0da"
dependencies = [ dependencies = [
"bindgen 0.72.1", "bindgen",
"bitflags 2.11.1", "bitflags 2.11.1",
"fslock", "fslock",
"gzip-header", "gzip-header",

54
crates/goose-cli/src/session/mod.rs
View file

@ -53,7 +53,7 @@ use std::collections::{HashMap, HashSet};
use std::io::IsTerminal; use std::io::IsTerminal;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::Arc; use std::sync::Arc;
use std::time::Instant; use std::time::{Duration, Instant};
use tokio; use tokio;
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
use tracing::warn; use tracing::warn;
@ -1769,25 +1769,54 @@ fn print_run_stats(
usage: Option<&ProviderUsage>, usage: Option<&ProviderUsage>,
) { ) {
let elapsed = run_started.elapsed(); let elapsed = run_started.elapsed();
let stats = usage.and_then(|usage| usage.stats.as_ref());
let generation_elapsed = stats
.and_then(|stats| stats.elapsed_ms)
.map(Duration::from_millis);
let output_tokens = usage let output_tokens = usage
.and_then(|usage| usage.usage.output_tokens) .and_then(|usage| usage.usage.output_tokens)
.and_then(|tokens| usize::try_from(tokens).ok()) .and_then(|tokens| usize::try_from(tokens).ok())
.or_else(|| usage.and_then(|usage| usage.stats.as_ref()?.output_tokens)); .or_else(|| stats.and_then(|stats| stats.output_tokens));
let tokens_per_second = output_tokens.map(|tokens| { let tokens_per_second = output_tokens.map(|tokens| {
if elapsed.as_secs_f64() > 0.0 { let rate_elapsed = generation_elapsed.unwrap_or(elapsed);
tokens as f64 / elapsed.as_secs_f64() if rate_elapsed.as_secs_f64() > 0.0 {
tokens as f64 / rate_elapsed.as_secs_f64()
} else { } else {
0.0 0.0
} }
}); });
let model_load_ms = stats.and_then(|stats| stats.model_load_ms);
let generation_time_to_first_token_ms = stats.and_then(|stats| stats.time_to_first_token_ms);
eprintln!("\nStats:"); eprintln!("\nStats:");
match first_token_at { if let Some(ms) = model_load_ms {
Some(first) => eprintln!( eprintln!(" Model load: {:.2}s", ms as f64 / 1000.0);
" Time to first token: {:.2}s", }
first.duration_since(run_started).as_secs_f64() if model_load_ms.is_some() {
), match generation_time_to_first_token_ms {
None => eprintln!(" Time to first token: unavailable"), Some(ms) => eprintln!(
" Generation time to first token: {:.2}s",
ms as f64 / 1000.0
),
None => eprintln!(" Generation time to first token: unavailable"),
}
match first_token_at {
Some(first) => eprintln!(
" End-to-end time to first token: {:.2}s",
first.duration_since(run_started).as_secs_f64()
),
None => eprintln!(" End-to-end time to first token: unavailable"),
}
} else if let Some(ms) = generation_time_to_first_token_ms {
eprintln!(" Time to first token: {:.2}s", ms as f64 / 1000.0);
} else {
match first_token_at {
Some(first) => eprintln!(
" Time to first token: {:.2}s",
first.duration_since(run_started).as_secs_f64()
),
None => eprintln!(" Time to first token: unavailable"),
}
} }
match tokens_per_second { match tokens_per_second {
Some(rate) => eprintln!(" Tokens/sec: {:.2}", rate), Some(rate) => eprintln!(" Tokens/sec: {:.2}", rate),
@ -1797,10 +1826,7 @@ fn print_run_stats(
eprintln!(" Output tokens: {tokens}"); eprintln!(" Output tokens: {tokens}");
} }
if let Some(draft) = usage if let Some(draft) = stats.and_then(|stats| stats.draft.as_ref()) {
.and_then(|usage| usage.stats.as_ref())
.and_then(|stats| stats.draft.as_ref())
{
eprintln!(" Draft accept rate: {:.1}%", draft.accept_rate * 100.0); eprintln!(" Draft accept rate: {:.1}%", draft.accept_rate * 100.0);
eprintln!( eprintln!(
" Draft tokens: {} accepted: {} target verified: {} rounds: {}", " Draft tokens: {} accepted: {} target verified: {} rounds: {}",

6
crates/goose-cli/src/session/output.rs
View file

@ -248,7 +248,8 @@ pub fn render_message(message: &Message, debug: bool) {
} }
MessageContent::SystemNotification(notification) => { MessageContent::SystemNotification(notification) => {
match notification.notification_type { match notification.notification_type {
SystemNotificationType::ThinkingMessage => { SystemNotificationType::ThinkingMessage
| SystemNotificationType::ProgressMessage => {
show_thinking(); show_thinking();
set_thinking_message(&notification.msg); set_thinking_message(&notification.msg);
} }
@ -330,7 +331,8 @@ pub fn render_message_streaming(
} }
MessageContent::SystemNotification(notification) => { MessageContent::SystemNotification(notification) => {
match notification.notification_type { match notification.notification_type {
SystemNotificationType::ThinkingMessage => { SystemNotificationType::ThinkingMessage
| SystemNotificationType::ProgressMessage => {
show_thinking(); show_thinking();
set_thinking_message(&notification.msg); set_thinking_message(&notification.msg);
} }

8
crates/goose-local-inference/Cargo.toml
View file

@ -15,7 +15,7 @@ workspace = true
default = [] default = []
cuda = ["llama-cpp-2/cuda"] cuda = ["llama-cpp-2/cuda"]
vulkan = ["llama-cpp-2/vulkan"] vulkan = ["llama-cpp-2/vulkan"]
mlx = ["dep:mlx-rs", "dep:mlx-lm", "dep:mlx-lm-utils"] mlx = ["dep:safemlx", "dep:safemlx-lm", "dep:safemlx-lm-utils"]
[dependencies] [dependencies]
anyhow = { workspace = true } anyhow = { workspace = true }
@ -46,9 +46,9 @@ utoipa = { workspace = true, features = ["chrono"] }
uuid = { workspace = true, features = ["v4", "std"] } uuid = { workspace = true, features = ["v4", "std"] }
tempfile = { workspace = true } tempfile = { workspace = true }
mlx-rs = { version = "0.25.3", default-features = false, features = ["accelerate", "metal", "safetensors"], optional = true } safemlx = { default-features = false, features = ["accelerate", "metal", "safetensors"], optional = true, version = "0.1.2" }
mlx-lm = { git = "https://github.com/jh-block/mlx-lm", optional = true } safemlx-lm = { optional = true, version = "0.1.5" }
mlx-lm-utils = { git = "https://github.com/jh-block/mlx-lm", optional = true } safemlx-lm-utils = { optional = true, version = "0.1.2" }
[target.'cfg(target_os = "macos")'.dependencies] [target.'cfg(target_os = "macos")'.dependencies]
llama-cpp-2 = { workspace = true, features = ["sampler", "metal", "mtmd"] } llama-cpp-2 = { workspace = true, features = ["sampler", "metal", "mtmd"] }

1
crates/goose-local-inference/src/backend.rs
View file

@ -22,6 +22,7 @@ pub(super) struct LocalGenerationRequest<'a> {
pub temperature: Option<f32>, pub temperature: Option<f32>,
pub max_tokens: Option<i32>, pub max_tokens: Option<i32>,
pub context_limit: usize, pub context_limit: usize,
pub model_load_ms: Option<u64>,
pub resolved_model: &'a ResolvedModelPaths, pub resolved_model: &'a ResolvedModelPaths,
pub draft_model_path: Option<std::path::PathBuf>, pub draft_model_path: Option<std::path::PathBuf>,
pub message_id: &'a str, pub message_id: &'a str,

333
crates/goose-local-inference/src/hf_models.rs
View file

@ -1085,20 +1085,13 @@ mod tests {
} }
#[test] #[test]
fn mlx_compatible_repo_accepts_supported_tokenizer_formats() { fn mlx_compatible_repo_accepts_tokenizer_json() {
let config = Some(serde_json::json!({ "model_type": "llama" })); let config = Some(serde_json::json!({ "model_type": "llama" }));
for tokenizer_files in [ assert!(is_mlx_compatible_repo(
vec!["tokenizer.json"], &config,
vec!["tokenizer.model"], &mlx_siblings(&["tokenizer.json"])
vec!["tokenizer.tiktoken"], ));
vec!["vocab.json", "merges.txt"],
] {
assert!(
is_mlx_compatible_repo(&config, &mlx_siblings(&tokenizer_files)),
"{tokenizer_files:?}"
);
}
} }
#[test] #[test]
@ -1107,8 +1100,11 @@ mod tests {
for tokenizer_files in [ for tokenizer_files in [
vec!["tokenizer_config.json"], vec!["tokenizer_config.json"],
vec!["tokenizer.model"],
vec!["tokenizer.tiktoken"],
vec!["vocab.json"], vec!["vocab.json"],
vec!["merges.txt"], vec!["merges.txt"],
vec!["vocab.json", "merges.txt"],
] { ] {
assert!( assert!(
!is_mlx_compatible_repo(&config, &mlx_siblings(&tokenizer_files)), !is_mlx_compatible_repo(&config, &mlx_siblings(&tokenizer_files)),
@ -1117,6 +1113,75 @@ mod tests {
} }
} }
#[test]
fn mlx_download_filenames_include_fp8_snapshot_metadata() {
let siblings = [
"config.json",
"configuration.json",
"generation_config.json",
"model.safetensors.index.json",
"layers-0.safetensors",
"outside.safetensors",
"tokenizer.json",
"tokenizer_config.json",
"chat_template.jinja",
"preprocessor_config.json",
"video_preprocessor_config.json",
"README.md",
]
.into_iter()
.map(sibling)
.collect::<Vec<_>>();
let filenames = mlx_download_filenames(&siblings)
.into_iter()
.collect::<std::collections::HashSet<_>>();
for filename in [
"configuration.json",
"chat_template.jinja",
"preprocessor_config.json",
"video_preprocessor_config.json",
] {
assert!(filenames.contains(filename), "{filename}");
}
assert!(!filenames.contains("README.md"));
}
#[test]
fn mlx_download_size_uses_safetensors_metadata_when_sibling_sizes_are_missing() {
let info: ModelInfo = serde_json::from_value(serde_json::json!({
"id": "owner/repo",
"safetensors": {
"parameters": {
"BF16": 10,
"F8_E4M3": 20
},
"total": 30
}
}))
.unwrap();
let siblings = vec![
RepoSibling {
rfilename: "config.json".to_string(),
size: None,
lfs: None,
},
RepoSibling {
rfilename: "model.safetensors".to_string(),
size: None,
lfs: None,
},
];
assert_eq!(mlx_download_size_bytes(&info, &siblings), 40);
}
#[test]
fn mlx_variant_id_detects_fp8_repo_name() {
assert_eq!(mlx_variant_id("Qwen/Qwen3.6-35B-A3B-FP8", &None), "fp8");
}
fn test_model(repo_id: &str) -> HfModelInfo { fn test_model(repo_id: &str) -> HfModelInfo {
HfModelInfo { HfModelInfo {
repo_id: repo_id.to_string(), repo_id: repo_id.to_string(),
@ -1157,6 +1222,13 @@ mod tests {
assert_eq!(results[0].repo_id, "gguf/repo"); assert_eq!(results[0].repo_id, "gguf/repo");
} }
#[test]
fn best_download_count_ignores_zero_primary() {
assert_eq!(best_download_count(Some(0), Some(42)), Some(42));
assert_eq!(best_download_count(Some(7), Some(42)), Some(7));
assert_eq!(best_download_count(Some(0), Some(0)), None);
}
#[test] #[test]
fn test_recommend_variant() { fn test_recommend_variant() {
let variants = vec![ let variants = vec![
@ -1468,7 +1540,13 @@ async fn search_mlx_models_with_query(query: &str, limit: usize) -> Result<Vec<H
if results.len() >= limit { if results.len() >= limit {
break; break;
} }
if let Some(model) = get_local_model_info_for_repo_with_client(&client, &info.id).await? { if let Some(model) = get_local_model_info_for_repo_with_client_and_downloads(
&client,
&info.id,
info.downloads,
)
.await?
{
results.push(model); results.push(model);
} }
} }
@ -1484,6 +1562,14 @@ async fn get_local_model_info_for_repo(repo_id: &str) -> Result<Option<HfModelIn
async fn get_local_model_info_for_repo_with_client( async fn get_local_model_info_for_repo_with_client(
client: &HFClient, client: &HFClient,
repo_id: &str, repo_id: &str,
) -> Result<Option<HfModelInfo>> {
get_local_model_info_for_repo_with_client_and_downloads(client, repo_id, None).await
}
async fn get_local_model_info_for_repo_with_client_and_downloads(
client: &HFClient,
repo_id: &str,
downloads_hint: Option<u64>,
) -> Result<Option<HfModelInfo>> { ) -> Result<Option<HfModelInfo>> {
let repo = model_repo(client, repo_id)?; let repo = model_repo(client, repo_id)?;
let info = repo let info = repo
@ -1495,7 +1581,7 @@ async fn get_local_model_info_for_repo_with_client(
]) ])
.send() .send()
.await?; .await?;
model_info_to_local_model_info(info).await model_info_to_local_model_info(&repo, info, downloads_hint).await
} }
async fn get_exact_name_local_model_info(model_name: &str) -> Result<Option<HfModelInfo>> { async fn get_exact_name_local_model_info(model_name: &str) -> Result<Option<HfModelInfo>> {
@ -1508,7 +1594,11 @@ async fn get_exact_name_local_model_info(model_name: &str) -> Result<Option<HfMo
Ok(None) Ok(None)
} }
async fn model_info_to_local_model_info(info: ModelInfo) -> Result<Option<HfModelInfo>> { async fn model_info_to_local_model_info(
repo: &HFRepository<RepoTypeModel>,
info: ModelInfo,
downloads_hint: Option<u64>,
) -> Result<Option<HfModelInfo>> {
let repo_id = info.id.clone(); let repo_id = info.id.clone();
let mut variants: Vec<HfModelVariant> = get_repo_gguf_variants(&repo_id) let mut variants: Vec<HfModelVariant> = get_repo_gguf_variants(&repo_id)
.await .await
@ -1516,7 +1606,13 @@ async fn model_info_to_local_model_info(info: ModelInfo) -> Result<Option<HfMode
.iter() .iter()
.map(|variant| variant.to_model_variant(&repo_id)) .map(|variant| variant.to_model_variant(&repo_id))
.collect(); .collect();
variants.extend(mlx_variants_from_model_info(&repo_id, &info)); if is_mlx_compatible_model_info(&info) {
let mlx_config = load_repo_config_json(repo).await.unwrap_or_else(|error| {
tracing::debug!(repo_id, %error, "Failed to load MLX config.json; falling back to API config");
info.config.clone()
});
variants.extend(mlx_variants_from_model_info(&repo_id, &info, &mlx_config));
}
if variants.is_empty() { if variants.is_empty() {
return Ok(None); return Ok(None);
@ -1530,17 +1626,49 @@ async fn model_info_to_local_model_info(info: ModelInfo) -> Result<Option<HfMode
.next_back() .next_back()
.unwrap_or(&repo_id) .unwrap_or(&repo_id)
.to_string(); .to_string();
let downloads = match best_download_count(info.downloads, downloads_hint) {
Some(downloads) => downloads,
None => get_repo_downloads(&repo_id).await?.unwrap_or(0),
};
Ok(Some(HfModelInfo { Ok(Some(HfModelInfo {
repo_id, repo_id,
author, author,
model_name, model_name,
downloads: info.downloads.unwrap_or(0), downloads,
gguf_files: Vec::new(), gguf_files: Vec::new(),
variants, variants,
})) }))
} }
fn best_download_count(primary: Option<u64>, hint: Option<u64>) -> Option<u64> {
primary
.filter(|downloads| *downloads > 0)
.or_else(|| hint.filter(|downloads| *downloads > 0))
}
async fn get_repo_downloads(repo_id: &str) -> Result<Option<u64>> {
let client = reqwest::Client::new();
let token = optional_hf_token(huggingface_auth::resolve_token_async()).await;
let url = format!("{}/{}", HF_API_BASE, repo_id);
let response = apply_hf_auth(client.get(&url), token.as_deref())
.header("User-Agent", "goose-ai-agent")
.send()
.await?;
if !response.status().is_success() {
bail!(
"HuggingFace API returned status {} for repo {}",
response.status(),
repo_id
);
}
let model: HfApiModel = response.json().await?;
Ok(model.downloads)
}
pub async fn get_repo_local_variants(repo_id: &str) -> Result<Vec<HfModelVariant>> { pub async fn get_repo_local_variants(repo_id: &str) -> Result<Vec<HfModelVariant>> {
let mut variants: Vec<HfModelVariant> = get_repo_gguf_variants(repo_id) let mut variants: Vec<HfModelVariant> = get_repo_gguf_variants(repo_id)
.await .await
@ -1570,25 +1698,39 @@ pub async fn get_repo_mlx_variants(repo_id: &str) -> Result<Vec<HfModelVariant>>
]) ])
.send() .send()
.await?; .await?;
Ok(mlx_variants_from_model_info(repo_id, &info)) if !is_mlx_compatible_model_info(&info) {
return Ok(Vec::new());
}
let mlx_config = load_repo_config_json(&repo)
.await
.unwrap_or_else(|_| info.config.clone());
Ok(mlx_variants_from_model_info(repo_id, &info, &mlx_config))
} }
fn mlx_variants_from_model_info(repo_id: &str, info: &ModelInfo) -> Vec<HfModelVariant> { async fn load_repo_config_json(
repo: &HFRepository<RepoTypeModel>,
) -> Result<Option<serde_json::Value>> {
let config_path = repo
.download_file()
.filename("config.json".to_string())
.send()
.await?;
let config_json = tokio::fs::read_to_string(config_path).await?;
Ok(Some(serde_json::from_str(&config_json)?))
}
fn mlx_variants_from_model_info(
repo_id: &str,
info: &ModelInfo,
mlx_config: &Option<serde_json::Value>,
) -> Vec<HfModelVariant> {
let siblings = info.siblings.as_deref().unwrap_or(&[]); let siblings = info.siblings.as_deref().unwrap_or(&[]);
if !is_mlx_compatible_repo(&info.config, siblings) { if !is_mlx_compatible_repo(&info.config, siblings) {
return Vec::new(); return Vec::new();
} }
let size_bytes = mlx_download_filenames(siblings) let size_bytes = mlx_download_size_bytes(info, siblings);
.into_iter()
.filter_map(|filename| {
siblings
.iter()
.find(|s| s.rfilename == filename)
.and_then(|s| s.size)
})
.sum();
let variant_id = mlx_variant_id(repo_id, &info.config); let variant_id = mlx_variant_id(repo_id, &info.config);
vec![HfModelVariant { vec![HfModelVariant {
@ -1601,20 +1743,24 @@ fn mlx_variants_from_model_info(repo_id: &str, info: &ModelInfo) -> Vec<HfModelV
size_bytes, size_bytes,
filename: None, filename: None,
download_url: None, download_url: None,
description: mlx_variant_description(&info.config), description: mlx_variant_description(mlx_config),
quality_rank: 91, quality_rank: 91,
sharded: siblings sharded: siblings
.iter() .iter()
.filter(|s| s.rfilename.ends_with(".safetensors")) .filter(|s| s.rfilename.ends_with(".safetensors"))
.count() .count()
> 1, > 1,
supported: is_mlx_runtime_supported(&info.config) supported: is_mlx_runtime_supported(mlx_config)
&& cfg!(target_os = "macos") && cfg!(target_os = "macos")
&& cfg!(feature = "mlx"), && cfg!(feature = "mlx"),
unsupported_reason: mlx_unsupported_reason(&info.config), unsupported_reason: mlx_unsupported_reason(mlx_config),
}] }]
} }
fn is_mlx_compatible_model_info(info: &ModelInfo) -> bool {
is_mlx_compatible_repo(&info.config, info.siblings.as_deref().unwrap_or_default())
}
fn is_mlx_compatible_repo(config: &Option<serde_json::Value>, siblings: &[RepoSibling]) -> bool { fn is_mlx_compatible_repo(config: &Option<serde_json::Value>, siblings: &[RepoSibling]) -> bool {
let has_config = siblings.iter().any(|s| s.rfilename == "config.json"); let has_config = siblings.iter().any(|s| s.rfilename == "config.json");
let has_tokenizer = has_mlx_tokenizer(siblings); let has_tokenizer = has_mlx_tokenizer(siblings);
@ -1625,16 +1771,50 @@ fn is_mlx_compatible_repo(config: &Option<serde_json::Value>, siblings: &[RepoSi
has_config && has_tokenizer && has_safetensors && mlx_model_type(config).is_some() has_config && has_tokenizer && has_safetensors && mlx_model_type(config).is_some()
} }
fn mlx_download_size_bytes(info: &ModelInfo, siblings: &[RepoSibling]) -> u64 {
let sibling_size: u64 = mlx_download_filenames(siblings)
.into_iter()
.filter_map(|filename| {
siblings
.iter()
.find(|s| s.rfilename == filename)
.and_then(|s| s.size)
})
.sum();
sibling_size.max(estimated_safetensors_size_bytes(info))
}
fn estimated_safetensors_size_bytes(info: &ModelInfo) -> u64 {
info.safetensors
.as_ref()
.map(|safetensors| {
safetensors
.parameters
.iter()
.map(|(dtype, count)| count.saturating_mul(dtype_size_bytes(dtype)))
.sum()
})
.unwrap_or(0)
}
fn dtype_size_bytes(dtype: &str) -> u64 {
match dtype.to_ascii_uppercase().as_str() {
"BOOL" | "I8" | "U8" | "F8_E4M3" | "F8_E4M3FN" | "F8_E5M2" | "F8_E5M2FNUZ" => 1,
"BF16" | "F16" | "I16" | "U16" => 2,
"F32" | "I32" | "U32" => 4,
"F64" | "I64" | "U64" => 8,
_ => 0,
}
}
fn has_mlx_tokenizer(siblings: &[RepoSibling]) -> bool { fn has_mlx_tokenizer(siblings: &[RepoSibling]) -> bool {
let has_file = |filename: &str| siblings.iter().any(|s| s.rfilename == filename);
siblings siblings
.iter() .iter()
.any(|s| is_standalone_mlx_tokenizer_file(&s.rfilename)) .any(|s| is_standalone_mlx_tokenizer_file(&s.rfilename))
|| (has_file("vocab.json") && has_file("merges.txt"))
} }
fn is_standalone_mlx_tokenizer_file(filename: &str) -> bool { fn is_standalone_mlx_tokenizer_file(filename: &str) -> bool {
filename == "tokenizer.json" || filename.ends_with(".model") || filename.ends_with(".tiktoken") filename == "tokenizer.json"
} }
fn mlx_model_type(config: &Option<serde_json::Value>) -> Option<&str> { fn mlx_model_type(config: &Option<serde_json::Value>) -> Option<&str> {
@ -1644,22 +1824,8 @@ fn mlx_model_type(config: &Option<serde_json::Value>) -> Option<&str> {
.and_then(|value| value.as_str()) .and_then(|value| value.as_str())
} }
fn is_mlx_runtime_supported_model_type(model_type: &str) -> bool {
matches!(model_type, "gemma4" | "gemma4_text" | "llama" | "qwen3")
}
fn is_mlx_moe_model(config: &Option<serde_json::Value>) -> bool {
config
.as_ref()
.and_then(|config| config.get("text_config"))
.and_then(|text_config| text_config.get("enable_moe_block"))
.and_then(|value| value.as_bool())
.unwrap_or(false)
}
fn is_mlx_runtime_supported(config: &Option<serde_json::Value>) -> bool { fn is_mlx_runtime_supported(config: &Option<serde_json::Value>) -> bool {
mlx_model_type(config).is_some_and(is_mlx_runtime_supported_model_type) mlx_config_support(config).is_none()
&& !is_mlx_moe_model(config)
} }
fn mlx_unsupported_reason(config: &Option<serde_json::Value>) -> Option<String> { fn mlx_unsupported_reason(config: &Option<serde_json::Value>) -> Option<String> {
@ -1670,16 +1836,23 @@ fn mlx_unsupported_reason(config: &Option<serde_json::Value>) -> Option<String>
return Some("MLX support was not compiled in".to_string()); return Some("MLX support was not compiled in".to_string());
} }
let model_type = mlx_model_type(config)?; mlx_config_support(config)
if !is_mlx_runtime_supported_model_type(model_type) { }
return Some(format!(
"MLX backend does not support '{}' models yet", fn mlx_config_support(config: &Option<serde_json::Value>) -> Option<String> {
model_type let config = config.as_ref()?;
)); mlx_config_support_for_value(config)
} }
if is_mlx_moe_model(config) {
return Some("MLX backend does not support Gemma 4 MoE models yet".to_string()); #[cfg(feature = "mlx")]
} fn mlx_config_support_for_value(config: &serde_json::Value) -> Option<String> {
safemlx_lm::check_model_config(config)
.unsupported_reason()
.map(str::to_string)
}
#[cfg(not(feature = "mlx"))]
fn mlx_config_support_for_value(_config: &serde_json::Value) -> Option<String> {
None None
} }
@ -1771,6 +1944,10 @@ fn should_download_for_mlx(filename: &str) -> bool {
|| is_standalone_mlx_tokenizer_file(filename) || is_standalone_mlx_tokenizer_file(filename)
|| filename == "tokenizer_config.json" || filename == "tokenizer_config.json"
|| filename == "generation_config.json" || filename == "generation_config.json"
|| filename == "configuration.json"
|| filename == "chat_template.jinja"
|| filename == "preprocessor_config.json"
|| filename == "video_preprocessor_config.json"
|| filename == "special_tokens_map.json" || filename == "special_tokens_map.json"
|| filename == "model.safetensors.index.json" || filename == "model.safetensors.index.json"
|| filename == "vocab.json" || filename == "vocab.json"
@ -1780,7 +1957,7 @@ fn should_download_for_mlx(filename: &str) -> bool {
fn mlx_variant_id(repo_id: &str, config: &Option<serde_json::Value>) -> String { fn mlx_variant_id(repo_id: &str, config: &Option<serde_json::Value>) -> String {
let repo_lower = repo_id.to_lowercase(); let repo_lower = repo_id.to_lowercase();
for marker in ["bf16", "f16", "fp16", "f32", "fp32", "4bit", "8bit"] { for marker in ["bf16", "f16", "fp16", "f32", "fp32", "fp8", "4bit", "8bit"] {
if repo_lower.contains(marker) { if repo_lower.contains(marker) {
return marker.to_string(); return marker.to_string();
} }
@ -1954,20 +2131,12 @@ async fn resolve_mlx_model(repo_id: &str, variant_id: &str) -> Result<ResolvedLo
let repo = client.model(owner.to_string(), name.to_string()); let repo = client.model(owner.to_string(), name.to_string());
let info = repo let info = repo
.info() .info()
.expand(vec!["siblings".to_string()]) .expand(vec!["siblings".to_string(), "safetensors".to_string()])
.send() .send()
.await?; .await?;
let siblings = info.siblings.as_deref().unwrap_or(&[]); let siblings = info.siblings.as_deref().unwrap_or(&[]);
let filenames = mlx_download_filenames(siblings); let filenames = mlx_download_filenames(siblings);
let total_size = filenames let total_size = mlx_download_size_bytes(&info, siblings);
.iter()
.filter_map(|filename| {
siblings
.iter()
.find(|s| s.rfilename == *filename)
.and_then(|s| s.size)
})
.sum();
let progress = HfDownloadProgress::new(repo_id.to_string(), total_size); let progress = HfDownloadProgress::new(repo_id.to_string(), total_size);
progress.init(); progress.init();
let mut snapshot_path = None; let mut snapshot_path = None;
@ -2103,8 +2272,13 @@ impl HfDownloadProgress {
} }
fn finish_file(&self, size_bytes: u64) { fn finish_file(&self, size_bytes: u64) {
let observed_size = self
.state
.lock()
.map(|state| state.current_file_total_bytes.max(state.bytes_downloaded))
.unwrap_or(0);
if let Ok(mut completed_bytes) = self.completed_bytes.lock() { if let Ok(mut completed_bytes) = self.completed_bytes.lock() {
*completed_bytes = completed_bytes.saturating_add(size_bytes); *completed_bytes = completed_bytes.saturating_add(size_bytes.max(observed_size));
} }
if let Ok(mut state) = self.state.lock() { if let Ok(mut state) = self.state.lock() {
state.bytes_downloaded = 0; state.bytes_downloaded = 0;
@ -2117,10 +2291,15 @@ impl HfDownloadProgress {
let completed_bytes = self.completed_bytes.lock().map(|value| *value).unwrap_or(0); let completed_bytes = self.completed_bytes.lock().map(|value| *value).unwrap_or(0);
if let Ok(state) = self.state.lock() { if let Ok(state) = self.state.lock() {
let bytes_downloaded = completed_bytes.saturating_add(state.bytes_downloaded); let bytes_downloaded = completed_bytes.saturating_add(state.bytes_downloaded);
let total_bytes =
self.total_bytes
.max(completed_bytes.saturating_add(
state.current_file_total_bytes.max(state.bytes_downloaded),
));
update_download_manager_progress( update_download_manager_progress(
&self.model_id, &self.model_id,
bytes_downloaded.min(self.total_bytes), bytes_downloaded.min(total_bytes),
self.total_bytes.max(state.current_file_total_bytes), total_bytes,
state.speed_bps, state.speed_bps,
); );
} }
@ -2152,10 +2331,12 @@ impl ProgressHandler for HfDownloadProgress {
file.bytes_completed file.bytes_completed
} }
}) })
.max() .sum();
.unwrap_or(0); let total_bytes = files.iter().map(|file| file.total_bytes).sum();
if let Ok(mut state) = self.state.lock() { if let Ok(mut state) = self.state.lock() {
state.bytes_downloaded = state.bytes_downloaded.max(bytes_downloaded); state.bytes_downloaded = state.bytes_downloaded.max(bytes_downloaded);
state.current_file_total_bytes =
state.current_file_total_bytes.max(total_bytes);
} }
self.update_progress_from_state(); self.update_progress_from_state();
} }
@ -2168,7 +2349,7 @@ impl ProgressHandler for HfDownloadProgress {
return; return;
} }
if let Ok(mut state) = self.state.lock() { if let Ok(mut state) = self.state.lock() {
state.bytes_downloaded = *bytes_completed; state.bytes_downloaded = state.bytes_downloaded.max(*bytes_completed);
state.current_file_total_bytes = state.current_file_total_bytes =
(*total_bytes).max(state.current_file_total_bytes); (*total_bytes).max(state.current_file_total_bytes);
state.speed_bps = bytes_per_sec.map(|speed| speed as u64); state.speed_bps = bytes_per_sec.map(|speed| speed as u64);

391
crates/goose-local-inference/src/lib.rs
View file

@ -14,6 +14,7 @@ mod mlx;
pub(crate) mod multimodal; pub(crate) mod multimodal;
#[cfg(feature = "mlx")] #[cfg(feature = "mlx")]
mod native_tool_parsing; mod native_tool_parsing;
pub(crate) mod thinking_output;
#[cfg(feature = "mlx")] #[cfg(feature = "mlx")]
mod tool_emulation; mod tool_emulation;
mod tool_parsing; mod tool_parsing;
@ -23,7 +24,9 @@ use async_stream::try_stream;
use async_trait::async_trait; use async_trait::async_trait;
use backend::{BackendLoadedModel, LocalInferenceBackend}; use backend::{BackendLoadedModel, LocalInferenceBackend};
use goose_provider_types::base::{MessageStream, Provider, ProviderDescriptor, ProviderMetadata}; use goose_provider_types::base::{MessageStream, Provider, ProviderDescriptor, ProviderMetadata};
use goose_provider_types::conversation::message::{Message, MessageContent}; use goose_provider_types::conversation::message::{
Message, MessageContent, SystemNotificationType,
};
use goose_provider_types::conversation::token_usage::{ProviderUsage, Usage}; use goose_provider_types::conversation::token_usage::{ProviderUsage, Usage};
use goose_provider_types::errors::ProviderError; use goose_provider_types::errors::ProviderError;
use goose_provider_types::images::ImageFormat; use goose_provider_types::images::ImageFormat;
@ -34,13 +37,33 @@ use local_model_registry::ChatTemplate;
use mlx::{MlxBackend, MLX_BACKEND_ID}; use mlx::{MlxBackend, MLX_BACKEND_ID};
use rmcp::model::Tool; use rmcp::model::Tool;
use serde_json::{json, Value}; use serde_json::{json, Value};
use std::collections::HashMap; use std::collections::{HashMap, HashSet};
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, Mutex as StdMutex, Weak}; use std::sync::{Arc, Mutex as StdMutex};
use tokio::sync::Mutex; use tokio::sync::{Mutex, Notify};
use uuid::Uuid; use uuid::Uuid;
type ModelSlot = Arc<Mutex<Option<Box<dyn BackendLoadedModel>>>>; type ModelSlotHandle = Arc<ModelSlot>;
struct ModelSlot {
state: Mutex<ModelSlotState>,
notify: Notify,
}
enum ModelSlotState {
Empty,
Loading,
Loaded(Box<dyn BackendLoadedModel>),
}
impl ModelSlot {
fn new() -> Self {
Self {
state: Mutex::new(ModelSlotState::Empty),
notify: Notify::new(),
}
}
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)] #[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct ModelCacheKey { struct ModelCacheKey {
@ -64,7 +87,8 @@ impl ModelCacheKey {
} }
pub struct InferenceRuntime { pub struct InferenceRuntime {
models: StdMutex<HashMap<ModelCacheKey, ModelSlot>>, models: StdMutex<HashMap<ModelCacheKey, ModelSlotHandle>>,
cold_load_lock: Mutex<()>,
backends: HashMap<&'static str, Arc<dyn LocalInferenceBackend>>, backends: HashMap<&'static str, Arc<dyn LocalInferenceBackend>>,
} }
@ -72,19 +96,17 @@ pub fn builtin_chat_template_names() -> Vec<String> {
llamacpp::builtin_chat_template_names() llamacpp::builtin_chat_template_names()
} }
/// Global weak reference used to share a single `InferenceRuntime` across static RUNTIME: StdMutex<Option<Arc<InferenceRuntime>>> = StdMutex::new(None);
/// 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 fn current_runtime() -> Option<Arc<InferenceRuntime>> {
/// strong refs drop (normal shutdown), the runtime is deallocated and the RUNTIME.lock().expect("runtime lock poisoned").clone()
/// backend freed. The `Weak` left behind is inert during `__cxa_finalize`, so no }
/// ggml statics race.
static RUNTIME: StdMutex<Weak<InferenceRuntime>> = StdMutex::new(Weak::new());
impl InferenceRuntime { impl InferenceRuntime {
pub fn get_or_init() -> Result<Arc<Self>> { pub fn get_or_init() -> Result<Arc<Self>> {
let mut guard = RUNTIME.lock().expect("runtime lock poisoned"); let mut guard = RUNTIME.lock().expect("runtime lock poisoned");
if let Some(runtime) = guard.upgrade() { if let Some(runtime) = guard.as_ref() {
return Ok(runtime); return Ok(runtime.clone());
} }
let llamacpp_backend: Arc<dyn LocalInferenceBackend> = Arc::new(LlamaCppBackend::new()?); let llamacpp_backend: Arc<dyn LocalInferenceBackend> = Arc::new(LlamaCppBackend::new()?);
let mlx_backend: Arc<dyn LocalInferenceBackend> = Arc::new(MlxBackend::new()); let mlx_backend: Arc<dyn LocalInferenceBackend> = Arc::new(MlxBackend::new());
@ -93,9 +115,10 @@ impl InferenceRuntime {
backends.insert(MLX_BACKEND_ID, mlx_backend); backends.insert(MLX_BACKEND_ID, mlx_backend);
let runtime = Arc::new(Self { let runtime = Arc::new(Self {
models: StdMutex::new(HashMap::new()), models: StdMutex::new(HashMap::new()),
cold_load_lock: Mutex::new(()),
backends, backends,
}); });
*guard = Arc::downgrade(&runtime); *guard = Some(runtime.clone());
Ok(runtime) Ok(runtime)
} }
@ -122,14 +145,19 @@ impl InferenceRuntime {
}) })
} }
fn get_or_create_model_slot(&self, key: ModelCacheKey) -> ModelSlot { fn get_or_create_model_slot(&self, key: ModelCacheKey) -> ModelSlotHandle {
let mut map = self.models.lock().expect("model cache lock poisoned"); let mut map = self.models.lock().expect("model cache lock poisoned");
map.entry(key) map.entry(key)
.or_insert_with(|| Arc::new(Mutex::new(None))) .or_insert_with(|| Arc::new(ModelSlot::new()))
.clone() .clone()
} }
fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec<ModelSlot> { fn model_slot(&self, key: &ModelCacheKey) -> Option<ModelSlotHandle> {
let map = self.models.lock().expect("model cache lock poisoned");
map.get(key).cloned()
}
fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec<ModelSlotHandle> {
let map = self.models.lock().expect("model cache lock poisoned"); let map = self.models.lock().expect("model cache lock poisoned");
map.iter() map.iter()
.filter(|(key, _)| *key != keep_key) .filter(|(key, _)| *key != keep_key)
@ -138,6 +166,76 @@ impl InferenceRuntime {
} }
} }
pub async fn is_model_loaded(model_name: &str) -> Result<bool, ProviderError> {
let resolved = match resolve_model_path(model_name) {
Some(resolved) => resolved,
None => return Ok(false),
};
let runtime = InferenceRuntime::get_or_init().map_err(|error| {
ProviderError::ExecutionError(format!("Failed to initialize local inference: {error}"))
})?;
let backend = runtime.backend_for_model(&resolved)?;
let key = ModelCacheKey::new(
backend.id(),
model_name.to_string(),
resolved.settings.chat_template,
);
let Some(slot) = runtime.model_slot(&key) else {
return Ok(false);
};
let state = slot.state.lock().await;
Ok(matches!(*state, ModelSlotState::Loaded(_)))
}
pub async fn loaded_model_ids() -> Result<HashSet<String>, ProviderError> {
let Some(runtime) = current_runtime() else {
return Ok(HashSet::new());
};
let slots = {
let map = runtime.models.lock().expect("model cache lock poisoned");
map.iter()
.map(|(key, slot)| (key.model_id.clone(), slot.clone()))
.collect::<Vec<_>>()
};
let mut loaded = HashSet::new();
for (model_id, slot) in slots {
if let Ok(state) = slot.state.try_lock() {
if matches!(*state, ModelSlotState::Loaded(_)) {
loaded.insert(model_id);
}
} else {
loaded.insert(model_id);
}
}
Ok(loaded)
}
pub async fn evict_model(model_name: &str) -> Result<bool, ProviderError> {
let Some(runtime) = current_runtime() else {
return Ok(false);
};
let slots = {
let map = runtime.models.lock().expect("model cache lock poisoned");
map.iter()
.filter(|(key, _)| key.model_id == model_name)
.map(|(_, slot)| slot.clone())
.collect::<Vec<_>>()
};
let mut evicted = false;
for slot in slots {
let mut state = slot.state.lock().await;
if matches!(*state, ModelSlotState::Loaded(_)) {
*state = ModelSlotState::Empty;
evicted = true;
slot.notify.notify_waiters();
}
}
Ok(evicted)
}
const PROVIDER_NAME: &str = "local"; const PROVIDER_NAME: &str = "local";
const DEFAULT_MODEL: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M"; const DEFAULT_MODEL: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M";
@ -565,41 +663,9 @@ impl Provider for LocalInferenceProvider {
})?; })?;
let backend = self.runtime.backend_for_model(&resolved)?; let backend = self.runtime.backend_for_model(&resolved)?;
let model_context_limit = resolved.context_limit; 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 // Allow request_params to override thinking
let mut model_settings = model_settings; let mut model_settings = resolved.settings.clone();
if let Some(false) = model_config if let Some(false) = model_config
.request_param::<bool>("enable_thinking") .request_param::<bool>("enable_thinking")
.or_else(|| { .or_else(|| {
@ -611,6 +677,15 @@ impl Provider for LocalInferenceProvider {
model_settings.enable_thinking = false; model_settings.enable_thinking = false;
} }
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());
let runtime = self.runtime.clone();
let cache_key = cache_key.clone();
let model_arc = model_slot.clone(); let model_arc = model_slot.clone();
let backend = backend.clone(); let backend = backend.clone();
let model_name = model_config.model_name.clone(); let model_name = model_config.model_name.clone();
@ -639,17 +714,172 @@ impl Provider for LocalInferenceProvider {
}, },
}); });
let mut log = start_log(model_config, &log_payload)?;
let (tx, mut rx) = tokio::sync::mpsc::channel::< let (tx, mut rx) = tokio::sync::mpsc::channel::<
Result<(Option<Message>, Option<ProviderUsage>), ProviderError>, Result<(Option<Message>, Option<ProviderUsage>), ProviderError>,
>(32); >(32);
let mut log = start_log(model_config, &log_payload)?;
tokio::task::spawn_blocking(move || { tokio::spawn(async move {
// Macro to log errors before sending them through the channel let mut model_load_ms = None;
macro_rules! send_err {
($err:expr) => {{ // Ensure model is loaded — unload any other models first to free memory.
let err = $err; loop {
let state = model_slot.state.lock().await;
match &*state {
ModelSlotState::Loaded(_) => break,
ModelSlotState::Loading => {
let notified = model_slot.notify.notified();
drop(state);
notified.await;
}
ModelSlotState::Empty => {
drop(state);
let cold_load_guard = runtime.cold_load_lock.lock().await;
let mut state = model_slot.state.lock().await;
match &*state {
ModelSlotState::Loaded(_) => break,
ModelSlotState::Loading => {
let notified = model_slot.notify.notified();
drop(state);
drop(cold_load_guard);
notified.await;
continue;
}
ModelSlotState::Empty => {}
}
*state = ModelSlotState::Loading;
drop(state);
let loading_message = Message::assistant().with_system_notification(
SystemNotificationType::ProgressMessage,
format!("Loading local model {model_name}..."),
);
if tx.send(Ok((Some(loading_message), None))).await.is_err() {
let mut state = model_slot.state.lock().await;
*state = ModelSlotState::Empty;
model_slot.notify.notify_waiters();
return;
}
let other_model_slots = runtime.other_model_slots(&cache_key);
for slot in other_model_slots {
let mut other = slot.state.lock().await;
if matches!(*other, ModelSlotState::Loaded(_)) {
tracing::info!("Unloading previous model to free memory");
*other = ModelSlotState::Empty;
}
}
let model_id = model_name.clone();
let resolved_for_load = resolved_model.clone();
let settings_for_load = settings.clone();
let backend_for_load = backend.clone();
let load_started = std::time::Instant::now();
let loaded = match tokio::task::spawn_blocking(move || {
backend_for_load.load_model(
&model_id,
&resolved_for_load,
&settings_for_load,
)
})
.await
{
Ok(Ok(loaded)) => loaded,
Ok(Err(err)) => {
let mut state = model_slot.state.lock().await;
*state = ModelSlotState::Empty;
model_slot.notify.notify_waiters();
let _ = log.error(&err);
let _ = tx.send(Err(err)).await;
return;
}
Err(err) => {
let mut state = model_slot.state.lock().await;
*state = ModelSlotState::Empty;
model_slot.notify.notify_waiters();
let err = ProviderError::ExecutionError(err.to_string());
let _ = log.error(&err);
let _ = tx.send(Err(err)).await;
return;
}
};
let elapsed_ms =
u64::try_from(load_started.elapsed().as_millis()).unwrap_or(u64::MAX);
model_load_ms = Some(elapsed_ms);
tracing::info!(
backend = backend.id(),
model = %model_name,
model_load_ms = elapsed_ms,
"Loaded local inference model"
);
let _ = log.write(
&json!({
"path": "model_load",
"backend": backend.id(),
"model": &model_name,
"model_load_ms": elapsed_ms,
}),
None,
);
let mut state = model_slot.state.lock().await;
*state = ModelSlotState::Loaded(loaded);
model_slot.notify.notify_waiters();
drop(cold_load_guard);
break;
}
}
}
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.state.blocking_lock();
let loaded = match &mut *model_guard {
ModelSlotState::Loaded(loaded) => loaded.as_mut(),
ModelSlotState::Empty | ModelSlotState::Loading => {
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,
model_load_ms,
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, request);
if let Err(err) = result {
let msg = match &err { let msg = match &err {
ProviderError::ExecutionError(s) => s.as_str(), ProviderError::ExecutionError(s) => s.as_str(),
ProviderError::ContextLengthExceeded(s) => s.as_str(), ProviderError::ContextLengthExceeded(s) => s.as_str(),
@ -657,49 +887,8 @@ impl Provider for LocalInferenceProvider {
}; };
let _ = log.error(msg); let _ = log.error(msg);
let _ = tx.blocking_send(Err(err)); 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! { Ok(Box::pin(try_stream! {

5
crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs
View file

@ -27,10 +27,9 @@ use serde_json::json;
use std::borrow::Cow; use std::borrow::Cow;
use uuid::Uuid; use uuid::Uuid;
use super::super::{finalize_usage, StreamSender}; use super::super::{finalize_usage, thinking_output::ThinkingOutputFilter, StreamSender};
use super::inference_engine::{ use super::inference_engine::{
generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, TokenAction,
ThinkingOutputFilter, TokenAction,
}; };
const SHELL_TOOL: &str = "developer__shell"; const SHELL_TOOL: &str = "developer__shell";

81
crates/goose-local-inference/src/llamacpp/inference_engine.rs
View file

@ -3,7 +3,6 @@ use crate::local_model_registry::ModelSettings;
use crate::multimodal::ExtractedImage; use crate::multimodal::ExtractedImage;
use goose_provider_types::errors::ProviderError; use goose_provider_types::errors::ProviderError;
use goose_provider_types::request_log::{LoggerHandleExt, RequestLogHandle}; 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::context::params::LlamaContextParams;
use llama_cpp_2::llama_batch::LlamaBatch; use llama_cpp_2::llama_batch::LlamaBatch;
use llama_cpp_2::model::{AddBos, ChatTemplateResult, LlamaChatTemplate, LlamaModel}; use llama_cpp_2::model::{AddBos, ChatTemplateResult, LlamaChatTemplate, LlamaModel};
@ -48,86 +47,6 @@ pub(super) struct PreparedGeneration<'model> {
pub effective_ctx: usize, pub effective_ctx: usize,
} }
pub(super) struct ThinkingOutputFilter {
enabled: bool,
saw_structured_reasoning: bool,
think_filter: ThinkFilter,
pending_inline_thinking: String,
accumulated_thinking: String,
}
impl ThinkingOutputFilter {
pub(super) fn new(enable_thinking: bool, generation_prompt: &str) -> Self {
let mut think_filter = ThinkFilter::new();
if enable_thinking && !generation_prompt.is_empty() {
let _ = think_filter.push(generation_prompt);
}
Self {
enabled: enable_thinking,
saw_structured_reasoning: false,
think_filter,
pending_inline_thinking: String::new(),
accumulated_thinking: String::new(),
}
}
pub(super) fn push_structured_reasoning(&mut self, reasoning: &str) -> Option<String> {
if reasoning.is_empty() {
return None;
}
self.saw_structured_reasoning = true;
self.pending_inline_thinking.clear();
self.think_filter = ThinkFilter::new();
self.accumulated_thinking.push_str(reasoning);
Some(reasoning.to_string())
}
pub(super) fn push_text(&mut self, text: &str) -> FilterOut {
if !self.enabled {
return FilterOut {
content: text.to_string(),
thinking: String::new(),
};
}
let mut filtered = self.think_filter.push(text);
if self.saw_structured_reasoning {
filtered.thinking.clear();
} else if !filtered.thinking.is_empty() {
self.pending_inline_thinking.push_str(&filtered.thinking);
filtered.thinking.clear();
}
filtered
}
pub(super) fn finish(&mut self) -> FilterOut {
let mut filtered = if self.enabled && !self.saw_structured_reasoning {
std::mem::take(&mut self.think_filter).finish()
} else {
FilterOut::default()
};
if !self.saw_structured_reasoning {
let mut thinking = std::mem::take(&mut self.pending_inline_thinking);
thinking.push_str(&filtered.thinking);
if !thinking.is_empty() {
self.accumulated_thinking.push_str(&thinking);
}
filtered.thinking = thinking;
} else {
filtered.thinking.clear();
}
filtered
}
pub(super) fn accumulated_thinking(&self) -> &str {
&self.accumulated_thinking
}
}
pub(super) struct StopSuffixTrimmer { pub(super) struct StopSuffixTrimmer {
pending: String, pending: String,
stops: Vec<String>, stops: Vec<String>,

4
crates/goose-local-inference/src/llamacpp/inference_native_tools.rs
View file

@ -6,9 +6,9 @@ use std::borrow::Cow;
use uuid::Uuid; use uuid::Uuid;
use super::super::finalize_usage; use super::super::finalize_usage;
use super::super::thinking_output::ThinkingOutputFilter;
use super::inference_engine::{ use super::inference_engine::{
generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, TokenAction,
ThinkingOutputFilter, TokenAction,
}; };
pub(super) fn generate_with_native_tools( pub(super) fn generate_with_native_tools(

27
crates/goose-local-inference/src/management.rs
View file

@ -26,6 +26,7 @@ use goose_sdk_types::custom_requests::{
LocalInferenceModelSettingsReadResponse, LocalInferenceModelSettingsUpdateResponse, LocalInferenceModelSettingsReadResponse, LocalInferenceModelSettingsUpdateResponse,
LocalInferenceModelsListResponse, LocalInferenceSamplingConfig, LocalInferenceToolCallingMode, LocalInferenceModelsListResponse, LocalInferenceSamplingConfig, LocalInferenceToolCallingMode,
}; };
use std::collections::HashSet;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, OnceLock}; use std::sync::{Arc, OnceLock};
@ -44,13 +45,16 @@ pub async fn list_models() -> Result<LocalInferenceModelsListResponse> {
let runtime = management_runtime()?; let runtime = management_runtime()?;
let recommended_id = recommend_local_model(&runtime); let recommended_id = recommend_local_model(&runtime);
let loaded_model_ids = crate::loaded_model_ids()
.await
.map_err(|error| anyhow!(error.to_string()))?;
let registry = get_registry() let registry = get_registry()
.lock() .lock()
.map_err(|_| anyhow!("Failed to acquire registry lock"))?; .map_err(|_| anyhow!("Failed to acquire registry lock"))?;
let mut models: Vec<LocalInferenceModelDto> = registry let mut models: Vec<LocalInferenceModelDto> = registry
.list_models() .list_models()
.iter() .iter()
.map(|entry| local_model_to_dto(entry, &recommended_id)) .map(|entry| local_model_to_dto(entry, &recommended_id, &loaded_model_ids))
.collect(); .collect();
models.sort_by(|a, b| { models.sort_by(|a, b| {
@ -204,6 +208,20 @@ pub fn delete_model(model_id: &str) -> Result<()> {
registry.delete_model(model_id) registry.delete_model(model_id)
} }
pub fn model_exists(model_id: &str) -> Result<bool> {
let registry = get_registry()
.lock()
.map_err(|_| anyhow!("Failed to acquire registry lock"))?;
Ok(registry.get_model(model_id).is_some())
}
pub async fn evict_model(model_id: &str) -> Result<()> {
crate::evict_model(model_id)
.await
.map(|_| ())
.map_err(|error| anyhow!(error.to_string()))
}
pub fn get_model_settings(model_id: &str) -> Result<LocalInferenceModelSettingsReadResponse> { pub fn get_model_settings(model_id: &str) -> Result<LocalInferenceModelSettingsReadResponse> {
let registry = get_registry() let registry = get_registry()
.lock() .lock()
@ -412,7 +430,11 @@ pub async fn ensure_featured_models_current() -> Result<()> {
Ok(()) Ok(())
} }
fn local_model_to_dto(entry: &LocalModelEntry, recommended_id: &str) -> LocalInferenceModelDto { fn local_model_to_dto(
entry: &LocalModelEntry,
recommended_id: &str,
loaded_model_ids: &HashSet<String>,
) -> LocalInferenceModelDto {
let vision_capable = entry.settings.vision_capable; let vision_capable = entry.settings.vision_capable;
LocalInferenceModelDto { LocalInferenceModelDto {
id: entry.id.clone(), id: entry.id.clone(),
@ -422,6 +444,7 @@ fn local_model_to_dto(entry: &LocalModelEntry, recommended_id: &str) -> LocalInf
size_bytes: entry.file_size(), size_bytes: entry.file_size(),
status: model_download_status_to_dto(entry.download_status()), status: model_download_status_to_dto(entry.download_status()),
recommended: recommended_id == entry.id, recommended: recommended_id == entry.id,
is_loaded: loaded_model_ids.contains(&entry.id),
settings: model_settings_to_dto(&entry.settings), settings: model_settings_to_dto(&entry.settings),
vision_capable, vision_capable,
mmproj_status: vision_capable mmproj_status: vision_capable

629
crates/goose-local-inference/src/mlx.rs
View file

@ -3,23 +3,24 @@ mod imp {
use std::any::Any; use std::any::Any;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use mlx_lm::cache::ConcatKeyValueCache; use safemlx::transforms::eval;
use mlx_lm::gemma4_mtp::generate_gemma4_mtp; use safemlx::{random, Array, Device, DeviceType, Stream};
use mlx_lm::models::{gemma4_assistant::load_gemma4_assistant_model, LoadedModel, Model}; use safemlx_lm::gemma4_mtp::generate_gemma4_mtp;
use mlx_lm_utils::tokenizer::{Chat, Conversation, Role}; use safemlx_lm::models::{gemma4_assistant::load_gemma4_assistant_model, LoadedModel, Model};
use mlx_rs::transforms::eval; use safemlx_lm_utils::tokenizer::{Chat, Conversation, Role, Tokenizer};
use serde_json::json; use serde_json::json;
use crate::backend::{BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend}; use crate::backend::{BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend};
use crate::local_model_registry::{ModelSettings, ToolCallingMode}; use crate::local_model_registry::{ModelSettings, ToolCallingMode};
use crate::native_tool_parsing::message_from_native_tool_text; use crate::native_tool_parsing::message_from_native_tool_text;
use crate::provider_utils::filter_extensions_from_system_prompt; use crate::provider_utils::filter_extensions_from_system_prompt;
use crate::thinking_output::ThinkingOutputFilter;
use crate::tool_emulation::{ use crate::tool_emulation::{
build_emulator_tool_description, load_tiny_model_prompt, message_for_emulator_action, build_emulator_tool_description, load_tiny_model_prompt, message_for_emulator_action,
StreamingEmulatorParser, CODE_EXECUTION_TOOL, StreamingEmulatorParser, CODE_EXECUTION_TOOL,
}; };
use crate::{extract_text_content, ResolvedModelPaths}; use crate::{extract_text_content, ResolvedModelPaths};
use goose_provider_types::conversation::message::Message; use goose_provider_types::conversation::message::{Message, MessageContent};
use goose_provider_types::conversation::token_usage::{ use goose_provider_types::conversation::token_usage::{
DraftStats, ProviderStats, ProviderUsage, Usage, DraftStats, ProviderStats, ProviderUsage, Usage,
}; };
@ -57,14 +58,25 @@ mod imp {
} }
let model_dir = model_dir_from_path(&resolved.model_path)?; let model_dir = model_dir_from_path(&resolved.model_path)?;
let model = LoadedModel::load(&model_dir).map_err(mlx_error)?; let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0));
let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0));
let model =
LoadedModel::load(&model_dir, &stream, &weights_stream).map_err(mlx_error)?;
let tokenizer =
Tokenizer::from_file(model_dir.join("tokenizer.json")).map_err(mlx_error)?;
tracing::info!( tracing::info!(
backend = self.id(), backend = self.id(),
model_id, model_id,
model_type = model.model_type(), model_type = model.model_type(),
"MLX model loaded successfully" "MLX model loaded successfully"
); );
Ok(Box::new(MlxLoadedModel { model, model_dir })) let stop_token_ids = mlx_stop_token_ids(&model, &model_dir);
Ok(Box::new(MlxLoadedModel {
model,
tokenizer,
model_dir,
stop_token_ids,
}))
} }
fn generate( fn generate(
@ -79,6 +91,7 @@ mod imp {
ProviderError::ExecutionError("Loaded model backend mismatch".to_string()) ProviderError::ExecutionError("Loaded model backend mismatch".to_string())
})?; })?;
let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0));
let tool_mode = if request.tools.is_empty() { let tool_mode = if request.tools.is_empty() {
ToolMode::None ToolMode::None
} else { } else {
@ -110,77 +123,114 @@ mod imp {
let prompt_array = loaded let prompt_array = loaded
.model .model
.encode_to_array(&prompt, false) .encode_to_array(&prompt, false, &stream)
.map_err(mlx_error)?; .map_err(mlx_error)?;
let max_tokens = request let max_tokens = mlx_max_tokens(
.settings request.settings,
.max_output_tokens request.max_tokens,
.or_else(|| { request.context_limit,
request prompt_tokens.len(),
.max_tokens );
.and_then(|tokens| usize::try_from(tokens).ok()) let (settings_temp, seed) = sampling(request.settings);
}) let temp = request.temperature.unwrap_or(settings_temp);
.unwrap_or(512); let prng_key = prng_key(temp, seed)?;
let temp = request let eos_token_ids = loaded.stop_token_ids.clone();
.temperature
.unwrap_or_else(|| temperature(request.settings));
let eos_token_ids = loaded.model.eos_token_ids().to_vec();
let generation_started = std::time::Instant::now(); let generation_started = std::time::Instant::now();
let (generated_ids, draft_stats) = let MlxGeneration {
if let Some(draft_model_path) = &request.draft_model_path { generated_ids,
if matches!(loaded.model.model_mut(), Model::Gemma4(_)) { generated_text,
let mut assistant = draft_stats,
load_gemma4_assistant_model(draft_model_path).map_err(|error| { time_to_first_token_ms,
streamed_response,
} = if let Some(draft_model_path) = &request.draft_model_path {
if matches!(loaded.model.model_mut(), Model::Gemma4(_)) {
let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0));
let mut assistant =
load_gemma4_assistant_model(draft_model_path, &stream, &weights_stream)
.map_err(|error| {
mlx_error(format!("failed to load MLX draft model: {error}")) mlx_error(format!("failed to load MLX draft model: {error}"))
})?; })?;
let target = match loaded.model.model_mut() { let target = match loaded.model.model_mut() {
Model::Gemma4(target) => target, Model::Gemma4(target) => target,
_ => unreachable!(), _ => unreachable!(),
}; };
let (ids, stats) = generate_gemma4_mtp( let (ids, stats) = generate_gemma4_mtp(
target, target,
&mut assistant, &mut assistant,
&prompt_array,
&eos_token_ids,
max_tokens,
temp,
)
.map_err(mlx_error)?;
(
ids,
Some(DraftStats {
model: Some(draft_model_path.display().to_string()),
draft_tokens: stats.draft_tokens,
accepted_tokens: stats.accepted_tokens,
target_tokens: stats.target_tokens,
rounds: stats.rounds,
accept_rate: stats.accept_rate(),
}),
)
} else {
generate_single_model(
&mut loaded.model,
&prompt_array,
&eos_token_ids,
max_tokens,
temp,
)?
}
} else {
generate_single_model(
&mut loaded.model,
&prompt_array, &prompt_array,
&eos_token_ids, &eos_token_ids,
max_tokens, max_tokens,
temp, temp,
prng_key,
&stream,
)
.map_err(mlx_error)?;
let generated_text = loaded.tokenizer.decode(&ids, true).map_err(mlx_error)?;
MlxGeneration {
generated_ids: ids,
generated_text,
draft_stats: Some(DraftStats {
model: Some(draft_model_path.display().to_string()),
draft_tokens: stats.draft_tokens,
accepted_tokens: stats.accepted_tokens,
target_tokens: stats.target_tokens,
rounds: stats.rounds,
accept_rate: stats.accept_rate(),
}),
time_to_first_token_ms: None,
streamed_response: false,
}
} else {
generate_single_model(
&mut loaded.model,
&loaded.tokenizer,
&prompt_array,
&eos_token_ids,
max_tokens,
temp,
prng_key,
&stream,
generation_started,
MlxStreamEmitter::new(
request.message_id,
tool_mode,
request.settings.enable_thinking,
&prompt,
request.tx,
),
)? )?
}; }
} else {
generate_single_model(
&mut loaded.model,
&loaded.tokenizer,
&prompt_array,
&eos_token_ids,
max_tokens,
temp,
prng_key,
&stream,
generation_started,
MlxStreamEmitter::new(
request.message_id,
tool_mode,
request.settings.enable_thinking,
&prompt,
request.tx,
),
)?
};
let generated_text = loaded if !streamed_response {
.model emit_generated_response(
.decode(&generated_ids, true) &generated_text,
.map_err(mlx_error)?; &prompt,
emit_generated_response(&generated_text, request.message_id, tool_mode, request.tx)?; request.settings.enable_thinking,
request.message_id,
tool_mode,
request.tx,
)?;
}
let output_tokens = generated_ids.len() as i32; let output_tokens = generated_ids.len() as i32;
let input_tokens = prompt_tokens.len() as i32; let input_tokens = prompt_tokens.len() as i32;
@ -194,12 +244,16 @@ mod imp {
"model_dir": loaded.model_dir, "model_dir": loaded.model_dir,
"prompt_tokens": input_tokens, "prompt_tokens": input_tokens,
"output_tokens": output_tokens, "output_tokens": output_tokens,
"model_load_ms": request.model_load_ms,
"time_to_first_token_ms": time_to_first_token_ms,
"elapsed_ms": generation_started.elapsed().as_millis() as u64,
"generated_text": generated_text, "generated_text": generated_text,
"draft": draft_stats, "draft": draft_stats,
}); });
let _ = request.log.write(&log_json, Some(&usage)); let _ = request.log.write(&log_json, Some(&usage));
let stats = ProviderStats { let stats = ProviderStats {
time_to_first_token_ms: None, time_to_first_token_ms,
model_load_ms: request.model_load_ms,
elapsed_ms: Some(generation_started.elapsed().as_millis() as u64), elapsed_ms: Some(generation_started.elapsed().as_millis() as u64),
output_tokens: Some(generated_ids.len()), output_tokens: Some(generated_ids.len()),
draft: draft_stats, draft: draft_stats,
@ -221,9 +275,19 @@ mod imp {
Emulated { code_mode_enabled: bool }, Emulated { code_mode_enabled: bool },
} }
struct MlxGeneration {
generated_ids: Vec<u32>,
generated_text: String,
draft_stats: Option<DraftStats>,
time_to_first_token_ms: Option<u64>,
streamed_response: bool,
}
struct MlxLoadedModel { struct MlxLoadedModel {
model: LoadedModel, model: LoadedModel,
tokenizer: Tokenizer,
model_dir: PathBuf, model_dir: PathBuf,
stop_token_ids: Vec<u32>,
} }
impl BackendLoadedModel for MlxLoadedModel { impl BackendLoadedModel for MlxLoadedModel {
@ -242,6 +306,44 @@ mod imp {
} }
} }
fn mlx_stop_token_ids(model: &LoadedModel, model_dir: &Path) -> Vec<u32> {
let mut ids = model.eos_token_ids().to_vec();
for id in generation_config_eos_token_ids(model_dir) {
if !ids.contains(&id) {
ids.push(id);
}
}
ids
}
fn generation_config_eos_token_ids(model_dir: &Path) -> Vec<u32> {
let Ok(config_json) = std::fs::read_to_string(model_dir.join("generation_config.json"))
else {
return Vec::new();
};
let Ok(config) = serde_json::from_str::<serde_json::Value>(&config_json) else {
return Vec::new();
};
match config.get("eos_token_id") {
Some(value) => token_id_or_ids(value),
None => Vec::new(),
}
}
fn token_id_or_ids(value: &serde_json::Value) -> Vec<u32> {
if let Some(id) = value.as_u64().and_then(|id| u32::try_from(id).ok()) {
return vec![id];
}
value
.as_array()
.map(|ids| {
ids.iter()
.filter_map(|id| id.as_u64().and_then(|id| u32::try_from(id).ok()))
.collect()
})
.unwrap_or_default()
}
fn build_prompt( fn build_prompt(
model: &mut LoadedModel, model: &mut LoadedModel,
model_name: &str, model_name: &str,
@ -316,28 +418,107 @@ mod imp {
fn generate_single_model( fn generate_single_model(
model: &mut LoadedModel, model: &mut LoadedModel,
prompt_array: &mlx_rs::Array, tokenizer: &Tokenizer,
prompt_array: &Array,
eos_token_ids: &[u32], eos_token_ids: &[u32],
max_tokens: usize, max_tokens: usize,
temp: f32, temp: f32,
) -> Result<(Vec<u32>, Option<DraftStats>), ProviderError> { prng_key: Option<Array>,
let mut cache: Vec<Option<ConcatKeyValueCache>> = Vec::new(); stream: &Stream,
generation_started: std::time::Instant,
mut emitter: MlxStreamEmitter<'_>,
) -> Result<MlxGeneration, ProviderError> {
let mut cache = model.new_cache();
let mut generated_ids = Vec::new(); let mut generated_ids = Vec::new();
let mut streamed_text = String::new();
let mut time_to_first_token_ms = None;
let stream_generation = emitter.can_stream();
let mut decode_stream = tokenizer.decode_stream(true);
{ {
let generator = model let generator = model
.generate(&mut cache, temp, prompt_array) .generate_with_cache(&mut cache, temp, prompt_array, prng_key, stream)
.take(max_tokens); .take(max_tokens);
for token in generator { for token in generator {
let token = token.map_err(mlx_error)?; let token = token.map_err(mlx_error)?;
eval([&token]).map_err(mlx_error)?; eval([&token]).map_err(mlx_error)?;
let token_id = token.item::<u32>(); let token_id = token.item::<u32>(stream);
time_to_first_token_ms.get_or_insert_with(|| {
u64::try_from(generation_started.elapsed().as_millis()).unwrap_or(u64::MAX)
});
if eos_token_ids.contains(&token_id) { if eos_token_ids.contains(&token_id) {
break; break;
} }
generated_ids.push(token_id); generated_ids.push(token_id);
if stream_generation {
if let Some(piece) = decode_stream.step(token_id).map_err(mlx_error)? {
if !piece.is_empty() {
let should_continue = emitter.push_text(&piece)?;
streamed_text.push_str(&piece);
if !should_continue {
break;
}
}
}
}
} }
} }
Ok((generated_ids, None)) let generated_text = tokenizer.decode(&generated_ids, true).map_err(mlx_error)?;
let streamed_response = if stream_generation {
match final_stream_suffix(&generated_text, &streamed_text)? {
Some(suffix) => {
if !suffix.is_empty() {
emitter.push_text(suffix)?;
}
true
}
None => false,
}
} else {
false
};
if streamed_response {
emitter.finish()?;
}
Ok(MlxGeneration {
generated_ids,
generated_text,
draft_stats: None,
time_to_first_token_ms,
streamed_response,
})
}
fn final_stream_suffix<'a>(
generated_text: &'a str,
streamed_text: &str,
) -> Result<Option<&'a str>, ProviderError> {
if streamed_text.is_empty() {
return Ok(None);
}
generated_text
.strip_prefix(streamed_text)
.map(Some)
.ok_or_else(|| mlx_error("streamed MLX decode did not match final tokenizer decode"))
}
fn mlx_max_tokens(
settings: &ModelSettings,
request_max_tokens: Option<i32>,
context_limit: usize,
prompt_tokens: usize,
) -> usize {
let configured_max = settings
.max_output_tokens
.or_else(|| request_max_tokens.and_then(|tokens| usize::try_from(tokens).ok()));
if context_limit == 0 {
return configured_max.unwrap_or(4096);
}
let context_headroom = context_limit.saturating_sub(prompt_tokens);
configured_max
.map(|max| max.min(context_headroom))
.unwrap_or(context_headroom)
} }
fn is_gemma4(model: &LoadedModel) -> bool { fn is_gemma4(model: &LoadedModel) -> bool {
@ -463,6 +644,8 @@ mod imp {
fn emit_generated_response( fn emit_generated_response(
generated_text: &str, generated_text: &str,
generation_prompt: &str,
enable_thinking: bool,
message_id: &str, message_id: &str,
tool_mode: ToolMode, tool_mode: ToolMode,
tx: &tokio::sync::mpsc::Sender< tx: &tokio::sync::mpsc::Sender<
@ -473,30 +656,27 @@ mod imp {
return Ok(()); return Ok(());
} }
let (content, thinking) =
split_generated_thinking(generated_text, generation_prompt, enable_thinking);
match tool_mode { match tool_mode {
ToolMode::None => { ToolMode::None => {
let mut msg = Message::assistant().with_text(generated_text); emit_assistant_message(message_id, &thinking, &content, tx)?;
msg.id = Some(message_id.to_string());
tx.blocking_send(Ok((Some(msg), None))).map_err(|_| {
ProviderError::ExecutionError("Failed to stream MLX response".to_string())
})?;
} }
ToolMode::Native => { ToolMode::Native => {
if let Some(message) = message_from_native_tool_text(generated_text, message_id)? { if let Some(mut message) = message_from_native_tool_text(&content, message_id)? {
prepend_thinking(&mut message, &thinking);
tx.blocking_send(Ok((Some(message), None))).map_err(|_| { tx.blocking_send(Ok((Some(message), None))).map_err(|_| {
ProviderError::ExecutionError("Failed to stream MLX response".to_string()) ProviderError::ExecutionError("Failed to stream MLX response".to_string())
})?; })?;
} else { } else {
let mut msg = Message::assistant().with_text(generated_text); emit_assistant_message(message_id, &thinking, &content, tx)?;
msg.id = Some(message_id.to_string());
tx.blocking_send(Ok((Some(msg), None))).map_err(|_| {
ProviderError::ExecutionError("Failed to stream MLX response".to_string())
})?;
} }
} }
ToolMode::Emulated { code_mode_enabled } => { ToolMode::Emulated { code_mode_enabled } => {
emit_assistant_message(message_id, &thinking, "", tx)?;
let mut parser = StreamingEmulatorParser::new(code_mode_enabled); let mut parser = StreamingEmulatorParser::new(code_mode_enabled);
let mut actions = parser.process_chunk(generated_text); let mut actions = parser.process_chunk(&content);
actions.extend(parser.flush()); actions.extend(parser.flush());
for action in actions { for action in actions {
@ -510,14 +690,191 @@ mod imp {
Ok(()) Ok(())
} }
fn temperature(settings: &ModelSettings) -> f32 { struct MlxStreamEmitter<'a> {
match &settings.sampling { message_id: &'a str,
crate::local_model_registry::SamplingConfig::Greedy => 0.0, tool_mode: ToolMode,
crate::local_model_registry::SamplingConfig::Temperature { temperature, .. } => { tx: &'a tokio::sync::mpsc::Sender<
*temperature Result<(Option<Message>, Option<ProviderUsage>), ProviderError>,
>,
output_filter: ThinkingOutputFilter,
emulator_parser: Option<StreamingEmulatorParser>,
stop_after_tool_call: bool,
}
impl<'a> MlxStreamEmitter<'a> {
fn new(
message_id: &'a str,
tool_mode: ToolMode,
enable_thinking: bool,
generation_prompt: &str,
tx: &'a tokio::sync::mpsc::Sender<
Result<(Option<Message>, Option<ProviderUsage>), ProviderError>,
>,
) -> Self {
let emulator_parser = match tool_mode {
ToolMode::Emulated { code_mode_enabled } => {
Some(StreamingEmulatorParser::new(code_mode_enabled))
}
ToolMode::None | ToolMode::Native => None,
};
Self {
message_id,
tool_mode,
tx,
output_filter: ThinkingOutputFilter::new(enable_thinking, generation_prompt),
emulator_parser,
stop_after_tool_call: false,
} }
crate::local_model_registry::SamplingConfig::MirostatV2 { .. } => 0.0,
} }
fn can_stream(&self) -> bool {
!matches!(self.tool_mode, ToolMode::Native)
}
fn push_text(&mut self, text: &str) -> Result<bool, ProviderError> {
let filtered = self.output_filter.push_text(text);
if !filtered.content.is_empty() {
self.emit_content(&filtered.content)?;
}
Ok(!self.stop_after_tool_call)
}
fn finish(&mut self) -> Result<(), ProviderError> {
self.flush_filtered_output()?;
let actions = self
.emulator_parser
.as_mut()
.map(StreamingEmulatorParser::flush)
.unwrap_or_default();
for action in actions {
let (message, is_tool) = message_for_emulator_action(&action, self.message_id);
if is_tool {
self.flush_filtered_output()?;
}
self.send(message)?;
self.stop_after_tool_call |= is_tool;
if is_tool {
break;
}
}
Ok(())
}
fn flush_filtered_output(&mut self) -> Result<(), ProviderError> {
let filtered = self.output_filter.finish();
if !filtered.thinking.is_empty() {
let mut message = Message::assistant().with_thinking(filtered.thinking, "");
message.id = Some(self.message_id.to_string());
self.send(message)?;
}
if !filtered.content.is_empty() {
self.emit_content(&filtered.content)?;
}
Ok(())
}
fn emit_content(&mut self, content: &str) -> Result<(), ProviderError> {
match self.tool_mode {
ToolMode::None => {
let mut message = Message::assistant().with_text(content);
message.id = Some(self.message_id.to_string());
self.send(message)
}
ToolMode::Emulated { .. } => {
let actions = self
.emulator_parser
.as_mut()
.map(|parser| parser.process_chunk(content))
.unwrap_or_default();
for action in actions {
let (message, is_tool) =
message_for_emulator_action(&action, self.message_id);
if is_tool {
self.flush_filtered_output()?;
}
self.send(message)?;
self.stop_after_tool_call |= is_tool;
if is_tool {
break;
}
}
Ok(())
}
ToolMode::Native => Ok(()),
}
}
fn send(&self, message: Message) -> Result<(), ProviderError> {
self.tx
.blocking_send(Ok((Some(message), None)))
.map_err(|_| {
ProviderError::ExecutionError("Failed to stream MLX response".to_string())
})
}
}
fn split_generated_thinking(
generated_text: &str,
generation_prompt: &str,
enable_thinking: bool,
) -> (String, String) {
let mut filter = ThinkingOutputFilter::new(enable_thinking, generation_prompt);
let mut filtered = filter.push_text(generated_text);
let final_filtered = filter.finish();
filtered.content.push_str(&final_filtered.content);
filtered.thinking.push_str(&final_filtered.thinking);
(filtered.content, filtered.thinking)
}
fn emit_assistant_message(
message_id: &str,
thinking: &str,
content: &str,
tx: &tokio::sync::mpsc::Sender<
Result<(Option<Message>, Option<ProviderUsage>), ProviderError>,
>,
) -> Result<(), ProviderError> {
if thinking.is_empty() && content.is_empty() {
return Ok(());
}
let mut message = Message::assistant();
if !thinking.is_empty() {
message = message.with_thinking(thinking, "");
}
if !content.is_empty() {
message = message.with_text(content);
}
message.id = Some(message_id.to_string());
tx.blocking_send(Ok((Some(message), None)))
.map_err(|_| ProviderError::ExecutionError("Failed to stream MLX response".to_string()))
}
fn prepend_thinking(message: &mut Message, thinking: &str) {
if !thinking.is_empty() {
message
.content
.insert(0, MessageContent::thinking(thinking, ""));
}
}
fn sampling(settings: &ModelSettings) -> (f32, Option<u32>) {
match &settings.sampling {
crate::local_model_registry::SamplingConfig::Greedy => (0.0, None),
crate::local_model_registry::SamplingConfig::Temperature {
temperature, seed, ..
} => (*temperature, *seed),
crate::local_model_registry::SamplingConfig::MirostatV2 { seed, .. } => (0.0, *seed),
}
}
fn prng_key(temp: f32, seed: Option<u32>) -> Result<Option<Array>, ProviderError> {
if temp == 0.0 {
return Ok(None);
}
random::key(seed.unwrap_or(0) as u64)
.map(Some)
.map_err(mlx_error)
} }
fn render_prompt(system: &str, messages: &[Message]) -> String { fn render_prompt(system: &str, messages: &[Message]) -> String {
@ -547,6 +904,90 @@ mod imp {
fn mlx_error(error: impl std::fmt::Display) -> ProviderError { fn mlx_error(error: impl std::fmt::Display) -> ProviderError {
ProviderError::ExecutionError(format!("MLX backend error: {}", error)) ProviderError::ExecutionError(format!("MLX backend error: {}", error))
} }
#[cfg(test)]
mod tests {
use super::{
final_stream_suffix, mlx_max_tokens, split_generated_thinking, token_id_or_ids,
};
use crate::local_model_registry::ModelSettings;
use serde_json::json;
#[test]
fn extracts_thinking_started_by_generation_prompt() {
let (content, thinking) = split_generated_thinking(
"hidden reasoning</think>visible answer",
"<|im_start|>assistant\n<think>\n",
true,
);
assert_eq!(thinking.trim(), "hidden reasoning");
assert_eq!(content, "visible answer");
}
#[test]
fn leaves_think_tags_as_content_when_thinking_disabled() {
let generated = "hidden reasoning</think>visible answer";
let (content, thinking) =
split_generated_thinking(generated, "<|im_start|>assistant\n<think>\n", false);
assert!(thinking.is_empty());
assert_eq!(content, generated);
}
#[test]
fn parses_single_and_multiple_eos_token_ids() {
assert_eq!(token_id_or_ids(&json!(248044)), vec![248044]);
assert_eq!(
token_id_or_ids(&json!([248046, 248044])),
vec![248046, 248044]
);
}
#[test]
fn final_stream_suffix_flushes_append_only_suffix() {
assert_eq!(
final_stream_suffix("hello world", "hello").unwrap(),
Some(" world")
);
}
#[test]
fn final_stream_suffix_does_not_replay_fully_streamed_text() {
assert_eq!(
final_stream_suffix("run tool", "run tool").unwrap(),
Some("")
);
}
#[test]
fn final_stream_suffix_allows_unstreamed_fallback() {
assert_eq!(final_stream_suffix("hello world", "").unwrap(), None);
}
#[test]
fn final_stream_suffix_rejects_rewritten_streamed_prefix() {
assert!(final_stream_suffix("corrected response", "stale prefix").is_err());
}
#[test]
fn max_tokens_defaults_to_context_headroom() {
let settings = ModelSettings::default();
assert_eq!(mlx_max_tokens(&settings, None, 128_000, 1_752), 126_248);
}
#[test]
fn max_tokens_respects_configured_caps() {
let mut settings = ModelSettings::default();
settings.max_output_tokens = Some(2048);
assert_eq!(mlx_max_tokens(&settings, None, 128_000, 1_752), 2048);
let settings = ModelSettings::default();
assert_eq!(mlx_max_tokens(&settings, Some(1024), 128_000, 1_752), 1024);
}
}
} }
#[cfg(not(feature = "mlx"))] #[cfg(not(feature = "mlx"))]

81
crates/goose-local-inference/src/thinking_output.rs Normal file
View file

@ -0,0 +1,81 @@
use goose_provider_types::thinking::{FilterOut, ThinkFilter};
pub(crate) struct ThinkingOutputFilter {
enabled: bool,
saw_structured_reasoning: bool,
think_filter: ThinkFilter,
pending_inline_thinking: String,
accumulated_thinking: String,
}
impl ThinkingOutputFilter {
pub(crate) fn new(enable_thinking: bool, generation_prompt: &str) -> Self {
let mut think_filter = ThinkFilter::new();
if enable_thinking && !generation_prompt.is_empty() {
let _ = think_filter.push(generation_prompt);
}
Self {
enabled: enable_thinking,
saw_structured_reasoning: false,
think_filter,
pending_inline_thinking: String::new(),
accumulated_thinking: String::new(),
}
}
pub(crate) fn push_structured_reasoning(&mut self, reasoning: &str) -> Option<String> {
if reasoning.is_empty() {
return None;
}
self.saw_structured_reasoning = true;
self.pending_inline_thinking.clear();
self.think_filter = ThinkFilter::new();
self.accumulated_thinking.push_str(reasoning);
Some(reasoning.to_string())
}
pub(crate) fn push_text(&mut self, text: &str) -> FilterOut {
if !self.enabled {
return FilterOut {
content: text.to_string(),
thinking: String::new(),
};
}
let mut filtered = self.think_filter.push(text);
if self.saw_structured_reasoning {
filtered.thinking.clear();
} else if !filtered.thinking.is_empty() {
self.pending_inline_thinking.push_str(&filtered.thinking);
filtered.thinking.clear();
}
filtered
}
pub(crate) fn finish(&mut self) -> FilterOut {
let mut filtered = if self.enabled && !self.saw_structured_reasoning {
std::mem::take(&mut self.think_filter).finish()
} else {
FilterOut::default()
};
if !self.saw_structured_reasoning {
let mut thinking = std::mem::take(&mut self.pending_inline_thinking);
thinking.push_str(&filtered.thinking);
if !thinking.is_empty() {
self.accumulated_thinking.push_str(&thinking);
}
filtered.thinking = thinking;
} else {
filtered.thinking.clear();
}
filtered
}
pub(crate) fn accumulated_thinking(&self) -> &str {
&self.accumulated_thinking
}
}

1
crates/goose-provider-types/src/conversation/message.rs
View file

@ -253,6 +253,7 @@ pub struct FrontendToolRequest {
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
pub enum SystemNotificationType { pub enum SystemNotificationType {
ThinkingMessage, ThinkingMessage,
ProgressMessage,
InlineMessage, InlineMessage,
CreditsExhausted, CreditsExhausted,
} }

1
crates/goose-provider-types/src/conversation/token_usage.rs
View file

@ -25,6 +25,7 @@ pub enum CostSource {
#[derive(Debug, Clone, Serialize, Deserialize, Default)] #[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ProviderStats { pub struct ProviderStats {
pub time_to_first_token_ms: Option<u64>, pub time_to_first_token_ms: Option<u64>,
pub model_load_ms: Option<u64>,
pub elapsed_ms: Option<u64>, pub elapsed_ms: Option<u64>,
pub output_tokens: Option<usize>, pub output_tokens: Option<usize>,
pub draft: Option<DraftStats>, pub draft: Option<DraftStats>,

11
crates/goose-sdk-types/src/custom_requests.rs
View file

@ -1910,6 +1910,7 @@ pub struct LocalInferenceModelDto {
pub size_bytes: u64, pub size_bytes: u64,
pub status: LocalInferenceModelDownloadStatusDto, pub status: LocalInferenceModelDownloadStatusDto,
pub recommended: bool, pub recommended: bool,
pub is_loaded: bool,
pub settings: LocalInferenceModelSettingsDto, pub settings: LocalInferenceModelSettingsDto,
pub vision_capable: bool, pub vision_capable: bool,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
@ -2031,6 +2032,16 @@ pub struct LocalInferenceModelDeleteRequest {
pub model_id: String, pub model_id: String,
} }
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request(
method = "_goose/unstable/local-inference/models/evict",
response = EmptyResponse
)]
#[serde(rename_all = "camelCase")]
pub struct LocalInferenceModelEvictRequest {
pub model_id: String,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)] #[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)]
#[request( #[request(
method = "_goose/unstable/local-inference/models/settings/read", method = "_goose/unstable/local-inference/models/settings/read",

5
crates/goose/acp-meta.json
View file

@ -530,6 +530,11 @@
"requestType": "LocalInferenceModelDeleteRequest_unstable", "requestType": "LocalInferenceModelDeleteRequest_unstable",
"responseType": "EmptyResponse" "responseType": "EmptyResponse"
}, },
{
"method": "_goose/unstable/local-inference/models/evict",
"requestType": "LocalInferenceModelEvictRequest_unstable",
"responseType": "EmptyResponse"
},
{ {
"method": "_goose/unstable/local-inference/models/settings/read", "method": "_goose/unstable/local-inference/models/settings/read",
"requestType": "LocalInferenceModelSettingsReadRequest_unstable", "requestType": "LocalInferenceModelSettingsReadRequest_unstable",

26
crates/goose/acp-schema.json
View file

@ -5734,6 +5734,9 @@
"recommended": { "recommended": {
"type": "boolean" "type": "boolean"
}, },
"isLoaded": {
"type": "boolean"
},
"settings": { "settings": {
"$ref": "#/$defs/LocalInferenceModelSettingsDto" "$ref": "#/$defs/LocalInferenceModelSettingsDto"
}, },
@ -5759,6 +5762,7 @@
"sizeBytes", "sizeBytes",
"status", "status",
"recommended", "recommended",
"isLoaded",
"settings", "settings",
"visionCapable" "visionCapable"
] ]
@ -6217,6 +6221,19 @@
"x-side": "agent", "x-side": "agent",
"x-method": "_goose/unstable/local-inference/models/delete" "x-method": "_goose/unstable/local-inference/models/delete"
}, },
"LocalInferenceModelEvictRequest_unstable": {
"type": "object",
"properties": {
"modelId": {
"type": "string"
}
},
"required": [
"modelId"
],
"x-side": "agent",
"x-method": "_goose/unstable/local-inference/models/evict"
},
"LocalInferenceModelSettingsReadRequest_unstable": { "LocalInferenceModelSettingsReadRequest_unstable": {
"type": "object", "type": "object",
"properties": { "properties": {
@ -7789,6 +7806,15 @@
"description": "Params for _goose/unstable/local-inference/models/delete", "description": "Params for _goose/unstable/local-inference/models/delete",
"title": "LocalInferenceModelDeleteRequest_unstable" "title": "LocalInferenceModelDeleteRequest_unstable"
}, },
{
"allOf": [
{
"$ref": "#/$defs/LocalInferenceModelEvictRequest_unstable"
}
],
"description": "Params for _goose/unstable/local-inference/models/evict",
"title": "LocalInferenceModelEvictRequest_unstable"
},
{ {
"allOf": [ "allOf": [
{ {

72
crates/goose/src/acp/server.rs
View file

@ -2072,6 +2072,23 @@ fn send_status_message_update(
Ok(()) Ok(())
} }
fn send_progress_message_update(
cx: &ConnectionTo<Client>,
supports_goose_custom_notifications: bool,
session_id: &str,
message: String,
) -> Result<(), agent_client_protocol::Error> {
if supports_goose_custom_notifications {
cx.send_notification(GooseSessionNotification {
session_id: session_id.to_string(),
update: GooseSessionUpdate::StatusMessage(StatusMessageUpdate {
status: StatusMessage::Progress { message },
}),
})?;
}
Ok(())
}
fn status_message_from_system_notification( fn status_message_from_system_notification(
notification: &SystemNotificationContent, notification: &SystemNotificationContent,
) -> Option<StatusMessage> { ) -> Option<StatusMessage> {
@ -2079,9 +2096,11 @@ fn status_message_from_system_notification(
SystemNotificationType::InlineMessage => Some(StatusMessage::Notice { SystemNotificationType::InlineMessage => Some(StatusMessage::Notice {
message: notification.msg.clone(), message: notification.msg.clone(),
}), }),
SystemNotificationType::ThinkingMessage => Some(StatusMessage::Progress { SystemNotificationType::ThinkingMessage | SystemNotificationType::ProgressMessage => {
message: notification.msg.clone(), Some(StatusMessage::Progress {
}), message: notification.msg.clone(),
})
}
SystemNotificationType::CreditsExhausted => None, SystemNotificationType::CreditsExhausted => None,
} }
} }
@ -2482,6 +2501,44 @@ impl GooseAcpAgent {
)) ))
} }
async fn send_local_inference_progress_update(
&self,
cx: &ConnectionTo<Client>,
acp_session_id: &SessionId,
session_id: &str,
agent: &Arc<Agent>,
) -> Result<(), agent_client_protocol::Error> {
let Ok(provider) = agent.provider().await else {
return Ok(());
};
if provider.get_name() != "local" {
return Ok(());
}
let model_config = agent.model_config_for_session(session_id).await.ok();
let model_name = model_config
.as_ref()
.map(|config| config.model_name.clone())
.unwrap_or_else(|| "local model".to_string());
#[cfg(feature = "local-inference")]
if let Some(model_config) = model_config.as_ref() {
if crate::providers::local_inference::is_model_loaded(&model_config.model_name)
.await
.unwrap_or(false)
{
return Ok(());
}
}
send_progress_message_update(
cx,
self.supports_goose_custom_notifications(),
acp_session_id.0.as_ref(),
format!("Loading local model {model_name}..."),
)
}
async fn on_load_session( async fn on_load_session(
&self, &self,
cx: &ConnectionTo<Client>, cx: &ConnectionTo<Client>,
@ -2524,6 +2581,15 @@ impl GooseAcpAgent {
return Err(error); return Err(error);
} }
if let Err(error) = self
.send_local_inference_progress_update(cx, &args.session_id, &session_id, &agent)
.await
{
self.clear_active_run(&session_id, &run_id).await;
let _ = Self::send_active_run_update(cx, &args.session_id, None);
return Err(error);
}
let user_message = Self::convert_acp_prompt_to_message(&args.prompt); let user_message = Self::convert_acp_prompt_to_message(&args.prompt);
let message_text = user_message.as_concat_text(); let message_text = user_message.as_concat_text();

17
crates/goose/src/acp/server/config.rs
View file

@ -216,7 +216,8 @@ impl GooseAcpAgent {
if let Some(model_id) = model_id.as_deref() { if let Some(model_id) = model_id.as_deref() {
let model_exists = entry.default_model == model_id let model_exists = entry.default_model == model_id
|| entry.models.iter().any(|model| model.id == model_id); || entry.models.iter().any(|model| model.id == model_id)
|| (provider_id == "local" && local_inference_model_exists(model_id)?);
if !model_exists { if !model_exists {
return Err(agent_client_protocol::Error::invalid_params().data(format!( return Err(agent_client_protocol::Error::invalid_params().data(format!(
"Model '{model_id}' is not available for provider '{provider_id}'" "Model '{model_id}' is not available for provider '{provider_id}'"
@ -254,6 +255,20 @@ impl GooseAcpAgent {
} }
} }
fn local_inference_model_exists(model_id: &str) -> Result<bool, agent_client_protocol::Error> {
#[cfg(feature = "local-inference")]
{
crate::providers::local_inference::management::model_exists(model_id)
.internal_err_ctx("Failed to read local inference models")
}
#[cfg(not(feature = "local-inference"))]
{
let _ = model_id;
Ok(false)
}
}
struct PreferenceDef { struct PreferenceDef {
key: PreferenceKey, key: PreferenceKey,
config_key: &'static str, config_key: &'static str,

8
crates/goose/src/acp/server/custom_dispatch.rs
View file

@ -875,6 +875,14 @@ impl GooseAcpAgent {
self.on_local_inference_model_delete(req).await self.on_local_inference_model_delete(req).await
} }
#[custom_method(LocalInferenceModelEvictRequest)]
async fn dispatch_local_inference_model_evict(
&self,
req: LocalInferenceModelEvictRequest,
) -> Result<EmptyResponse, agent_client_protocol::Error> {
self.on_local_inference_model_evict(req).await
}
#[custom_method(LocalInferenceModelSettingsReadRequest)] #[custom_method(LocalInferenceModelSettingsReadRequest)]
async fn dispatch_local_inference_model_settings_read( async fn dispatch_local_inference_model_settings_read(
&self, &self,

20
crates/goose/src/acp/server/local_inference.rs
View file

@ -98,6 +98,26 @@ impl GooseAcpAgent {
} }
} }
pub(super) async fn on_local_inference_model_evict(
&self,
req: LocalInferenceModelEvictRequest,
) -> Result<EmptyResponse, agent_client_protocol::Error> {
#[cfg(feature = "local-inference")]
{
crate::providers::local_inference::configure_huggingface_auth();
crate::providers::local_inference::management::evict_model(&req.model_id)
.await
.invalid_params_err()?;
Ok(EmptyResponse {})
}
#[cfg(not(feature = "local-inference"))]
{
let _ = req;
Err(local_inference_unavailable())
}
}
pub(super) async fn on_local_inference_model_settings_read( pub(super) async fn on_local_inference_model_settings_read(
&self, &self,
req: LocalInferenceModelSettingsReadRequest, req: LocalInferenceModelSettingsReadRequest,

21
crates/goose/src/agents/agent.rs
View file

@ -68,7 +68,7 @@ use tracing::{debug, error, info, instrument, warn};
const DEFAULT_MAX_TURNS: u32 = 1000; const DEFAULT_MAX_TURNS: u32 = 1000;
const DEFAULT_STOP_HOOK_BLOCK_CAP: u32 = 8; const DEFAULT_STOP_HOOK_BLOCK_CAP: u32 = 8;
const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation..."; const COMPACTION_PROGRESS_TEXT: &str = "goose is compacting the conversation...";
const MAX_TURNS_MESSAGE: &str = "I've reached the maximum number of actions I can do without user input. Would you like me to continue?"; const MAX_TURNS_MESSAGE: &str = "I've reached the maximum number of actions I can do without user input. Would you like me to continue?";
const DEFAULT_FRONTEND_INSTRUCTIONS: &str = "The following tools are provided directly by the frontend and will be executed by the frontend when called."; const DEFAULT_FRONTEND_INSTRUCTIONS: &str = "The following tools are provided directly by the frontend and will be executed by the frontend when called.";
@ -1769,8 +1769,8 @@ impl Agent {
yield AgentEvent::Message( yield AgentEvent::Message(
Message::assistant().with_system_notification( Message::assistant().with_system_notification(
SystemNotificationType::ThinkingMessage, SystemNotificationType::ProgressMessage,
COMPACTION_THINKING_TEXT, COMPACTION_PROGRESS_TEXT,
) )
); );
@ -2060,6 +2060,17 @@ impl Agent {
} }
if let Some(response) = response { if let Some(response) = response {
if !response.content.is_empty()
&& response
.content
.iter()
.all(|content| matches!(content, MessageContent::SystemNotification(_)))
{
yield AgentEvent::Message(response);
tokio::task::yield_now().await;
continue;
}
let ToolCategorizeResult { let ToolCategorizeResult {
frontend_requests, frontend_requests,
remaining_requests, remaining_requests,
@ -2427,8 +2438,8 @@ impl Agent {
); );
yield AgentEvent::Message( yield AgentEvent::Message(
Message::assistant().with_system_notification( Message::assistant().with_system_notification(
SystemNotificationType::ThinkingMessage, SystemNotificationType::ProgressMessage,
COMPACTION_THINKING_TEXT, COMPACTION_PROGRESS_TEXT,
) )
); );

36
crates/goose/src/agents/reply_parts.rs
View file

@ -161,6 +161,13 @@ fn fill_stream_timing(
} }
} }
fn message_has_timing_content(message: &Message) -> bool {
message
.content
.iter()
.any(|content| !matches!(content, MessageContent::SystemNotification(_)))
}
impl Agent { impl Agent {
pub async fn prepare_tools_and_prompt( pub async fn prepare_tools_and_prompt(
&self, &self,
@ -345,7 +352,7 @@ impl Agent {
let (msg_opt, usage_opt) = result?; let (msg_opt, usage_opt) = result?;
if let Some(msg) = msg_opt { if let Some(msg) = msg_opt {
if first_content_at.is_none() { if first_content_at.is_none() && message_has_timing_content(&msg) {
first_content_at = Some(std::time::Instant::now()); first_content_at = Some(std::time::Instant::now());
} }
accumulated_message = Some(match accumulated_message { accumulated_message = Some(match accumulated_message {
@ -394,7 +401,9 @@ impl Agent {
while let Some(result) = stream.next().await { while let Some(result) = stream.next().await {
let (message, mut usage) = result?; let (message, mut usage) = result?;
if message.is_some() && first_content_at.is_none() { if first_content_at.is_none()
&& message.as_ref().is_some_and(message_has_timing_content)
{
first_content_at = Some(std::time::Instant::now()); first_content_at = Some(std::time::Instant::now());
} }
if let Some(usage) = usage.as_mut() { if let Some(usage) = usage.as_mut() {
@ -656,7 +665,7 @@ pub fn is_tool_visible_to_model(tool: &Tool) -> bool {
mod tests { mod tests {
use super::*; use super::*;
use crate::config::GooseMode; use crate::config::GooseMode;
use crate::conversation::message::Message; use crate::conversation::message::{Message, SystemNotificationType};
use crate::providers::base::Provider; use crate::providers::base::Provider;
use crate::session::session_manager::SessionType; use crate::session::session_manager::SessionType;
use async_trait::async_trait; use async_trait::async_trait;
@ -1036,6 +1045,27 @@ mod tests {
usage usage
} }
#[test]
fn message_has_timing_content_ignores_system_notification_only_messages() {
let message = Message::assistant().with_system_notification(
SystemNotificationType::ProgressMessage,
"Loading local model test-model...",
);
assert!(!message_has_timing_content(&message));
}
#[test]
fn message_has_timing_content_counts_user_visible_messages() {
let text_message = Message::assistant().with_text("hello");
let mixed_message = Message::assistant()
.with_system_notification(SystemNotificationType::ProgressMessage, "Loading...")
.with_text("ready");
assert!(message_has_timing_content(&text_message));
assert!(message_has_timing_content(&mixed_message));
}
#[test] #[test]
fn fill_stream_timing_fills_both_fields_when_stats_absent() { fn fill_stream_timing_fills_both_fields_when_stats_absent() {
let request_started = Instant::now() - Duration::from_millis(100); let request_started = Instant::now() - Duration::from_millis(100);

2
ui/desktop/src/acp/__tests__/chatNotifications.test.ts
View file

@ -81,6 +81,7 @@ function snapshotWithName(name: string): AcpChatSessionSnapshot {
accumulatedTotalTokens: 0, accumulatedTotalTokens: 0,
}, },
notifications: [], notifications: [],
progressMessage: undefined,
chatState: ChatState.Idle, chatState: ChatState.Idle,
sessionLoadError: undefined, sessionLoadError: undefined,
activePromptAttemptId: null, activePromptAttemptId: null,
@ -102,6 +103,7 @@ function snapshotWithoutSession(): AcpChatSessionSnapshot {
accumulatedTotalTokens: 0, accumulatedTotalTokens: 0,
}, },
notifications: [], notifications: [],
progressMessage: undefined,
chatState: ChatState.Idle, chatState: ChatState.Idle,
sessionLoadError: undefined, sessionLoadError: undefined,
activePromptAttemptId: null, activePromptAttemptId: null,

1
ui/desktop/src/acp/__tests__/chatSessionController.test.ts
View file

@ -108,6 +108,7 @@ function snapshotWithActivePrompt(activePromptAttemptId: string | null): AcpChat
accumulatedTotalTokens: 0, accumulatedTotalTokens: 0,
}, },
notifications: [], notifications: [],
progressMessage: undefined,
chatState: activePromptAttemptId ? ChatState.Streaming : ChatState.Idle, chatState: activePromptAttemptId ? ChatState.Streaming : ChatState.Idle,
sessionLoadError: undefined, sessionLoadError: undefined,
activePromptAttemptId, activePromptAttemptId,

10
ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts
View file

@ -583,13 +583,9 @@ describe('createAcpSessionNotificationAdapter', () => {
status: { type: 'progress', message: 'Still working' }, status: { type: 'progress', message: 'Still working' },
}) })
); );
messages = expectOnlyMessagesChange(progressStateChanges); expect(progressStateChanges).toEqual([
{ type: 'progressMessage', message: 'Still working' },
expect(firstContent(messages[2])).toMatchObject({ ]);
type: 'systemNotification',
notificationType: 'thinkingMessage',
msg: 'Still working',
});
}); });
}); });

6
ui/desktop/src/acp/adapter/gooseSessionNotifications.ts
View file

@ -38,7 +38,9 @@ function applyStatusMessage(
sessionId: string, sessionId: string,
update: Extract<GooseSessionNotification_unstable['update'], { sessionUpdate: 'status_message' }> update: Extract<GooseSessionNotification_unstable['update'], { sessionUpdate: 'status_message' }>
): AcpChatStateChange[] { ): AcpChatStateChange[] {
const notificationType = update.status.type === 'notice' ? 'inlineMessage' : 'thinkingMessage'; if (update.status.type === 'progress') {
return [{ type: 'progressMessage', message: update.status.message }];
}
state.messages.push({ state.messages.push({
id: `acp_status_${sessionId}_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`, id: `acp_status_${sessionId}_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`,
@ -47,7 +49,7 @@ function applyStatusMessage(
content: [ content: [
{ {
type: 'systemNotification', type: 'systemNotification',
notificationType, notificationType: 'inlineMessage',
msg: update.status.message, msg: update.status.message,
}, },
], ],

1
ui/desktop/src/acp/adapter/shared.ts
View file

@ -5,6 +5,7 @@ import type { Message, NotificationEvent } from '../../types/message';
export type AcpChatStateChange = export type AcpChatStateChange =
| { type: 'messages'; messages: Message[] } | { type: 'messages'; messages: Message[] }
| { type: 'tokenState'; tokenState: Partial<TokenState> } | { type: 'tokenState'; tokenState: Partial<TokenState> }
| { type: 'progressMessage'; message: string | undefined }
| { | {
type: 'sessionInfo'; type: 'sessionInfo';
name?: string; name?: string;

30
ui/desktop/src/acp/chatSessionStore.ts
View file

@ -19,6 +19,7 @@ export interface AcpChatSessionSnapshot {
messages: Message[]; messages: Message[];
tokenState: TokenState; tokenState: TokenState;
notifications: NotificationEvent[]; notifications: NotificationEvent[];
progressMessage: string | undefined;
chatState: ChatState; chatState: ChatState;
sessionLoadError: string | undefined; sessionLoadError: string | undefined;
activePromptAttemptId: string | null; activePromptAttemptId: string | null;
@ -155,6 +156,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
messages: [], messages: [],
tokenState: { ...initialTokenState }, tokenState: { ...initialTokenState },
notifications: [], notifications: [],
progressMessage: undefined,
chatState: ChatState.Idle, chatState: ChatState.Idle,
sessionLoadError: undefined, sessionLoadError: undefined,
activePromptAttemptId: null, activePromptAttemptId: null,
@ -190,6 +192,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
const entry = getOrCreateEntry(sessionId); const entry = getOrCreateEntry(sessionId);
resetReplayState(entry); resetReplayState(entry);
entry.sessionLoadError = undefined; entry.sessionLoadError = undefined;
entry.progressMessage = undefined;
entry.chatState = ChatState.LoadingConversation; entry.chatState = ChatState.LoadingConversation;
return notify(sessionId, entry); return notify(sessionId, entry);
}; };
@ -198,6 +201,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
const entry = getOrCreateEntry(sessionId); const entry = getOrCreateEntry(sessionId);
entry.session = session; entry.session = session;
entry.sessionLoadError = undefined; entry.sessionLoadError = undefined;
entry.progressMessage = undefined;
entry.chatState = entry.activePromptAttemptId ? ChatState.Streaming : ChatState.Idle; entry.chatState = entry.activePromptAttemptId ? ChatState.Streaming : ChatState.Idle;
return notify(sessionId, entry); return notify(sessionId, entry);
}; };
@ -208,6 +212,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
) => { ) => {
const entry = getOrCreateEntry(sessionId); const entry = getOrCreateEntry(sessionId);
entry.sessionLoadError = sessionLoadError; entry.sessionLoadError = sessionLoadError;
entry.progressMessage = undefined;
entry.chatState = ChatState.Idle; entry.chatState = ChatState.Idle;
return notify(sessionId, entry); return notify(sessionId, entry);
}; };
@ -237,6 +242,9 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
const setChatState: AcpChatSessionActions['setChatState'] = (sessionId, chatState) => { const setChatState: AcpChatSessionActions['setChatState'] = (sessionId, chatState) => {
const entry = getOrCreateEntry(sessionId); const entry = getOrCreateEntry(sessionId);
if (chatState === ChatState.Idle) {
entry.progressMessage = undefined;
}
entry.chatState = chatState; entry.chatState = chatState;
return notify(sessionId, entry); return notify(sessionId, entry);
}; };
@ -287,6 +295,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
entry.chatState = ChatState.Streaming; entry.chatState = ChatState.Streaming;
entry.sessionLoadError = undefined; entry.sessionLoadError = undefined;
entry.notifications = []; entry.notifications = [];
entry.progressMessage = undefined;
return notify(sessionId, entry); return notify(sessionId, entry);
}; };
@ -309,6 +318,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
entry.pendingCancelPromptAttemptId = promptAttemptId; entry.pendingCancelPromptAttemptId = promptAttemptId;
entry.pendingUserInputRequestIds.clear(); entry.pendingUserInputRequestIds.clear();
discardPendingLocalSteerMessages(entry); discardPendingLocalSteerMessages(entry);
entry.progressMessage = undefined;
entry.chatState = ChatState.Idle; entry.chatState = ChatState.Idle;
return notify(sessionId, entry); return notify(sessionId, entry);
}; };
@ -385,6 +395,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
entry.promptCancellationRestoreState = null; entry.promptCancellationRestoreState = null;
entry.pendingUserInputRequestIds.clear(); entry.pendingUserInputRequestIds.clear();
discardPendingLocalSteerMessages(entry); discardPendingLocalSteerMessages(entry);
entry.progressMessage = undefined;
entry.chatState = ChatState.Idle; entry.chatState = ChatState.Idle;
entry.sessionLoadError = error; entry.sessionLoadError = error;
notify(sessionId, entry); notify(sessionId, entry);
@ -416,6 +427,9 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal {
notification notification
) => { ) => {
const entry = getOrCreateEntry(notification.sessionId); const entry = getOrCreateEntry(notification.sessionId);
if (shouldClearProgressMessage(notification)) {
entry.progressMessage = undefined;
}
const changes = entry.adapter.apply(notification); const changes = entry.adapter.apply(notification);
applyChatStateChanges(entry, changes); applyChatStateChanges(entry, changes);
return notify(notification.sessionId, entry); return notify(notification.sessionId, entry);
@ -580,6 +594,9 @@ function applyChatStateChanges(entry: StoreEntry, changes: AcpChatStateChange[])
case 'tokenState': case 'tokenState':
entry.tokenState = { ...entry.tokenState, ...change.tokenState }; entry.tokenState = { ...entry.tokenState, ...change.tokenState };
break; break;
case 'progressMessage':
entry.progressMessage = change.message;
break;
case 'sessionInfo': case 'sessionInfo':
if (change.name && entry.session) { if (change.name && entry.session) {
entry.session = { ...entry.session, name: change.name }; entry.session = { ...entry.session, name: change.name };
@ -598,10 +615,22 @@ function applyChatStateChanges(entry: StoreEntry, changes: AcpChatStateChange[])
} }
} }
function shouldClearProgressMessage(notification: SessionNotification): boolean {
switch (notification.update.sessionUpdate) {
case 'agent_message_chunk':
case 'agent_thought_chunk':
case 'tool_call':
return true;
default:
return false;
}
}
function resetReplayState(entry: StoreEntry): void { function resetReplayState(entry: StoreEntry): void {
entry.messages = []; entry.messages = [];
entry.tokenState = { ...initialTokenState }; entry.tokenState = { ...initialTokenState };
entry.notifications = []; entry.notifications = [];
entry.progressMessage = undefined;
entry.activeRunId = null; entry.activeRunId = null;
entry.pendingCancelPromptAttemptId = null; entry.pendingCancelPromptAttemptId = null;
entry.promptCancellationRestoreState = null; entry.promptCancellationRestoreState = null;
@ -675,6 +704,7 @@ function snapshotFromEntry(entry: StoreEntry): AcpChatSessionSnapshot {
messages: cloneMessages(entry.messages), messages: cloneMessages(entry.messages),
tokenState: { ...entry.tokenState }, tokenState: { ...entry.tokenState },
notifications: [...entry.notifications], notifications: [...entry.notifications],
progressMessage: entry.progressMessage,
chatState: entry.chatState, chatState: entry.chatState,
sessionLoadError: entry.sessionLoadError, sessionLoadError: entry.sessionLoadError,
activePromptAttemptId: entry.activePromptAttemptId, activePromptAttemptId: entry.activePromptAttemptId,

5
ui/desktop/src/acp/local-inference.ts
View file

@ -56,6 +56,11 @@ export async function deleteLocalModel(modelId: string): Promise<void> {
await client.goose.localInferenceModelsDelete_unstable({ modelId }); await client.goose.localInferenceModelsDelete_unstable({ modelId });
} }
export async function evictLocalModel(modelId: string): Promise<void> {
const client = await getAcpClient();
await client.goose.localInferenceModelsEvict_unstable({ modelId });
}
export async function getModelSettings(modelId: string): Promise<ModelSettings> { export async function getModelSettings(modelId: string): Promise<ModelSettings> {
const client = await getAcpClient(); const client = await getAcpClient();
const response = await client.goose.localInferenceModelsSettingsRead_unstable({ modelId }); const response = await client.goose.localInferenceModelsSettingsRead_unstable({ modelId });

17
ui/desktop/src/components/BaseChat.tsx
View file

@ -23,12 +23,7 @@ import { RecipeWarningModal } from './ui/RecipeWarningModal';
import { scanRecipe } from '../recipe'; import { scanRecipe } from '../recipe';
import type { Recipe } from '../recipe'; import type { Recipe } from '../recipe';
import RecipeActivities from './recipes/RecipeActivities'; import RecipeActivities from './recipes/RecipeActivities';
import { import { getTextAndImageContent, type Message, type UserInput } from '../types/message';
getThinkingMessage,
getTextAndImageContent,
type Message,
type UserInput,
} from '../types/message';
import { substituteParameters } from '../utils/parameterSubstitution'; import { substituteParameters } from '../utils/parameterSubstitution';
import { useAutoSubmit } from '../hooks/useAutoSubmit'; import { useAutoSubmit } from '../hooks/useAutoSubmit';
import { Goose } from './icons'; import { Goose } from './icons';
@ -92,6 +87,7 @@ export default function BaseChat({
session, session,
messages, messages,
chatState, chatState,
progressMessage,
updateSession, updateSession,
handleSubmit, handleSubmit,
onSteerQueuedMessage, onSteerQueuedMessage,
@ -469,14 +465,7 @@ export default function BaseChat({
{chatState !== ChatState.Idle && ( {chatState !== ChatState.Idle && (
<div className="absolute bottom-1 left-4 z-20 pointer-events-none"> <div className="absolute bottom-1 left-4 z-20 pointer-events-none">
<LoadingGoose <LoadingGoose chatState={chatState} message={progressMessage} />
chatState={chatState}
message={
messages.length > 0
? getThinkingMessage(messages[messages.length - 1])
: undefined
}
/>
</div> </div>
)} )}
</div> </div>

60
ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx
View file

@ -8,6 +8,8 @@ import {
Settings2, Settings2,
Eye, Eye,
RefreshCw, RefreshCw,
Cpu,
PowerOff,
} from 'lucide-react'; } from 'lucide-react';
import { Button } from '../../ui/button'; import { Button } from '../../ui/button';
import { useModelAndProvider } from '../../ModelAndProviderContext'; import { useModelAndProvider } from '../../ModelAndProviderContext';
@ -18,6 +20,7 @@ import {
getLocalModelDownloadProgress, getLocalModelDownloadProgress,
cancelLocalModelDownload, cancelLocalModelDownload,
deleteLocalModel, deleteLocalModel,
evictLocalModel,
type DownloadProgress, type DownloadProgress,
type DownloadModelRequest, type DownloadModelRequest,
type LocalModelResponse, type LocalModelResponse,
@ -106,6 +109,14 @@ const i18n = defineMessages({
id: 'localInferenceSettings.modelSettingsTitle', id: 'localInferenceSettings.modelSettingsTitle',
defaultMessage: 'Model settings', defaultMessage: 'Model settings',
}, },
loadedInMemory: {
id: 'localInferenceSettings.loadedInMemory',
defaultMessage: 'Loaded in memory',
},
evictFromMemory: {
id: 'localInferenceSettings.evictFromMemory',
defaultMessage: 'Evict from memory',
},
vision: { vision: {
id: 'localInferenceSettings.vision', id: 'localInferenceSettings.vision',
defaultMessage: 'Vision', defaultMessage: 'Vision',
@ -181,6 +192,7 @@ export const LocalInferenceSettings = () => {
const [downloadRequests, setDownloadRequests] = useState<Map<string, DownloadModelRequest>>( const [downloadRequests, setDownloadRequests] = useState<Map<string, DownloadModelRequest>>(
new Map() new Map()
); );
const [evictingModelId, setEvictingModelId] = useState<string | null>(null);
const [showAllFeatured, setShowAllFeatured] = useState(false); const [showAllFeatured, setShowAllFeatured] = useState(false);
const [settingsOpenFor, setSettingsOpenFor] = useState<string | null>(null); const [settingsOpenFor, setSettingsOpenFor] = useState<string | null>(null);
const { currentModel, currentProvider, refreshCurrentModelAndProvider } = useModelAndProvider(); const { currentModel, currentProvider, refreshCurrentModelAndProvider } = useModelAndProvider();
@ -371,6 +383,18 @@ export const LocalInferenceSettings = () => {
} }
}; };
const handleEvictModel = async (modelId: string) => {
setEvictingModelId(modelId);
try {
await evictLocalModel(modelId);
await loadModels();
} catch (error) {
console.error('Failed to evict model:', error);
} finally {
setEvictingModelId(null);
}
};
const handleHfDownloadStarted = (modelId: string, request: DownloadModelRequest) => { const handleHfDownloadStarted = (modelId: string, request: DownloadModelRequest) => {
setDownloadRequests((prev) => new Map(prev).set(modelId, request)); setDownloadRequests((prev) => new Map(prev).set(modelId, request));
pollDownloadProgress(modelId); pollDownloadProgress(modelId);
@ -393,6 +417,15 @@ export const LocalInferenceSettings = () => {
.map(([modelId]) => modelId) .map(([modelId]) => modelId)
); );
useEffect(() => {
if (downloadedModels.length === 0) return;
const interval = setInterval(() => {
loadModels();
}, 3000);
return () => clearInterval(interval);
}, [downloadedModels.length, loadModels]);
return ( return (
<div className="space-y-6"> <div className="space-y-6">
<div> <div>
@ -526,18 +559,26 @@ export const LocalInferenceSettings = () => {
: 'border-border-subtle bg-background-default hover:border-border-default' : 'border-border-subtle bg-background-default hover:border-border-default'
}`} }`}
> >
<div className="flex items-center justify-between"> <div className="flex items-center justify-between gap-3">
<div className="flex items-center gap-2"> <div className="flex min-w-0 items-center gap-2 flex-wrap">
<input <input
type="radio" type="radio"
checked={isSelected} checked={isSelected}
onChange={() => selectModel(model.id)} onChange={() => selectModel(model.id)}
className="cursor-pointer" className="cursor-pointer"
/> />
<span className="text-sm font-medium text-text-default">{model.id}</span> <span className="text-sm font-medium text-text-default break-all">
{model.id}
</span>
<span className="text-xs text-text-muted"> <span className="text-xs text-text-muted">
{formatBytes(model.sizeBytes)} {formatBytes(model.sizeBytes)}
</span> </span>
{model.isLoaded && (
<span className="inline-flex items-center gap-1 text-xs text-green-400 bg-green-500/10 px-2 py-0.5 rounded">
<Cpu className="w-3 h-3" />
{intl.formatMessage(i18n.loadedInMemory)}
</span>
)}
{model.recommended && ( {model.recommended && (
<span className="text-xs bg-blue-500 text-white px-2 py-0.5 rounded"> <span className="text-xs bg-blue-500 text-white px-2 py-0.5 rounded">
{intl.formatMessage(i18n.recommended)} {intl.formatMessage(i18n.recommended)}
@ -545,7 +586,7 @@ export const LocalInferenceSettings = () => {
)} )}
<VisionBadge model={model} intl={intl} /> <VisionBadge model={model} intl={intl} />
</div> </div>
<div className="flex items-center gap-1"> <div className="flex shrink-0 items-center gap-1">
<Button <Button
variant="ghost" variant="ghost"
size="sm" size="sm"
@ -554,6 +595,17 @@ export const LocalInferenceSettings = () => {
> >
<Settings2 className="w-4 h-4" /> <Settings2 className="w-4 h-4" />
</Button> </Button>
{model.isLoaded && (
<Button
variant="ghost"
size="sm"
onClick={() => handleEvictModel(model.id)}
disabled={evictingModelId === model.id}
title={intl.formatMessage(i18n.evictFromMemory)}
>
<PowerOff className="w-4 h-4" />
</Button>
)}
<Button <Button
variant="ghost" variant="ghost"
size="sm" size="sm"

2
ui/desktop/src/hooks/useChatSession.ts
View file

@ -61,6 +61,7 @@ export function useChatSession({
const messages = acpSnapshot?.messages ?? []; const messages = acpSnapshot?.messages ?? [];
const session = acpSnapshot?.session; const session = acpSnapshot?.session;
const chatState = acpSnapshot?.chatState ?? ChatState.LoadingConversation; const chatState = acpSnapshot?.chatState ?? ChatState.LoadingConversation;
const progressMessage = acpSnapshot?.progressMessage;
const sessionLoadError = acpSnapshot?.sessionLoadError; const sessionLoadError = acpSnapshot?.sessionLoadError;
const tokenState = acpSnapshot?.tokenState ?? initialTokenState; const tokenState = acpSnapshot?.tokenState ?? initialTokenState;
const queueProcessingBlocked = acpSnapshot?.pendingCancelPromptAttemptId != null; const queueProcessingBlocked = acpSnapshot?.pendingCancelPromptAttemptId != null;
@ -319,6 +320,7 @@ export function useChatSession({
messages, messages,
session, session,
chatState, chatState,
progressMessage,
updateSession, updateSession,
handleSubmit, handleSubmit,
onSteerQueuedMessage, onSteerQueuedMessage,

1
ui/desktop/src/hooks/useChatSessionTypes.ts
View file

@ -13,6 +13,7 @@ export interface UseChatSessionResult {
session?: Session; session?: Session;
messages: Message[]; messages: Message[];
chatState: ChatState; chatState: ChatState;
progressMessage?: string;
updateSession: (updater: (session: Session) => Session) => void; updateSession: (updater: (session: Session) => Session) => void;
handleSubmit: (input: UserInput) => Promise<void>; handleSubmit: (input: UserInput) => Promise<void>;
onSteerQueuedMessage?: (input: UserInput) => Promise<boolean>; onSteerQueuedMessage?: (input: UserInput) => Promise<boolean>;

6
ui/desktop/src/i18n/messages/de.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Wird heruntergeladen" "defaultMessage": "Wird heruntergeladen"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Empfohlene Modelle" "defaultMessage": "Empfohlene Modelle"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Melden Sie sich an, um die Ratenbegrenzungen beim Suchen und Herunterladen von Modellen zu erhöhen und auf private oder zugriffsbeschränkte Hugging-Face-Repositories zuzugreifen." "defaultMessage": "Melden Sie sich an, um die Ratenbegrenzungen beim Suchen und Herunterladen von Modellen zu erhöhen und auf private oder zugriffsbeschränkte Hugging-Face-Repositories zuzugreifen."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Modelleinstellungen" "defaultMessage": "Modelleinstellungen"
}, },

6
ui/desktop/src/i18n/messages/en.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Downloading" "defaultMessage": "Downloading"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Featured Models" "defaultMessage": "Featured Models"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Sign in to increase rate limits when searching and downloading models, and to access private or gated Hugging Face repositories." "defaultMessage": "Sign in to increase rate limits when searching and downloading models, and to access private or gated Hugging Face repositories."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Model Settings" "defaultMessage": "Model Settings"
}, },

6
ui/desktop/src/i18n/messages/es.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": { "localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)" "defaultMessage": "{downloaded} / {total} ({percent}%)"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Modelos destacados" "defaultMessage": "Modelos destacados"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Inicia sesión para aumentar los límites de uso al buscar y descargar modelos, y para acceder a repositorios de Hugging Face privados o restringidos." "defaultMessage": "Inicia sesión para aumentar los límites de uso al buscar y descargar modelos, y para acceder a repositorios de Hugging Face privados o restringidos."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Ajustes del modelo" "defaultMessage": "Ajustes del modelo"
}, },

6
ui/desktop/src/i18n/messages/fr.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Téléchargement en cours" "defaultMessage": "Téléchargement en cours"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Modèles en vedette" "defaultMessage": "Modèles en vedette"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Connectez-vous pour augmenter les limites de débit lors de la recherche et du téléchargement de modèles, et pour accéder aux dépôts Hugging Face privés ou à accès restreint." "defaultMessage": "Connectez-vous pour augmenter les limites de débit lors de la recherche et du téléchargement de modèles, et pour accéder aux dépôts Hugging Face privés ou à accès restreint."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Paramètres du modèle" "defaultMessage": "Paramètres du modèle"
}, },

6
ui/desktop/src/i18n/messages/hi.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": { "localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)" "defaultMessage": "{downloaded} / {total} ({percent}%)"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "विशेष रुप से प्रदर्शित मॉडल" "defaultMessage": "विशेष रुप से प्रदर्शित मॉडल"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "मॉडल खोजते और डाउनलोड करते समय दर सीमा बढ़ाने के लिए और निजी या गेटेड हगिंग फेस रिपॉजिटरी तक पहुंचने के लिए साइन इन करें।" "defaultMessage": "मॉडल खोजते और डाउनलोड करते समय दर सीमा बढ़ाने के लिए और निजी या गेटेड हगिंग फेस रिपॉजिटरी तक पहुंचने के लिए साइन इन करें।"
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "मॉडल Settings" "defaultMessage": "मॉडल Settings"
}, },

6
ui/desktop/src/i18n/messages/id.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Mengunduh" "defaultMessage": "Mengunduh"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Model Unggulan" "defaultMessage": "Model Unggulan"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Masuk untuk meningkatkan batas laju saat mencari dan mengunduh model, serta untuk mengakses repositori Hugging Face yang privat atau terbatas." "defaultMessage": "Masuk untuk meningkatkan batas laju saat mencari dan mengunduh model, serta untuk mengakses repositori Hugging Face yang privat atau terbatas."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Pengaturan Model" "defaultMessage": "Pengaturan Model"
}, },

6
ui/desktop/src/i18n/messages/it.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Download in corso" "defaultMessage": "Download in corso"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Modelli in evidenza" "defaultMessage": "Modelli in evidenza"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Accedi per aumentare i limiti di frequenza durante la ricerca e il download dei modelli e per accedere ai repository Hugging Face privati o ad accesso limitato." "defaultMessage": "Accedi per aumentare i limiti di frequenza durante la ricerca e il download dei modelli e per accedere ai repository Hugging Face privati o ad accesso limitato."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Impostazioni modello" "defaultMessage": "Impostazioni modello"
}, },

6
ui/desktop/src/i18n/messages/ja.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": { "localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)" "defaultMessage": "{downloaded} / {total} ({percent}%)"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "注目モデル" "defaultMessage": "注目モデル"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "モデルの検索・ダウンロード時のレート制限を緩和し、非公開またはゲート付きHugging Faceリポジトリにアクセスするにはサインインしてください。" "defaultMessage": "モデルの検索・ダウンロード時のレート制限を緩和し、非公開またはゲート付きHugging Faceリポジトリにアクセスするにはサインインしてください。"
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "モデル設定" "defaultMessage": "モデル設定"
}, },

6
ui/desktop/src/i18n/messages/ko.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": { "localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)" "defaultMessage": "{downloaded} / {total} ({percent}%)"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "주요 모델" "defaultMessage": "주요 모델"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "모델 검색 및 다운로드 시 속도 제한 한도를 높이고 비공개 또는 접근 제한된 Hugging Face 리포지토리에 액세스하려면 로그인하세요." "defaultMessage": "모델 검색 및 다운로드 시 속도 제한 한도를 높이고 비공개 또는 접근 제한된 Hugging Face 리포지토리에 액세스하려면 로그인하세요."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "모델 설정" "defaultMessage": "모델 설정"
}, },

6
ui/desktop/src/i18n/messages/ms.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Memuat turun" "defaultMessage": "Memuat turun"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Model Pilihan" "defaultMessage": "Model Pilihan"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Log masuk untuk meningkatkan had kadar semasa mencari dan memuat turun model, dan untuk mengakses repositori Hugging Face peribadi atau berpagar." "defaultMessage": "Log masuk untuk meningkatkan had kadar semasa mencari dan memuat turun model, dan untuk mengakses repositori Hugging Face peribadi atau berpagar."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Tetapan Model" "defaultMessage": "Tetapan Model"
}, },

6
ui/desktop/src/i18n/messages/pt.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Baixando" "defaultMessage": "Baixando"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Modelos em Destaque" "defaultMessage": "Modelos em Destaque"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Faça login para aumentar os limites de taxa ao pesquisar e baixar modelos, e para acessar repositórios privados ou restritos do Hugging Face." "defaultMessage": "Faça login para aumentar os limites de taxa ao pesquisar e baixar modelos, e para acessar repositórios privados ou restritos do Hugging Face."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Configurações do Modelo" "defaultMessage": "Configurações do Modelo"
}, },

6
ui/desktop/src/i18n/messages/ru.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": { "localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)" "defaultMessage": "{downloaded} / {total} ({percent}%)"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Рекомендуемые модели" "defaultMessage": "Рекомендуемые модели"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Войдите, чтобы увеличить лимиты запросов при поиске и загрузке моделей, а также получить доступ к приватным или ограниченным репозиториям Hugging Face." "defaultMessage": "Войдите, чтобы увеличить лимиты запросов при поиске и загрузке моделей, а также получить доступ к приватным или ограниченным репозиториям Hugging Face."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Настройки модели" "defaultMessage": "Настройки модели"
}, },

6
ui/desktop/src/i18n/messages/tr.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": { "localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} (%{percent})" "defaultMessage": "{downloaded} / {total} (%{percent})"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Öne Çıkan Modeller" "defaultMessage": "Öne Çıkan Modeller"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Model ararken ve indirirken hız sınırlarını artırmak ve özel ya da erişimi kısıtlı Hugging Face depolarına erişmek için oturum açın." "defaultMessage": "Model ararken ve indirirken hız sınırlarını artırmak ve özel ya da erişimi kısıtlı Hugging Face depolarına erişmek için oturum açın."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Modeli Ayarları" "defaultMessage": "Modeli Ayarları"
}, },

6
ui/desktop/src/i18n/messages/vi.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "Đang tải xuống" "defaultMessage": "Đang tải xuống"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "Mô hình nổi bật" "defaultMessage": "Mô hình nổi bật"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Đăng nhập để tăng giới hạn tốc độ khi tìm kiếm và tải xuống mô hình, và để truy cập các kho lưu trữ Hugging Face riêng tư hoặc bị giới hạn." "defaultMessage": "Đăng nhập để tăng giới hạn tốc độ khi tìm kiếm và tải xuống mô hình, và để truy cập các kho lưu trữ Hugging Face riêng tư hoặc bị giới hạn."
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "Cài đặt mô hình" "defaultMessage": "Cài đặt mô hình"
}, },

6
ui/desktop/src/i18n/messages/zh-CN.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": { "localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total}（{percent}%）" "defaultMessage": "{downloaded} / {total}（{percent}%）"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "推荐模型" "defaultMessage": "推荐模型"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "登录后可在搜索和下载模型时提高速率限制，并访问私有或受限的 Hugging Face 仓库。" "defaultMessage": "登录后可在搜索和下载模型时提高速率限制，并访问私有或受限的 Hugging Face 仓库。"
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "模型设置" "defaultMessage": "模型设置"
}, },

6
ui/desktop/src/i18n/messages/zh-TW.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": { "localInferenceSettings.downloading": {
"defaultMessage": "下載中" "defaultMessage": "下載中"
}, },
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": { "localInferenceSettings.featuredModels": {
"defaultMessage": "精選模型" "defaultMessage": "精選模型"
}, },
"localInferenceSettings.huggingFaceSignInNote": { "localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "登入以在搜尋與下載模型時提高速率上限，並存取私人或受限的 Hugging Face 儲存庫。" "defaultMessage": "登入以在搜尋與下載模型時提高速率上限，並存取私人或受限的 Hugging Face 儲存庫。"
}, },
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": { "localInferenceSettings.modelSettings": {
"defaultMessage": "模型設定" "defaultMessage": "模型設定"
}, },

20
ui/desktop/src/types/message.ts
View file

@ -93,7 +93,11 @@ type ContentIcon = {
theme?: 'light' | 'dark' | JsonObject; theme?: 'light' | 'dark' | JsonObject;
}; };
export type SystemNotificationType = 'thinkingMessage' | 'inlineMessage' | 'creditsExhausted'; export type SystemNotificationType =
| 'thinkingMessage'
| 'progressMessage'
| 'inlineMessage'
| 'creditsExhausted';
export type SystemNotificationContent = { export type SystemNotificationContent = {
data?: unknown; data?: unknown;
@ -443,17 +447,3 @@ export function hasCompletedToolCalls(message: Message): boolean {
const toolRequests = getToolRequests(message); const toolRequests = getToolRequests(message);
return toolRequests.length > 0; return toolRequests.length > 0;
} }
export function getThinkingMessage(message: Message | undefined): string | undefined {
if (!message || message.role !== 'assistant') {
return undefined;
}
for (const content of message.content) {
if (content.type === 'systemNotification' && content.notificationType === 'thinkingMessage') {
return content.msg;
}
}
return undefined;
}

26
ui/desktop/src/utils/conversionUtils.test.ts Normal file
View file

@ -0,0 +1,26 @@
import { describe, expect, it } from 'vitest';
import { errorMessage } from './conversionUtils';
describe('errorMessage', () => {
it('prefers ACP JSON-RPC error data over generic messages', () => {
expect(
errorMessage({
error: {
message: 'Invalid params',
data: 'MLX backend error: failed to load model',
},
})
).toBe('MLX backend error: failed to load model');
});
it('prefers ACP JSON-RPC error data from Error instances', () => {
const error = Object.assign(new Error('Invalid params'), {
error: {
message: 'Invalid params',
data: 'MLX backend error: failed to load model',
},
});
expect(errorMessage(error)).toBe('MLX backend error: failed to load model');
});
});

18
ui/desktop/src/utils/conversionUtils.ts
View file

@ -13,6 +13,11 @@ export async function safeJsonParse<T>(
} }
export function errorMessage(err: Error | unknown, default_value?: string) { export function errorMessage(err: Error | unknown, default_value?: string) {
const acpData = acpErrorData(err);
if (typeof acpData === 'string') {
return acpData;
}
if (err instanceof Error) { if (err instanceof Error) {
return err.message; return err.message;
} else if (typeof err === 'object' && err !== null && 'message' in err) { } else if (typeof err === 'object' && err !== null && 'message' in err) {
@ -22,6 +27,19 @@ export function errorMessage(err: Error | unknown, default_value?: string) {
} }
} }
function acpErrorData(err: unknown): unknown {
if (typeof err !== 'object' || err === null) {
return undefined;
}
const candidate = 'error' in err && isRecord(err.error) ? err.error : err;
return isRecord(candidate) ? candidate.data : undefined;
}
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === 'object' && value !== null;
}
export function formatErrorForLogging(error: unknown): string { export function formatErrorForLogging(error: unknown): string {
if (error instanceof Error) { if (error instanceof Error) {
return `${error.name}: ${error.message}${error.stack ? `\n${error.stack}` : ''}`; return `${error.name}: ${error.message}${error.stack ? `\n${error.stack}` : ''}`;

10
ui/sdk/src/generated/client.gen.ts
View file

@ -121,6 +121,7 @@ import type {
LocalInferenceModelDownloadProgressResponse_unstable, LocalInferenceModelDownloadProgressResponse_unstable,
LocalInferenceModelDownloadRequest_unstable, LocalInferenceModelDownloadRequest_unstable,
LocalInferenceModelDownloadResponse_unstable, LocalInferenceModelDownloadResponse_unstable,
LocalInferenceModelEvictRequest_unstable,
LocalInferenceModelSettingsReadRequest_unstable, LocalInferenceModelSettingsReadRequest_unstable,
LocalInferenceModelSettingsReadResponse_unstable, LocalInferenceModelSettingsReadResponse_unstable,
LocalInferenceModelSettingsUpdateRequest_unstable, LocalInferenceModelSettingsUpdateRequest_unstable,
@ -1377,6 +1378,15 @@ export class GooseExtClient {
); );
} }
async localInferenceModelsEvict_unstable(
params: LocalInferenceModelEvictRequest_unstable,
): Promise<void> {
await this.conn.extMethod(
"_goose/unstable/local-inference/models/evict",
params,
);
}
async localInferenceModelsSettingsRead_unstable( async localInferenceModelsSettingsRead_unstable(
params: LocalInferenceModelSettingsReadRequest_unstable, params: LocalInferenceModelSettingsReadRequest_unstable,
): Promise<LocalInferenceModelSettingsReadResponse_unstable> { ): Promise<LocalInferenceModelSettingsReadResponse_unstable> {

7
ui/sdk/src/generated/index.ts
View file

File diff suppressed because one or more lines are too long

7
ui/sdk/src/generated/types.gen.ts
View file

@ -2329,6 +2329,7 @@ export type LocalInferenceModelDto = {
sizeBytes: number; sizeBytes: number;
status: LocalInferenceModelDownloadStatusDto; status: LocalInferenceModelDownloadStatusDto;
recommended: boolean; recommended: boolean;
isLoaded: boolean;
settings: LocalInferenceModelSettingsDto; settings: LocalInferenceModelSettingsDto;
visionCapable: boolean; visionCapable: boolean;
mmprojStatus?: LocalInferenceModelDownloadStatusDto | null; mmprojStatus?: LocalInferenceModelDownloadStatusDto | null;
@ -2433,6 +2434,10 @@ export type LocalInferenceModelDeleteRequest_unstable = {
modelId: string; modelId: string;
}; };
export type LocalInferenceModelEvictRequest_unstable = {
modelId: string;
};
export type LocalInferenceModelSettingsReadRequest_unstable = { export type LocalInferenceModelSettingsReadRequest_unstable = {
modelId: string; modelId: string;
}; };
@ -2618,7 +2623,7 @@ export type RecipeParamsAction = 'submit' | 'cancel';
export type ExtRequest = { export type ExtRequest = {
id: string; id: string;
method: string; method: string;
params?: AddSessionExtensionRequest_unstable | RemoveSessionExtensionRequest_unstable | GetToolsRequest_unstable | SetToolPermissionsRequest_unstable | GooseToolCallRequest_unstable | ReadResourceRequest_unstable | AppsListRequest_unstable | AppsExportRequest_unstable | AppsImportRequest_unstable | UpdateWorkingDirRequest_unstable | SetSessionSystemPromptRequest_unstable | SteerSessionRequest_unstable | DiagnosticsGetRequest_unstable | ListPromptsRequest_unstable | GetPromptRequest_unstable | SavePromptRequest_unstable | ResetPromptRequest_unstable | DeleteSessionRequest | GetConfigExtensionsRequest_unstable | GetAvailableExtensionsRequest_unstable | AddConfigExtensionRequest_unstable | RemoveConfigExtensionRequest_unstable | SetConfigExtensionEnabledRequest_unstable | GetSessionExtensionsRequest_unstable | ListProvidersRequest_unstable | ProviderSupportedModelsListRequest_unstable | ProviderCatalogListRequest_unstable | ProviderSetupCatalogListRequest_unstable | ProviderCatalogTemplateRequest_unstable | CustomProviderCreateRequest_unstable | CustomProviderReadRequest_unstable | CustomProviderUpdateRequest_unstable | CustomProviderDeleteRequest_unstable | RefreshProviderInventoryRequest_unstable | ProviderConfigReadRequest_unstable | ProviderConfigStatusRequest_unstable | ProviderConfigSaveRequest_unstable | ProviderConfigDeleteRequest_unstable | ProviderConfigAuthenticateRequest_unstable | ProviderSecretsListRequest_unstable | ProviderSecretDeleteRequest_unstable | CanonicalModelInfoRequest_unstable | PreferencesReadRequest_unstable | PreferencesSaveRequest_unstable | PreferencesRemoveRequest_unstable | ConfigReadRequest_unstable | ConfigUpsertRequest_unstable | ConfigRemoveRequest_unstable | ConfigReadAllRequest_unstable | DefaultsReadRequest_unstable | DefaultsSaveRequest_unstable | DefaultsClearRequest_unstable | OnboardingImportScanRequest_unstable | OnboardingImportApplyRequest_unstable | ExportSessionRequest_unstable | ImportSessionRequest_unstable | ShareSessionNostrRequest_unstable | EncodeRecipeRequest_unstable | DecodeRecipeRequest_unstable | ScanRecipeRequest_unstable | ListRecipesRequest_unstable | DeleteRecipeRequest_unstable | ScheduleRecipeRequest_unstable | SetRecipeSlashCommandRequest_unstable | SaveRecipeRequest_unstable | ParseRecipeRequest_unstable | RecipeToYamlRequest_unstable | ListSchedulesRequest_unstable | ListScheduleSessionsRequest_unstable | CreateScheduleRequest_unstable | DeleteScheduleRequest_unstable | PauseScheduleRequest_unstable | UnpauseScheduleRequest_unstable | UpdateScheduleRequest_unstable | RunScheduleNowRequest_unstable | KillRunningJobRequest_unstable | InspectRunningJobRequest_unstable | GetSessionInfoRequest_unstable | TruncateSessionConversationRequest_unstable | UpdateSessionProjectRequest_unstable | RenameSessionRequest_unstable | ArchiveSessionRequest_unstable | UnarchiveSessionRequest_unstable | CreateSourceRequest_unstable | ListSourcesRequest_unstable | ListAgentMentionsRequest_unstable | ListSlashCommandsRequest_unstable | UpdateSourceRequest_unstable | DeleteSourceRequest_unstable | ExportSourceRequest_unstable | ImportSourcesRequest_unstable | DictationTranscribeRequest_unstable | DictationConfigRequest_unstable | DictationSecretSaveRequest_unstable | DictationSecretDeleteRequest_unstable | DictationModelsListRequest_unstable | DictationModelDownloadRequest_unstable | DictationModelDownloadProgressRequest_unstable | DictationModelCancelRequest_unstable | DictationModelDeleteRequest_unstable | DictationModelSelectRequest_unstable | LocalInferenceModelsListRequest_unstable | LocalInferenceModelDownloadRequest_unstable | LocalInferenceModelDownloadProgressRequest_unstable | LocalInferenceModelDownloadCancelRequest_unstable | LocalInferenceModelDeleteRequest_unstable | LocalInferenceModelSettingsReadRequest_unstable | LocalInferenceModelSettingsUpdateRequest_unstable | LocalInferenceHuggingFaceSearchRequest_unstable | LocalInferenceHuggingFaceRepoVariantsRequest_unstable | LocalInferenceBuiltinChatTemplatesListRequest_unstable | { params?: AddSessionExtensionRequest_unstable | RemoveSessionExtensionRequest_unstable | GetToolsRequest_unstable | SetToolPermissionsRequest_unstable | GooseToolCallRequest_unstable | ReadResourceRequest_unstable | AppsListRequest_unstable | AppsExportRequest_unstable | AppsImportRequest_unstable | UpdateWorkingDirRequest_unstable | SetSessionSystemPromptRequest_unstable | SteerSessionRequest_unstable | DiagnosticsGetRequest_unstable | ListPromptsRequest_unstable | GetPromptRequest_unstable | SavePromptRequest_unstable | ResetPromptRequest_unstable | DeleteSessionRequest | GetConfigExtensionsRequest_unstable | GetAvailableExtensionsRequest_unstable | AddConfigExtensionRequest_unstable | RemoveConfigExtensionRequest_unstable | SetConfigExtensionEnabledRequest_unstable | GetSessionExtensionsRequest_unstable | ListProvidersRequest_unstable | ProviderSupportedModelsListRequest_unstable | ProviderCatalogListRequest_unstable | ProviderSetupCatalogListRequest_unstable | ProviderCatalogTemplateRequest_unstable | CustomProviderCreateRequest_unstable | CustomProviderReadRequest_unstable | CustomProviderUpdateRequest_unstable | CustomProviderDeleteRequest_unstable | RefreshProviderInventoryRequest_unstable | ProviderConfigReadRequest_unstable | ProviderConfigStatusRequest_unstable | ProviderConfigSaveRequest_unstable | ProviderConfigDeleteRequest_unstable | ProviderConfigAuthenticateRequest_unstable | ProviderSecretsListRequest_unstable | ProviderSecretDeleteRequest_unstable | CanonicalModelInfoRequest_unstable | PreferencesReadRequest_unstable | PreferencesSaveRequest_unstable | PreferencesRemoveRequest_unstable | ConfigReadRequest_unstable | ConfigUpsertRequest_unstable | ConfigRemoveRequest_unstable | ConfigReadAllRequest_unstable | DefaultsReadRequest_unstable | DefaultsSaveRequest_unstable | DefaultsClearRequest_unstable | OnboardingImportScanRequest_unstable | OnboardingImportApplyRequest_unstable | ExportSessionRequest_unstable | ImportSessionRequest_unstable | ShareSessionNostrRequest_unstable | EncodeRecipeRequest_unstable | DecodeRecipeRequest_unstable | ScanRecipeRequest_unstable | ListRecipesRequest_unstable | DeleteRecipeRequest_unstable | ScheduleRecipeRequest_unstable | SetRecipeSlashCommandRequest_unstable | SaveRecipeRequest_unstable | ParseRecipeRequest_unstable | RecipeToYamlRequest_unstable | ListSchedulesRequest_unstable | ListScheduleSessionsRequest_unstable | CreateScheduleRequest_unstable | DeleteScheduleRequest_unstable | PauseScheduleRequest_unstable | UnpauseScheduleRequest_unstable | UpdateScheduleRequest_unstable | RunScheduleNowRequest_unstable | KillRunningJobRequest_unstable | InspectRunningJobRequest_unstable | GetSessionInfoRequest_unstable | TruncateSessionConversationRequest_unstable | UpdateSessionProjectRequest_unstable | RenameSessionRequest_unstable | ArchiveSessionRequest_unstable | UnarchiveSessionRequest_unstable | CreateSourceRequest_unstable | ListSourcesRequest_unstable | ListAgentMentionsRequest_unstable | ListSlashCommandsRequest_unstable | UpdateSourceRequest_unstable | DeleteSourceRequest_unstable | ExportSourceRequest_unstable | ImportSourcesRequest_unstable | DictationTranscribeRequest_unstable | DictationConfigRequest_unstable | DictationSecretSaveRequest_unstable | DictationSecretDeleteRequest_unstable | DictationModelsListRequest_unstable | DictationModelDownloadRequest_unstable | DictationModelDownloadProgressRequest_unstable | DictationModelCancelRequest_unstable | DictationModelDeleteRequest_unstable | DictationModelSelectRequest_unstable | LocalInferenceModelsListRequest_unstable | LocalInferenceModelDownloadRequest_unstable | LocalInferenceModelDownloadProgressRequest_unstable | LocalInferenceModelDownloadCancelRequest_unstable | LocalInferenceModelDeleteRequest_unstable | LocalInferenceModelEvictRequest_unstable | LocalInferenceModelSettingsReadRequest_unstable | LocalInferenceModelSettingsUpdateRequest_unstable | LocalInferenceHuggingFaceSearchRequest_unstable | LocalInferenceHuggingFaceRepoVariantsRequest_unstable | LocalInferenceBuiltinChatTemplatesListRequest_unstable | {
[key: string]: unknown; [key: string]: unknown;
} | null; } | null;
}; };

6
ui/sdk/src/generated/zod.gen.ts
View file

@ -2505,6 +2505,7 @@ export const zLocalInferenceModelDto = z.object({
sizeBytes: z.number().int().gte(0), sizeBytes: z.number().int().gte(0),
status: zLocalInferenceModelDownloadStatusDto, status: zLocalInferenceModelDownloadStatusDto,
recommended: z.boolean(), recommended: z.boolean(),
isLoaded: z.boolean(),
settings: zLocalInferenceModelSettingsDto, settings: zLocalInferenceModelSettingsDto,
visionCapable: z.boolean(), visionCapable: z.boolean(),
mmprojStatus: z.union([ mmprojStatus: z.union([
@ -2573,6 +2574,10 @@ export const zLocalInferenceModelDeleteRequest_unstable = z.object({
modelId: z.string() modelId: z.string()
}); });
export const zLocalInferenceModelEvictRequest_unstable = z.object({
modelId: z.string()
});
export const zLocalInferenceModelSettingsReadRequest_unstable = z.object({ export const zLocalInferenceModelSettingsReadRequest_unstable = z.object({
modelId: z.string() modelId: z.string()
}); });
@ -2914,6 +2919,7 @@ export const zExtRequest = z.object({
zLocalInferenceModelDownloadProgressRequest_unstable, zLocalInferenceModelDownloadProgressRequest_unstable,
zLocalInferenceModelDownloadCancelRequest_unstable, zLocalInferenceModelDownloadCancelRequest_unstable,
zLocalInferenceModelDeleteRequest_unstable, zLocalInferenceModelDeleteRequest_unstable,
zLocalInferenceModelEvictRequest_unstable,
zLocalInferenceModelSettingsReadRequest_unstable, zLocalInferenceModelSettingsReadRequest_unstable,
zLocalInferenceModelSettingsUpdateRequest_unstable, zLocalInferenceModelSettingsUpdateRequest_unstable,
zLocalInferenceHuggingFaceSearchRequest_unstable, zLocalInferenceHuggingFaceSearchRequest_unstable,