mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
Merge
This commit is contained in:
Douwe Osinga
Generated
+175
-14
@@ -2093,7 +2093,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "961b955a666e25ee5a1091d219128d6e6401e3dab84efb1a2bf6b4035d797b39"
|
||||
dependencies = [
|
||||
"crmf",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"spki",
|
||||
"x509-cert",
|
||||
]
|
||||
@@ -2105,7 +2105,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7b77c319abfd5219629c45c34c89ba945ed3c5e49fcde9d16b6c3885f118a730"
|
||||
dependencies = [
|
||||
"const-oid 0.9.6",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"spki",
|
||||
"x509-cert",
|
||||
]
|
||||
@@ -2395,7 +2395,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "36fe21b96d5b87f5de4b5b7202ec41c00110ac817ce6728fe75fb2fe5962ed92"
|
||||
dependencies = [
|
||||
"cms",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"spki",
|
||||
"x509-cert",
|
||||
]
|
||||
@@ -3103,7 +3103,17 @@ dependencies = [
|
||||
"const-oid 0.9.6",
|
||||
"der_derive",
|
||||
"flagset",
|
||||
"pem-rfc7468",
|
||||
"pem-rfc7468 0.7.0",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "der"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "71fd89660b2dc699704064e59e9dba0147b903e85319429e131620d022be411b"
|
||||
dependencies = [
|
||||
"pem-rfc7468 1.0.0",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
@@ -3454,7 +3464,7 @@ version = "0.16.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca"
|
||||
dependencies = [
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"digest 0.10.7",
|
||||
"elliptic-curve",
|
||||
"rfc6979",
|
||||
@@ -3517,7 +3527,7 @@ dependencies = [
|
||||
"generic-array",
|
||||
"group",
|
||||
"hkdf",
|
||||
"pem-rfc7468",
|
||||
"pem-rfc7468 0.7.0",
|
||||
"pkcs8",
|
||||
"rand_core 0.6.4",
|
||||
"sec1",
|
||||
@@ -4502,6 +4512,7 @@ dependencies = [
|
||||
"minijinja",
|
||||
"mockall",
|
||||
"nanoid",
|
||||
"ndarray",
|
||||
"nostr",
|
||||
"nostr-sdk",
|
||||
"oauth2",
|
||||
@@ -4510,7 +4521,12 @@ dependencies = [
|
||||
"opentelemetry-appender-tracing",
|
||||
"opentelemetry-otlp 0.32.0",
|
||||
"opentelemetry-stdout",
|
||||
<<<<<<< Updated upstream
|
||||
"opentelemetry_sdk 0.32.0",
|
||||
=======
|
||||
"opentelemetry_sdk",
|
||||
"ort",
|
||||
>>>>>>> Stashed changes
|
||||
"pastey",
|
||||
"pctx_code_mode",
|
||||
"pem",
|
||||
@@ -4525,6 +4541,7 @@ dependencies = [
|
||||
"rmcp",
|
||||
"rubato",
|
||||
"rustls",
|
||||
"safetensors 0.7.0",
|
||||
"schemars 1.2.1",
|
||||
"sec1",
|
||||
"serde",
|
||||
@@ -4939,6 +4956,12 @@ dependencies = [
|
||||
"digest 0.11.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hmac-sha256"
|
||||
version = "1.1.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f"
|
||||
|
||||
[[package]]
|
||||
name = "home"
|
||||
version = "0.5.12"
|
||||
@@ -5591,7 +5614,11 @@ dependencies = [
|
||||
"portable-atomic",
|
||||
"portable-atomic-util",
|
||||
"serde_core",
|
||||
<<<<<<< Updated upstream
|
||||
"windows-link",
|
||||
=======
|
||||
"windows-sys 0.61.2",
|
||||
>>>>>>> Stashed changes
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6025,6 +6052,12 @@ version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
|
||||
|
||||
[[package]]
|
||||
name = "lzma-rust2"
|
||||
version = "0.15.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e20f57f9918e5bd7bc58c22cdd70a6afc7375d4dd9683af5f2b34bd3d2bba619"
|
||||
|
||||
[[package]]
|
||||
name = "macro_rules_attribute"
|
||||
version = "0.2.2"
|
||||
@@ -6065,6 +6098,16 @@ version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "matrixmultiply"
|
||||
version = "0.3.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
"rawpointer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "md-5"
|
||||
version = "0.10.6"
|
||||
@@ -6264,6 +6307,21 @@ dependencies = [
|
||||
"tempfile",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ndarray"
|
||||
version = "0.17.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
|
||||
dependencies = [
|
||||
"matrixmultiply",
|
||||
"num-complex",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
"portable-atomic",
|
||||
"portable-atomic-util",
|
||||
"rawpointer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ndk-context"
|
||||
version = "0.1.1"
|
||||
@@ -7006,6 +7064,44 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
|
||||
|
||||
[[package]]
|
||||
<<<<<<< Updated upstream
|
||||
=======
|
||||
name = "ordered-multimap"
|
||||
version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "49203cdcae0030493bad186b28da2fa25645fa276a51b6fec8010d281e02ef79"
|
||||
dependencies = [
|
||||
"dlv-list",
|
||||
"hashbrown 0.14.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ort"
|
||||
version = "2.0.0-rc.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133"
|
||||
dependencies = [
|
||||
"libloading 0.9.0",
|
||||
"ndarray",
|
||||
"ort-sys",
|
||||
"smallvec",
|
||||
"tracing",
|
||||
"ureq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ort-sys"
|
||||
version = "2.0.0-rc.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90"
|
||||
dependencies = [
|
||||
"hmac-sha256",
|
||||
"lzma-rust2",
|
||||
"ureq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
>>>>>>> Stashed changes
|
||||
name = "outref"
|
||||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
@@ -7299,6 +7395,15 @@ dependencies = [
|
||||
"base64ct",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pem-rfc7468"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "percent-encoding"
|
||||
version = "2.3.2"
|
||||
@@ -7446,7 +7551,7 @@ version = "0.7.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f"
|
||||
dependencies = [
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"pkcs8",
|
||||
"spki",
|
||||
]
|
||||
@@ -7457,7 +7562,7 @@ version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
|
||||
dependencies = [
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"spki",
|
||||
]
|
||||
|
||||
@@ -7998,6 +8103,12 @@ dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rawpointer"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
||||
|
||||
[[package]]
|
||||
name = "rayon"
|
||||
version = "1.12.0"
|
||||
@@ -8715,7 +8826,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc"
|
||||
dependencies = [
|
||||
"base16ct",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"generic-array",
|
||||
"pkcs8",
|
||||
"subtle",
|
||||
@@ -9099,7 +9210,7 @@ dependencies = [
|
||||
"aws-lc-rs",
|
||||
"base64 0.22.1",
|
||||
"const-oid 0.9.6",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"digest 0.10.7",
|
||||
"pem",
|
||||
"rand_core 0.9.5",
|
||||
@@ -9172,7 +9283,7 @@ dependencies = [
|
||||
"cmpv2",
|
||||
"cms",
|
||||
"const-oid 0.9.6",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"hex",
|
||||
"jiff",
|
||||
"rand 0.9.4",
|
||||
@@ -9340,6 +9451,20 @@ checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.61.2",
|
||||
<<<<<<< Updated upstream
|
||||
=======
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "socks"
|
||||
version = "0.3.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"libc",
|
||||
"winapi",
|
||||
>>>>>>> Stashed changes
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -9376,7 +9501,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -11626,6 +11751,36 @@ version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
|
||||
|
||||
[[package]]
|
||||
name = "ureq"
|
||||
version = "3.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"der 0.8.0",
|
||||
"log",
|
||||
"native-tls",
|
||||
"percent-encoding",
|
||||
"rustls-pki-types",
|
||||
"socks",
|
||||
"ureq-proto",
|
||||
"utf8-zero",
|
||||
"webpki-root-certs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ureq-proto"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"http 1.4.0",
|
||||
"httparse",
|
||||
"log",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "url"
|
||||
version = "2.5.8"
|
||||
@@ -11651,6 +11806,12 @@ version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
|
||||
|
||||
[[package]]
|
||||
name = "utf8-zero"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e"
|
||||
|
||||
[[package]]
|
||||
name = "utf8_iter"
|
||||
version = "1.0.4"
|
||||
@@ -12692,7 +12853,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94"
|
||||
dependencies = [
|
||||
"const-oid 0.9.6",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
"sha1",
|
||||
"signature",
|
||||
"spki",
|
||||
@@ -12725,7 +12886,7 @@ checksum = "f5ceece934a21607055b7ac5c25adb56a2ff559804b10705dc674d1d838c15e1"
|
||||
dependencies = [
|
||||
"cmpv2",
|
||||
"cms",
|
||||
"der",
|
||||
"der 0.7.10",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -206,6 +206,9 @@ process-wrap = { version = "9", default-features = false, features = ["std"] }
|
||||
nostr = { version = "0.44", default-features = false, features = ["nip44", "std"], optional = true }
|
||||
nostr-sdk = { version = "0.44", default-features = false, features = ["nip44"], optional = true }
|
||||
rustls = { workspace = true, optional = true }
|
||||
ort = { version = "2.0.0-rc.12", features = ["load-dynamic"] }
|
||||
safetensors = "0.7.0"
|
||||
ndarray = "0.17.2"
|
||||
|
||||
pctx_code_mode = { version = "0.3", default-features = false, optional = true }
|
||||
|
||||
|
||||
@@ -214,6 +214,7 @@ pub struct Agent {
|
||||
container: Mutex<Option<Container>>,
|
||||
goal: Mutex<Option<String>>,
|
||||
grind: Mutex<Option<String>>,
|
||||
complexity_router: Option<Arc<super::complexity_router::ComplexityModel>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -333,6 +334,8 @@ impl Agent {
|
||||
container: Mutex::new(None),
|
||||
goal: Mutex::new(None),
|
||||
grind: Mutex::new(None),
|
||||
complexity_router: super::complexity_router::ComplexityModel::try_load_default()
|
||||
.map(Arc::new),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1586,17 +1589,48 @@ impl Agent {
|
||||
|
||||
let provider = self.provider().await?;
|
||||
let provider_name = provider.get_name().to_string();
|
||||
let requested_model = provider.get_model_config().model_name;
|
||||
let inference = provider
|
||||
.fetch_model_info(&requested_model)
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|model_info| model_info.resolved_model)
|
||||
.map(|resolved_model| InferenceMetadata {
|
||||
provider: provider_name,
|
||||
requested_model,
|
||||
resolved_model: Some(resolved_model),
|
||||
let base_model_config = provider.get_model_config();
|
||||
let requested_model = base_model_config.model_name.clone();
|
||||
|
||||
let route_decision = self
|
||||
.complexity_router
|
||||
.as_ref()
|
||||
.and_then(|router| {
|
||||
super::complexity_router::route(router.as_ref(), &conversation)
|
||||
});
|
||||
|
||||
// Decide once per `reply_internal` call which model config to use.
|
||||
// If the complexity router says "use_fast" and the provider has a
|
||||
// `fast_model_config`, swap. Otherwise main config.
|
||||
let effective_model_config = match route_decision {
|
||||
Some(d) if d.use_fast && base_model_config.fast_model_config.is_some() => {
|
||||
tracing::info!(
|
||||
target: "goose::complexity_router",
|
||||
complexity = d.complexity,
|
||||
elapsed_ms = d.elapsed_ms,
|
||||
"routing to fast model",
|
||||
);
|
||||
base_model_config.use_fast_model()
|
||||
}
|
||||
Some(d) => {
|
||||
tracing::info!(
|
||||
target: "goose::complexity_router",
|
||||
complexity = d.complexity,
|
||||
elapsed_ms = d.elapsed_ms,
|
||||
use_fast = d.use_fast,
|
||||
has_fast_config = base_model_config.fast_model_config.is_some(),
|
||||
"routing to main model",
|
||||
);
|
||||
base_model_config.clone()
|
||||
}
|
||||
None => base_model_config.clone(),
|
||||
};
|
||||
|
||||
let inference = Some(InferenceMetadata {
|
||||
provider: provider_name,
|
||||
requested_model,
|
||||
resolved_model: Some(effective_model_config.model_name.clone()),
|
||||
});
|
||||
let session_manager = self.config.session_manager.clone();
|
||||
let session_id = session_config.id.clone();
|
||||
if !self.config.disable_session_naming {
|
||||
@@ -1676,6 +1710,7 @@ impl Agent {
|
||||
|
||||
let mut stream = Self::stream_response_from_provider(
|
||||
self.provider().await?,
|
||||
effective_model_config.clone(),
|
||||
&session_config.id,
|
||||
&system_prompt,
|
||||
conversation_with_moim.messages(),
|
||||
|
||||
@@ -0,0 +1,458 @@
|
||||
//! Pre-flight router that scores a conversation's cognitive complexity in [0, 1]
|
||||
//! so the agent can decide whether to send a turn to the fast model.
|
||||
//!
|
||||
//! `ComplexityModel` loads a self-contained bundle from `~/.goose/complexity_model/`
|
||||
//! containing a fastembed-style ONNX embedder, an HF tokenizer, an MLP head
|
||||
//! exported as safetensors, and a `config.json` describing the architecture.
|
||||
//!
|
||||
//! On top of the model this module exposes the small bits of policy the agent
|
||||
//! needs: rendering a `Conversation` into the same text format the model was
|
||||
//! trained on, and the threshold logic for "is this a fast-model turn?"
|
||||
//!
|
||||
//! See `GOOSE_INTEGRATION_PLAN.md` in the llm-router repo for the design;
|
||||
//! the bundle is produced by `train_complexity.py`.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use anyhow::{anyhow, bail, Context, Result};
|
||||
use ndarray::{Array1, Array2};
|
||||
use ort::session::Session;
|
||||
use ort::value::TensorRef;
|
||||
use safetensors::SafeTensors;
|
||||
use serde::Deserialize;
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
use crate::conversation::message::{Message, Role};
|
||||
use crate::conversation::Conversation;
|
||||
|
||||
const DEFAULT_BUNDLE_SUBDIR: &str = "complexity_model";
|
||||
const MAX_SEQ_LEN: usize = 512;
|
||||
const ANCHOR_MARKER: &str = ">>>";
|
||||
|
||||
/// Complexity threshold under which the agent routes the turn to the fast model.
|
||||
/// Same dial we used at training time; tuned empirically by eyeballing the
|
||||
/// WildChat distribution. Demo-friendly default.
|
||||
pub const FAST_MODEL_THRESHOLD: f32 = 0.5;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct BundleConfig {
|
||||
format_version: u32,
|
||||
embedder: EmbedderConfig,
|
||||
head: HeadConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct EmbedderConfig {
|
||||
repo_id: String,
|
||||
output_dim: usize,
|
||||
onnx_file: String,
|
||||
tokenizer_file: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct HeadConfig {
|
||||
input_dim: usize,
|
||||
hidden_dims: Vec<usize>,
|
||||
output_dim: usize,
|
||||
#[serde(default)]
|
||||
activation: String,
|
||||
#[serde(default)]
|
||||
output_activation: String,
|
||||
}
|
||||
|
||||
/// One linear layer worth of weights, as owned `ndarray` matrices.
|
||||
struct LinearLayer {
|
||||
weight: Array2<f32>, // (out_dim, in_dim) — matches torch.nn.Linear convention
|
||||
bias: Array1<f32>, // (out_dim,)
|
||||
}
|
||||
|
||||
impl LinearLayer {
|
||||
fn apply(&self, x: &Array1<f32>) -> Array1<f32> {
|
||||
self.weight.dot(x) + &self.bias
|
||||
}
|
||||
}
|
||||
|
||||
/// Output of a single complexity scoring call.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct ComplexityScore {
|
||||
pub complexity: f32,
|
||||
pub tool_calls_norm: f32,
|
||||
pub elapsed_ms: u64,
|
||||
}
|
||||
|
||||
/// Loaded complexity model. Cheap to clone (everything inside an `Arc`).
|
||||
pub struct ComplexityModel {
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
struct Inner {
|
||||
embedder_dim: usize,
|
||||
tokenizer: Tokenizer,
|
||||
session: Session,
|
||||
head: Vec<LinearLayer>,
|
||||
head_out_dim: usize,
|
||||
repo_id: String,
|
||||
}
|
||||
|
||||
impl ComplexityModel {
|
||||
/// Default location: `~/.goose/complexity_model/`. Returns `None` (not an
|
||||
/// error) if the bundle is missing — callers should treat that as
|
||||
/// "complexity routing disabled."
|
||||
pub fn try_load_default() -> Option<Self> {
|
||||
let dir = default_bundle_dir()?;
|
||||
if !dir.join("config.json").exists() {
|
||||
return None;
|
||||
}
|
||||
match Self::load_from_dir(&dir) {
|
||||
Ok(m) => Some(m),
|
||||
Err(e) => {
|
||||
tracing::warn!("failed to load complexity model from {:?}: {:#}", dir, e);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_from_dir(dir: &Path) -> Result<Self> {
|
||||
let cfg_path = dir.join("config.json");
|
||||
let cfg_text = std::fs::read_to_string(&cfg_path)
|
||||
.with_context(|| format!("reading {}", cfg_path.display()))?;
|
||||
let cfg: BundleConfig = serde_json::from_str(&cfg_text)
|
||||
.with_context(|| format!("parsing {}", cfg_path.display()))?;
|
||||
|
||||
if cfg.format_version != 1 {
|
||||
bail!("unsupported bundle format_version {}", cfg.format_version);
|
||||
}
|
||||
if cfg.head.input_dim != cfg.embedder.output_dim {
|
||||
bail!(
|
||||
"head input_dim {} != embedder output_dim {}",
|
||||
cfg.head.input_dim,
|
||||
cfg.embedder.output_dim
|
||||
);
|
||||
}
|
||||
|
||||
let tokenizer_path = dir.join(&cfg.embedder.tokenizer_file);
|
||||
let tokenizer = Tokenizer::from_file(&tokenizer_path)
|
||||
.map_err(|e| anyhow!("loading tokenizer at {}: {}", tokenizer_path.display(), e))?;
|
||||
|
||||
let onnx_path = dir.join(&cfg.embedder.onnx_file);
|
||||
let session = Session::builder()?
|
||||
.commit_from_file(&onnx_path)
|
||||
.with_context(|| format!("loading ONNX from {}", onnx_path.display()))?;
|
||||
|
||||
let weights_path = dir.join("weights.safetensors");
|
||||
let weights_bytes = std::fs::read(&weights_path)
|
||||
.with_context(|| format!("reading {}", weights_path.display()))?;
|
||||
let head = load_head_weights(&weights_bytes, &cfg.head)?;
|
||||
|
||||
Ok(Self {
|
||||
inner: Arc::new(Inner {
|
||||
embedder_dim: cfg.embedder.output_dim,
|
||||
tokenizer,
|
||||
session,
|
||||
head,
|
||||
head_out_dim: cfg.head.output_dim,
|
||||
repo_id: cfg.embedder.repo_id,
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
/// Score one rendered conversation. ~25-50ms on CPU for typical inputs.
|
||||
pub fn score(&self, text: &str) -> Result<ComplexityScore> {
|
||||
let started = Instant::now();
|
||||
let embedding = self.embed(text)?;
|
||||
let mut activations = embedding;
|
||||
let last = self.inner.head.len() - 1;
|
||||
for (i, layer) in self.inner.head.iter().enumerate() {
|
||||
activations = layer.apply(&activations);
|
||||
if i != last {
|
||||
activations.mapv_inplace(|x| x.max(0.0)); // ReLU
|
||||
}
|
||||
}
|
||||
activations.mapv_inplace(sigmoid);
|
||||
|
||||
if activations.len() < self.inner.head_out_dim {
|
||||
bail!(
|
||||
"head output has {} values, expected {}",
|
||||
activations.len(),
|
||||
self.inner.head_out_dim
|
||||
);
|
||||
}
|
||||
|
||||
Ok(ComplexityScore {
|
||||
complexity: activations[0],
|
||||
tool_calls_norm: if self.inner.head_out_dim >= 2 {
|
||||
activations[1]
|
||||
} else {
|
||||
0.0
|
||||
},
|
||||
elapsed_ms: started.elapsed().as_millis() as u64,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn embedder_repo(&self) -> &str {
|
||||
&self.inner.repo_id
|
||||
}
|
||||
|
||||
fn embed(&self, text: &str) -> Result<Array1<f32>> {
|
||||
let encoding = self
|
||||
.inner
|
||||
.tokenizer
|
||||
.encode(text, true)
|
||||
.map_err(|e| anyhow!("tokenize: {}", e))?;
|
||||
let mut ids: Vec<i64> = encoding.get_ids().iter().map(|&x| x as i64).collect();
|
||||
let mut mask: Vec<i64> = encoding
|
||||
.get_attention_mask()
|
||||
.iter()
|
||||
.map(|&x| x as i64)
|
||||
.collect();
|
||||
if ids.len() > MAX_SEQ_LEN {
|
||||
ids.truncate(MAX_SEQ_LEN);
|
||||
mask.truncate(MAX_SEQ_LEN);
|
||||
}
|
||||
let seq_len = ids.len();
|
||||
let ids_arr = Array2::from_shape_vec((1, seq_len), ids)?;
|
||||
let mask_arr = Array2::from_shape_vec((1, seq_len), mask)?;
|
||||
// bge-m3 is XLM-RoBERTa under the hood — single segment, so token_type_ids
|
||||
// is all zeros. The ONNX input is required nonetheless.
|
||||
let type_arr = Array2::<i64>::zeros((1, seq_len));
|
||||
|
||||
let outputs = self.inner.session.run(ort::inputs![
|
||||
"input_ids" => TensorRef::from_array_view(ids_arr.view())?,
|
||||
"attention_mask" => TensorRef::from_array_view(mask_arr.view())?,
|
||||
"token_type_ids" => TensorRef::from_array_view(type_arr.view())?,
|
||||
])?;
|
||||
|
||||
// bge-m3 ONNX outputs `last_hidden_state` of shape (1, seq_len, hidden).
|
||||
// We CLS-pool: take position 0.
|
||||
let last_hidden = outputs
|
||||
.get("last_hidden_state")
|
||||
.ok_or_else(|| anyhow!("ONNX output is missing 'last_hidden_state'"))?;
|
||||
let tensor_view = last_hidden.try_extract_array::<f32>()?;
|
||||
let shape = tensor_view.shape();
|
||||
if shape.len() != 3 {
|
||||
bail!("expected 3D embedder output, got shape {:?}", shape);
|
||||
}
|
||||
let hidden = shape[2];
|
||||
if hidden != self.inner.embedder_dim {
|
||||
bail!(
|
||||
"embedder hidden dim {} != config output_dim {}",
|
||||
hidden,
|
||||
self.inner.embedder_dim
|
||||
);
|
||||
}
|
||||
// CLS token = position 0 along the sequence axis.
|
||||
let cls_view = tensor_view.slice(ndarray::s![0, 0, ..]);
|
||||
let mut emb = Array1::from_iter(cls_view.iter().copied());
|
||||
|
||||
// bge-m3 outputs are unit-normalized in fastembed; do the same here.
|
||||
let norm = emb.dot(&emb).sqrt();
|
||||
if norm > 1e-12 {
|
||||
emb.mapv_inplace(|v| v / norm);
|
||||
}
|
||||
Ok(emb)
|
||||
}
|
||||
}
|
||||
|
||||
fn default_bundle_dir() -> Option<PathBuf> {
|
||||
let home = dirs::home_dir()?;
|
||||
Some(home.join(".goose").join(DEFAULT_BUNDLE_SUBDIR))
|
||||
}
|
||||
|
||||
fn sigmoid(x: f32) -> f32 {
|
||||
1.0 / (1.0 + (-x).exp())
|
||||
}
|
||||
|
||||
fn load_head_weights(bytes: &[u8], cfg: &HeadConfig) -> Result<Vec<LinearLayer>> {
|
||||
let st = SafeTensors::deserialize(bytes).context("parse safetensors")?;
|
||||
|
||||
// The Python head is `Sequential(Linear, ReLU, Dropout, Linear, ReLU, Dropout, ...)`
|
||||
// followed by `Linear` named `out`. Inside `Sequential` the linears are at
|
||||
// indices 0, 3, 6, … which corresponds to `trunk.{0,3,6,…}` in the
|
||||
// state_dict. We materialize them in order plus the final `out` linear.
|
||||
let mut layers = Vec::with_capacity(cfg.hidden_dims.len() + 1);
|
||||
let mut sequential_idx = 0usize;
|
||||
let mut prev_dim = cfg.input_dim;
|
||||
for &h in &cfg.hidden_dims {
|
||||
let w_name = format!("trunk.{}.weight", sequential_idx);
|
||||
let b_name = format!("trunk.{}.bias", sequential_idx);
|
||||
layers.push(read_linear(&st, &w_name, &b_name, h, prev_dim)?);
|
||||
prev_dim = h;
|
||||
sequential_idx += 3; // skip ReLU + Dropout
|
||||
}
|
||||
layers.push(read_linear(
|
||||
&st,
|
||||
"out.weight",
|
||||
"out.bias",
|
||||
cfg.output_dim,
|
||||
prev_dim,
|
||||
)?);
|
||||
Ok(layers)
|
||||
}
|
||||
|
||||
fn read_linear(
|
||||
st: &SafeTensors,
|
||||
weight_name: &str,
|
||||
bias_name: &str,
|
||||
expected_out: usize,
|
||||
expected_in: usize,
|
||||
) -> Result<LinearLayer> {
|
||||
let w_tensor = st
|
||||
.tensor(weight_name)
|
||||
.with_context(|| format!("missing tensor {}", weight_name))?;
|
||||
let b_tensor = st
|
||||
.tensor(bias_name)
|
||||
.with_context(|| format!("missing tensor {}", bias_name))?;
|
||||
|
||||
let w_shape = w_tensor.shape();
|
||||
if w_shape != [expected_out, expected_in] {
|
||||
bail!(
|
||||
"{}: shape {:?} != expected [{}, {}]",
|
||||
weight_name,
|
||||
w_shape,
|
||||
expected_out,
|
||||
expected_in
|
||||
);
|
||||
}
|
||||
let b_shape = b_tensor.shape();
|
||||
if b_shape != [expected_out] {
|
||||
bail!(
|
||||
"{}: shape {:?} != expected [{}]",
|
||||
bias_name,
|
||||
b_shape,
|
||||
expected_out
|
||||
);
|
||||
}
|
||||
|
||||
let w_data = bytes_to_f32(w_tensor.data())?;
|
||||
let b_data = bytes_to_f32(b_tensor.data())?;
|
||||
let weight = Array2::from_shape_vec((expected_out, expected_in), w_data)?;
|
||||
let bias = Array1::from_vec(b_data);
|
||||
Ok(LinearLayer { weight, bias })
|
||||
}
|
||||
|
||||
fn bytes_to_f32(bytes: &[u8]) -> Result<Vec<f32>> {
|
||||
if bytes.len() % 4 != 0 {
|
||||
bail!("tensor byte length {} is not a multiple of 4", bytes.len());
|
||||
}
|
||||
let mut out = Vec::with_capacity(bytes.len() / 4);
|
||||
for chunk in bytes.chunks_exact(4) {
|
||||
out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Render `conversation` into the `user:/assistant:` form used at training
|
||||
/// time. The most recent user message is the anchor, marked with `>>>`;
|
||||
/// everything before it is the context.
|
||||
///
|
||||
/// Returns `None` if no user message exists. We don't enforce a token budget
|
||||
/// here — the embedder truncates to its max sequence length, and the anchor
|
||||
/// is at the end so truncation drops the oldest context first.
|
||||
pub fn render_for_routing(conversation: &Conversation) -> Option<String> {
|
||||
let messages = conversation.messages();
|
||||
let anchor_idx = messages
|
||||
.iter()
|
||||
.enumerate()
|
||||
.rev()
|
||||
.find(|(_, m)| matches!(m.role, Role::User))
|
||||
.map(|(i, _)| i)?;
|
||||
|
||||
let mut lines = Vec::with_capacity(anchor_idx + 1);
|
||||
for (i, msg) in messages.iter().take(anchor_idx + 1).enumerate() {
|
||||
let role = match msg.role {
|
||||
Role::User => "user",
|
||||
Role::Assistant => "assistant",
|
||||
};
|
||||
let text = msg.as_concat_text();
|
||||
let text = text.trim();
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if i == anchor_idx {
|
||||
lines.push(format!("{} {}: {}", ANCHOR_MARKER, role, text));
|
||||
} else {
|
||||
lines.push(format!("{}: {}", role, text));
|
||||
}
|
||||
}
|
||||
if lines.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(lines.join("\n"))
|
||||
}
|
||||
}
|
||||
|
||||
/// All-in-one: render a conversation, score it, decide whether to use the
|
||||
/// fast model. Returns `None` if the conversation has no user turn (which
|
||||
/// would be weird).
|
||||
pub fn route(model: &ComplexityModel, conversation: &Conversation) -> Option<RouteDecision> {
|
||||
let rendered = render_for_routing(conversation)?;
|
||||
match model.score(&rendered) {
|
||||
Ok(score) => Some(RouteDecision {
|
||||
complexity: score.complexity,
|
||||
use_fast: score.complexity < FAST_MODEL_THRESHOLD,
|
||||
elapsed_ms: score.elapsed_ms,
|
||||
}),
|
||||
Err(e) => {
|
||||
tracing::warn!("complexity scoring failed, defaulting to smart model: {:#}", e);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the agent loop needs to know.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct RouteDecision {
|
||||
pub complexity: f32,
|
||||
pub use_fast: bool,
|
||||
pub elapsed_ms: u64,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::conversation::message::Message;
|
||||
|
||||
fn convo(msgs: Vec<Message>) -> Conversation {
|
||||
Conversation::new_unvalidated(msgs)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_empty_returns_none() {
|
||||
let c = convo(vec![]);
|
||||
assert!(render_for_routing(&c).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_single_user_message() {
|
||||
let c = convo(vec![Message::user().with_text("hi there")]);
|
||||
let r = render_for_routing(&c).expect("some");
|
||||
assert_eq!(r, ">>> user: hi there");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_multi_turn_marks_last_user() {
|
||||
let c = convo(vec![
|
||||
Message::user().with_text("what is 2+2"),
|
||||
Message::assistant().with_text("4"),
|
||||
Message::user().with_text("now squared"),
|
||||
]);
|
||||
let r = render_for_routing(&c).expect("some");
|
||||
assert_eq!(
|
||||
r,
|
||||
"user: what is 2+2\nassistant: 4\n>>> user: now squared"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_skips_empty_messages() {
|
||||
let c = convo(vec![
|
||||
Message::user().with_text("real question"),
|
||||
Message::assistant().with_text(""),
|
||||
Message::user().with_text("follow up"),
|
||||
]);
|
||||
let r = render_for_routing(&c).expect("some");
|
||||
assert_eq!(r, "user: real question\n>>> user: follow up");
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
mod agent;
|
||||
pub mod complexity_router;
|
||||
pub mod container;
|
||||
pub mod execute_commands;
|
||||
pub mod extension;
|
||||
|
||||
@@ -250,14 +250,13 @@ impl Agent {
|
||||
)]
|
||||
pub(crate) async fn stream_response_from_provider(
|
||||
provider: Arc<dyn Provider>,
|
||||
model_config: crate::model::ModelConfig,
|
||||
session_id: &str,
|
||||
system_prompt: &str,
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
toolshim_tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
let config = provider.get_model_config();
|
||||
|
||||
let filtered_messages: Vec<Message> = messages
|
||||
.iter()
|
||||
.filter(|m| m.is_agent_visible())
|
||||
@@ -265,7 +264,7 @@ impl Agent {
|
||||
.collect();
|
||||
|
||||
// Convert tool messages to text if toolshim is enabled
|
||||
let messages_for_provider = if config.toolshim {
|
||||
let messages_for_provider = if model_config.toolshim {
|
||||
convert_tool_messages_to_text(&filtered_messages)
|
||||
} else {
|
||||
Conversation::new_unvalidated(filtered_messages)
|
||||
@@ -276,10 +275,8 @@ impl Agent {
|
||||
let tools = tools.to_owned();
|
||||
let toolshim_tools = toolshim_tools.to_owned();
|
||||
let provider = provider.clone();
|
||||
let config = model_config.clone();
|
||||
|
||||
// Capture errors during stream creation and return them as part of the stream
|
||||
// so they can be handled by the existing error handling logic in the agent
|
||||
let model_config = provider.get_model_config();
|
||||
debug!("WAITING_LLM_STREAM_START");
|
||||
let stream_result = provider
|
||||
.stream(
|
||||
|
||||
@@ -47,9 +47,12 @@ pub struct OpenRouterProvider {
|
||||
|
||||
impl OpenRouterProvider {
|
||||
pub async fn from_env(model: ModelConfig) -> Result<Self> {
|
||||
let model = model.with_fast(OPENROUTER_DEFAULT_FAST_MODEL, OPENROUTER_PROVIDER_NAME)?;
|
||||
|
||||
let config = crate::config::Config::global();
|
||||
let fast_model_name: String = config
|
||||
.get_param("OPENROUTER_FAST_MODEL")
|
||||
.unwrap_or_else(|_| OPENROUTER_DEFAULT_FAST_MODEL.to_string());
|
||||
let model = model.with_fast(&fast_model_name, OPENROUTER_PROVIDER_NAME)?;
|
||||
|
||||
let api_key: String = config.get_secret("OPENROUTER_API_KEY")?;
|
||||
let host: String = config
|
||||
.get_param("OPENROUTER_HOST")
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
//! Parity test for the complexity model. Loads `~/.goose/complexity_router/`
|
||||
//! and `parity_fixture.jsonl` produced by `dump_parity_fixture.py`, then
|
||||
//! asserts the Rust pipeline matches Python within tolerance on every row.
|
||||
//!
|
||||
//! Skipped (not failed) if the bundle is absent — CI without local weights
|
||||
//! just won't exercise this.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use goose::agents::complexity_router::ComplexityModel;
|
||||
use serde::Deserialize;
|
||||
|
||||
const TOLERANCE: f32 = 1e-3;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ParityRow {
|
||||
text: String,
|
||||
#[serde(default)]
|
||||
lang: Option<String>,
|
||||
expected_complexity: f32,
|
||||
expected_tool_calls_norm: f32,
|
||||
}
|
||||
|
||||
fn bundle_dir() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.expect("home dir")
|
||||
.join(".goose")
|
||||
.join("complexity_router")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parity_against_python_reference() {
|
||||
let dir = bundle_dir();
|
||||
let cfg_path = dir.join("config.json");
|
||||
let fixture_path = dir.join("parity_fixture.jsonl");
|
||||
if !cfg_path.exists() || !fixture_path.exists() {
|
||||
eprintln!(
|
||||
"skipping: missing {} or {}",
|
||||
cfg_path.display(),
|
||||
fixture_path.display()
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let model = ComplexityModel::load_from_dir(&dir).expect("load complexity model");
|
||||
|
||||
let fixture_text = std::fs::read_to_string(&fixture_path).expect("read fixture");
|
||||
let rows: Vec<ParityRow> = fixture_text
|
||||
.lines()
|
||||
.filter(|l| !l.trim().is_empty())
|
||||
.map(|l| serde_json::from_str(l).expect("parse fixture row"))
|
||||
.collect();
|
||||
|
||||
assert!(!rows.is_empty(), "fixture has no rows");
|
||||
|
||||
let mut failures = Vec::new();
|
||||
for (i, row) in rows.iter().enumerate() {
|
||||
let out = model.score(&row.text).expect("score");
|
||||
let dc = (out.complexity - row.expected_complexity).abs();
|
||||
let dt = (out.tool_calls_norm - row.expected_tool_calls_norm).abs();
|
||||
let preview: String = row.text.chars().take(60).collect();
|
||||
eprintln!(
|
||||
"[{:>2}] lang={:<8} expected=({:.4}, {:.4}) got=({:.4}, {:.4}) Δ=({:.4}, {:.4}) {:?}",
|
||||
i,
|
||||
row.lang.as_deref().unwrap_or("?"),
|
||||
row.expected_complexity,
|
||||
row.expected_tool_calls_norm,
|
||||
out.complexity,
|
||||
out.tool_calls_norm,
|
||||
dc,
|
||||
dt,
|
||||
preview,
|
||||
);
|
||||
if dc > TOLERANCE || dt > TOLERANCE {
|
||||
failures.push((i, dc, dt));
|
||||
}
|
||||
}
|
||||
|
||||
if !failures.is_empty() {
|
||||
panic!(
|
||||
"{}/{} rows exceeded tolerance {} — see stderr",
|
||||
failures.len(),
|
||||
rows.len(),
|
||||
TOLERANCE,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -187,7 +187,16 @@ export default function GooseMessage({
|
||||
})}
|
||||
</div>
|
||||
<div className="text-xs text-text-secondary transition-all duration-200 group-hover:-translate-y-4 group-hover:opacity-0 pt-1">
|
||||
{!isStreaming && !hideTimestamp && timestamp}
|
||||
{!isStreaming && !hideTimestamp && (
|
||||
<>
|
||||
{timestamp}
|
||||
{message.metadata?.inference?.resolvedModel && (
|
||||
<span className="ml-2 opacity-70">
|
||||
@ {message.metadata.inference.resolvedModel}
|
||||
</span>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
Reference in New Issue
Block a user