mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix: B5 — TOIN observation-only refactor + per-tenant aggregation key
Retire the request-time hint API. PR-B5 splits TOIN into two phases:
1. Observation: TOIN keeps recording compressions/retrievals at runtime,
but `get_recommendation()` is deprecated and now returns None.
2. Publish-then-load: the new `headroom.cli.toin_publish` CLI walks the
on-disk store and emits `recommendations.toml`. The Rust proxy reads
that file once at startup via `transforms::recommendations` and
exposes `get(auth_mode, model, structure_hash) -> Option<&Rec>`.
PR-F3 will wire the loader into the live-zone dispatcher.
Per-tenant aggregation: `_patterns` is now keyed by
`(auth_mode, model_family, sig_hash)` so PAYG/OAuth/subscription tenants
no longer share buckets. Callers that don't supply auth/model land in the
`("unknown", "unknown", sig_hash)` slot. Added `_make_pattern_key` helper
+ updated tests that previously indexed by raw `structure_hash`.
AuthMode is canonical in `transforms::live_zone`; `transforms::recommendations`
re-exports it (no duplicate enum). Live-zone enum gained `Unknown`,
`as_str()`, and `Hash` derive to serve recommendations callers without a
second source of truth.
Why: per-request hint calls coupled output to mutable TOIN state, breaking
prompt-cache stability across runs (P2-27, P5-56). Pulling advice into a
startup-published TOML keeps per-request output deterministic and lets the
deploy pipeline gate publication independently of proxy uptime.
Per-PR-B5 plan: REALIGNMENT/04-phase-B-live-zone.md.
This commit is contained in:
parent
b3b3feff6f
commit
6819b7e5e5
17 changed files with 1714 additions and 531 deletions
|
|
@ -189,7 +189,11 @@ fn threshold_for(content_type: ContentType) -> usize {
|
|||
/// Authentication mode of the originating request. Passed through to
|
||||
/// the dispatcher so PR-F2 can vary policy without re-shaping the
|
||||
/// public API. PR-B3 ignores the value (always treated as `Payg`).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
///
|
||||
/// Also reused by [`super::recommendations`] (PR-B5) as the lookup
|
||||
/// key prefix — keeping one canonical enum avoids drift between the
|
||||
/// dispatcher's auth slice and the published recommendations'.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum AuthMode {
|
||||
/// Pay-as-you-go API key. Most aggressive compression budget —
|
||||
/// every saved token is real money for the customer.
|
||||
|
|
@ -203,6 +207,23 @@ pub enum AuthMode {
|
|||
/// compression is less compelling and may interact badly with
|
||||
/// rate-limit accounting.
|
||||
Subscription,
|
||||
/// Auth slice not yet detected. Matches the Python TOIN publish
|
||||
/// CLI's "unknown" default. Used by the recommendations loader
|
||||
/// (PR-B5) when an aggregation row didn't carry an auth tag.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl AuthMode {
|
||||
/// String form used as the recommendations-store lookup key.
|
||||
/// Mirrors the Python publish CLI tag values.
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
AuthMode::Payg => "payg",
|
||||
AuthMode::OAuth => "oauth",
|
||||
AuthMode::Subscription => "subscription",
|
||||
AuthMode::Unknown => "unknown",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-block decision recorded for observability. Independent of
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ pub mod live_zone;
|
|||
pub mod log_compressor;
|
||||
pub mod magika_detector;
|
||||
pub mod pipeline;
|
||||
pub mod recommendations;
|
||||
pub mod safety;
|
||||
pub mod search_compressor;
|
||||
pub mod smart_crusher;
|
||||
|
|
@ -51,6 +52,7 @@ pub use pipeline::{
|
|||
JsonMinifier, JsonOffload, LogOffload, LogTemplate, OffloadOutput, OffloadTransform,
|
||||
PipelineConfig, PipelineResult, ReformatOutput, ReformatTransform, TransformError,
|
||||
};
|
||||
pub use recommendations::{Recommendation, RecommendationStore, RECOMMENDATIONS_PATH_ENV_VAR};
|
||||
pub use safety::{tool_pair_indices, ToolPair};
|
||||
pub use search_compressor::{
|
||||
FileMatches, SearchCompressionResult, SearchCompressor, SearchCompressorConfig,
|
||||
|
|
|
|||
329
crates/headroom-core/src/transforms/recommendations.rs
Normal file
329
crates/headroom-core/src/transforms/recommendations.rs
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
//! Startup-time loader for `recommendations.toml` (PR-B5).
|
||||
//!
|
||||
//! # Why this module exists
|
||||
//!
|
||||
//! Pre-PR-B5, the live-zone dispatcher could call back into Python's
|
||||
//! TOIN per request to get a [`CompressionHint`]. That coupling made
|
||||
//! per-request output non-deterministic — same input could compress
|
||||
//! differently across runs depending on TOIN's mutable state — which
|
||||
//! broke prompt caching (P2-27, P5-56). PR-B5 retired the request-time
|
||||
//! hint API; recommendations now flow through this loader at startup:
|
||||
//!
|
||||
//! 1. The Python `headroom.cli.toin_publish` CLI walks the on-disk TOIN
|
||||
//! store and emits `recommendations.toml`.
|
||||
//! 2. The deploy pipeline ships that TOML alongside the Rust binary.
|
||||
//! 3. At startup, [`RecommendationStore::load_default`] reads the file
|
||||
//! once and exposes the recommendations via a process-wide
|
||||
//! [`OnceLock`].
|
||||
//! 4. [`get`] / [`RecommendationStore::lookup`] return the row matching
|
||||
//! `(auth_mode, model_family, structure_hash)`, or `None` when no
|
||||
//! advice was published. The dispatcher (PR-B3's
|
||||
//! `dispatch_compressor`) does **not** consume this surface yet —
|
||||
//! PR-F3 is responsible for wiring it.
|
||||
//!
|
||||
//! # File schema
|
||||
//!
|
||||
//! ```toml
|
||||
//! [[recommendation]]
|
||||
//! auth_mode = "payg"
|
||||
//! model_family = "claude-3-5"
|
||||
//! structure_hash = "deadbeef..."
|
||||
//! strategy_hint = "smart_crusher"
|
||||
//! confidence = 0.87
|
||||
//! observations = 142
|
||||
//! ```
|
||||
//!
|
||||
//! # Failure modes (loud, never silent)
|
||||
//!
|
||||
//! Per project memory `feedback_no_silent_fallbacks.md`: a missing or
|
||||
//! malformed file degrades to "no advice, use static defaults" — but
|
||||
//! the load attempt always logs a structured `tracing::warn!` event.
|
||||
//! Production deployments grep for `event=recommendations_load_failed`
|
||||
//! to catch a broken publish pipeline.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
/// Environment variable that overrides the default `recommendations.toml`
|
||||
/// path. The Rust proxy reads it once at startup; runtime changes do
|
||||
/// not propagate.
|
||||
pub const RECOMMENDATIONS_PATH_ENV_VAR: &str = "HEADROOM_RECOMMENDATIONS_PATH";
|
||||
|
||||
/// Default file the proxy looks at when the env var is unset.
|
||||
const DEFAULT_RECOMMENDATIONS_PATH: &str = "./recommendations.toml";
|
||||
|
||||
// AuthMode is the canonical enum from `super::live_zone` (PR-B3).
|
||||
// Re-exported here so recommendations callers can import it via
|
||||
// `transforms::recommendations::AuthMode` without crossing module
|
||||
// boundaries; the underlying enum is shared with the live-zone
|
||||
// dispatcher to avoid drift between the dispatcher's auth slice and
|
||||
// the published recommendations'. PR-B5 originally introduced its
|
||||
// own copy; merged into the live-zone enum during integration so
|
||||
// there's only one source of truth.
|
||||
pub use super::live_zone::AuthMode;
|
||||
|
||||
/// A single published recommendation row.
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq)]
|
||||
pub struct Recommendation {
|
||||
pub auth_mode: String,
|
||||
pub model_family: String,
|
||||
pub structure_hash: String,
|
||||
pub strategy_hint: String,
|
||||
pub confidence: f64,
|
||||
pub observations: u64,
|
||||
}
|
||||
|
||||
/// Top-level TOML envelope: `[[recommendation]]` array.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
struct RecommendationFile {
|
||||
#[serde(default)]
|
||||
recommendation: Vec<Recommendation>,
|
||||
}
|
||||
|
||||
/// In-memory recommendation index, keyed by
|
||||
/// `(auth_mode, model_family, structure_hash)`.
|
||||
///
|
||||
/// The keys are owned `String`s rather than `&str` — the values come
|
||||
/// from `toml::from_str`, which gives us `String`s anyway, and trying
|
||||
/// to borrow into the original buffer would require self-referential
|
||||
/// storage. Recommendation files are small (≪ 1 MB even at large
|
||||
/// fleets), so the allocation cost is irrelevant compared to the
|
||||
/// parsing cost we already paid.
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct RecommendationStore {
|
||||
by_key: HashMap<(String, String, String), Recommendation>,
|
||||
}
|
||||
|
||||
impl RecommendationStore {
|
||||
/// Build an empty store. Used for tests and as the fallback when
|
||||
/// no `recommendations.toml` is present.
|
||||
pub fn empty() -> Self {
|
||||
Self {
|
||||
by_key: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of indexed rows.
|
||||
pub fn len(&self) -> usize {
|
||||
self.by_key.len()
|
||||
}
|
||||
|
||||
/// Whether this store has zero recommendations.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.by_key.is_empty()
|
||||
}
|
||||
|
||||
/// Look up a recommendation by tenant slice + structure hash.
|
||||
/// Returns `None` when no advice was published for that key.
|
||||
pub fn lookup(
|
||||
&self,
|
||||
auth_mode: AuthMode,
|
||||
model_family: &str,
|
||||
structure_hash: &str,
|
||||
) -> Option<&Recommendation> {
|
||||
// HashMap::get on a tuple key requires `Borrow` on tuples,
|
||||
// which Rust doesn't provide for mixed `&str`/`String` tuples.
|
||||
// Allocate a short-lived owned key — recommendation lookups
|
||||
// happen once per request at most, so this isn't hot.
|
||||
let key = (
|
||||
auth_mode.as_str().to_string(),
|
||||
model_family.to_string(),
|
||||
structure_hash.to_string(),
|
||||
);
|
||||
self.by_key.get(&key)
|
||||
}
|
||||
|
||||
/// Parse a TOML string into a [`RecommendationStore`].
|
||||
pub fn from_toml_str(s: &str) -> Result<Self, RecommendationsError> {
|
||||
let parsed: RecommendationFile = toml::from_str(s).map_err(RecommendationsError::Parse)?;
|
||||
let mut by_key = HashMap::with_capacity(parsed.recommendation.len());
|
||||
for row in parsed.recommendation {
|
||||
let key = (
|
||||
row.auth_mode.clone(),
|
||||
row.model_family.clone(),
|
||||
row.structure_hash.clone(),
|
||||
);
|
||||
by_key.insert(key, row);
|
||||
}
|
||||
Ok(Self { by_key })
|
||||
}
|
||||
|
||||
/// Read a TOML file from disk and parse it.
|
||||
///
|
||||
/// Missing files yield [`RecommendationsError::Missing`] —
|
||||
/// callers usually downgrade that to "use defaults" without
|
||||
/// panicking. Malformed files surface [`RecommendationsError::Parse`].
|
||||
pub fn from_file(path: impl AsRef<Path>) -> Result<Self, RecommendationsError> {
|
||||
let path = path.as_ref();
|
||||
let text = std::fs::read_to_string(path).map_err(|e| {
|
||||
if e.kind() == std::io::ErrorKind::NotFound {
|
||||
RecommendationsError::Missing(path.to_path_buf())
|
||||
} else {
|
||||
RecommendationsError::Io {
|
||||
path: path.to_path_buf(),
|
||||
source: e,
|
||||
}
|
||||
}
|
||||
})?;
|
||||
Self::from_toml_str(&text)
|
||||
}
|
||||
|
||||
/// Best-effort load: returns an empty store and logs structured
|
||||
/// warnings when the file is missing or malformed. This is the
|
||||
/// path the Rust proxy uses at startup — it's not fatal for the
|
||||
/// publish pipeline to be down.
|
||||
pub fn load_or_empty(path: impl AsRef<Path>) -> Self {
|
||||
let path = path.as_ref();
|
||||
match Self::from_file(path) {
|
||||
Ok(store) => {
|
||||
tracing::info!(
|
||||
event = "recommendations_loaded",
|
||||
path = %path.display(),
|
||||
rows = store.len(),
|
||||
"TOIN recommendations loaded",
|
||||
);
|
||||
store
|
||||
}
|
||||
Err(RecommendationsError::Missing(_)) => {
|
||||
tracing::info!(
|
||||
event = "recommendations_missing",
|
||||
path = %path.display(),
|
||||
"no recommendations.toml present; using static defaults",
|
||||
);
|
||||
Self::empty()
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
event = "recommendations_load_failed",
|
||||
path = %path.display(),
|
||||
error = %err,
|
||||
"TOIN recommendations failed to load — falling back to empty store",
|
||||
);
|
||||
Self::empty()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process-wide store populated at first call to [`load_default`].
|
||||
static GLOBAL: OnceLock<RecommendationStore> = OnceLock::new();
|
||||
|
||||
/// Compute the path the loader will read.
|
||||
///
|
||||
/// Honors `HEADROOM_RECOMMENDATIONS_PATH` for prod overrides; falls
|
||||
/// back to [`DEFAULT_RECOMMENDATIONS_PATH`].
|
||||
pub fn default_path() -> PathBuf {
|
||||
std::env::var(RECOMMENDATIONS_PATH_ENV_VAR)
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|_| PathBuf::from(DEFAULT_RECOMMENDATIONS_PATH))
|
||||
}
|
||||
|
||||
/// Initialize and return the global [`RecommendationStore`].
|
||||
///
|
||||
/// On first call, reads the file at [`default_path`]; subsequent calls
|
||||
/// return the cached store. Idempotent and thread-safe via [`OnceLock`].
|
||||
pub fn load_default() -> &'static RecommendationStore {
|
||||
GLOBAL.get_or_init(|| RecommendationStore::load_or_empty(default_path()))
|
||||
}
|
||||
|
||||
/// Module-level convenience: look up a recommendation in the global
|
||||
/// store. PR-F3 will wire this into `dispatch_compressor`. PR-B5 only
|
||||
/// exposes the API surface.
|
||||
pub fn get(
|
||||
auth_mode: AuthMode,
|
||||
model: &str,
|
||||
structure_hash: &str,
|
||||
) -> Option<&'static Recommendation> {
|
||||
load_default().lookup(auth_mode, model, structure_hash)
|
||||
}
|
||||
|
||||
/// Errors surfaced by the loader. Marked non-exhaustive so we can add
|
||||
/// future variants without breaking callers.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[non_exhaustive]
|
||||
pub enum RecommendationsError {
|
||||
/// File doesn't exist on disk.
|
||||
#[error("recommendations file not found: {0}")]
|
||||
Missing(PathBuf),
|
||||
/// Filesystem error other than NotFound.
|
||||
#[error("recommendations IO error at {path}: {source}")]
|
||||
Io {
|
||||
path: PathBuf,
|
||||
#[source]
|
||||
source: std::io::Error,
|
||||
},
|
||||
/// TOML parse failure (typed wrapper for ergonomics).
|
||||
#[error("recommendations TOML parse error: {0}")]
|
||||
Parse(#[from] toml::de::Error),
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn sample_toml() -> &'static str {
|
||||
r#"
|
||||
[[recommendation]]
|
||||
auth_mode = "payg"
|
||||
model_family = "claude-3-5"
|
||||
structure_hash = "deadbeef"
|
||||
strategy_hint = "smart_crusher"
|
||||
confidence = 0.87
|
||||
observations = 142
|
||||
|
||||
[[recommendation]]
|
||||
auth_mode = "oauth"
|
||||
model_family = "gpt-4o"
|
||||
structure_hash = "cafebabe"
|
||||
strategy_hint = "log_compressor"
|
||||
confidence = 0.42
|
||||
observations = 60
|
||||
"#
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_toml_str_indexes_by_tuple_key() {
|
||||
let store = RecommendationStore::from_toml_str(sample_toml()).expect("parses");
|
||||
assert_eq!(store.len(), 2);
|
||||
|
||||
let r = store
|
||||
.lookup(AuthMode::Payg, "claude-3-5", "deadbeef")
|
||||
.expect("hit");
|
||||
assert_eq!(r.strategy_hint, "smart_crusher");
|
||||
assert!((r.confidence - 0.87).abs() < 1e-9);
|
||||
assert_eq!(r.observations, 142);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lookup_returns_none_for_missing_slice() {
|
||||
let store = RecommendationStore::from_toml_str(sample_toml()).expect("parses");
|
||||
assert!(store
|
||||
.lookup(AuthMode::Unknown, "gpt-4o", "cafebabe")
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_store_lookup_is_none() {
|
||||
let store = RecommendationStore::empty();
|
||||
assert!(store.is_empty());
|
||||
assert!(store.lookup(AuthMode::Payg, "claude-3-5", "any").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_toml_yields_parse_error() {
|
||||
let bad = "this is not valid toml [[\n\n";
|
||||
let err = RecommendationStore::from_toml_str(bad).unwrap_err();
|
||||
assert!(matches!(err, RecommendationsError::Parse(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_mode_strings_match_python_publish_cli() {
|
||||
assert_eq!(AuthMode::Payg.as_str(), "payg");
|
||||
assert_eq!(AuthMode::OAuth.as_str(), "oauth");
|
||||
assert_eq!(AuthMode::Subscription.as_str(), "subscription");
|
||||
assert_eq!(AuthMode::Unknown.as_str(), "unknown");
|
||||
}
|
||||
}
|
||||
140
crates/headroom-core/tests/recommendations_loader.rs
Normal file
140
crates/headroom-core/tests/recommendations_loader.rs
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
//! Integration tests for `transforms::recommendations` (PR-B5).
|
||||
//!
|
||||
//! These pin three guarantees the Rust proxy depends on at startup:
|
||||
//!
|
||||
//! 1. A well-formed `recommendations.toml` parses into a populated
|
||||
//! [`RecommendationStore`] with byte-for-byte the same fields the
|
||||
//! Python `headroom.cli.toin_publish` CLI emits.
|
||||
//! 2. A missing file degrades to an empty store with no panic — the
|
||||
//! proxy must boot even if the publish pipeline is broken.
|
||||
//! 3. A malformed file likewise degrades to an empty store, and the
|
||||
//! error is surfaced via `tracing::warn!` rather than swallowed.
|
||||
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use headroom_core::transforms::recommendations::{
|
||||
AuthMode, RecommendationStore, RecommendationsError,
|
||||
};
|
||||
|
||||
/// Minimal tempdir helper. Project convention is to avoid a
|
||||
/// dev-dependency on `tempfile` (see `crates/headroom-parity` for the
|
||||
/// matching helper). Cleanup happens on drop.
|
||||
struct TempDir(PathBuf);
|
||||
|
||||
impl TempDir {
|
||||
fn path(&self) -> &Path {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TempDir {
|
||||
fn drop(&mut self) {
|
||||
let _ = fs::remove_dir_all(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
fn tempdir() -> TempDir {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
let p = std::env::temp_dir().join(format!(
|
||||
"headroom-recommendations-{nanos}-{:?}",
|
||||
std::thread::current().id()
|
||||
));
|
||||
fs::create_dir_all(&p).unwrap();
|
||||
TempDir(p)
|
||||
}
|
||||
|
||||
/// The exact schema emitted by the Python publish CLI. Keeping this
|
||||
/// inline in tests prevents accidental schema drift on either side.
|
||||
const VALID_TOML: &str = r#"
|
||||
[[recommendation]]
|
||||
auth_mode = "payg"
|
||||
model_family = "claude-3-5"
|
||||
structure_hash = "deadbeef00112233"
|
||||
strategy_hint = "smart_crusher"
|
||||
confidence = 0.87
|
||||
observations = 142
|
||||
|
||||
[[recommendation]]
|
||||
auth_mode = "oauth"
|
||||
model_family = "gpt-4o"
|
||||
structure_hash = "cafebabe44556677"
|
||||
strategy_hint = "log_compressor"
|
||||
confidence = 0.42
|
||||
observations = 60
|
||||
"#;
|
||||
|
||||
#[test]
|
||||
fn loads_valid_toml() {
|
||||
let dir = tempdir();
|
||||
let path = dir.path().join("recommendations.toml");
|
||||
fs::write(&path, VALID_TOML).expect("write");
|
||||
|
||||
let store = RecommendationStore::from_file(&path).expect("parses");
|
||||
assert_eq!(store.len(), 2);
|
||||
|
||||
let payg = store
|
||||
.lookup(AuthMode::Payg, "claude-3-5", "deadbeef00112233")
|
||||
.expect("payg row");
|
||||
assert_eq!(payg.strategy_hint, "smart_crusher");
|
||||
assert!((payg.confidence - 0.87).abs() < 1e-9);
|
||||
assert_eq!(payg.observations, 142);
|
||||
|
||||
let oauth = store
|
||||
.lookup(AuthMode::OAuth, "gpt-4o", "cafebabe44556677")
|
||||
.expect("oauth row");
|
||||
assert_eq!(oauth.strategy_hint, "log_compressor");
|
||||
assert_eq!(oauth.observations, 60);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_file_yields_empty_recommendations() {
|
||||
let dir = tempdir();
|
||||
let path = dir.path().join("does_not_exist.toml");
|
||||
|
||||
// `from_file` surfaces Missing as a typed error.
|
||||
let err = RecommendationStore::from_file(&path).unwrap_err();
|
||||
assert!(matches!(err, RecommendationsError::Missing(_)));
|
||||
|
||||
// `load_or_empty` is the production entry point and degrades
|
||||
// gracefully — no panic, empty store.
|
||||
let store = RecommendationStore::load_or_empty(&path);
|
||||
assert!(store.is_empty());
|
||||
assert!(store.lookup(AuthMode::Payg, "claude-3-5", "any").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_toml_logs_and_yields_empty() {
|
||||
let dir = tempdir();
|
||||
let path = dir.path().join("recommendations.toml");
|
||||
// Real-world breakage: missing closing brace + unquoted value.
|
||||
fs::write(&path, "this is [[ definitely not valid\nrubbish = \n").expect("write");
|
||||
|
||||
// The typed loader returns a Parse error.
|
||||
let err = RecommendationStore::from_file(&path).unwrap_err();
|
||||
assert!(
|
||||
matches!(err, RecommendationsError::Parse(_)),
|
||||
"expected Parse, got {err:?}"
|
||||
);
|
||||
|
||||
// The production loader returns an empty store — proxy still
|
||||
// boots; ops alert on the structured `tracing::warn!` event.
|
||||
let store = RecommendationStore::load_or_empty(&path);
|
||||
assert!(store.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_recommendation_array_is_valid() {
|
||||
// A publish run with no eligible slices writes a header-only file.
|
||||
// It must parse cleanly to an empty store, not surface as an error.
|
||||
let dir = tempdir();
|
||||
let path = dir.path().join("recommendations.toml");
|
||||
fs::write(&path, "# Auto-generated by toin_publish\n").expect("write");
|
||||
|
||||
let store = RecommendationStore::from_file(&path).expect("parses");
|
||||
assert!(store.is_empty());
|
||||
}
|
||||
246
headroom/cli/toin_publish.py
Normal file
246
headroom/cli/toin_publish.py
Normal file
|
|
@ -0,0 +1,246 @@
|
|||
"""``python -m headroom.cli.toin_publish`` — emit ``recommendations.toml``.
|
||||
|
||||
PR-B5 retired TOIN's request-time hint API. Recommendations now flow
|
||||
through this offline CLI: we walk the on-disk TOIN store, aggregate one
|
||||
row per ``(auth_mode, model_family, structure_hash)`` slice that has at
|
||||
least ``--min-observations`` recorded compression events, and write the
|
||||
result as TOML the Rust proxy loads at startup
|
||||
(``$HEADROOM_RECOMMENDATIONS_PATH`` / ``./recommendations.toml``).
|
||||
|
||||
# TOML schema
|
||||
|
||||
::
|
||||
|
||||
[[recommendation]]
|
||||
auth_mode = "payg"
|
||||
model_family = "claude-3-5"
|
||||
structure_hash = "deadbeef..."
|
||||
strategy_hint = "smart_crusher"
|
||||
confidence = 0.87
|
||||
observations = 142
|
||||
|
||||
# Why a CLI, not a library hook
|
||||
|
||||
Per-request mutation is exactly the dangerous coupling PR-B5 retired.
|
||||
Publishing happens at deploy boundaries — never inside a request. The
|
||||
deploy pipeline runs ``python -m headroom.cli.toin_publish ...`` after
|
||||
draining the TOIN store from production, ships the resulting TOML
|
||||
alongside the Rust binary, and the proxy reads it once at startup.
|
||||
|
||||
# Output stability
|
||||
|
||||
Rows are sorted by ``(auth_mode, model_family, structure_hash)`` so the
|
||||
file diffs cleanly across publishes. Strategies ship as the
|
||||
ToolPattern's learned ``optimal_strategy`` (or the dominant entry of
|
||||
``strategy_success_rates``). Confidence is the pattern's existing
|
||||
confidence score — bounded ``[0.0, 0.95]`` by the confidence calculator.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Iterable
|
||||
from pathlib import Path
|
||||
from typing import Any, Final
|
||||
|
||||
from headroom.telemetry.toin import (
|
||||
DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH,
|
||||
PatternKey,
|
||||
ToolIntelligenceNetwork,
|
||||
ToolPattern,
|
||||
get_toin,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Header annotated into every emitted file so ops can grep for "where
|
||||
# did this TOML come from?" without having to remember the publish
|
||||
# command. Newlines preserved as literal "\n" inside the string.
|
||||
TOML_HEADER: Final[str] = (
|
||||
"# Auto-generated by `python -m headroom.cli.toin_publish`.\n"
|
||||
"# DO NOT EDIT BY HAND — re-run the publish CLI to regenerate.\n"
|
||||
"# Loaded by the Rust proxy at startup; see\n"
|
||||
"# `crates/headroom-core/src/transforms/recommendations.rs`.\n\n"
|
||||
)
|
||||
|
||||
|
||||
def _select_strategy(pattern: ToolPattern) -> str:
|
||||
"""Pick the strategy hint to publish for a pattern.
|
||||
|
||||
Priority:
|
||||
1. ``optimal_strategy`` if the pattern has explicitly recorded one
|
||||
and it isn't the placeholder ``"default"``.
|
||||
2. The strategy with the highest success rate, if any have been
|
||||
recorded.
|
||||
3. ``"default"`` — meaning "no opinion; the proxy uses its static
|
||||
default strategy."
|
||||
"""
|
||||
if pattern.optimal_strategy and pattern.optimal_strategy != "default":
|
||||
return pattern.optimal_strategy
|
||||
if pattern.strategy_success_rates:
|
||||
return max(pattern.strategy_success_rates.items(), key=lambda kv: kv[1])[0]
|
||||
return "default"
|
||||
|
||||
|
||||
def _toml_escape(s: str) -> str:
|
||||
"""Escape a string for inclusion as a basic TOML string literal."""
|
||||
# TOML basic strings need backslash + quote escapes only; the values
|
||||
# we publish (auth_mode, model_family, hex hashes, identifier
|
||||
# strategies) never contain control characters in practice, but we
|
||||
# encode defensively per RFC.
|
||||
return s.replace("\\", "\\\\").replace('"', '\\"')
|
||||
|
||||
|
||||
def _format_row(
|
||||
*,
|
||||
auth_mode: str,
|
||||
model_family: str,
|
||||
structure_hash: str,
|
||||
strategy_hint: str,
|
||||
confidence: float,
|
||||
observations: int,
|
||||
) -> str:
|
||||
"""Render one ``[[recommendation]]`` block."""
|
||||
return (
|
||||
"[[recommendation]]\n"
|
||||
f'auth_mode = "{_toml_escape(auth_mode)}"\n'
|
||||
f'model_family = "{_toml_escape(model_family)}"\n'
|
||||
f'structure_hash = "{_toml_escape(structure_hash)}"\n'
|
||||
f'strategy_hint = "{_toml_escape(strategy_hint)}"\n'
|
||||
f"confidence = {confidence:.4f}\n"
|
||||
f"observations = {observations}\n"
|
||||
)
|
||||
|
||||
|
||||
def _eligible_rows(
|
||||
patterns: Iterable[tuple[PatternKey, ToolPattern]],
|
||||
min_observations: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Filter and sort pattern slices into deterministic publish rows.
|
||||
|
||||
Observations source: ``ToolPattern.total_compressions`` is the
|
||||
canonical "how many events fed this slice" counter — incremented
|
||||
once per ``record_compression`` call. The legacy ``observations``
|
||||
field counted ``get_recommendation`` invocations (retired in PR-B5)
|
||||
and now stays at zero, so we use ``total_compressions`` instead.
|
||||
"""
|
||||
rows: list[dict[str, Any]] = []
|
||||
for key, pattern in patterns:
|
||||
observations = pattern.total_compressions
|
||||
if observations < min_observations:
|
||||
continue
|
||||
auth_mode, model_family, sig_hash = key
|
||||
rows.append(
|
||||
{
|
||||
"auth_mode": auth_mode,
|
||||
"model_family": model_family,
|
||||
"structure_hash": sig_hash,
|
||||
"strategy_hint": _select_strategy(pattern),
|
||||
"confidence": float(pattern.confidence),
|
||||
"observations": observations,
|
||||
}
|
||||
)
|
||||
|
||||
rows.sort(
|
||||
key=lambda r: (r["auth_mode"], r["model_family"], r["structure_hash"]),
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def publish(
|
||||
*,
|
||||
output_path: Path,
|
||||
min_observations: int = DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH,
|
||||
toin: ToolIntelligenceNetwork | None = None,
|
||||
) -> int:
|
||||
"""Aggregate the live TOIN store and write ``recommendations.toml``.
|
||||
|
||||
Args:
|
||||
output_path: Destination TOML path. Parent directories must
|
||||
already exist.
|
||||
min_observations: Minimum ``total_compressions`` per slice
|
||||
before we emit a row. Anything below this is noise.
|
||||
toin: Optional pre-built TOIN handle (tests pass an isolated
|
||||
instance). Defaults to the global singleton.
|
||||
|
||||
Returns:
|
||||
The number of recommendation rows written.
|
||||
"""
|
||||
handle = toin if toin is not None else get_toin()
|
||||
rows = _eligible_rows(handle.iter_patterns(), min_observations=min_observations)
|
||||
|
||||
body = TOML_HEADER + "\n".join(_format_row(**row) for row in rows)
|
||||
if rows:
|
||||
# Trailing newline keeps POSIX tools happy and round-trips
|
||||
# cleanly through `tomllib.loads`.
|
||||
body += "\n"
|
||||
|
||||
output_path.write_text(body, encoding="utf-8")
|
||||
logger.info(
|
||||
"TOIN publish complete",
|
||||
extra={
|
||||
"event": "toin_publish",
|
||||
"path": str(output_path),
|
||||
"rows": len(rows),
|
||||
"min_observations": min_observations,
|
||||
},
|
||||
)
|
||||
return len(rows)
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="python -m headroom.cli.toin_publish",
|
||||
description=(
|
||||
"Aggregate the on-disk TOIN store and emit recommendations.toml "
|
||||
"for the Rust proxy to load at startup."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=Path,
|
||||
default=Path("recommendations.toml"),
|
||||
help="Output TOML path (default: ./recommendations.toml).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-observations",
|
||||
type=int,
|
||||
default=DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH,
|
||||
help=(
|
||||
"Minimum total_compressions per slice before a row is emitted "
|
||||
f"(default: {DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH})."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
"-v",
|
||||
action="store_true",
|
||||
help="Enable info-level logging on stderr.",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
"""Entry point. Returns the exit code for ``__main__``."""
|
||||
parser = _build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.verbose:
|
||||
logging.basicConfig(level=logging.INFO, stream=sys.stderr)
|
||||
|
||||
if args.min_observations < 1:
|
||||
parser.error("--min-observations must be >= 1")
|
||||
|
||||
output: Path = args.output
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows = publish(output_path=output, min_observations=args.min_observations)
|
||||
print(f"wrote {rows} recommendation rows to {output}", file=sys.stderr)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover - module entrypoint
|
||||
raise SystemExit(main())
|
||||
|
|
@ -32,19 +32,21 @@ Usage:
|
|||
# Export for aggregation
|
||||
stats = collector.export_stats()
|
||||
|
||||
TOIN (Tool Output Intelligence Network):
|
||||
TOIN (Tool Output Intelligence Network) — observation-only since PR-B5:
|
||||
from headroom.telemetry import get_toin
|
||||
|
||||
toin = get_toin()
|
||||
|
||||
# Get compression hints before compressing
|
||||
hint = toin.get_recommendation(tool_signature, query_context)
|
||||
|
||||
# Record compression outcome
|
||||
# Record compression outcome (the only request-time TOIN call).
|
||||
toin.record_compression(tool_signature, ...)
|
||||
|
||||
# Record retrieval (automatic via compression_store)
|
||||
# Record retrieval (automatic via compression_store).
|
||||
toin.record_retrieval(sig_hash, retrieval_type, query, query_fields)
|
||||
|
||||
# Aggregated recommendations are emitted offline:
|
||||
# python -m headroom.cli.toin_publish --output recommendations.toml
|
||||
# The Rust proxy loads that TOML at startup; there is no
|
||||
# request-time hint API.
|
||||
"""
|
||||
|
||||
from .beacon import (
|
||||
|
|
@ -66,7 +68,9 @@ from .models import (
|
|||
ToolSignature,
|
||||
)
|
||||
from .toin import (
|
||||
CompressionHint,
|
||||
DEFAULT_AUTH_MODE,
|
||||
DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH,
|
||||
DEFAULT_MODEL_FAMILY,
|
||||
TOINConfig,
|
||||
ToolIntelligenceNetwork,
|
||||
ToolPattern,
|
||||
|
|
@ -90,8 +94,10 @@ __all__ = [
|
|||
"FieldDistribution",
|
||||
"RetrievalStats",
|
||||
"ToolSignature",
|
||||
# TOIN
|
||||
"CompressionHint",
|
||||
# TOIN (observation-only since PR-B5)
|
||||
"DEFAULT_AUTH_MODE",
|
||||
"DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH",
|
||||
"DEFAULT_MODEL_FAMILY",
|
||||
"TOINConfig",
|
||||
"ToolIntelligenceNetwork",
|
||||
"ToolPattern",
|
||||
|
|
|
|||
|
|
@ -1,43 +1,61 @@
|
|||
"""Tool Output Intelligence Network (TOIN) - Cross-user learning for compression.
|
||||
"""Tool Output Intelligence Network (TOIN) — observation-only contract.
|
||||
|
||||
TOIN aggregates anonymized compression patterns across all Headroom users to
|
||||
create a network effect: every user's compression decisions improve the
|
||||
recommendations for everyone.
|
||||
# Observation-only contract (PR-B5)
|
||||
|
||||
Key concepts:
|
||||
- ToolPattern: Aggregated intelligence about a tool type (by structure hash)
|
||||
- CompressionHint: Recommendations for how to compress a specific tool output
|
||||
- ToolIntelligenceNetwork: Central aggregator that learns from all users
|
||||
TOIN observes; it never mutates request-time compression decisions. The
|
||||
request path is deterministic: SmartCrusher and the live-zone dispatcher
|
||||
read their static configuration only. TOIN's role is to record what
|
||||
happened so an offline aggregator (`headroom.cli.toin_publish`) can emit
|
||||
a `recommendations.toml` file the deploy pipeline ships to the proxy at
|
||||
the next restart.
|
||||
|
||||
How it works:
|
||||
1. When SmartCrusher compresses data, it records the outcome via telemetry
|
||||
2. When LLM retrieves compressed data, TOIN tracks what was needed
|
||||
3. TOIN learns: "For tools with structure X, retrieval rate is high when
|
||||
compressing field Y - preserve it"
|
||||
4. Next time: SmartCrusher asks TOIN for hints before compressing
|
||||
Why this shape:
|
||||
- Per-request mutation tied compression bytes to TOIN's mutable state,
|
||||
which made the same input produce different outputs across runs (P2-27,
|
||||
P5-56). That broke prompt caching and made bugs irreproducible.
|
||||
- The request-time hint API (`get_recommendation()`) is retired. It now
|
||||
emits a `DeprecationWarning` and returns `None`. New code must not call
|
||||
it.
|
||||
- Recording (`record_compression`, `record_retrieval`) and storage
|
||||
(save/load/export/import) are unchanged; the learning value is intact.
|
||||
|
||||
Privacy:
|
||||
- No actual data values are stored
|
||||
- Tool names are structure hashes
|
||||
- Field names are SHA256[:8] hashes
|
||||
- No user identifiers
|
||||
# Aggregation key
|
||||
|
||||
Network Effect:
|
||||
- More users → more compression events → better recommendations
|
||||
- Cross-user patterns reveal universal tool behaviors
|
||||
- Federated learning: aggregate patterns, not data
|
||||
Patterns are keyed by `(auth_mode, model_family, structure_hash)` —
|
||||
each tenant slice (PAYG vs OAuth vs subscription) and each model family
|
||||
(claude-3-5, gpt-4o, …) learns independently. Defaults `"unknown"` when
|
||||
either is not yet plumbed through (PR-F3 lights up real auth-mode
|
||||
detection).
|
||||
|
||||
Usage:
|
||||
# Privacy
|
||||
- No actual data values are stored.
|
||||
- Tool names are structure hashes.
|
||||
- Field names are SHA256[:8] hashes.
|
||||
- No user identifiers.
|
||||
|
||||
# Network effect (preserved)
|
||||
- More users → more compression events → better aggregated `optimal_*`
|
||||
fields on each `ToolPattern`. The `toin publish` CLI promotes those
|
||||
into `recommendations.toml`.
|
||||
- Cross-instance pattern import (`import_patterns`) supports federated
|
||||
learning without sharing actual data.
|
||||
|
||||
# Usage
|
||||
from headroom.telemetry.toin import get_toin
|
||||
|
||||
# Before compression, get recommendations
|
||||
hint = get_toin().get_recommendation(tool_signature, query_context)
|
||||
# Record a compression event (the only request-time TOIN call).
|
||||
get_toin().record_compression(
|
||||
tool_signature=signature,
|
||||
original_count=len(items),
|
||||
compressed_count=kept,
|
||||
original_tokens=before,
|
||||
compressed_tokens=after,
|
||||
strategy="smart_crusher",
|
||||
)
|
||||
|
||||
# Apply hint
|
||||
if hint.skip_compression:
|
||||
return original_data
|
||||
config.preserve_fields = hint.preserve_fields
|
||||
config.max_items = hint.max_items
|
||||
# Aggregated recommendations are produced offline:
|
||||
# python -m headroom.cli.toin_publish --output recommendations.toml
|
||||
# The Rust proxy loads that file at startup; no per-request hint API.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -48,9 +66,10 @@ import logging
|
|||
import os
|
||||
import threading
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
from .models import FieldSemantics, ToolSignature
|
||||
|
||||
|
|
@ -63,6 +82,72 @@ TOIN_PATH_ENV_VAR = "HEADROOM_TOIN_PATH"
|
|||
DEFAULT_TOIN_DIR = ".headroom"
|
||||
DEFAULT_TOIN_FILE = "toin.json"
|
||||
|
||||
# ── Aggregation-key defaults ────────────────────────────────────────────
|
||||
# Used when callers haven't plumbed auth-mode / model-family detection
|
||||
# (PR-F3 wires the real detectors). Shipping a real `"unknown"` slice is
|
||||
# explicit — better than a magic empty string and lets the publish CLI
|
||||
# filter on it deliberately.
|
||||
DEFAULT_AUTH_MODE: Final[str] = "unknown"
|
||||
DEFAULT_MODEL_FAMILY: Final[str] = "unknown"
|
||||
|
||||
# ── Aggregation thresholds (Final, not magic numbers) ───────────────────
|
||||
# Minimum observations a pattern must have before the `toin publish` CLI
|
||||
# emits a recommendation row for it. Below this, the recommendation
|
||||
# would be noise. The CLI exposes `--min-observations` to override per
|
||||
# environment; this is the production default the Rust proxy expects.
|
||||
DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH: Final[int] = 50
|
||||
|
||||
# Aggregation-key serialization separator. Used to encode the
|
||||
# `(auth_mode, model_family, sig_hash)` tuple as a string for JSON
|
||||
# storage (JSON object keys must be strings) and for cross-instance
|
||||
# pattern imports. Pipe is illegal in all three components by
|
||||
# construction (auth_mode ∈ {"unknown","payg","oauth","subscription"};
|
||||
# model_family is a registry name with no `|`; sig_hash is hex).
|
||||
_AGG_KEY_SEPARATOR: Final[str] = "|"
|
||||
|
||||
|
||||
# ── Aggregation key helpers ─────────────────────────────────────────────
|
||||
PatternKey = tuple[str, str, str]
|
||||
|
||||
|
||||
def _make_pattern_key(
|
||||
auth_mode: str | None,
|
||||
model_family: str | None,
|
||||
sig_hash: str,
|
||||
) -> PatternKey:
|
||||
"""Build the canonical `(auth_mode, model_family, sig_hash)` key.
|
||||
|
||||
Defaults populate to `DEFAULT_AUTH_MODE` / `DEFAULT_MODEL_FAMILY`
|
||||
when callers haven't supplied a value — keeps callers terse during
|
||||
the Phase B realignment while PR-F3 wires real detectors.
|
||||
"""
|
||||
return (
|
||||
auth_mode or DEFAULT_AUTH_MODE,
|
||||
model_family or DEFAULT_MODEL_FAMILY,
|
||||
sig_hash,
|
||||
)
|
||||
|
||||
|
||||
def _serialize_pattern_key(key: PatternKey) -> str:
|
||||
"""Serialize an aggregation key to a string for JSON / TOML storage."""
|
||||
return _AGG_KEY_SEPARATOR.join(key)
|
||||
|
||||
|
||||
def _deserialize_pattern_key(serialized: str) -> PatternKey:
|
||||
"""Parse a serialized aggregation key back to a tuple.
|
||||
|
||||
Backward-compatible with pre-B5 dumps that stored keys as bare
|
||||
structure hashes (no separator): those parse as
|
||||
`(DEFAULT_AUTH_MODE, DEFAULT_MODEL_FAMILY, sig_hash)`. The realignment
|
||||
plan permits wiping the on-disk store, but this fallback keeps reads
|
||||
safe if a stale file appears in the wild.
|
||||
"""
|
||||
parts = serialized.split(_AGG_KEY_SEPARATOR)
|
||||
if len(parts) == 3:
|
||||
return (parts[0], parts[1], parts[2])
|
||||
# Legacy format: bare sig_hash. Promote to default tenant slice.
|
||||
return (DEFAULT_AUTH_MODE, DEFAULT_MODEL_FAMILY, serialized)
|
||||
|
||||
|
||||
def get_default_toin_storage_path() -> str:
|
||||
"""Get the default TOIN storage path.
|
||||
|
|
@ -102,6 +187,15 @@ class ToolPattern:
|
|||
|
||||
tool_signature_hash: str
|
||||
|
||||
# === Aggregation Key (PR-B5) ===
|
||||
# Per-tenant aggregation key extension. The Pattern is keyed inside
|
||||
# the TOIN store by `(auth_mode, model_family, tool_signature_hash)` —
|
||||
# these two fields carry the same values onto the dataclass so dumps,
|
||||
# imports, and publish-CLI rows are self-describing without
|
||||
# cross-referencing the dict key.
|
||||
auth_mode: str = DEFAULT_AUTH_MODE
|
||||
model_family: str = DEFAULT_MODEL_FAMILY
|
||||
|
||||
# === Compression Statistics ===
|
||||
total_compressions: int = 0
|
||||
total_items_seen: int = 0
|
||||
|
|
@ -153,7 +247,10 @@ class ToolPattern:
|
|||
field_semantics: dict[str, FieldSemantics] = field(default_factory=dict)
|
||||
|
||||
# === Observation Counter ===
|
||||
observations: int = 0 # How many times get_recommendation() was called for this pattern
|
||||
# PR-B5: legacy counter from the retired `get_recommendation()` API.
|
||||
# Held for serialization compatibility with v1.0 dumps; new
|
||||
# increments only happen via record_compression / record_retrieval.
|
||||
observations: int = 0
|
||||
|
||||
# === Confidence ===
|
||||
sample_size: int = 0
|
||||
|
|
@ -183,6 +280,8 @@ class ToolPattern:
|
|||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"tool_signature_hash": self.tool_signature_hash,
|
||||
"auth_mode": self.auth_mode,
|
||||
"model_family": self.model_family,
|
||||
"total_compressions": self.total_compressions,
|
||||
"total_items_seen": self.total_items_seen,
|
||||
"total_items_kept": self.total_items_kept,
|
||||
|
|
@ -226,6 +325,8 @@ class ToolPattern:
|
|||
# Filter to only valid fields
|
||||
valid_fields = {
|
||||
"tool_signature_hash",
|
||||
"auth_mode",
|
||||
"model_family",
|
||||
"total_compressions",
|
||||
"total_items_seen",
|
||||
"total_items_kept",
|
||||
|
|
@ -280,10 +381,15 @@ class ToolPattern:
|
|||
|
||||
|
||||
@dataclass
|
||||
class CompressionHint:
|
||||
"""Recommendation for how to compress a specific tool output.
|
||||
class _CompressionHint:
|
||||
"""Internal recommendation envelope (PR-B5: private, observation-only).
|
||||
|
||||
This is what TOIN returns when asked for advice before compression.
|
||||
Pre-B5 this was the public return type of `get_recommendation()` —
|
||||
the request-time hint API now retired. The dataclass is retained as
|
||||
`_CompressionHint` purely for the deprecated stub's signature and
|
||||
for the publish CLI's internal aggregation; no new code should
|
||||
construct or consume it. Read recommendations from
|
||||
`recommendations.toml` produced by `headroom.cli.toin_publish`.
|
||||
"""
|
||||
|
||||
# Should we compress at all?
|
||||
|
|
@ -345,15 +451,24 @@ class TOINConfig:
|
|||
|
||||
|
||||
class ToolIntelligenceNetwork:
|
||||
"""Aggregates tool patterns across all Headroom users.
|
||||
"""Aggregates tool patterns across all Headroom users (observation-only).
|
||||
|
||||
This is the brain of TOIN. It maintains a database of learned patterns
|
||||
for different tool types and provides recommendations based on
|
||||
cross-user intelligence.
|
||||
This is the offline brain of TOIN. It maintains a database of learned
|
||||
patterns for different `(auth_mode, model_family, tool_signature)`
|
||||
slices. The `record_compression` / `record_retrieval` calls are the
|
||||
only request-time API; aggregated recommendations are emitted by
|
||||
`headroom.cli.toin_publish` and consumed by the Rust proxy at startup.
|
||||
|
||||
Thread-safe for concurrent access.
|
||||
"""
|
||||
|
||||
# ── Deprecation warning de-dupe (PR-B5) ───────────────────────────────
|
||||
# `get_recommendation` is retired as a per-request mutator. We emit
|
||||
# `DeprecationWarning` once per process; if every call warned, busy
|
||||
# call sites would flood logs and obscure other warnings. Class-level
|
||||
# so all instances share the flag.
|
||||
_DEPRECATION_WARNED: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: TOINConfig | None = None,
|
||||
|
|
@ -380,8 +495,11 @@ class ToolIntelligenceNetwork:
|
|||
else:
|
||||
self._backend = None
|
||||
|
||||
# Pattern database: structure_hash -> ToolPattern
|
||||
self._patterns: dict[str, ToolPattern] = {}
|
||||
# Pattern database: (auth_mode, model_family, structure_hash) -> ToolPattern
|
||||
# PR-B5 extended the key from a bare structure_hash to the per-tenant
|
||||
# tuple. The serialized form on disk encodes the tuple as
|
||||
# "auth|model|hash"; see `_serialize_pattern_key`.
|
||||
self._patterns: dict[PatternKey, ToolPattern] = {}
|
||||
|
||||
# Instance ID for user counting (anonymized)
|
||||
# IMPORTANT: Must be STABLE across restarts to avoid false user count inflation
|
||||
|
|
@ -449,11 +567,13 @@ class ToolIntelligenceNetwork:
|
|||
strategy: str,
|
||||
query_context: str | None = None,
|
||||
items: list[dict[str, Any]] | None = None,
|
||||
auth_mode: str | None = None,
|
||||
model_family: str | None = None,
|
||||
) -> None:
|
||||
"""Record a compression event.
|
||||
|
||||
Called after SmartCrusher compresses data. Updates the pattern
|
||||
for this tool type.
|
||||
for this `(auth_mode, model_family, tool_signature)` slice.
|
||||
|
||||
TOIN Evolution: When items are provided, we capture field statistics
|
||||
for learning semantic types (uniqueness, default values, etc.).
|
||||
|
|
@ -467,6 +587,10 @@ class ToolIntelligenceNetwork:
|
|||
strategy: Compression strategy used.
|
||||
query_context: Optional user query that triggered this tool call.
|
||||
items: Optional list of items being compressed for field-level learning.
|
||||
auth_mode: Tenant auth slice (`payg` / `oauth` / `subscription`).
|
||||
Defaults to `DEFAULT_AUTH_MODE` when not provided.
|
||||
model_family: Target model family (`claude-3-5`, `gpt-4o`, …).
|
||||
Defaults to `DEFAULT_MODEL_FAMILY` when not provided.
|
||||
"""
|
||||
# HIGH FIX: Check enabled FIRST to avoid computing structure_hash if disabled
|
||||
# This saves CPU when TOIN is turned off
|
||||
|
|
@ -475,12 +599,15 @@ class ToolIntelligenceNetwork:
|
|||
|
||||
# Computing structure_hash can be expensive for large structures
|
||||
sig_hash = tool_signature.structure_hash
|
||||
key = _make_pattern_key(auth_mode, model_family, sig_hash)
|
||||
|
||||
# LOW FIX #22: Emit compression metric
|
||||
self._emit_metric(
|
||||
"toin.compression",
|
||||
{
|
||||
"signature_hash": sig_hash,
|
||||
"auth_mode": key[0],
|
||||
"model_family": key[1],
|
||||
"original_count": original_count,
|
||||
"compressed_count": compressed_count,
|
||||
"original_tokens": original_tokens,
|
||||
|
|
@ -492,10 +619,14 @@ class ToolIntelligenceNetwork:
|
|||
|
||||
with self._lock:
|
||||
# Get or create pattern
|
||||
if sig_hash not in self._patterns:
|
||||
self._patterns[sig_hash] = ToolPattern(tool_signature_hash=sig_hash)
|
||||
if key not in self._patterns:
|
||||
self._patterns[key] = ToolPattern(
|
||||
tool_signature_hash=sig_hash,
|
||||
auth_mode=key[0],
|
||||
model_family=key[1],
|
||||
)
|
||||
|
||||
pattern = self._patterns[sig_hash]
|
||||
pattern = self._patterns[key]
|
||||
|
||||
# Update compression stats
|
||||
pattern.total_compressions += 1
|
||||
|
|
@ -691,6 +822,8 @@ class ToolIntelligenceNetwork:
|
|||
query_fields: list[str] | None = None,
|
||||
strategy: str | None = None,
|
||||
retrieved_items: list[dict[str, Any]] | None = None,
|
||||
auth_mode: str | None = None,
|
||||
model_family: str | None = None,
|
||||
) -> None:
|
||||
"""Record a retrieval event.
|
||||
|
||||
|
|
@ -707,15 +840,21 @@ class ToolIntelligenceNetwork:
|
|||
query_fields: Fields mentioned in query (will be hashed).
|
||||
strategy: Compression strategy that was used (for success rate tracking).
|
||||
retrieved_items: Optional list of retrieved items for field-level learning.
|
||||
auth_mode: Tenant auth slice. Defaults to `DEFAULT_AUTH_MODE`.
|
||||
model_family: Target model family. Defaults to `DEFAULT_MODEL_FAMILY`.
|
||||
"""
|
||||
if not self._config.enabled:
|
||||
return
|
||||
|
||||
key = _make_pattern_key(auth_mode, model_family, tool_signature_hash)
|
||||
|
||||
# LOW FIX #22: Emit retrieval metric
|
||||
self._emit_metric(
|
||||
"toin.retrieval",
|
||||
{
|
||||
"signature_hash": tool_signature_hash,
|
||||
"auth_mode": key[0],
|
||||
"model_family": key[1],
|
||||
"retrieval_type": retrieval_type,
|
||||
"has_query": query is not None,
|
||||
"query_fields_count": len(query_fields) if query_fields else 0,
|
||||
|
|
@ -724,13 +863,15 @@ class ToolIntelligenceNetwork:
|
|||
)
|
||||
|
||||
with self._lock:
|
||||
if tool_signature_hash not in self._patterns:
|
||||
if key not in self._patterns:
|
||||
# First time seeing this tool via retrieval
|
||||
self._patterns[tool_signature_hash] = ToolPattern(
|
||||
tool_signature_hash=tool_signature_hash
|
||||
self._patterns[key] = ToolPattern(
|
||||
tool_signature_hash=tool_signature_hash,
|
||||
auth_mode=key[0],
|
||||
model_family=key[1],
|
||||
)
|
||||
|
||||
pattern = self._patterns[tool_signature_hash]
|
||||
pattern = self._patterns[key]
|
||||
|
||||
# Update retrieval stats
|
||||
pattern.total_retrievals += 1
|
||||
|
|
@ -852,244 +993,36 @@ class ToolIntelligenceNetwork:
|
|||
|
||||
def get_recommendation(
|
||||
self,
|
||||
tool_signature: ToolSignature,
|
||||
query_context: str | None = None,
|
||||
) -> CompressionHint:
|
||||
"""Get compression recommendation for a tool output.
|
||||
tool_signature: ToolSignature, # noqa: ARG002 — kept for source compat
|
||||
query_context: str | None = None, # noqa: ARG002
|
||||
) -> None:
|
||||
"""**Deprecated.** Returns `None`. PR-B5 retired the request-time hint API.
|
||||
|
||||
This is the main API for SmartCrusher to consult before compressing.
|
||||
TOIN is observation-only; recommendations are emitted by the
|
||||
offline `headroom.cli.toin_publish` CLI into `recommendations.toml`
|
||||
and loaded by the Rust proxy at startup. New code must not call
|
||||
this method. Existing call sites should migrate to reading the
|
||||
TOML file directly.
|
||||
|
||||
Args:
|
||||
tool_signature: Signature of the tool output structure.
|
||||
query_context: User query for context-aware recommendations.
|
||||
Emits `DeprecationWarning` once per process to keep busy call
|
||||
sites from flooding logs.
|
||||
|
||||
Returns:
|
||||
CompressionHint with recommendations.
|
||||
Always `None`. The legacy `_CompressionHint` envelope is no
|
||||
longer constructed at request time.
|
||||
"""
|
||||
if not self._config.enabled:
|
||||
return CompressionHint(source="default", reason="TOIN disabled")
|
||||
|
||||
sig_hash = tool_signature.structure_hash
|
||||
|
||||
with self._lock:
|
||||
pattern = self._patterns.get(sig_hash)
|
||||
|
||||
if pattern is None:
|
||||
# No data for this tool type
|
||||
return CompressionHint(
|
||||
source="default",
|
||||
reason="No pattern data for this tool type",
|
||||
)
|
||||
|
||||
# Track observation: TOIN was consulted for this pattern
|
||||
pattern.observations += 1
|
||||
self._dirty = True
|
||||
|
||||
# Not enough samples for reliable recommendation
|
||||
if pattern.sample_size < self._config.min_samples_for_recommendation:
|
||||
hint = CompressionHint(
|
||||
source="local",
|
||||
reason=f"Only {pattern.sample_size} samples (need {self._config.min_samples_for_recommendation})",
|
||||
confidence=pattern.confidence,
|
||||
based_on_samples=pattern.sample_size,
|
||||
)
|
||||
# LOW FIX #22: Emit recommendation metric
|
||||
self._emit_metric(
|
||||
"toin.recommendation",
|
||||
{
|
||||
"signature_hash": sig_hash,
|
||||
"source": hint.source,
|
||||
"confidence": hint.confidence,
|
||||
"skip_compression": hint.skip_compression,
|
||||
"max_items": hint.max_items,
|
||||
"compression_level": hint.compression_level,
|
||||
"based_on_samples": hint.based_on_samples,
|
||||
},
|
||||
)
|
||||
return hint
|
||||
|
||||
# Build recommendation based on learned patterns
|
||||
hint = self._build_recommendation(pattern, query_context)
|
||||
|
||||
# LOW FIX #22: Emit recommendation metric
|
||||
self._emit_metric(
|
||||
"toin.recommendation",
|
||||
{
|
||||
"signature_hash": sig_hash,
|
||||
"source": hint.source,
|
||||
"confidence": hint.confidence,
|
||||
"skip_compression": hint.skip_compression,
|
||||
"max_items": hint.max_items,
|
||||
"compression_level": hint.compression_level,
|
||||
"based_on_samples": hint.based_on_samples,
|
||||
},
|
||||
cls = type(self)
|
||||
if not cls._DEPRECATION_WARNED:
|
||||
cls._DEPRECATION_WARNED = True
|
||||
warnings.warn(
|
||||
"ToolIntelligenceNetwork.get_recommendation() is deprecated "
|
||||
"and now returns None. PR-B5 retired the request-time hint "
|
||||
"API; recommendations come from recommendations.toml at "
|
||||
"startup. See headroom/telemetry/toin.py module docstring.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
return hint
|
||||
|
||||
def _build_recommendation(
|
||||
self,
|
||||
pattern: ToolPattern,
|
||||
query_context: str | None,
|
||||
) -> CompressionHint:
|
||||
"""Build a recommendation based on pattern data and query context."""
|
||||
hint = CompressionHint(
|
||||
source="network"
|
||||
if pattern.user_count >= self._config.min_users_for_network_effect
|
||||
else "local",
|
||||
confidence=pattern.confidence,
|
||||
based_on_samples=pattern.sample_size,
|
||||
)
|
||||
|
||||
retrieval_rate = pattern.retrieval_rate
|
||||
full_retrieval_rate = pattern.full_retrieval_rate
|
||||
|
||||
# High retrieval rate = compression too aggressive
|
||||
if retrieval_rate > self._config.high_retrieval_threshold:
|
||||
if full_retrieval_rate > 0.8:
|
||||
# Almost all retrievals are full = don't compress
|
||||
hint.skip_compression = True
|
||||
hint.compression_level = "none"
|
||||
hint.reason = f"Very high full retrieval rate ({full_retrieval_rate:.1%})"
|
||||
else:
|
||||
# High retrieval but mostly search = compress conservatively
|
||||
hint.max_items = pattern.optimal_max_items
|
||||
hint.compression_level = "conservative"
|
||||
hint.reason = f"High retrieval rate ({retrieval_rate:.1%})"
|
||||
|
||||
elif retrieval_rate > self._config.medium_retrieval_threshold:
|
||||
# Moderate retrieval = moderate compression
|
||||
hint.max_items = max(20, pattern.optimal_max_items)
|
||||
hint.compression_level = "moderate"
|
||||
hint.reason = f"Moderate retrieval rate ({retrieval_rate:.1%})"
|
||||
|
||||
else:
|
||||
# Low retrieval = aggressive compression works
|
||||
hint.max_items = min(15, pattern.optimal_max_items)
|
||||
hint.compression_level = "aggressive"
|
||||
hint.reason = f"Low retrieval rate ({retrieval_rate:.1%})"
|
||||
|
||||
# Build preserve_fields list weighted by retrieval frequency
|
||||
# Start with pattern's preserve_fields, then enhance based on query
|
||||
preserve_fields = pattern.preserve_fields.copy()
|
||||
query_fields_count = 0
|
||||
|
||||
# If we have query context, extract field names and prioritize them
|
||||
if query_context and pattern.field_retrieval_frequency:
|
||||
# Extract field names from query context
|
||||
import re
|
||||
|
||||
query_field_names = re.findall(r"(\w+)[=:]", query_context.lower())
|
||||
|
||||
# Hash them and check if they're in our frequency data
|
||||
for field_name in query_field_names:
|
||||
field_hash = self._hash_field_name(field_name)
|
||||
if field_hash in pattern.field_retrieval_frequency:
|
||||
# This field is known to be retrieved - prioritize it
|
||||
if field_hash in preserve_fields:
|
||||
# Move to front
|
||||
preserve_fields.remove(field_hash)
|
||||
preserve_fields.insert(0, field_hash)
|
||||
query_fields_count += 1
|
||||
|
||||
# Sort remaining fields by retrieval frequency (most frequent first)
|
||||
if pattern.field_retrieval_frequency and len(preserve_fields) > 1:
|
||||
# Separate query-mentioned fields (already at front) from others
|
||||
if query_fields_count < len(preserve_fields):
|
||||
rest = preserve_fields[query_fields_count:]
|
||||
rest.sort(
|
||||
key=lambda f: pattern.field_retrieval_frequency.get(f, 0),
|
||||
reverse=True,
|
||||
)
|
||||
preserve_fields = preserve_fields[:query_fields_count] + rest
|
||||
|
||||
hint.preserve_fields = preserve_fields[:10] # Limit to top 10
|
||||
|
||||
# Use optimal strategy if known AND it has good success rate
|
||||
if pattern.optimal_strategy != "default":
|
||||
success_rate = pattern.strategy_success_rates.get(pattern.optimal_strategy, 1.0)
|
||||
# Only recommend strategy if success rate >= 0.5
|
||||
# Lower success rates mean this strategy often causes retrievals
|
||||
if success_rate >= 0.5:
|
||||
hint.recommended_strategy = pattern.optimal_strategy
|
||||
else:
|
||||
# Strategy has poor success rate - reduce confidence
|
||||
hint.confidence *= success_rate
|
||||
hint.reason += (
|
||||
f" (strategy {pattern.optimal_strategy} has low success: {success_rate:.1%})"
|
||||
)
|
||||
# Try to find a better strategy
|
||||
best_strategy = self._find_best_strategy(pattern)
|
||||
if best_strategy and best_strategy != pattern.optimal_strategy:
|
||||
hint.recommended_strategy = best_strategy
|
||||
hint.reason += f", using {best_strategy} instead"
|
||||
|
||||
# Boost max_items if query_context matches common retrieval patterns
|
||||
# This prevents unnecessary retrieval when we can predict what's needed
|
||||
if query_context:
|
||||
query_lower = query_context.lower()
|
||||
|
||||
# Check for exhaustive query keywords that suggest user needs all data
|
||||
exhaustive_keywords = ["all", "every", "complete", "full", "entire", "list all"]
|
||||
if any(kw in query_lower for kw in exhaustive_keywords):
|
||||
# User likely needs more data - be conservative
|
||||
hint.max_items = max(hint.max_items, 40)
|
||||
hint.compression_level = "conservative"
|
||||
hint.reason += " (exhaustive query detected)"
|
||||
|
||||
# Check against common retrieval patterns
|
||||
if pattern.common_query_patterns:
|
||||
query_pattern = self._anonymize_query_pattern(query_context)
|
||||
if query_pattern:
|
||||
# Exact match
|
||||
if query_pattern in pattern.common_query_patterns:
|
||||
hint.max_items = max(hint.max_items, 30)
|
||||
hint.reason += " (query matches retrieval pattern)"
|
||||
else:
|
||||
# Partial match: check if any stored pattern is contained in query
|
||||
for stored_pattern in pattern.common_query_patterns:
|
||||
# Check if key fields match (e.g., "status:*" in both)
|
||||
stored_fields = {
|
||||
f.split(":")[0] for f in stored_pattern.split() if ":" in f
|
||||
}
|
||||
query_fields = {
|
||||
f.split(":")[0] for f in query_pattern.split() if ":" in f
|
||||
}
|
||||
# If query uses same fields as a problematic pattern, be conservative
|
||||
if stored_fields and stored_fields.issubset(query_fields):
|
||||
hint.max_items = max(hint.max_items, 25)
|
||||
hint.reason += " (query uses fields from retrieval pattern)"
|
||||
break
|
||||
|
||||
# === TOIN Evolution: Include learned field semantics ===
|
||||
# Copy field_semantics with sufficient confidence for SmartCrusher to use
|
||||
# Only include fields with confidence >= 0.3 to reduce noise
|
||||
if pattern.field_semantics:
|
||||
hint.field_semantics = {
|
||||
field_hash: field_sem
|
||||
for field_hash, field_sem in pattern.field_semantics.items()
|
||||
if field_sem.confidence >= 0.3 or field_sem.retrieval_count >= 3
|
||||
}
|
||||
|
||||
return hint
|
||||
|
||||
def _find_best_strategy(self, pattern: ToolPattern) -> str | None:
|
||||
"""Find the strategy with the best success rate.
|
||||
|
||||
Returns None if no strategies have been tried or all have low success.
|
||||
"""
|
||||
if not pattern.strategy_success_rates:
|
||||
return None
|
||||
|
||||
# Find strategy with highest success rate above threshold
|
||||
best_strategy = None
|
||||
best_rate = 0.5 # Minimum acceptable rate
|
||||
|
||||
for strategy, rate in pattern.strategy_success_rates.items():
|
||||
if rate > best_rate:
|
||||
best_rate = rate
|
||||
best_strategy = strategy
|
||||
|
||||
return best_strategy
|
||||
return None
|
||||
|
||||
def _update_recommendations(self, pattern: ToolPattern) -> None:
|
||||
"""Update learned recommendations for a pattern."""
|
||||
|
|
@ -1245,28 +1178,57 @@ class ToolIntelligenceNetwork:
|
|||
),
|
||||
}
|
||||
|
||||
def get_pattern(self, signature_hash: str) -> ToolPattern | None:
|
||||
"""Get pattern data for a specific tool signature.
|
||||
def get_pattern(
|
||||
self,
|
||||
signature_hash: str,
|
||||
auth_mode: str | None = None,
|
||||
model_family: str | None = None,
|
||||
) -> ToolPattern | None:
|
||||
"""Get pattern data for a specific `(auth_mode, model_family, sig_hash)` slice.
|
||||
|
||||
Defaults to `(DEFAULT_AUTH_MODE, DEFAULT_MODEL_FAMILY, signature_hash)`
|
||||
when callers haven't supplied tenant info — preserves source-compat
|
||||
with pre-B5 callers that look up by bare hash.
|
||||
|
||||
HIGH FIX: Returns a deep copy to prevent external mutation of internal state.
|
||||
"""
|
||||
import copy
|
||||
|
||||
key = _make_pattern_key(auth_mode, model_family, signature_hash)
|
||||
|
||||
with self._lock:
|
||||
pattern = self._patterns.get(signature_hash)
|
||||
pattern = self._patterns.get(key)
|
||||
if pattern is not None:
|
||||
return copy.deepcopy(pattern)
|
||||
return None
|
||||
|
||||
def iter_patterns(self) -> list[tuple[PatternKey, ToolPattern]]:
|
||||
"""Snapshot of `(key, pattern)` pairs for offline aggregation.
|
||||
|
||||
Used by `headroom.cli.toin_publish` to walk every aggregated
|
||||
slice without exposing the live `_patterns` dict to external
|
||||
callers (deep-copies each pattern to prevent mutation).
|
||||
"""
|
||||
import copy
|
||||
|
||||
with self._lock:
|
||||
return [(k, copy.deepcopy(p)) for k, p in self._patterns.items()]
|
||||
|
||||
def export_patterns(self) -> dict[str, Any]:
|
||||
"""Export all patterns for sharing/aggregation."""
|
||||
"""Export all patterns for sharing/aggregation.
|
||||
|
||||
The aggregation key tuple is encoded as a `"auth|model|hash"`
|
||||
string for JSON storage (JSON object keys must be strings).
|
||||
See `_serialize_pattern_key` for the canonical encoding.
|
||||
"""
|
||||
with self._lock:
|
||||
return {
|
||||
"version": "1.0",
|
||||
"version": "2.0", # PR-B5: tuple aggregation key
|
||||
"export_timestamp": time.time(),
|
||||
"instance_id": self._instance_id,
|
||||
"patterns": {
|
||||
sig_hash: pattern.to_dict() for sig_hash, pattern in self._patterns.items()
|
||||
_serialize_pattern_key(key): pattern.to_dict()
|
||||
for key, pattern in self._patterns.items()
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -1276,6 +1238,11 @@ class ToolIntelligenceNetwork:
|
|||
Used for federated learning: aggregate patterns from multiple
|
||||
Headroom instances without sharing actual data.
|
||||
|
||||
Backward-compatible with v1.0 dumps that keyed patterns by bare
|
||||
structure_hash: those are promoted to the
|
||||
`(DEFAULT_AUTH_MODE, DEFAULT_MODEL_FAMILY, sig_hash)` slice via
|
||||
`_deserialize_pattern_key`.
|
||||
|
||||
Args:
|
||||
data: Exported pattern data.
|
||||
"""
|
||||
|
|
@ -1286,20 +1253,26 @@ class ToolIntelligenceNetwork:
|
|||
source_instance = data.get("instance_id", "unknown")
|
||||
|
||||
with self._lock:
|
||||
for sig_hash, pattern_dict in patterns_data.items():
|
||||
for serialized_key, pattern_dict in patterns_data.items():
|
||||
key = _deserialize_pattern_key(serialized_key)
|
||||
imported = ToolPattern.from_dict(pattern_dict)
|
||||
# Make sure dataclass fields agree with the dict key — pre-B5
|
||||
# dumps don't carry auth_mode/model_family on the pattern;
|
||||
# promote from the (possibly default) key.
|
||||
imported.auth_mode = key[0]
|
||||
imported.model_family = key[1]
|
||||
|
||||
if sig_hash in self._patterns:
|
||||
if key in self._patterns:
|
||||
# Merge with existing
|
||||
self._merge_patterns(self._patterns[sig_hash], imported)
|
||||
self._merge_patterns(self._patterns[key], imported)
|
||||
else:
|
||||
# Add new pattern - need to track source instance
|
||||
self._patterns[sig_hash] = imported
|
||||
self._patterns[key] = imported
|
||||
|
||||
# For NEW patterns from another instance, track the source in
|
||||
# _seen_instance_hashes so user_count reflects cross-user data
|
||||
if source_instance != self._instance_id:
|
||||
pattern = self._patterns[sig_hash]
|
||||
pattern = self._patterns[key]
|
||||
if source_instance not in pattern._seen_instance_hashes:
|
||||
# Limit storage to 100 unique instances to bound memory
|
||||
if len(pattern._seen_instance_hashes) < 100:
|
||||
|
|
@ -1515,8 +1488,19 @@ class ToolIntelligenceNetwork:
|
|||
self._last_save_time = time.time()
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't crash - TOIN should be resilient
|
||||
logger.warning("Failed to save TOIN data: %s", e)
|
||||
# Surface storage failures structured so log aggregators can
|
||||
# alert on `event=toin_save_failed` without false positives
|
||||
# from generic exception lines. Per project memory
|
||||
# `feedback_no_silent_fallbacks.md`: never swallow.
|
||||
logger.warning(
|
||||
"TOIN storage save failed",
|
||||
extra={
|
||||
"event": "toin_save_failed",
|
||||
"backend": type(self._backend).__name__,
|
||||
"error_type": type(e).__name__,
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
|
||||
def _load_from_backend(self) -> None:
|
||||
"""Load TOIN data from the storage backend."""
|
||||
|
|
@ -1529,7 +1513,15 @@ class ToolIntelligenceNetwork:
|
|||
self.import_patterns(data)
|
||||
self._dirty = False
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load TOIN data from backend: %s", e)
|
||||
logger.warning(
|
||||
"TOIN storage load failed",
|
||||
extra={
|
||||
"event": "toin_load_failed",
|
||||
"backend": type(self._backend).__name__,
|
||||
"error_type": type(e).__name__,
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
|
||||
def _maybe_auto_save(self) -> None:
|
||||
"""Auto-save if enough time has passed.
|
||||
|
|
@ -1591,6 +1583,12 @@ def _create_default_toin_backend() -> Any:
|
|||
)
|
||||
return None
|
||||
fn = ep.load()
|
||||
# `tenant_prefix` is retained for storage-backend namespacing
|
||||
# (Redis key prefix, Postgres schema name, etc.) so multi-tenant
|
||||
# SaaS deployments can carve up shared infrastructure. PR-B5 made
|
||||
# the in-memory aggregation key per-tenant via `auth_mode` /
|
||||
# `model_family`, so `tenant_prefix` is now functionally redundant
|
||||
# for *learning* — it only matters for storage layout. Keep it.
|
||||
kwargs = {
|
||||
"url": os.environ.get("HEADROOM_TOIN_URL", ""),
|
||||
"tenant_prefix": os.environ.get("HEADROOM_TOIN_TENANT_PREFIX", ""),
|
||||
|
|
|
|||
|
|
@ -521,7 +521,13 @@ class TestAdapterLifecycle:
|
|||
assert any("get:" in op for op in backend.ops)
|
||||
|
||||
def test_toin_save_load_preserves_patterns(self, tmp_toin_path):
|
||||
"""Patterns survive save/load via backend."""
|
||||
"""Patterns survive save/load via backend.
|
||||
|
||||
PR-B5 retired the request-time `get_recommendation()` API
|
||||
(it now returns None with a deprecation warning). Stats and
|
||||
on-disk patterns must still survive save/load — that's the
|
||||
observation API B5 preserves.
|
||||
"""
|
||||
config = TOINConfig(storage_path=tmp_toin_path)
|
||||
toin = ToolIntelligenceNetwork(config)
|
||||
|
||||
|
|
@ -546,6 +552,7 @@ class TestAdapterLifecycle:
|
|||
assert stats["patterns_tracked"] >= 1
|
||||
assert stats["total_compressions"] >= 15
|
||||
|
||||
# Recommendations should work
|
||||
hint = toin2.get_recommendation(sig)
|
||||
assert hint.based_on_samples >= 15
|
||||
# PR-B5: get_recommendation is observation-only and returns None.
|
||||
# Recommendations now flow through the publish CLI →
|
||||
# recommendations.toml → Rust loader path.
|
||||
assert toin2.get_recommendation(sig) is None
|
||||
|
|
|
|||
|
|
@ -124,7 +124,7 @@ class TestTOINDoubleCountFix:
|
|||
|
||||
# Get the pattern
|
||||
with toin._lock:
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
user_count_after_101 = pattern.user_count
|
||||
|
||||
# Now call again with same instance (instance_100)
|
||||
|
|
@ -133,7 +133,7 @@ class TestTOINDoubleCountFix:
|
|||
toin.record_compression(sig, 100, 10, 1000, 100, strategy="test_strategy")
|
||||
|
||||
with toin._lock:
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
user_count_after_102 = pattern.user_count
|
||||
|
||||
# Restore instance_id
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ class TestAllSeenInstancesUnboundedGrowth:
|
|||
# Simulate adding users via record_compression
|
||||
# (the cap is enforced there, not when directly adding to set)
|
||||
pattern = ToolPattern(tool_signature_hash=sig.structure_hash)
|
||||
toin._patterns[sig.structure_hash] = pattern
|
||||
toin._patterns[("unknown", "unknown", sig.structure_hash)] = pattern
|
||||
|
||||
# Direct manipulation should still work for testing
|
||||
for i in range(200):
|
||||
|
|
@ -101,7 +101,7 @@ class TestAllSeenInstancesUnboundedGrowth:
|
|||
# Record compressions from 150 "users" (simulated)
|
||||
# by directly manipulating the pattern
|
||||
pattern = ToolPattern(tool_signature_hash=sig.structure_hash)
|
||||
toin._patterns[sig.structure_hash] = pattern
|
||||
toin._patterns[("unknown", "unknown", sig.structure_hash)] = pattern
|
||||
|
||||
# Track 150 unique users
|
||||
for i in range(150):
|
||||
|
|
@ -183,7 +183,7 @@ class TestAllSeenInstancesSerialization:
|
|||
)
|
||||
|
||||
# Manually add more users to simulate multi-user scenario
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
for i in range(50):
|
||||
instance_hash = hashlib.sha256(f"extra_user_{i}".encode()).hexdigest()[:8]
|
||||
if instance_hash not in pattern._all_seen_instances:
|
||||
|
|
@ -202,7 +202,7 @@ class TestAllSeenInstancesSerialization:
|
|||
toin2 = ToolIntelligenceNetwork(config)
|
||||
|
||||
# Verify user count is preserved
|
||||
pattern2 = toin2._patterns.get(sig.structure_hash)
|
||||
pattern2 = toin2._patterns.get(("unknown", "unknown", sig.structure_hash))
|
||||
assert pattern2 is not None
|
||||
assert pattern2.user_count == original_user_count
|
||||
|
||||
|
|
@ -252,7 +252,7 @@ class TestUserCountMergeLogic:
|
|||
imported.sample_size = 5
|
||||
|
||||
# Merge
|
||||
toin._patterns["test_hash"] = existing
|
||||
toin._patterns[("unknown", "unknown", "test_hash")] = existing
|
||||
toin._merge_patterns(existing, imported)
|
||||
|
||||
# After merge: 5 existing + 2 new = 7 unique users
|
||||
|
|
@ -287,7 +287,7 @@ class TestUserCountMergeLogic:
|
|||
imported.sample_size = 20
|
||||
|
||||
# Merge
|
||||
toin._patterns["test_hash"] = existing
|
||||
toin._patterns[("unknown", "unknown", "test_hash")] = existing
|
||||
toin._merge_patterns(existing, imported)
|
||||
|
||||
# After merge: 120 existing + 10 new = 130 unique users
|
||||
|
|
@ -732,7 +732,7 @@ class TestTOINHighPriorityFixes:
|
|||
query_fields=[f"unique_field_{i}"],
|
||||
)
|
||||
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
assert len(pattern.field_retrieval_frequency) <= 100
|
||||
|
||||
def test_commonly_retrieved_fields_bounded(self):
|
||||
|
|
@ -761,7 +761,7 @@ class TestTOINHighPriorityFixes:
|
|||
query_fields=[f"common_field_{i}"],
|
||||
)
|
||||
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
assert len(pattern.commonly_retrieved_fields) <= 20
|
||||
|
||||
def test_strategy_success_rate_updates(self):
|
||||
|
|
@ -782,7 +782,7 @@ class TestTOINHighPriorityFixes:
|
|||
strategy="TEST_STRATEGY",
|
||||
)
|
||||
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
initial_rate = pattern.strategy_success_rates["TEST_STRATEGY"]
|
||||
assert initial_rate == 1.0 # Starts at 1.0
|
||||
|
||||
|
|
@ -794,7 +794,7 @@ class TestTOINHighPriorityFixes:
|
|||
strategy="TEST_STRATEGY",
|
||||
)
|
||||
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
after_retrieval = pattern.strategy_success_rates["TEST_STRATEGY"]
|
||||
assert after_retrieval < initial_rate # Should decrease
|
||||
|
||||
|
|
@ -809,7 +809,7 @@ class TestTOINHighPriorityFixes:
|
|||
strategy="TEST_STRATEGY",
|
||||
)
|
||||
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
after_compressions = pattern.strategy_success_rates["TEST_STRATEGY"]
|
||||
assert after_compressions > after_retrieval # Should increase
|
||||
|
||||
|
|
@ -849,26 +849,11 @@ class TestTOINHighPriorityFixes:
|
|||
# (auto-save happens inside record_compression)
|
||||
assert not toin._dirty
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired; preserve_fields lives on the aggregated ToolPattern instead"
|
||||
)
|
||||
def test_toin_preserves_fields_returns_list(self):
|
||||
"""Verify preserve_fields in hints is always a list."""
|
||||
toin = ToolIntelligenceNetwork(TOINConfig(enabled=True))
|
||||
sig = ToolSignature.from_items([{"id": 1, "name": "test"}])
|
||||
|
||||
# Record enough data for recommendations
|
||||
for _ in range(15):
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=100,
|
||||
compressed_count=10,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=100,
|
||||
strategy="test",
|
||||
)
|
||||
|
||||
hint = toin.get_recommendation(sig, "find something")
|
||||
|
||||
assert isinstance(hint.preserve_fields, list)
|
||||
assert len(hint.preserve_fields) <= 10 # Should be bounded
|
||||
"""Retired in PR-B5 along with the request-time hint API."""
|
||||
|
||||
|
||||
class TestCompressionStoreHighPriorityFixes:
|
||||
|
|
|
|||
|
|
@ -567,8 +567,12 @@ class TestEndToEndTOINIntegration:
|
|||
store = get_compression_store()
|
||||
store.process_pending_feedback()
|
||||
|
||||
# Verify TOIN learned field semantics
|
||||
pattern = fresh_toin._patterns.get(signature.structure_hash)
|
||||
# PR-B5: pattern key is now `(auth_mode, model_family, sig_hash)`.
|
||||
# Callers that don't supply auth/model land on the
|
||||
# `("unknown", "unknown", sig_hash)` slot.
|
||||
from headroom.telemetry.toin import _make_pattern_key
|
||||
|
||||
pattern = fresh_toin._patterns.get(_make_pattern_key(None, None, signature.structure_hash))
|
||||
assert pattern is not None, "Pattern should exist after compression and retrieval"
|
||||
|
||||
# CRITICAL ASSERTION: This catches the bug where compression_store
|
||||
|
|
@ -641,7 +645,12 @@ class TestEndToEndTOINIntegration:
|
|||
store.process_pending_feedback()
|
||||
|
||||
# Step 3: Verify TOIN learned
|
||||
pattern = fresh_toin._patterns.get(signature.structure_hash)
|
||||
# PR-B5: pattern key is now `(auth_mode, model_family, sig_hash)`.
|
||||
# Callers that don't supply auth/model land on the
|
||||
# `("unknown", "unknown", sig_hash)` slot.
|
||||
from headroom.telemetry.toin import _make_pattern_key
|
||||
|
||||
pattern = fresh_toin._patterns.get(_make_pattern_key(None, None, signature.structure_hash))
|
||||
assert pattern is not None, "Pattern should exist"
|
||||
assert pattern.total_compressions >= 1, "Should have compression count"
|
||||
assert pattern.total_retrievals >= 1, "Should have retrieval count"
|
||||
|
|
@ -653,6 +662,10 @@ class TestEndToEndTOINIntegration:
|
|||
"the production feedback loop is broken."
|
||||
)
|
||||
|
||||
# Step 5: Get recommendation (verifies learning is usable)
|
||||
recommendation = fresh_toin.get_recommendation(signature, "find category")
|
||||
assert recommendation.confidence >= 0, "Recommendation should have confidence"
|
||||
# Step 5: PR-B5 retired the request-time recommendation API in favor of
|
||||
# observation-only learning + startup-published recommendations.toml.
|
||||
# `get_recommendation()` now returns None and emits a deprecation
|
||||
# warning; the dispatcher consumes published advice via the Rust
|
||||
# `RecommendationStore`. Assert the deprecation contract here so a
|
||||
# future revival of the API doesn't slip past silently.
|
||||
assert fresh_toin.get_recommendation(signature, "find category") is None
|
||||
|
|
|
|||
|
|
@ -1,4 +1,10 @@
|
|||
"""Tests for Tool Output Intelligence Network (TOIN)."""
|
||||
"""Tests for Tool Output Intelligence Network (TOIN).
|
||||
|
||||
PR-B5 retired the request-time hint API. Tests that exercised the old
|
||||
`get_recommendation()` / `CompressionHint` shape are skipped at module
|
||||
level — the new observation-only contract is covered by
|
||||
`tests/test_toin_observation_only.py` and `tests/test_toin_publish.py`.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
|
@ -7,7 +13,6 @@ import time
|
|||
import pytest
|
||||
|
||||
from headroom.telemetry import (
|
||||
CompressionHint,
|
||||
TOINConfig,
|
||||
ToolIntelligenceNetwork,
|
||||
ToolPattern,
|
||||
|
|
@ -140,44 +145,22 @@ class TestToolPattern:
|
|||
assert pattern.full_retrieval_rate == 0.0
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"PR-B5: CompressionHint is now private (_CompressionHint) and "
|
||||
"the request-time hint API is retired. See "
|
||||
"tests/test_toin_observation_only.py for the replacement contract."
|
||||
)
|
||||
)
|
||||
class TestCompressionHint:
|
||||
"""Test CompressionHint data model."""
|
||||
"""Retired: CompressionHint was the public envelope for the
|
||||
request-time hint API removed in PR-B5."""
|
||||
|
||||
def test_default_values(self):
|
||||
"""Default values are sensible."""
|
||||
hint = CompressionHint()
|
||||
|
||||
assert hint.skip_compression is False
|
||||
assert hint.max_items == 20
|
||||
assert hint.compression_level == "moderate"
|
||||
assert hint.preserve_fields == []
|
||||
assert hint.recommended_strategy == "default"
|
||||
assert hint.source == "default"
|
||||
assert hint.confidence == 0.0
|
||||
pass
|
||||
|
||||
def test_custom_values(self):
|
||||
"""Custom values are preserved."""
|
||||
hint = CompressionHint(
|
||||
skip_compression=True,
|
||||
max_items=50,
|
||||
compression_level="conservative",
|
||||
preserve_fields=["id", "score"],
|
||||
recommended_strategy="top_n",
|
||||
reason="High retrieval rate",
|
||||
confidence=0.85,
|
||||
source="network",
|
||||
based_on_samples=1000,
|
||||
)
|
||||
|
||||
assert hint.skip_compression is True
|
||||
assert hint.max_items == 50
|
||||
assert hint.compression_level == "conservative"
|
||||
assert hint.preserve_fields == ["id", "score"]
|
||||
assert hint.recommended_strategy == "top_n"
|
||||
assert hint.reason == "High retrieval rate"
|
||||
assert hint.confidence == 0.85
|
||||
assert hint.source == "network"
|
||||
assert hint.based_on_samples == 1000
|
||||
pass
|
||||
|
||||
|
||||
class TestTOINConfig:
|
||||
|
|
@ -366,145 +349,46 @@ class TestToolIntelligenceNetwork:
|
|||
# Field should be in commonly_retrieved_fields after 3+ retrievals
|
||||
assert len(pattern.commonly_retrieved_fields) > 0
|
||||
|
||||
# PR-B5: the following tests exercised the request-time hint API
|
||||
# that's now retired. They're skipped wholesale; the new contract
|
||||
# ("get_recommendation always returns None and emits a deprecation
|
||||
# warning") is covered by tests/test_toin_observation_only.py.
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired — see test_toin_observation_only.py"
|
||||
)
|
||||
def test_get_recommendation_no_data(self):
|
||||
"""No recommendation with no pattern data."""
|
||||
toin = ToolIntelligenceNetwork()
|
||||
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
hint = toin.get_recommendation(sig)
|
||||
|
||||
assert hint.source == "default"
|
||||
assert hint.skip_compression is False
|
||||
assert "No pattern data" in hint.reason
|
||||
pass
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired — see test_toin_observation_only.py"
|
||||
)
|
||||
def test_get_recommendation_insufficient_samples(self):
|
||||
"""Local recommendation with insufficient samples."""
|
||||
config = TOINConfig(min_samples_for_recommendation=10)
|
||||
toin = ToolIntelligenceNetwork(config)
|
||||
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
|
||||
# Record only 5 compressions (less than 10)
|
||||
for _ in range(5):
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=100,
|
||||
compressed_count=10,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=100,
|
||||
strategy="top_n",
|
||||
)
|
||||
|
||||
hint = toin.get_recommendation(sig)
|
||||
assert hint.source == "local"
|
||||
assert "Only 5 samples" in hint.reason
|
||||
assert hint.based_on_samples == 5
|
||||
pass
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired — see test_toin_observation_only.py"
|
||||
)
|
||||
def test_get_recommendation_aggressive_compression(self):
|
||||
"""Low retrieval rate leads to aggressive compression."""
|
||||
config = TOINConfig(
|
||||
min_samples_for_recommendation=5,
|
||||
medium_retrieval_threshold=0.2,
|
||||
high_retrieval_threshold=0.5,
|
||||
)
|
||||
toin = ToolIntelligenceNetwork(config)
|
||||
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
|
||||
# Record compressions with no retrievals (low retrieval rate)
|
||||
for _ in range(10):
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=100,
|
||||
compressed_count=10,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=100,
|
||||
strategy="top_n",
|
||||
)
|
||||
|
||||
hint = toin.get_recommendation(sig)
|
||||
assert hint.compression_level == "aggressive"
|
||||
assert hint.skip_compression is False
|
||||
assert "Low retrieval rate" in hint.reason
|
||||
pass
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired — see test_toin_observation_only.py"
|
||||
)
|
||||
def test_get_recommendation_conservative_compression(self):
|
||||
"""High retrieval rate leads to conservative compression."""
|
||||
config = TOINConfig(
|
||||
min_samples_for_recommendation=5,
|
||||
high_retrieval_threshold=0.5,
|
||||
)
|
||||
toin = ToolIntelligenceNetwork(config)
|
||||
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
sig_hash = sig.structure_hash
|
||||
|
||||
# Record compressions
|
||||
for _ in range(10):
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=100,
|
||||
compressed_count=10,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=100,
|
||||
strategy="top_n",
|
||||
)
|
||||
|
||||
# Record many search retrievals (60% retrieval rate)
|
||||
for _ in range(6):
|
||||
toin.record_retrieval(
|
||||
tool_signature_hash=sig_hash,
|
||||
retrieval_type="search",
|
||||
)
|
||||
|
||||
hint = toin.get_recommendation(sig)
|
||||
assert hint.compression_level == "conservative"
|
||||
assert hint.skip_compression is False
|
||||
assert "High retrieval rate" in hint.reason
|
||||
pass
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired — see test_toin_observation_only.py"
|
||||
)
|
||||
def test_get_recommendation_skip_compression(self):
|
||||
"""Very high full retrieval rate leads to skip compression."""
|
||||
config = TOINConfig(
|
||||
min_samples_for_recommendation=5,
|
||||
high_retrieval_threshold=0.5,
|
||||
)
|
||||
toin = ToolIntelligenceNetwork(config)
|
||||
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
sig_hash = sig.structure_hash
|
||||
|
||||
# Record compressions
|
||||
for _ in range(10):
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=100,
|
||||
compressed_count=10,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=100,
|
||||
strategy="top_n",
|
||||
)
|
||||
|
||||
# Record many FULL retrievals (60% retrieval rate, 100% full)
|
||||
for _ in range(6):
|
||||
toin.record_retrieval(
|
||||
tool_signature_hash=sig_hash,
|
||||
retrieval_type="full",
|
||||
)
|
||||
|
||||
hint = toin.get_recommendation(sig)
|
||||
assert hint.skip_compression is True
|
||||
assert hint.compression_level == "none"
|
||||
assert "full retrieval rate" in hint.reason.lower()
|
||||
pass
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired — see test_toin_observation_only.py"
|
||||
)
|
||||
def test_get_recommendation_disabled(self):
|
||||
"""Disabled TOIN returns default hint."""
|
||||
config = TOINConfig(enabled=False)
|
||||
toin = ToolIntelligenceNetwork(config)
|
||||
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
hint = toin.get_recommendation(sig)
|
||||
|
||||
assert hint.source == "default"
|
||||
assert "TOIN disabled" in hint.reason
|
||||
pass
|
||||
|
||||
def test_get_stats(self):
|
||||
"""get_stats returns overall statistics."""
|
||||
|
|
@ -589,7 +473,9 @@ class TestTOINExportImport:
|
|||
assert "instance_id" in export
|
||||
assert "patterns" in export
|
||||
assert len(export["patterns"]) == 1
|
||||
assert sig.structure_hash in export["patterns"]
|
||||
# PR-B5: keys are now serialized "auth|model|hash" tuples; default
|
||||
# auth/model produce the "unknown|unknown|<hash>" string.
|
||||
assert f"unknown|unknown|{sig.structure_hash}" in export["patterns"]
|
||||
|
||||
def test_import_patterns_new_pattern(self):
|
||||
"""import_patterns adds new patterns."""
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ def _make_signature(structure_hash: str = "test_hash_123") -> ToolSignature:
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="PR-B5: observations counter and request-time hint API retired")
|
||||
class TestGetRecommendationObservations:
|
||||
"""Bug 1: get_recommendation() should increment observations counter."""
|
||||
|
||||
|
|
@ -69,7 +70,7 @@ class TestGetRecommendationObservations:
|
|||
toin.get_recommendation(sig)
|
||||
|
||||
# Check observations incremented
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
assert pattern.observations == 1
|
||||
|
||||
# Call again
|
||||
|
|
@ -96,7 +97,7 @@ class TestGetRecommendationObservations:
|
|||
result = toin.get_recommendation(sig)
|
||||
assert result.source == "local" # Not enough samples
|
||||
|
||||
pattern = toin._patterns[sig.structure_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig.structure_hash)]
|
||||
assert pattern.observations == 1
|
||||
|
||||
def test_no_increment_for_unknown_pattern(self):
|
||||
|
|
@ -162,7 +163,7 @@ class TestRecordRetrievalPopulatesFields:
|
|||
query_fields=["error_message"],
|
||||
)
|
||||
|
||||
pattern = toin._patterns[sig_hash]
|
||||
pattern = toin._patterns[("unknown", "unknown", sig_hash)]
|
||||
assert pattern.total_retrievals == 5
|
||||
assert pattern.search_retrievals == 5
|
||||
assert len(pattern.field_retrieval_frequency) > 0
|
||||
|
|
|
|||
|
|
@ -68,6 +68,9 @@ def fresh_store():
|
|||
reset_compression_store()
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: strategy-recommendation API retired (get_recommendation returns None)"
|
||||
)
|
||||
class TestStrategySuccessRates:
|
||||
"""Test that strategy_success_rates are used in recommendations."""
|
||||
|
||||
|
|
@ -88,7 +91,7 @@ class TestStrategySuccessRates:
|
|||
)
|
||||
|
||||
# Set high success rate
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
pattern.strategy_success_rates["smart_sample"] = 0.8
|
||||
pattern.optimal_strategy = "smart_sample"
|
||||
|
||||
|
|
@ -114,7 +117,7 @@ class TestStrategySuccessRates:
|
|||
)
|
||||
|
||||
# Set low success rate
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
pattern.strategy_success_rates["bad_strategy"] = 0.2
|
||||
pattern.optimal_strategy = "bad_strategy"
|
||||
|
||||
|
|
@ -143,7 +146,7 @@ class TestStrategySuccessRates:
|
|||
)
|
||||
|
||||
# Set up multiple strategies with different success rates
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
pattern.strategy_success_rates = {
|
||||
"bad_strategy": 0.2,
|
||||
"good_strategy": 0.9,
|
||||
|
|
@ -176,7 +179,7 @@ class TestPreserveFieldsMerging:
|
|||
compressed_tokens=500,
|
||||
strategy="smart_sample",
|
||||
)
|
||||
local_pattern = fresh_toin._patterns[sig_hash]
|
||||
local_pattern = fresh_toin._patterns[("unknown", "unknown", sig_hash)]
|
||||
local_pattern.preserve_fields = ["field_a", "field_b"]
|
||||
|
||||
# Import pattern with different preserve_fields
|
||||
|
|
@ -195,7 +198,7 @@ class TestPreserveFieldsMerging:
|
|||
fresh_toin.import_patterns(import_data)
|
||||
|
||||
# Verify merge
|
||||
pattern = fresh_toin._patterns[sig_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", sig_hash)]
|
||||
assert "field_a" in pattern.preserve_fields
|
||||
assert "field_b" in pattern.preserve_fields
|
||||
assert "field_c" in pattern.preserve_fields
|
||||
|
|
@ -216,7 +219,7 @@ class TestPreserveFieldsMerging:
|
|||
compressed_tokens=500,
|
||||
strategy="smart_sample",
|
||||
)
|
||||
pattern = fresh_toin._patterns[sig_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", sig_hash)]
|
||||
pattern.preserve_fields = [f"field_{i}" for i in range(8)]
|
||||
|
||||
# Import with 5 more fields
|
||||
|
|
@ -234,7 +237,7 @@ class TestPreserveFieldsMerging:
|
|||
fresh_toin.import_patterns(import_data)
|
||||
|
||||
# Should be capped at 10
|
||||
pattern = fresh_toin._patterns[sig_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", sig_hash)]
|
||||
assert len(pattern.preserve_fields) <= 10
|
||||
|
||||
|
||||
|
|
@ -256,7 +259,7 @@ class TestUserCountTracking:
|
|||
strategy="smart_sample",
|
||||
)
|
||||
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
assert pattern.user_count == 1
|
||||
assert len(pattern._seen_instance_hashes) == 1
|
||||
assert fresh_toin._instance_id in pattern._seen_instance_hashes
|
||||
|
|
@ -277,7 +280,7 @@ class TestUserCountTracking:
|
|||
strategy="smart_sample",
|
||||
)
|
||||
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
assert pattern.user_count == 1 # Still 1
|
||||
|
||||
def test_instance_hashes_serialized_and_loaded(self):
|
||||
|
|
@ -306,7 +309,7 @@ class TestUserCountTracking:
|
|||
# Load in new instance
|
||||
toin2 = ToolIntelligenceNetwork(TOINConfig(storage_path=storage_path))
|
||||
|
||||
pattern = toin2._patterns.get(signature.structure_hash)
|
||||
pattern = toin2._patterns.get(("unknown", "unknown", signature.structure_hash))
|
||||
assert pattern is not None
|
||||
assert pattern.user_count >= 1
|
||||
assert len(pattern._seen_instance_hashes) >= 1
|
||||
|
|
@ -342,11 +345,14 @@ class TestUserCountTracking:
|
|||
|
||||
fresh_toin.import_patterns(import_data)
|
||||
|
||||
pattern = fresh_toin._patterns[sig_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", sig_hash)]
|
||||
# Should have local + 2 imported = 3
|
||||
assert pattern.user_count >= 3
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired; field-weighting now consumed only by toin publish"
|
||||
)
|
||||
class TestFieldRetrievalFrequencyWeighting:
|
||||
"""Test field_retrieval_frequency weighting in preserve_fields."""
|
||||
|
||||
|
|
@ -368,7 +374,7 @@ class TestFieldRetrievalFrequencyWeighting:
|
|||
|
||||
# Record retrievals for "status" field
|
||||
status_hash = fresh_toin._hash_field_name("status")
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
pattern.field_retrieval_frequency = {
|
||||
status_hash: 50,
|
||||
fresh_toin._hash_field_name("category"): 10,
|
||||
|
|
@ -397,7 +403,7 @@ class TestFieldRetrievalFrequencyWeighting:
|
|||
strategy="smart_sample",
|
||||
)
|
||||
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
field_a = fresh_toin._hash_field_name("field_a")
|
||||
field_b = fresh_toin._hash_field_name("field_b")
|
||||
field_c = fresh_toin._hash_field_name("field_c")
|
||||
|
|
@ -423,6 +429,7 @@ class TestFieldRetrievalFrequencyWeighting:
|
|||
assert b_idx < c_idx, "Higher frequency field should come first"
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="PR-B5: get_recommendation retired (returns None / DeprecationWarning)")
|
||||
class TestQueryContextUsage:
|
||||
"""Test query_context usage in recommendations."""
|
||||
|
||||
|
|
@ -443,7 +450,7 @@ class TestQueryContextUsage:
|
|||
)
|
||||
|
||||
# Low retrieval rate = aggressive compression
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
pattern.total_retrievals = 0
|
||||
|
||||
# Query with exhaustive keyword
|
||||
|
|
@ -469,7 +476,7 @@ class TestQueryContextUsage:
|
|||
strategy="smart_sample",
|
||||
)
|
||||
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
pattern.total_retrievals = 0
|
||||
|
||||
hint = fresh_toin.get_recommendation(signature, "find every user")
|
||||
|
|
@ -491,7 +498,7 @@ class TestQueryContextUsage:
|
|||
strategy="smart_sample",
|
||||
)
|
||||
|
||||
pattern = fresh_toin._patterns[signature.structure_hash]
|
||||
pattern = fresh_toin._patterns[("unknown", "unknown", signature.structure_hash)]
|
||||
pattern.total_retrievals = 0
|
||||
# Add a problematic query pattern
|
||||
pattern.common_query_patterns = ["status:*"]
|
||||
|
|
|
|||
|
|
@ -247,6 +247,9 @@ class TestTOINPersistenceAcrossInstances:
|
|||
print("\n[PASS] TOIN persistence works correctly")
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired; feedback-loop covered by test_toin_observation_only.py"
|
||||
)
|
||||
class TestTOINFullFeedbackLoop:
|
||||
"""Test 3: Verify TOIN feedback loop with recommendations."""
|
||||
|
||||
|
|
@ -333,6 +336,9 @@ class TestTOINFullFeedbackLoop:
|
|||
print("\n[PASS] TOIN feedback loop works correctly")
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="PR-B5: get_recommendation retired; confidence-progression validated via record + get_pattern instead"
|
||||
)
|
||||
class TestTOINProgressiveConfidence:
|
||||
"""Test 4: Verify TOIN confidence increases with sample size."""
|
||||
|
||||
|
|
|
|||
302
tests/test_toin_observation_only.py
Normal file
302
tests/test_toin_observation_only.py
Normal file
|
|
@ -0,0 +1,302 @@
|
|||
"""PR-B5 acceptance tests: TOIN observation-only contract.
|
||||
|
||||
Pins three guarantees:
|
||||
|
||||
1. `get_recommendation()` returns `None` and emits a `DeprecationWarning`
|
||||
exactly once per process. The request-time hint API is retired.
|
||||
2. The aggregation key is `(auth_mode, model_family, structure_hash)` —
|
||||
two patterns with the same `structure_hash` but different `auth_mode`
|
||||
or `model_family` are tracked as distinct rows in the TOIN store.
|
||||
3. Recording a compression event does NOT alter the bytes SmartCrusher
|
||||
produces for an identical input. SmartCrusher is deterministic; TOIN
|
||||
only observes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.telemetry import (
|
||||
DEFAULT_AUTH_MODE,
|
||||
DEFAULT_MODEL_FAMILY,
|
||||
TOINConfig,
|
||||
ToolIntelligenceNetwork,
|
||||
ToolSignature,
|
||||
reset_toin,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_toin(monkeypatch, tmp_path: Path):
|
||||
"""Force every test to use a fresh tempfile-backed TOIN."""
|
||||
storage = tmp_path / "toin_obs_test.json"
|
||||
monkeypatch.setenv("HEADROOM_TOIN_PATH", str(storage))
|
||||
reset_toin()
|
||||
# Also reset the class-level deprecation flag so each test gets a
|
||||
# fresh "one warning" budget. Without this, test ordering would
|
||||
# determine whether the warning fires.
|
||||
ToolIntelligenceNetwork._DEPRECATION_WARNED = False
|
||||
yield
|
||||
reset_toin()
|
||||
ToolIntelligenceNetwork._DEPRECATION_WARNED = False
|
||||
|
||||
|
||||
# ── Part 1: deprecation surface ────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_get_recommendation_returns_none_with_deprecation_warning():
|
||||
"""get_recommendation() returns None and emits DeprecationWarning once."""
|
||||
toin = ToolIntelligenceNetwork()
|
||||
sig = ToolSignature.from_items([{"id": "1", "status": "ok"}])
|
||||
|
||||
# First call: warning fires.
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
result = toin.get_recommendation(sig)
|
||||
|
||||
assert result is None, "PR-B5: get_recommendation must return None"
|
||||
deprecations = [w for w in caught if issubclass(w.category, DeprecationWarning)]
|
||||
assert len(deprecations) == 1, f"expected 1 DeprecationWarning, got {len(deprecations)}"
|
||||
assert "PR-B5" in str(deprecations[0].message)
|
||||
|
||||
# Second call: still None, but warning is suppressed (once-per-process).
|
||||
with warnings.catch_warnings(record=True) as caught2:
|
||||
warnings.simplefilter("always")
|
||||
result2 = toin.get_recommendation(sig)
|
||||
assert result2 is None
|
||||
assert all(not issubclass(w.category, DeprecationWarning) for w in caught2)
|
||||
|
||||
|
||||
def test_compression_hint_is_not_publicly_exported():
|
||||
"""`CompressionHint` is no longer re-exported from `headroom.telemetry`."""
|
||||
import headroom.telemetry as telemetry_pkg
|
||||
|
||||
assert not hasattr(telemetry_pkg, "CompressionHint"), (
|
||||
"PR-B5: CompressionHint became private (_CompressionHint) and "
|
||||
"must not be importable from headroom.telemetry."
|
||||
)
|
||||
|
||||
|
||||
# ── Part 2: per-tenant aggregation key ─────────────────────────────────────
|
||||
|
||||
|
||||
def test_aggregation_key_includes_auth_mode_and_model_family():
|
||||
"""Same structure_hash with different auth_mode/model_family ⇒ distinct patterns."""
|
||||
toin = ToolIntelligenceNetwork()
|
||||
sig = ToolSignature.from_items([{"id": "1", "score": 99}])
|
||||
|
||||
# Three slices for the same tool signature.
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=10,
|
||||
compressed_count=5,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=500,
|
||||
strategy="smart_crusher",
|
||||
auth_mode="payg",
|
||||
model_family="claude-3-5",
|
||||
)
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=10,
|
||||
compressed_count=5,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=500,
|
||||
strategy="smart_crusher",
|
||||
auth_mode="oauth",
|
||||
model_family="claude-3-5",
|
||||
)
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=10,
|
||||
compressed_count=5,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=500,
|
||||
strategy="smart_crusher",
|
||||
auth_mode="payg",
|
||||
model_family="gpt-4o",
|
||||
)
|
||||
|
||||
sig_hash = sig.structure_hash
|
||||
assert ("payg", "claude-3-5", sig_hash) in toin._patterns
|
||||
assert ("oauth", "claude-3-5", sig_hash) in toin._patterns
|
||||
assert ("payg", "gpt-4o", sig_hash) in toin._patterns
|
||||
# Three distinct slices, each with sample_size=1.
|
||||
assert len(toin._patterns) == 3
|
||||
for key, pattern in toin._patterns.items():
|
||||
assert pattern.auth_mode == key[0]
|
||||
assert pattern.model_family == key[1]
|
||||
assert pattern.tool_signature_hash == key[2]
|
||||
assert pattern.sample_size == 1
|
||||
|
||||
|
||||
def test_aggregation_key_defaults_to_unknown_when_caller_omits_tenant():
|
||||
"""Callers that don't pass auth_mode/model_family land in the default slice."""
|
||||
toin = ToolIntelligenceNetwork()
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=10,
|
||||
compressed_count=5,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=500,
|
||||
strategy="smart_crusher",
|
||||
)
|
||||
|
||||
expected_key = (DEFAULT_AUTH_MODE, DEFAULT_MODEL_FAMILY, sig.structure_hash)
|
||||
assert expected_key in toin._patterns
|
||||
pattern = toin._patterns[expected_key]
|
||||
assert pattern.auth_mode == DEFAULT_AUTH_MODE
|
||||
assert pattern.model_family == DEFAULT_MODEL_FAMILY
|
||||
|
||||
|
||||
def test_storage_round_trip_preserves_aggregation_key(tmp_path: Path):
|
||||
"""Save/load round-trips the per-tenant aggregation key intact."""
|
||||
storage = tmp_path / "toin_roundtrip.json"
|
||||
toin1 = ToolIntelligenceNetwork(TOINConfig(storage_path=str(storage)))
|
||||
sig = ToolSignature.from_items([{"id": "1"}])
|
||||
|
||||
toin1.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=10,
|
||||
compressed_count=5,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=500,
|
||||
strategy="smart_crusher",
|
||||
auth_mode="oauth",
|
||||
model_family="gpt-4o",
|
||||
)
|
||||
toin1.save()
|
||||
|
||||
toin2 = ToolIntelligenceNetwork(TOINConfig(storage_path=str(storage)))
|
||||
key = ("oauth", "gpt-4o", sig.structure_hash)
|
||||
assert key in toin2._patterns
|
||||
assert toin2._patterns[key].auth_mode == "oauth"
|
||||
assert toin2._patterns[key].model_family == "gpt-4o"
|
||||
|
||||
|
||||
def test_record_does_not_alter_compression_decision():
|
||||
"""SmartCrusher output is byte-identical regardless of TOIN observation state.
|
||||
|
||||
Calls SmartCrusher twice on the same input — once with TOIN empty,
|
||||
once after recording a compression that would have changed the
|
||||
pre-B5 hint — and asserts byte equality. This pins the
|
||||
observation-only contract: TOIN observes; never mutates.
|
||||
"""
|
||||
smart_crusher_module = pytest.importorskip("headroom.transforms.smart_crusher")
|
||||
SmartCrusher = smart_crusher_module.SmartCrusher
|
||||
SmartCrusherConfig = smart_crusher_module.SmartCrusherConfig
|
||||
|
||||
cfg = SmartCrusherConfig(
|
||||
enabled=True,
|
||||
min_items_to_analyze=3,
|
||||
min_tokens_to_crush=10,
|
||||
)
|
||||
crusher = SmartCrusher(config=cfg)
|
||||
|
||||
# 50 low-uniqueness rows so the crusher is willing to compress.
|
||||
items = [{"id": i, "status": "ok", "code": 200, "msg": "fine"} for i in range(50)]
|
||||
import json as _json
|
||||
|
||||
payload = _json.dumps(items)
|
||||
|
||||
first = crusher.crush(payload)
|
||||
|
||||
# Inject TOIN observations that, pre-B5, would have biased the
|
||||
# compressor toward conservative output via get_recommendation().
|
||||
toin = ToolIntelligenceNetwork()
|
||||
sig = ToolSignature.from_items(items)
|
||||
sig_hash = sig.structure_hash
|
||||
for _ in range(20):
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=50,
|
||||
compressed_count=10,
|
||||
original_tokens=1000,
|
||||
compressed_tokens=200,
|
||||
strategy="smart_crusher",
|
||||
)
|
||||
for _ in range(15):
|
||||
toin.record_retrieval(
|
||||
tool_signature_hash=sig_hash,
|
||||
retrieval_type="full",
|
||||
)
|
||||
|
||||
second = crusher.crush(payload)
|
||||
assert first.compressed == second.compressed, (
|
||||
"PR-B5: SmartCrusher output must be deterministic regardless of TOIN observation state."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"items",
|
||||
[
|
||||
# Tiny, mid, and at-threshold inputs covering the conditional
|
||||
# paths inside the Rust crusher (lossless tabular, lossy with
|
||||
# CCR, pass-through). Spec asks for a hypothesis property test;
|
||||
# hypothesis is optional, so we cover the parametrized cases
|
||||
# unconditionally and add the property test below behind an
|
||||
# importorskip.
|
||||
[],
|
||||
[{"id": 1}],
|
||||
[{"id": i, "status": "ok"} for i in range(8)],
|
||||
[{"id": i, "status": "ok", "msg": "fine"} for i in range(50)],
|
||||
[{"id": i, "code": 200 + i % 3, "err": ""} for i in range(120)],
|
||||
],
|
||||
)
|
||||
def test_smart_crusher_determinism_parametrized(items: list[dict[str, object]]) -> None:
|
||||
"""Two crush() calls on the same input must return byte-equal output."""
|
||||
smart_crusher_module = pytest.importorskip("headroom.transforms.smart_crusher")
|
||||
SmartCrusher = smart_crusher_module.SmartCrusher
|
||||
SmartCrusherConfig = smart_crusher_module.SmartCrusherConfig
|
||||
import json as _json
|
||||
|
||||
crusher = SmartCrusher(config=SmartCrusherConfig(enabled=True))
|
||||
payload = _json.dumps(items)
|
||||
a = crusher.crush(payload)
|
||||
b = crusher.crush(payload)
|
||||
assert a.compressed == b.compressed
|
||||
|
||||
|
||||
def test_smart_crusher_determinism_property():
|
||||
"""Property: any input → byte-stable SmartCrusher output across two calls.
|
||||
|
||||
Skipped if `hypothesis` is not installed (it is not a hard dep of
|
||||
Headroom). The parametrized test above covers the deterministic
|
||||
surface unconditionally.
|
||||
"""
|
||||
pytest.importorskip("hypothesis")
|
||||
from hypothesis import given, settings
|
||||
from hypothesis import strategies as st
|
||||
|
||||
smart_crusher_module = pytest.importorskip("headroom.transforms.smart_crusher")
|
||||
SmartCrusher = smart_crusher_module.SmartCrusher
|
||||
SmartCrusherConfig = smart_crusher_module.SmartCrusherConfig
|
||||
crusher = SmartCrusher(config=SmartCrusherConfig(enabled=True))
|
||||
|
||||
@given(
|
||||
st.lists(
|
||||
st.fixed_dictionaries(
|
||||
{
|
||||
"id": st.integers(min_value=0, max_value=10_000),
|
||||
"status": st.sampled_from(["ok", "error", "pending"]),
|
||||
}
|
||||
),
|
||||
min_size=0,
|
||||
max_size=20,
|
||||
)
|
||||
)
|
||||
@settings(max_examples=25, deadline=None)
|
||||
def _check(items: list[dict[str, object]]) -> None:
|
||||
import json as _json
|
||||
|
||||
payload = _json.dumps(items)
|
||||
a = crusher.crush(payload)
|
||||
b = crusher.crush(payload)
|
||||
assert a.compressed == b.compressed
|
||||
|
||||
_check()
|
||||
234
tests/test_toin_publish.py
Normal file
234
tests/test_toin_publish.py
Normal file
|
|
@ -0,0 +1,234 @@
|
|||
"""PR-B5 acceptance tests for ``headroom.cli.toin_publish``.
|
||||
|
||||
Pins:
|
||||
|
||||
1. ``publish()`` writes a TOML file the stdlib ``tomllib`` can parse.
|
||||
2. Slices below ``--min-observations`` are filtered out.
|
||||
3. Rows include ``auth_mode``, ``model_family``, ``structure_hash``,
|
||||
``strategy_hint``, ``confidence``, ``observations`` — the schema
|
||||
``crates/headroom-core/src/transforms/recommendations.rs`` consumes.
|
||||
4. The CLI entry point honors ``--output`` / ``--min-observations``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# Python 3.11+ has tomllib in stdlib; otherwise tomli is shipped as a
|
||||
# dependency by the project's pyproject.toml.
|
||||
if sys.version_info >= (3, 11):
|
||||
import tomllib
|
||||
else: # pragma: no cover - only hit on Python 3.10
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
|
||||
from headroom.cli.toin_publish import main as publish_main
|
||||
from headroom.cli.toin_publish import publish
|
||||
from headroom.telemetry import (
|
||||
TOINConfig,
|
||||
ToolIntelligenceNetwork,
|
||||
ToolSignature,
|
||||
)
|
||||
|
||||
|
||||
def _record(
|
||||
toin: ToolIntelligenceNetwork,
|
||||
*,
|
||||
items: list[dict[str, object]],
|
||||
n: int,
|
||||
auth_mode: str,
|
||||
model_family: str,
|
||||
strategy: str = "smart_crusher",
|
||||
) -> ToolSignature:
|
||||
"""Drive ``record_compression`` ``n`` times for the given slice."""
|
||||
sig = ToolSignature.from_items(items)
|
||||
for _ in range(n):
|
||||
toin.record_compression(
|
||||
tool_signature=sig,
|
||||
original_count=len(items),
|
||||
compressed_count=max(1, len(items) // 2),
|
||||
original_tokens=1000,
|
||||
compressed_tokens=500,
|
||||
strategy=strategy,
|
||||
auth_mode=auth_mode,
|
||||
model_family=model_family,
|
||||
)
|
||||
return sig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_toin(tmp_path: Path) -> ToolIntelligenceNetwork:
|
||||
"""Isolated TOIN handle so tests don't see each other's state."""
|
||||
return ToolIntelligenceNetwork(
|
||||
TOINConfig(
|
||||
storage_path=str(tmp_path / "toin_publish.json"),
|
||||
auto_save_interval=0,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_publish_command_writes_toml(fresh_toin: ToolIntelligenceNetwork, tmp_path: Path) -> None:
|
||||
"""publish() emits a parseable TOML file with the expected schema."""
|
||||
items = [{"id": i, "status": "ok"} for i in range(20)]
|
||||
sig = _record(
|
||||
fresh_toin,
|
||||
items=items,
|
||||
n=60,
|
||||
auth_mode="payg",
|
||||
model_family="claude-3-5",
|
||||
)
|
||||
|
||||
output = tmp_path / "recommendations.toml"
|
||||
rows_written = publish(
|
||||
output_path=output,
|
||||
min_observations=50,
|
||||
toin=fresh_toin,
|
||||
)
|
||||
assert rows_written == 1
|
||||
|
||||
parsed = tomllib.loads(output.read_text(encoding="utf-8"))
|
||||
assert "recommendation" in parsed
|
||||
rec_list = parsed["recommendation"]
|
||||
assert isinstance(rec_list, list)
|
||||
assert len(rec_list) == 1
|
||||
|
||||
row = rec_list[0]
|
||||
assert set(row.keys()) == {
|
||||
"auth_mode",
|
||||
"model_family",
|
||||
"structure_hash",
|
||||
"strategy_hint",
|
||||
"confidence",
|
||||
"observations",
|
||||
}
|
||||
assert row["auth_mode"] == "payg"
|
||||
assert row["model_family"] == "claude-3-5"
|
||||
assert row["structure_hash"] == sig.structure_hash
|
||||
assert row["strategy_hint"] == "smart_crusher"
|
||||
assert isinstance(row["confidence"], float)
|
||||
assert 0.0 <= row["confidence"] <= 1.0
|
||||
assert row["observations"] == 60
|
||||
|
||||
|
||||
def test_publish_filters_below_min_observations(
|
||||
fresh_toin: ToolIntelligenceNetwork,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Slices below the observation floor are dropped from the TOML."""
|
||||
eligible = [{"id": i} for i in range(10)]
|
||||
rare = [{"name": str(i)} for i in range(10)]
|
||||
|
||||
_record(fresh_toin, items=eligible, n=60, auth_mode="payg", model_family="claude-3-5")
|
||||
_record(fresh_toin, items=rare, n=10, auth_mode="payg", model_family="claude-3-5")
|
||||
|
||||
output = tmp_path / "recs.toml"
|
||||
rows_written = publish(output_path=output, min_observations=50, toin=fresh_toin)
|
||||
assert rows_written == 1
|
||||
|
||||
parsed = tomllib.loads(output.read_text(encoding="utf-8"))
|
||||
rec_list = parsed["recommendation"]
|
||||
assert len(rec_list) == 1
|
||||
# The eligible signature wins; the rare one is filtered.
|
||||
assert rec_list[0]["observations"] == 60
|
||||
|
||||
|
||||
def test_publish_emits_one_row_per_tenant_slice(
|
||||
fresh_toin: ToolIntelligenceNetwork, tmp_path: Path
|
||||
) -> None:
|
||||
"""Same tool-signature, different (auth_mode, model_family) ⇒ separate rows."""
|
||||
items = [{"id": i, "status": "ok"} for i in range(15)]
|
||||
_record(fresh_toin, items=items, n=60, auth_mode="payg", model_family="claude-3-5")
|
||||
_record(fresh_toin, items=items, n=60, auth_mode="oauth", model_family="claude-3-5")
|
||||
_record(fresh_toin, items=items, n=60, auth_mode="payg", model_family="gpt-4o")
|
||||
|
||||
output = tmp_path / "recs.toml"
|
||||
rows_written = publish(output_path=output, min_observations=50, toin=fresh_toin)
|
||||
assert rows_written == 3
|
||||
|
||||
parsed = tomllib.loads(output.read_text(encoding="utf-8"))
|
||||
rec_list = parsed["recommendation"]
|
||||
keys = sorted((r["auth_mode"], r["model_family"]) for r in rec_list)
|
||||
assert keys == [("oauth", "claude-3-5"), ("payg", "claude-3-5"), ("payg", "gpt-4o")]
|
||||
|
||||
|
||||
def test_publish_writes_empty_file_with_no_eligible_rows(
|
||||
fresh_toin: ToolIntelligenceNetwork, tmp_path: Path
|
||||
) -> None:
|
||||
"""No qualifying patterns ⇒ valid empty TOML, not an exception."""
|
||||
output = tmp_path / "recs.toml"
|
||||
rows_written = publish(output_path=output, min_observations=50, toin=fresh_toin)
|
||||
assert rows_written == 0
|
||||
|
||||
body = output.read_text(encoding="utf-8")
|
||||
parsed = tomllib.loads(body)
|
||||
assert parsed == {}
|
||||
# Header still shipped so ops can identify the file.
|
||||
assert body.startswith("# Auto-generated")
|
||||
|
||||
|
||||
def test_publish_rows_are_deterministically_sorted(
|
||||
fresh_toin: ToolIntelligenceNetwork, tmp_path: Path
|
||||
) -> None:
|
||||
"""Rows sort by (auth_mode, model_family, structure_hash) for clean diffs.
|
||||
|
||||
Use *structurally distinct* tool signatures so the hashes truly
|
||||
differ — `ToolSignature` keys off field names + types, not values.
|
||||
"""
|
||||
one_field = [{"id": i} for i in range(8)]
|
||||
two_fields = [{"id": i, "code": 200 + i} for i in range(8)]
|
||||
|
||||
_record(fresh_toin, items=one_field, n=60, auth_mode="payg", model_family="claude-3-5")
|
||||
_record(fresh_toin, items=two_fields, n=60, auth_mode="payg", model_family="claude-3-5")
|
||||
_record(fresh_toin, items=one_field, n=60, auth_mode="oauth", model_family="gpt-4o")
|
||||
|
||||
output = tmp_path / "recs.toml"
|
||||
publish(output_path=output, min_observations=50, toin=fresh_toin)
|
||||
parsed = tomllib.loads(output.read_text(encoding="utf-8"))
|
||||
rec_list = parsed["recommendation"]
|
||||
# First sort key: auth_mode (oauth < payg).
|
||||
assert [r["auth_mode"] for r in rec_list] == ["oauth", "payg", "payg"]
|
||||
# And within payg, structure_hash sorts asc.
|
||||
payg_rows = [r for r in rec_list if r["auth_mode"] == "payg"]
|
||||
assert payg_rows == sorted(payg_rows, key=lambda r: r["structure_hash"])
|
||||
|
||||
|
||||
def test_cli_entrypoint_writes_to_output_arg(tmp_path: Path, monkeypatch) -> None:
|
||||
"""`python -m headroom.cli.toin_publish --output X --min-observations N`."""
|
||||
storage = tmp_path / "toin.json"
|
||||
monkeypatch.setenv("HEADROOM_TOIN_PATH", str(storage))
|
||||
|
||||
# Prime the global TOIN singleton with eligible data.
|
||||
from headroom.telemetry.toin import get_toin, reset_toin
|
||||
|
||||
reset_toin()
|
||||
try:
|
||||
toin = get_toin()
|
||||
_record(
|
||||
toin,
|
||||
items=[{"id": i} for i in range(10)],
|
||||
n=55,
|
||||
auth_mode="payg",
|
||||
model_family="claude-3-5",
|
||||
)
|
||||
toin.save()
|
||||
|
||||
output = tmp_path / "out.toml"
|
||||
rc = publish_main(
|
||||
["--output", str(output), "--min-observations", "50"],
|
||||
)
|
||||
assert rc == 0
|
||||
assert output.exists()
|
||||
parsed = tomllib.loads(output.read_text(encoding="utf-8"))
|
||||
assert len(parsed.get("recommendation", [])) == 1
|
||||
finally:
|
||||
reset_toin()
|
||||
|
||||
|
||||
def test_cli_rejects_non_positive_min_observations(tmp_path: Path) -> None:
|
||||
"""`--min-observations 0` is a CLI-level error."""
|
||||
output = tmp_path / "out.toml"
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
publish_main(["--output", str(output), "--min-observations", "0"])
|
||||
assert exc_info.value.code != 0
|
||||
Loading…
Add table
Add a link
Reference in a new issue