diff --git a/Cargo.lock b/Cargo.lock index acb14a9..93e26fb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -169,6 +169,16 @@ dependencies = [ "tracing", ] +[[package]] +name = "tinytools-jev" +version = "0.3.0" +dependencies = [ + "async-trait", + "tinytools", + "tokio", + "tracing", +] + [[package]] name = "tokio" version = "1.53.1" diff --git a/Cargo.toml b/Cargo.toml index c8b237f..95366f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,8 +38,8 @@ tracing = { version = "0.1", default-features = false } # `Tool::execute` returns `anyhow::Result` because a tool body calls arbitrary # host code and has no useful closed error set of its own. anyhow = "1" -# Unit tests for the async trait defaults drive a real executor. -tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +# Unit tests for async traits drive a real executor. +tokio = { version = "1", features = ["macros", "rt-multi-thread", "time"] } # Lints apply to every member that opts in with `[lints] workspace = true`, and # to every target of that member. CI runs clippy with `-D warnings`, so anything diff --git a/README.md b/README.md index 866171b..408d764 100644 --- a/README.md +++ b/README.md @@ -65,11 +65,15 @@ compiles neither the harness nor the host. | `context` | `ToolRunContext` — the narrow seam onto a live run | | `workspace` | `WorkspaceDescriptor`, `SandboxMode` — the root a tool may touch, and how strictly it is sandboxed | | `naming` | `humanize_tool_name`, `context_detail_from_args` — rendering a call for a human | +| `rank` | `ToolRanker`, `RankCandidate`, `RankHit`, `Bm25Ranker` — ranking a catalogue of tools against an intent, with the lexical ranker built in | The workspace also contains `tinytools-agent`, a separate crate for model-facing tool-call parsing, dialects, catalogue/result rendering, and transcript replay. It builds on `ToolSpec` without adding agent-loop or provider -dependencies to the base `tinytools` vocabulary crate. +dependencies to the base `tinytools` vocabulary crate. `tinytools-jev` is a +second sibling: a `ToolRanker` backed by TypeSafe's Jev decision model through +`tinyjevclient`, kept out of the vocabulary crate because it carries an HTTP +transport. ## What is deliberately not here diff --git a/crates/tinytools-jev/Cargo.toml b/crates/tinytools-jev/Cargo.toml new file mode 100644 index 0000000..e029391 --- /dev/null +++ b/crates/tinytools-jev/Cargo.toml @@ -0,0 +1,26 @@ +[package] +publish = false +name = "tinytools-jev" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +repository.workspace = true +description = "A tinytools ToolRanker backed by TypeSafe's Jev decision model: retrieve a shortlist lexically, let Jev decide." +documentation = "https://docs.rs/tinytools-jev" +readme = "README.md" + +[dependencies] +async-trait = { workspace = true } +tinytools = { path = "../tinytools", version = "0.3.0" } +tracing = { workspace = true, optional = true } + +[dev-dependencies] +tokio = { workspace = true } + +[features] +default = [] +tracing = ["dep:tracing"] + +[lints] +workspace = true diff --git a/crates/tinytools-jev/README.md b/crates/tinytools-jev/README.md new file mode 100644 index 0000000..fda20a6 --- /dev/null +++ b/crates/tinytools-jev/README.md @@ -0,0 +1,48 @@ +# tinytools-jev + +A `tinytools::ToolRanker` backed by a host-provided Jev evaluator. The host +owns its HTTP client, credentials, retry policy, and deadline. + +## Retrieve, then decide + +Jev answers a `Choice` question with a calibrated probability for every +option in ~150 ms, accepts at most 255 options, and loses accuracy as the +option list fills with entries unrelated to the request. So `JevRanker` +never shows it a whole catalogue: + +1. A retriever (`Bm25Ranker` unless the host supplies one) narrows the + catalogue to `retrieval_k` candidates (20 by default). Skipped when the + catalogue already fits. When the retriever finds nothing and the catalogue + fits one Choice, Jev sees all of it — a lexical miss on a paraphrase is + exactly the case a decision model is for. +2. One request: a `Choice` over the shortlist plus a `none` option, and a + `Noul` asking whether the request needs a tool at all. +3. Hits are the options by probability, `none` removed, anything below + `min_probability` dropped. Each hit's `confidence` is its probability; + `rank_detailed` also returns the Choice confidence, the `needs_tool` + probability, the `none` probability, tokens, latency and attempts. + +Every evaluator failure is a `RankError` the caller falls back from. + +## Limits that shape the design + +| Limit | Value | Consequence | +|------------------------------|--------|-----------------------------------------------| +| Options per Choice | 255 | `retrieval_k` is clamped to it | +| Summary shown per option | 240 ch | clipped, with the family named after | +| Context per request | 64k | `RankContext` stays to a few recent turns | +| Pricing (jev-1.13) | $0.042 / M input, output free | a search is ~1–2k tokens | + +## Wording + +Jev reads literally. The instructions name the user's `request` and ask which +tool accomplishes it "by what each tool does, not by shared words"; each +option is `name: first sentence (from family)`; the state carries the request +and at most the caller's few recent turns. + +## Building a ranker + +Implement `JevEvaluator` in the host by translating `JevRequest` into the +client's wire request and translating its answer into `JevDecision`. Pass that +implementation to `JevRanker::new`. This keeps transport and runtime choices +at the host boundary where their policy belongs. diff --git a/crates/tinytools-jev/src/lib.rs b/crates/tinytools-jev/src/lib.rs new file mode 100644 index 0000000..18bb197 --- /dev/null +++ b/crates/tinytools-jev/src/lib.rs @@ -0,0 +1,214 @@ +//! Dependency-free Jev-backed tool ranking. +//! +//! A host supplies [`JevEvaluator`], retaining ownership of transport, +//! authentication, retry, and deadline policy. + +#[cfg(test)] +mod test; +mod types; +pub use types::{JevDecision, JevEvaluator, JevOption, JevRankerConfig, JevRanking, JevRequest}; + +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::Arc, + time::Duration, +}; +use tinytools::{RankCandidate, RankContext, RankError, RankHit, ToolRanker}; +const NONE_OPTION: &str = "none"; +const MAX_SUMMARY_CHARS: usize = 240; + +/// Ranks tools using a host-provided evaluator. +#[derive(Debug, Clone)] +pub struct JevRanker { + evaluator: Arc, + config: JevRankerConfig, +} +impl JevRanker { + /// Stable ranker kind. + pub const KIND: &'static str = "jev"; + /// Creates a ranker. + #[must_use] + pub fn new(evaluator: Arc, config: JevRankerConfig) -> Self { + Self { evaluator, config } + } + /// Returns the active configuration. + #[must_use] + pub fn config(&self) -> &JevRankerConfig { + &self.config + } + /// Ranks and returns all decision metadata. + /// + /// # Errors + /// Returns invalid-input errors and forwards evaluator/retriever failures. + pub async fn rank_detailed( + &self, + intent: &str, + context: &RankContext, + candidates: &[RankCandidate], + limit: usize, + ) -> Result { + let intent = intent.trim(); + if intent.is_empty() { + return Err(RankError::invalid_input("intent is empty")); + } + if candidates.is_empty() || limit == 0 { + return Ok(JevRanking::empty()); + } + let shortlist = self.shortlist(intent, context, candidates).await?; + if shortlist.is_empty() { + return Ok(JevRanking::empty()); + } + let request = self.build_request(intent, context, &shortlist)?; + let started = std::time::Instant::now(); + let decision = self.evaluator.evaluate(&request).await?; + let mut ranking = decode(&decision, &shortlist, self.config.min_probability); + ranking.hits.truncate(limit); + ranking.latency = started.elapsed(); + Ok(ranking) + } + async fn shortlist<'a>( + &self, + intent: &str, + context: &RankContext, + candidates: &'a [RankCandidate], + ) -> Result, RankError> { + validate_candidates(candidates)?; + let k = self.config.retrieval_k.min(JevRankerConfig::MAX_CANDIDATES); + if candidates.len() <= k { + return Ok(candidates.iter().collect()); + } + let hits = self + .config + .retriever + .rank(intent, context, candidates, k) + .await?; + let by_key: BTreeMap<&str, &RankCandidate> = + candidates.iter().map(|c| (c.key.as_str(), c)).collect(); + let shortlist: Vec<_> = hits + .iter() + .filter_map(|h| by_key.get(h.key.as_str()).copied()) + .take(JevRankerConfig::MAX_CANDIDATES) + .collect(); + if shortlist.is_empty() && candidates.len() <= JevRankerConfig::MAX_CANDIDATES { + return Ok(candidates.iter().collect()); + } + Ok(shortlist) + } + fn build_request( + &self, + intent: &str, + context: &RankContext, + shortlist: &[&RankCandidate], + ) -> Result { + if shortlist.len() > JevRankerConfig::MAX_CANDIDATES { + return Err(RankError::invalid_input("too many shortlisted candidates")); + } + let mut options: Vec<_> = shortlist + .iter() + .map(|c| JevOption { + key: c.key.clone(), + description: option_text(c), + }) + .collect(); + options.push(JevOption { + key: NONE_OPTION.into(), + description: "No listed tool accomplishes the request.".into(), + }); + Ok(JevRequest { + intent: intent.into(), + recent_turns: context.recent_turns.clone(), + options, + model: self.config.model.clone(), + }) + } +} +#[async_trait::async_trait] +impl ToolRanker for JevRanker { + fn kind(&self) -> &'static str { + Self::KIND + } + async fn rank( + &self, + intent: &str, + context: &RankContext, + candidates: &[RankCandidate], + limit: usize, + ) -> Result, RankError> { + self.rank_detailed(intent, context, candidates, limit) + .await + .map(|r| r.hits) + } +} +impl JevRanking { + fn empty() -> Self { + Self { + hits: vec![], + choice_confidence: 0.0, + needs_tool: None, + none_probability: 0.0, + shortlisted: 0, + input_tokens: None, + latency: Duration::ZERO, + attempts: 0, + } + } +} +fn validate_candidates(candidates: &[RankCandidate]) -> Result<(), RankError> { + let mut keys = BTreeSet::new(); + for c in candidates { + if c.key == NONE_OPTION { + return Err(RankError::invalid_input("candidate key `none` is reserved")); + } + if !keys.insert(c.key.as_str()) { + return Err(RankError::invalid_input("duplicate candidate key")); + } + } + Ok(()) +} +fn option_text(candidate: &RankCandidate) -> String { + let mut summary: String = candidate.summary.chars().take(MAX_SUMMARY_CHARS).collect(); + if candidate.summary.chars().count() > MAX_SUMMARY_CHARS { + summary.push('…'); + } + candidate.family.as_ref().map_or(summary.clone(), |family| { + format!("{summary} (from {family})") + }) +} +fn decode(decision: &JevDecision, shortlist: &[&RankCandidate], floor: f64) -> JevRanking { + let none = decision + .probabilities + .get(NONE_OPTION) + .copied() + .unwrap_or(0.0); + let best = shortlist + .iter() + .filter_map(|c| decision.probabilities.get(&c.key).copied()) + .fold(0.0_f64, f64::max); + let abstained = none >= best || decision.needs_tool.is_some_and(|p| p < 0.5); + let mut hits: Vec<_> = if abstained { + vec![] + } else { + shortlist + .iter() + .filter_map(|c| { + let p = *decision.probabilities.get(&c.key)?; + (p >= floor).then(|| RankHit { + key: c.key.clone(), + score: p, + confidence: Some(p), + }) + }) + .collect() + }; + hits.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.key.cmp(&b.key))); + JevRanking { + hits, + choice_confidence: decision.choice_confidence, + needs_tool: decision.needs_tool, + none_probability: none, + shortlisted: shortlist.len(), + input_tokens: decision.input_tokens, + latency: Duration::ZERO, + attempts: decision.attempts, + } +} diff --git a/crates/tinytools-jev/src/test.rs b/crates/tinytools-jev/src/test.rs new file mode 100644 index 0000000..5ba6355 --- /dev/null +++ b/crates/tinytools-jev/src/test.rs @@ -0,0 +1,209 @@ +//! Tests for provider-neutral Jev ranking. + +use super::*; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, +}; + +#[derive(Debug)] +struct FakeEvaluator { + decision: JevDecision, + seen: Mutex>, +} +#[async_trait::async_trait] +impl JevEvaluator for FakeEvaluator { + async fn evaluate(&self, request: &JevRequest) -> Result { + if let Ok(mut seen) = self.seen.lock() { + seen.push(request.clone()); + } + Ok(self.decision.clone()) + } +} + +fn candidates() -> Vec { + vec![ + RankCandidate::new("slack", "Send a Slack message").with_family("chat"), + RankCandidate::new("gmail", "Send an email"), + ] +} + +fn ranker( + probabilities: [(&str, f64); 3], + needs_tool: Option, +) -> (JevRanker, Arc) { + let evaluator = Arc::new(FakeEvaluator { + decision: JevDecision { + probabilities: probabilities + .into_iter() + .map(|(key, value)| (key.to_owned(), value)) + .collect::>(), + choice_confidence: 0.8, + needs_tool, + input_tokens: Some(10), + attempts: 1, + }, + seen: Mutex::new(vec![]), + }); + ( + JevRanker::new(evaluator.clone(), JevRankerConfig::new()), + evaluator, + ) +} + +#[tokio::test] +async fn ranks_candidates_and_builds_none_option() { + let (ranker, evaluator) = ranker([("slack", 0.8), ("gmail", 0.1), ("none", 0.1)], Some(0.9)); + let result = ranker + .rank_detailed("message Alex", &RankContext::empty(), &candidates(), 2) + .await; + assert_eq!( + result + .as_ref() + .ok() + .and_then(|ranking| ranking.hits.first()) + .map(|hit| hit.key.as_str()), + Some("slack") + ); + let options = evaluator + .seen + .lock() + .ok() + .and_then(|seen| seen.first().map(|request| request.options.len())); + assert_eq!(options, Some(3)); +} + +#[tokio::test] +async fn suppresses_hits_when_none_wins_or_no_tool_is_needed() { + let (none_ranker, _) = ranker([("slack", 0.2), ("gmail", 0.1), ("none", 0.7)], Some(0.9)); + let none_hits = none_ranker + .rank("question", &RankContext::empty(), &candidates(), 2) + .await + .unwrap_or_default(); + assert!(none_hits.is_empty()); + let (no_tool_ranker, _) = ranker([("slack", 0.8), ("gmail", 0.1), ("none", 0.1)], Some(0.2)); + let no_tool_hits = no_tool_ranker + .rank("question", &RankContext::empty(), &candidates(), 2) + .await + .unwrap_or_default(); + assert!(no_tool_hits.is_empty()); +} + +#[test] +fn configuration_reserves_none_slot_and_rejects_nan() { + let config = JevRankerConfig::new() + .with_retrieval_k(usize::MAX) + .with_min_probability(f64::NAN); + assert_eq!(config.retrieval_k, JevRankerConfig::MAX_CANDIDATES); + assert!((config.min_probability - 0.05).abs() < f64::EPSILON); +} + +#[test] +fn option_text_clips_by_characters() { + let candidate = RankCandidate::new("key", "é".repeat(300)).with_family("family"); + let text = option_text(&candidate); + assert_eq!( + text.chars().filter(|character| *character == 'é').count(), + MAX_SUMMARY_CHARS + ); + assert!(text.ends_with("… (from family)")); +} + +#[tokio::test] +async fn validates_inputs_and_short_circuits_empty_work() { + let (ranker, _) = ranker([("slack", 0.8), ("gmail", 0.1), ("none", 0.1)], Some(0.9)); + let empty = ranker + .rank("request", &RankContext::empty(), &[], 2) + .await + .unwrap_or_default(); + assert!(empty.is_empty()); + let zero = ranker + .rank("request", &RankContext::empty(), &candidates(), 0) + .await + .unwrap_or_default(); + assert!(zero.is_empty()); + assert!(matches!( + ranker + .rank(" ", &RankContext::empty(), &candidates(), 2) + .await, + Err(RankError::InvalidInput { .. }) + )); + let reserved = vec![RankCandidate::new("none", "reserved")]; + assert!(matches!( + ranker + .rank("request", &RankContext::empty(), &reserved, 2) + .await, + Err(RankError::InvalidInput { .. }) + )); + let duplicate = vec![ + RankCandidate::new("x", "one"), + RankCandidate::new("x", "two"), + ]; + assert!(matches!( + ranker + .rank("request", &RankContext::empty(), &duplicate, 2) + .await, + Err(RankError::InvalidInput { .. }) + )); +} + +#[tokio::test] +async fn retrieves_large_catalogues_and_handles_a_retrieval_miss() { + let (_initial_ranker, evaluator) = + ranker([("slack", 0.8), ("gmail", 0.1), ("none", 0.1)], Some(0.9)); + let configured_ranker = JevRanker::new(evaluator, JevRankerConfig::new().with_retrieval_k(1)); + let result = configured_ranker + .rank("Slack message", &RankContext::empty(), &candidates(), 2) + .await + .unwrap_or_default(); + assert_eq!(result.first().map(|hit| hit.key.as_str()), Some("slack")); + + let (_unused_ranker, evaluator) = + ranker([("slack", 0.8), ("gmail", 0.1), ("none", 0.1)], Some(0.9)); + let miss_ranker = JevRanker::new( + evaluator.clone(), + JevRankerConfig::new().with_retrieval_k(1), + ); + let _ = miss_ranker + .rank("unrelated", &RankContext::empty(), &candidates(), 2) + .await; + let shown = evaluator + .seen + .lock() + .ok() + .and_then(|seen| seen.first().map(|request| request.options.len())); + assert_eq!(shown, Some(3)); +} + +#[derive(Debug)] +struct FailingEvaluator; +#[async_trait::async_trait] +impl JevEvaluator for FailingEvaluator { + async fn evaluate(&self, _request: &JevRequest) -> Result { + Err(RankError::backend("provider unavailable")) + } +} + +#[tokio::test] +async fn forwards_evaluator_errors() { + let ranker = JevRanker::new(Arc::new(FailingEvaluator), JevRankerConfig::default()); + let result = ranker + .rank("message", &RankContext::empty(), &candidates(), 2) + .await; + assert!(matches!(result, Err(RankError::Backend { .. }))); + assert_eq!(ranker.kind(), JevRanker::KIND); + assert!(format!("{:?}", ranker.config()).contains("bm25")); +} + +#[test] +fn configuration_builders_are_observable() { + let retriever: Arc = Arc::new(tinytools::Bm25Ranker); + let config = JevRankerConfig::new() + .with_retriever(retriever) + .with_retrieval_k(0) + .with_min_probability(2.0) + .with_model("jev-pinned"); + assert_eq!(config.retrieval_k, 1); + assert!((config.min_probability - 1.0).abs() < f64::EPSILON); + assert_eq!(config.model, "jev-pinned"); +} diff --git a/crates/tinytools-jev/src/types.rs b/crates/tinytools-jev/src/types.rs new file mode 100644 index 0000000..bd5b3d3 --- /dev/null +++ b/crates/tinytools-jev/src/types.rs @@ -0,0 +1,139 @@ +//! Provider-neutral request, response, and configuration types. + +use std::{collections::BTreeMap, fmt, sync::Arc}; +use tinytools::{Bm25Ranker, RankHit, ToolRanker}; + +/// One candidate presented to an evaluator. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JevOption { + /// Opaque candidate key. + pub key: String, + /// Concise description the evaluator judges. + pub description: String, +} + +/// A provider-neutral tool-selection request. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JevRequest { + /// The user's request. + pub intent: String, + /// Recent user turns, oldest first. + pub recent_turns: Vec, + /// Candidate tools, including `none`. + pub options: Vec, + /// Model identifier configured by the host. + pub model: String, +} + +/// The decision returned by a host-provided evaluator. +#[derive(Clone, Debug, PartialEq)] +pub struct JevDecision { + /// Probability for each option key. + pub probabilities: BTreeMap, + /// Confidence in the overall choice. + pub choice_confidence: f64, + /// Probability that the request needs a tool. + pub needs_tool: Option, + /// Input tokens billed, when reported. + pub input_tokens: Option, + /// Attempts made by the host client. + pub attempts: u32, +} + +/// Evaluates a request without coupling this workspace to a transport. +#[async_trait::async_trait] +pub trait JevEvaluator: Send + Sync + fmt::Debug { + /// Evaluates one tool-selection request. + /// + /// # Errors + /// Returns a ranking error for invalid requests, timeouts, or provider failures. + async fn evaluate(&self, request: &JevRequest) -> Result; +} + +/// How [`JevRanker`][crate::JevRanker] retrieves and decides. +#[derive(Clone)] +pub struct JevRankerConfig { + pub(crate) retriever: Arc, + pub(crate) retrieval_k: usize, + pub(crate) min_probability: f64, + pub(crate) model: String, +} + +impl JevRankerConfig { + /// Maximum options accepted, including `none`. + pub const MAX_OPTIONS: usize = 255; + /// Maximum real candidates, reserving one slot for `none`. + pub const MAX_CANDIDATES: usize = Self::MAX_OPTIONS - 1; + /// Returns the default configuration. + #[must_use] + pub fn new() -> Self { + Self { + retriever: Arc::new(Bm25Ranker), + retrieval_k: 20, + min_probability: 0.05, + model: "jev-latest".into(), + } + } + /// Replaces the retriever. + #[must_use] + pub fn with_retriever(mut self, value: Arc) -> Self { + self.retriever = value; + self + } + /// Sets the shortlist size, clamped to the valid candidate range. + #[must_use] + pub fn with_retrieval_k(mut self, value: usize) -> Self { + self.retrieval_k = value.clamp(1, Self::MAX_CANDIDATES); + self + } + /// Sets a finite probability floor, clamped to `0.0..=1.0`. + #[must_use] + pub fn with_min_probability(mut self, value: f64) -> Self { + if value.is_finite() { + self.min_probability = value.clamp(0.0, 1.0); + } + self + } + /// Sets the model id. + #[must_use] + pub fn with_model(mut self, value: impl Into) -> Self { + self.model = value.into(); + self + } +} +impl Default for JevRankerConfig { + fn default() -> Self { + Self::new() + } +} +impl fmt::Debug for JevRankerConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JevRankerConfig") + .field("retriever", &self.retriever.kind()) + .field("retrieval_k", &self.retrieval_k) + .field("min_probability", &self.min_probability) + .field("model", &self.model) + .finish() + } +} + +/// Everything one ranking learned, beyond its hits. +#[derive(Clone, Debug, PartialEq)] +pub struct JevRanking { + /// Ranked hits. + pub hits: Vec, + /// Confidence in the choice. + pub choice_confidence: f64, + /// Probability that the request needs a tool. + pub needs_tool: Option, + /// Probability assigned to `none`. + pub none_probability: f64, + /// Candidate count shown. + pub shortlisted: usize, + /// Input tokens billed, when reported. + pub input_tokens: Option, + /// Evaluator wall time. + pub latency: std::time::Duration, + /// Host-client attempt count. + pub attempts: u32, +} diff --git a/crates/tinytools/src/lib.rs b/crates/tinytools/src/lib.rs index 22e1d5a..7564d34 100644 --- a/crates/tinytools/src/lib.rs +++ b/crates/tinytools/src/lib.rs @@ -32,6 +32,8 @@ //! - [`context`] — [`ToolRunContext`], the narrow seam onto a live run. //! - [`workspace`] — [`WorkspaceDescriptor`], the root a tool may touch. //! - [`naming`] — rendering a call for a human. +//! - [`rank`] — [`ToolRanker`], ranking a catalogue of tools against an +//! intent, and the lexical [`Bm25Ranker`] every host gets for free. //! //! # What is deliberately not here //! @@ -104,6 +106,7 @@ pub mod context; pub mod naming; pub mod permission; pub mod policy; +pub mod rank; pub mod result; pub mod spec; pub mod tool; @@ -124,6 +127,9 @@ pub use permission::PermissionLevel; pub use policy::{ ToolAccess, ToolDisplay, ToolPolicy, ToolReplay, ToolRuntime, ToolSideEffects, WorkspaceAccess, }; +pub use rank::{ + Bm25Index, Bm25Ranker, RankCandidate, RankContext, RankError, RankHit, ToolRanker, tokenize, +}; pub use result::{FileData, ImageData, ToolContent, ToolControl, ToolErrorKind, ToolResult}; pub use spec::ToolSpec; pub use tool::{Tool, ToolExposure}; diff --git a/crates/tinytools/src/rank/README.md b/crates/tinytools/src/rank/README.md new file mode 100644 index 0000000..3fa8e7f --- /dev/null +++ b/crates/tinytools/src/rank/README.md @@ -0,0 +1,40 @@ +# `rank` + +Ranking a catalogue of tools against an intent: the `ToolRanker` trait, the +`RankCandidate` / `RankHit` / `RankContext` vocabulary, and `Bm25Ranker`, the +lexical implementation every host gets without a network. + +## Design + +A host that registers more tools than a model should see on every request +advertises a few and lets the model search for the rest. That search is a +ranking problem, and the interesting rankers (a decision model such as +TypeSafe's Jev, an embedding index) talk to a service this crate must never +depend on. So the crate owns the *question* — here are candidates, here is an +intent, order them — and one free answer, BM25. A harness asks through the +trait and never learns which kind answered; a host composes them, typically +BM25 to retrieve a shortlist and a model to decide (see `tinytools-jev`). + +`Bm25Index` and `tokenize` moved here from the `tinyagents` harness's +discovery module so both crates rank with one arithmetic. + +## Contract + +- `rank` returns at most `limit` hits, best first, and only candidates it + considers relevant. Empty means "nothing fits", never "padded to `limit`". +- Every returned key names a candidate the caller passed. +- Failure is a `RankError`; the caller falls back. Never a panic on input. +- `RankHit::confidence` is a calibrated probability or `None`. BM25 returns + `None`: a BM25 score is not a probability, and a caller gating on + confidence must treat `None` as unknown rather than zero. +- `RankContext` is deliberately small. A model-backed ranker pays for every + byte on every search, and unrelated context is a distractor. + +## Files + +| File | Owns | +|------------|-------------------------------------------------------------| +| `mod.rs` | `ToolRanker`, the `Arc` blanket impl, re-exports | +| `types.rs` | `RankCandidate`, `RankHit`, `RankContext`, `RankError` | +| `bm25.rs` | `tokenize`, `Bm25Index`, `Bm25Ranker` | +| `test.rs` | Unit tests | diff --git a/crates/tinytools/src/rank/bm25.rs b/crates/tinytools/src/rank/bm25.rs new file mode 100644 index 0000000..4fd4441 --- /dev/null +++ b/crates/tinytools/src/rank/bm25.rs @@ -0,0 +1,283 @@ +//! BM25 ranking over short documents, and the [`ToolRanker`] built on it. +//! +//! Takes `(sort_key, text)` pairs and returns ranked indices, so the +//! tool-specific half (what text is searchable, what a hit looks like) stays +//! with the caller. Hand-rolled rather than the `bm25` crate: the arithmetic +//! is ~80 lines, and this crate is the dependency floor of every tool author. +//! +//! Moved here from the `tinyagents` harness's discovery module so a host can +//! rank with the same arithmetic the harness uses, and so a model-backed +//! ranker can retrieve a shortlist with it before deciding. + +use std::collections::HashMap; + +use super::{RankCandidate, RankContext, RankError, RankHit, ToolRanker}; + +/// Words that carry no capability meaning, dropped from a query before ranking. +/// +/// A document-frequency threshold alone cannot do this on a small corpus: +/// with three deferred tools, "a" may appear in exactly one description and +/// so rank as the *most* distinguishing term in "send a calendar invite". +/// Deliberately short and English-only — it can only remove terms, so a +/// description in another language ranks exactly as it would without it. +/// Words that could name a capability ("up", as in "look up") are left in. +const STOPWORDS: &[&str] = &[ + "a", "an", "and", "are", "as", "at", "be", "but", "by", "can", "do", "for", "from", "how", "i", + "if", "in", "into", "is", "it", "its", "me", "my", "of", "on", "or", "our", "so", "that", + "the", "their", "them", "then", "there", "these", "they", "this", "to", "was", "we", "were", + "what", "when", "which", "who", "will", "with", "would", "you", "your", +]; + +/// BM25 term-frequency saturation (the standard default). +const K1: f64 = 1.2; +/// BM25 length normalisation (the standard default). +const B: f64 = 0.75; + +/// Splits text into search terms. +/// +/// Splits on non-alphanumerics **and** on a lower→upper transition, so +/// `memory_hybrid_search` and `readWorkflowResource` both yield the words a +/// person would type. Without the camel-case rule a query for "workflow" +/// misses a tool whose only mention of it is inside an identifier. +#[must_use] +pub fn tokenize(text: &str) -> Vec { + let mut out = Vec::new(); + let mut current = String::new(); + let mut previous_lower = false; + for ch in text.chars() { + if ch.is_alphanumeric() { + if ch.is_uppercase() && previous_lower && !current.is_empty() { + out.push(std::mem::take(&mut current)); + } + current.extend(ch.to_lowercase()); + previous_lower = ch.is_lowercase() || ch.is_numeric(); + } else if !current.is_empty() { + out.push(std::mem::take(&mut current)); + previous_lower = false; + } + } + if !current.is_empty() { + out.push(current); + } + out +} + +/// A ranked corpus of short documents, identified by their index into the +/// slice the caller built the index from. +#[derive(Default, Debug, Clone)] +pub struct Bm25Index { + documents: Vec, + document_frequency: HashMap, + average_length: f64, +} + +#[derive(Debug, Clone)] +struct Document { + /// Used only to break score ties deterministically. + sort_key: String, + tokens: Vec, +} + +impl Bm25Index { + /// Builds from `(sort_key, searchable_text)` pairs, in caller order. + /// + /// `sort_key` breaks ties; make it the id the caller would print, so two + /// identical queries produce identical output. An unstable order would make + /// a model's transcript non-reproducible for no benefit. + pub fn build<'a>(documents: impl IntoIterator) -> Self { + let documents: Vec = documents + .into_iter() + .map(|(sort_key, text)| Document { + sort_key: sort_key.to_string(), + tokens: tokenize(text), + }) + .collect(); + + let mut document_frequency: HashMap = HashMap::new(); + for doc in &documents { + let mut seen: Vec<&str> = Vec::new(); + for token in &doc.tokens { + if !seen.contains(&token.as_str()) { + seen.push(token); + *document_frequency.entry(token.clone()).or_insert(0) += 1; + } + } + } + + let total: usize = documents.iter().map(|d| d.tokens.len()).sum(); + #[allow(clippy::cast_precision_loss)] + let average_length = if documents.is_empty() { + 0.0 + } else { + total as f64 / documents.len() as f64 + }; + + Self { + documents, + document_frequency, + average_length, + } + } + + /// `true` when the corpus holds no documents. + #[must_use] + pub fn is_empty(&self) -> bool { + self.documents.is_empty() + } + + /// Number of documents in the corpus. + #[must_use] + pub fn len(&self) -> usize { + self.documents.len() + } + + /// Ranks against `query`, best first, returning **indices** into the corpus. + /// + /// Only documents scoring above zero are returned. Padding the list out to + /// `limit` with unrelated entries would spend exactly the tokens deferral + /// exists to save, and would invite the model to call something unrelated + /// to what it asked for. + #[must_use] + pub fn search(&self, query: &str, limit: usize) -> Vec { + self.search_scored(query, limit) + .into_iter() + .map(|(_, index)| index) + .collect() + } + + /// [`Self::search`], keeping each hit's score. + #[must_use] + pub fn search_scored(&self, query: &str, limit: usize) -> Vec<(f64, usize)> { + let terms = self.significant(tokenize(query)); + if terms.is_empty() || self.documents.is_empty() { + return Vec::new(); + } + let mut scored: Vec<(f64, usize)> = self + .documents + .iter() + .enumerate() + .map(|(index, doc)| (self.score(doc, &terms), index)) + .filter(|(score, _)| *score > 0.0) + .collect(); + + scored.sort_by(|a, b| { + b.0.partial_cmp(&a.0) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| { + self.documents[a.1] + .sort_key + .cmp(&self.documents[b.1].sort_key) + }) + }); + scored.truncate(limit); + scored + } + + /// Drops query terms too common in this corpus to mean anything. + /// + /// The IDF below carries the standard `+ 1`, which keeps a term present in + /// every document at a small **positive** weight — without it a corpus of + /// one document scores every term at zero and nothing is ever findable. + /// The cost is that "a" and "the" score, so a document-frequency filter + /// runs first: drop the term when `df >= max(2, ceil(0.8 * n))`. The floor + /// of 2 makes it inert on a one-document corpus; [`STOPWORDS`] covers the + /// small-corpus case the threshold cannot. + fn significant(&self, terms: Vec) -> Vec { + let n = self.documents.len(); + if n == 0 { + return Vec::new(); + } + #[allow( + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss + )] + let threshold = std::cmp::max(2, (0.8 * n as f64).ceil() as usize); + terms + .into_iter() + .filter(|term| !STOPWORDS.contains(&term.as_str())) + .filter(|term| *self.document_frequency.get(term).unwrap_or(&0) < threshold) + .collect() + } + + #[allow(clippy::cast_precision_loss)] + fn score(&self, doc: &Document, terms: &[String]) -> f64 { + let length = doc.tokens.len() as f64; + let count = self.documents.len() as f64; + terms + .iter() + .map(|term| { + let frequency = doc.tokens.iter().filter(|t| *t == term).count() as f64; + if frequency == 0.0 { + return 0.0; + } + let df = *self.document_frequency.get(term).unwrap_or(&0) as f64; + // Standard BM25 IDF with the +1 that keeps a term present in + // every document at a small positive weight (see `significant`). + let idf = ((count - df + 0.5) / (df + 0.5) + 1.0).ln(); + let normalised = frequency * (K1 + 1.0) + / (frequency + K1 * (1.0 - B + B * length / self.average_length.max(1.0))); + idf * normalised + }) + .sum() + } +} + +/// The lexical [`ToolRanker`]: BM25 over each candidate's summary and family. +/// +/// Builds a fresh [`Bm25Index`] per call. A catalogue of a few hundred short +/// summaries indexes in microseconds, and a per-call index means the ranker +/// holds no state to invalidate when the caller's catalogue changes. Returns +/// no [`RankHit::confidence`]: a BM25 score is not a probability. +#[derive(Debug, Default, Clone, Copy)] +pub struct Bm25Ranker; + +impl Bm25Ranker { + /// The stable [`ToolRanker::kind`] of this ranker. + pub const KIND: &'static str = "bm25"; + + /// Ranks synchronously; [`ToolRanker::rank`] delegates here. + #[must_use] + pub fn rank_sync(candidates: &[RankCandidate], intent: &str, limit: usize) -> Vec { + let texts: Vec = candidates + .iter() + .map(|candidate| match &candidate.family { + Some(family) => format!("{} {}", candidate.summary, family), + None => candidate.summary.clone(), + }) + .collect(); + let index = Bm25Index::build( + candidates + .iter() + .zip(&texts) + .map(|(candidate, text)| (candidate.key.as_str(), text.as_str())), + ); + index + .search_scored(intent, limit) + .into_iter() + .map(|(score, i)| RankHit::new(candidates[i].key.clone(), score)) + .collect() + } +} + +#[async_trait::async_trait] +impl ToolRanker for Bm25Ranker { + fn kind(&self) -> &'static str { + Self::KIND + } + + async fn rank( + &self, + intent: &str, + _context: &RankContext, + candidates: &[RankCandidate], + limit: usize, + ) -> Result, RankError> { + if intent.trim().is_empty() { + return Err(RankError::InvalidInput { + reason: "intent is empty".to_owned(), + }); + } + Ok(Self::rank_sync(candidates, intent, limit)) + } +} diff --git a/crates/tinytools/src/rank/mod.rs b/crates/tinytools/src/rank/mod.rs new file mode 100644 index 0000000..7b82213 --- /dev/null +++ b/crates/tinytools/src/rank/mod.rs @@ -0,0 +1,81 @@ +//! Ranking a catalogue of tools against an intent. +//! +//! A host that registers more tools than a model should see on every request +//! advertises a few and lets the model *search* for the rest. The search is a +//! ranking problem — "which of these tools is the one for *send Alex a note +//! that I'm late*?" — and this module is the vocabulary for it: the +//! [`ToolRanker`] trait, the [`RankCandidate`] a ranker reads, the +//! [`RankHit`] it returns, and one implementation every host can use without +//! a network, [`Bm25Ranker`]. +//! +//! # Why a trait +//! +//! Lexical ranking is free and answers most queries whose wording overlaps a +//! tool's description. It misses the paraphrase — "ping" for a tool described +//! as "send a message" — and it has no notion of confidence, so a host cannot +//! tell a strong match from the least-bad one. A decision model (`TypeSafe`'s +//! Jev, or an embedding index) answers both, at the cost of a network call the +//! vocabulary crate must not make. The trait lets a harness ask "rank these" +//! without knowing which kind is answering, and lets a host compose them: +//! retrieve a shortlist lexically, then let the model decide. +//! +//! # Contract +//! +//! - `rank` returns at most `limit` hits, best first, and only candidates it +//! considers relevant: an empty result means "nothing here fits", never +//! "I padded to `limit`". +//! - Every returned key names a candidate the caller passed. A ranker never +//! invents one. +//! - A ranker that cannot answer returns [`RankError`]; the caller decides +//! what to fall back to. It never panics on caller input. +//! - `rank` is `async` because the interesting implementations talk to a +//! service; [`Bm25Ranker`] completes without yielding. + +mod bm25; +#[cfg(test)] +mod test; +mod types; + +pub use bm25::{Bm25Index, Bm25Ranker, tokenize}; +pub use types::{RankCandidate, RankContext, RankError, RankHit}; + +/// Ranks tool candidates against an intent. +/// +/// See the [module docs](self) for the contract every implementation keeps. +#[async_trait::async_trait] +pub trait ToolRanker: Send + Sync { + /// A short stable name for logs and telemetry: `"bm25"`, `"jev"`. + fn kind(&self) -> &'static str; + + /// Ranks `candidates` against `intent`, best first, at most `limit` hits. + /// + /// # Errors + /// + /// Returns [`RankError`] when the ranker cannot produce an answer. A + /// caller should treat every variant as "fall back", and read the variant + /// only to say why in a log line. + async fn rank( + &self, + intent: &str, + context: &RankContext, + candidates: &[RankCandidate], + limit: usize, + ) -> Result, RankError>; +} + +#[async_trait::async_trait] +impl ToolRanker for std::sync::Arc { + fn kind(&self) -> &'static str { + (**self).kind() + } + + async fn rank( + &self, + intent: &str, + context: &RankContext, + candidates: &[RankCandidate], + limit: usize, + ) -> Result, RankError> { + (**self).rank(intent, context, candidates, limit).await + } +} diff --git a/crates/tinytools/src/rank/test.rs b/crates/tinytools/src/rank/test.rs new file mode 100644 index 0000000..25ad1cf --- /dev/null +++ b/crates/tinytools/src/rank/test.rs @@ -0,0 +1,123 @@ +use super::*; + +fn candidates() -> Vec { + vec![ + RankCandidate::new( + "calendar_invite", + "calendar_invite Send a calendar invite to attendees. title start attendees", + ) + .with_family("calendar"), + RankCandidate::new( + "pdf_read", + "pdf_read Read the text of a PDF document. path pages", + ) + .with_family("documents"), + RankCandidate::new( + "stock_quote", + "stock_quote Fetch the latest price for a ticker symbol. symbol", + ) + .with_family("finance"), + ] +} + +#[test] +fn tokenize_splits_identifiers_and_camel_case() { + assert_eq!( + tokenize("memory_hybrid_search readWorkflowResource v2"), + vec![ + "memory", "hybrid", "search", "read", "workflow", "resource", "v2" + ] + ); +} + +#[test] +fn index_ranks_by_description_and_breaks_ties_by_key() { + let index = Bm25Index::build([ + ("b_tool", "send a message"), + ("a_tool", "send a message"), + ("c_tool", "read a file"), + ]); + assert_eq!(index.len(), 3); + assert_eq!(index.search("send message", 5), vec![1, 0]); +} + +#[test] +fn index_on_one_document_corpus_still_finds_it() { + let index = Bm25Index::build([("only", "fetch the latest price for a ticker symbol")]); + assert_eq!(index.search("ticker price", 5), vec![0]); + assert!(index.search("calendar", 5).is_empty()); +} + +#[test] +fn index_returns_nothing_for_stopword_only_queries() { + let index = Bm25Index::build([("a", "send a message"), ("b", "read a file")]); + assert!(index.search("the a of", 5).is_empty()); + assert!(Bm25Index::default().search("anything", 5).is_empty()); + assert!(Bm25Index::default().is_empty()); +} + +#[test] +fn bm25_ranker_returns_positive_hits_only_best_first() { + let hits = Bm25Ranker::rank_sync(&candidates(), "read the text of a pdf", 5); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].key, "pdf_read"); + assert!(hits[0].score > 0.0); + assert!(hits[0].confidence.is_none()); +} + +#[test] +fn bm25_ranker_matches_on_family() { + let hits = Bm25Ranker::rank_sync(&candidates(), "finance", 5); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].key, "stock_quote"); +} + +#[test] +fn bm25_ranker_honours_limit() { + let hits = Bm25Ranker::rank_sync(&candidates(), "send read fetch", 2); + assert_eq!(hits.len(), 2); +} + +#[tokio::test] +async fn bm25_ranker_rejects_empty_intent() { + let err = Bm25Ranker + .rank(" ", &RankContext::empty(), &candidates(), 3) + .await + .err() + .map(|e| e.to_string()); + assert_eq!( + err.as_deref(), + Some("invalid ranking input: intent is empty") + ); +} + +#[tokio::test] +async fn ranker_is_object_safe_behind_an_arc() { + let ranker: std::sync::Arc = std::sync::Arc::new(Bm25Ranker); + assert_eq!(ranker.kind(), "bm25"); + let hits = ranker + .rank("calendar invite", &RankContext::empty(), &candidates(), 3) + .await + .unwrap_or_default(); + assert_eq!( + hits.first().map(|hit| hit.key.as_str()), + Some("calendar_invite") + ); +} + +#[test] +fn rank_error_displays_without_credentials() { + let err = RankError::Backend { + reason: "status 401".to_owned(), + }; + assert_eq!(err.to_string(), "ranker backend failed: status 401"); + assert_eq!(RankError::Timeout.to_string(), "ranker timed out"); + assert_eq!( + RankError::backend("offline").to_string(), + "ranker backend failed: offline" + ); + assert_eq!( + RankError::invalid_input("bad key").to_string(), + "invalid ranking input: bad key" + ); +} diff --git a/crates/tinytools/src/rank/types.rs b/crates/tinytools/src/rank/types.rs new file mode 100644 index 0000000..352ad27 --- /dev/null +++ b/crates/tinytools/src/rank/types.rs @@ -0,0 +1,146 @@ +//! The vocabulary of tool ranking: what a ranker is given and what it returns. + +use std::fmt; + +/// One tool as a ranker sees it: an opaque key the caller maps back to the +/// tool, an optional family (the pack, toolkit, or server it belongs to), and +/// a short summary — the text a ranker reads. +/// +/// The summary is what the ranker judges *on*. Callers should keep it to the +/// tool's name, its first sentence or two, and its argument names: a decision +/// model degrades as its input fills with content unrelated to the choice. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RankCandidate { + /// Caller-owned identifier, returned verbatim in [`RankHit::key`]. Must be + /// unique within one `rank` call. + pub key: String, + /// The group this tool belongs to, when the caller has one — a toolpack, + /// a connector toolkit, an MCP server. Used as ranking text and echoed + /// back so a caller can render it; a ranker never requires it. + pub family: Option, + /// Name plus a one-line description plus argument names: the searchable + /// text. + pub summary: String, +} + +impl RankCandidate { + /// A candidate with no family. + #[must_use] + pub fn new(key: impl Into, summary: impl Into) -> Self { + Self { + key: key.into(), + family: None, + summary: summary.into(), + } + } + + /// Sets the family. + #[must_use] + pub fn with_family(mut self, family: impl Into) -> Self { + self.family = Some(family.into()); + self + } +} + +/// One ranked tool. +#[derive(Clone, Debug, PartialEq)] +pub struct RankHit { + /// The [`RankCandidate::key`] this hit refers to. + pub key: String, + /// Ranker-specific relevance, higher is better. Comparable only within one + /// `rank` call and one ranker: BM25 scores and calibrated probabilities do + /// not share a scale. + pub score: f64, + /// A calibrated probability that this is the right tool, in `0.0..=1.0`, + /// when the ranker produces one. Lexical rankers return `None`; a caller + /// gating on confidence treats `None` as "unknown", not as zero. + pub confidence: Option, +} + +impl RankHit { + /// A hit with no confidence estimate. + #[must_use] + pub fn new(key: impl Into, score: f64) -> Self { + Self { + key: key.into(), + score, + confidence: None, + } + } +} + +/// Context a ranker may use beyond the intent itself. +/// +/// Deliberately small. Rankers that consult a decision model pay for every +/// byte of this on every search, and unrelated context is a distractor, not +/// a help. The intent is always passed separately and is always what the +/// candidates are judged against. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RankContext { + /// The most recent user turns, oldest first, each already clipped by the + /// caller. Lets a ranker resolve "do it again for Bob" against the turn + /// before. Empty when the caller has nothing worth adding. + pub recent_turns: Vec, +} + +impl RankContext { + /// A context with nothing beyond the intent. + #[must_use] + pub fn empty() -> Self { + Self::default() + } +} + +/// Why a ranker could not rank. +/// +/// A caller treats every variant the same way — fall back to a cheaper ranker +/// or to the candidates' declared order — so the variants exist for logs, not +/// for control flow. +#[derive(Debug)] +#[non_exhaustive] +pub enum RankError { + /// The ranker's backing service refused or failed the request. The + /// message never carries a credential. + Backend { + /// Stable, log-safe description. + reason: String, + }, + /// The request could not be built: too many candidates for the ranker, + /// an empty intent, a duplicate key. + InvalidInput { + /// Stable description of the rejected input. + reason: String, + }, + /// The ranker did not answer within its deadline. + Timeout, +} + +impl RankError { + /// Creates a backend failure with a log-safe description. + #[must_use] + pub fn backend(reason: impl Into) -> Self { + Self::Backend { + reason: reason.into(), + } + } + + /// Creates an invalid-input failure. + #[must_use] + pub fn invalid_input(reason: impl Into) -> Self { + Self::InvalidInput { + reason: reason.into(), + } + } +} + +impl fmt::Display for RankError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Backend { reason } => write!(f, "ranker backend failed: {reason}"), + Self::InvalidInput { reason } => write!(f, "invalid ranking input: {reason}"), + Self::Timeout => f.write_str("ranker timed out"), + } + } +} + +impl std::error::Error for RankError {} diff --git a/crates/tinytools/src/tool/types.rs b/crates/tinytools/src/tool/types.rs index 9185fc8..cce2f4f 100644 --- a/crates/tinytools/src/tool/types.rs +++ b/crates/tinytools/src/tool/types.rs @@ -189,6 +189,17 @@ pub trait Tool: Send + Sync { ToolExposure::Direct } + /// The group this tool belongs to, when it has one: a toolpack, a + /// connector toolkit, an MCP server. + /// + /// Read by a tool-search index so a model can find "the Slack one" and so + /// a hit can say where it came from. Purely descriptive — a host neither + /// gates nor routes on it. Most tools have no family and keep the + /// default. + fn family(&self) -> Option<&str> { + None + } + /// Whether two concurrent invocations are safe to run in parallel within a /// single model turn. ///