car-topology 0.55.0

//! Execution records — what a topology actually cost and actually achieved.
//!
//! The whole point of this crate is that topology design should be scored on
//! *measured* outcomes rather than a structural surrogate. A record is one
//! (query, topology) pair executed for real, carrying the utility it earned and
//! the tokens it burned.
//!
//! ## Per-task token normalization
//!
//! Raw token counts are dominated by task difficulty: a hard query costs more
//! under every topology, so a model regressed on raw `τ` mostly learns to
//! predict difficulty. The paper's fix, reproduced here exactly, is to divide
//! by the mean token count of the *same* task:
//!
//! ```text
//! τ̃_j = τ_j / mean{ τ_j' : task(j') == task(j) }        (Eq. 2)
//! R_j  = u_j − λ · τ̃_j,   λ = 0.1
//! ```
//!
//! Difficulty cancels and the within-task ranking survives, which is the only
//! ranking a topology selector can act on.

use std::collections::{HashMap, HashSet};

use serde::{Deserialize, Serialize};

use crate::error::TopologyError;
use crate::topology::Topology;

/// The default cost weight `λ` in `R = u − λ·τ̃`.
///
/// 0.1 is the paper's operating point, chosen — per its ablation (d) — because
/// it sits on the accuracy/token frontier rather than at an unconstrained
/// cost minimum: pushing `λ` more negative keeps cutting tokens but starts
/// trading accuracy away.
pub const DEFAULT_COST_WEIGHT: f32 = 0.1;

/// The utility above which a record counts as having solved its task, and so is
/// eligible to enter the codebook.
pub const DEFAULT_SURVIVOR_THRESHOLD: f32 = 0.5;

/// One (query, topology) pair executed for real.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ExecutionRecord {
    /// Identity of the task this record was collected on. Records sharing a
    /// `task_id` form the group Eq. (2) normalizes within, so this must be the
    /// task — not the run — or the normalizer divides every record by itself.
    pub task_id: String,
    /// The query embedding `c`. Any frozen sentence encoder will do; the crate
    /// only requires that every record in a set uses the same one, at the same
    /// dimension.
    pub query: Vec<f32>,
    /// The topology the team ran under.
    pub topology: Topology,
    /// Measured utility in `[0, 1]` — 1 for a solved task, 0 for a failed one,
    /// or a graded score where the caller has one.
    pub utility: f32,
    /// Measured total LLM tokens for the run. Not an estimate, not `|E|`.
    pub tokens: u64,
}

impl ExecutionRecord {
    /// Construct a record, clamping utility into `[0, 1]`.
    pub fn new(
        task_id: impl Into<String>,
        query: Vec<f32>,
        topology: Topology,
        utility: f32,
        tokens: u64,
    ) -> Self {
        Self {
            task_id: task_id.into(),
            query,
            topology,
            utility: utility.clamp(0.0, 1.0),
            tokens,
        }
    }
}

/// A validated collection of execution records, plus the per-task normalization
/// Eq. (2) derives from it.
///
/// Validation happens once, at construction, so every downstream stage can read
/// `queries` and `topologies` without re-checking dimensions.
#[derive(Debug, Clone)]
pub struct RecordSet {
    records: Vec<ExecutionRecord>,
    /// `τ̃_j` for each record, parallel to `records`.
    normalized_cost: Vec<f32>,
    query_dim: usize,
    n: usize,
    /// Which encoder produced every `query` in this set, when the caller knows.
    embedder: Option<String>,
}

impl RecordSet {
    /// Validate and index a set of records.
    ///
    /// Rejects a set whose records disagree on team size or query dimension:
    /// those are the two shapes every later stage indexes by, and a mixed set
    /// would otherwise fail deep inside a matrix solve with no useful message.
    pub fn new(records: Vec<ExecutionRecord>) -> Result<Self, TopologyError> {
        if records.is_empty() {
            return Err(TopologyError::NoRecords { kind: "execution" });
        }
        let query_dim = records[0].query.len();
        let n = records[0].topology.n();
        for (index, r) in records.iter().enumerate() {
            if r.query.len() != query_dim {
                return Err(TopologyError::QueryDimMismatch {
                    expected: query_dim,
                    found: r.query.len(),
                });
            }
            if r.topology.n() != n {
                return Err(TopologyError::SizeMismatch {
                    expected: n,
                    found: r.topology.n(),
                });
            }
            if !r.utility.is_finite() {
                return Err(TopologyError::NonFinite {
                    index,
                    field: "utility",
                });
            }
            if r.query.iter().any(|v| !v.is_finite()) {
                return Err(TopologyError::NonFinite {
                    index,
                    field: "query embedding",
                });
            }
        }

        let normalized_cost = normalize_per_task(&records);
        Ok(Self {
            records,
            normalized_cost,
            query_dim,
            n,
            embedder: None,
        })
    }

    /// Validate and index a set of records, recording which encoder produced
    /// their query embeddings.
    ///
    /// Worth carrying even though nothing here reads the embeddings' *meaning*.
    /// Every downstream stage compares queries by cosine similarity, so the
    /// whole pipeline is defined over directions in one embedding space; mixing
    /// two encoders in a set, or querying a fitted selector with a third, does
    /// not fail — it silently answers from arbitrary similarities.
    /// [`crate::journal`] refuses to build a mixed set, and
    /// [`crate::TopologySelector::select_with`] refuses a mismatched query.
    pub fn with_embedder(
        records: Vec<ExecutionRecord>,
        embedder: impl Into<String>,
    ) -> Result<Self, TopologyError> {
        let mut set = Self::new(records)?;
        set.embedder = Some(embedder.into());
        Ok(set)
    }

    /// The encoder that produced this set's query embeddings, when known.
    pub fn embedder(&self) -> Option<&str> {
        self.embedder.as_deref()
    }

    /// The records, in the order they were supplied.
    pub fn records(&self) -> &[ExecutionRecord] {
        &self.records
    }

    /// Number of records.
    pub fn len(&self) -> usize {
        self.records.len()
    }

    /// Whether the set holds no records. Always false for a constructed
    /// [`RecordSet`] — [`RecordSet::new`] rejects the empty case — but clippy
    /// asks for it alongside `len`, and a caller holding one by reference
    /// shouldn't have to know that.
    pub fn is_empty(&self) -> bool {
        self.records.is_empty()
    }

    /// Team size shared by every record.
    pub fn team_size(&self) -> usize {
        self.n
    }

    /// Query-embedding dimension shared by every record.
    pub fn query_dim(&self) -> usize {
        self.query_dim
    }

    /// `τ̃_j` — the record's tokens divided by the mean tokens of its task.
    ///
    /// A task whose records all cost zero tokens normalizes to 1.0 rather than
    /// dividing by zero: "no measured cost" should read as "average", not as a
    /// NaN that poisons every downstream fit.
    pub fn normalized_cost(&self, index: usize) -> f32 {
        self.normalized_cost.get(index).copied().unwrap_or(1.0)
    }

    /// All normalized costs, parallel to [`RecordSet::records`].
    pub fn normalized_costs(&self) -> &[f32] {
        &self.normalized_cost
    }

    /// `R_j = u_j − λ·τ̃_j` (Eq. 2).
    pub fn reward(&self, index: usize, cost_weight: f32) -> f32 {
        self.records[index].utility - cost_weight * self.normalized_cost(index)
    }

    /// Distinct task ids, in first-seen order.
    ///
    /// Membership goes through a set rather than a linear scan of what has been
    /// seen: a journal that a daemon appends to indefinitely has one task per
    /// coordination run, so the scan version was quadratic in exactly the
    /// dimension that grows without bound.
    pub fn task_ids(&self) -> Vec<String> {
        let mut seen = HashSet::new();
        let mut order = Vec::new();
        for r in &self.records {
            if seen.insert(r.task_id.as_str()) {
                order.push(r.task_id.clone());
            }
        }
        order
    }

    /// Indices of the records whose utility clears `threshold` — the
    /// "reward-surviving" set the codebook is fitted on.
    pub fn survivors(&self, threshold: f32) -> Vec<usize> {
        (0..self.records.len())
            .filter(|&i| self.records[i].utility > threshold)
            .collect()
    }
}

/// Compute `τ̃` for every record by dividing by its task's mean token count.
fn normalize_per_task(records: &[ExecutionRecord]) -> Vec<f32> {
    let mut sums: HashMap<&str, (f64, usize)> = HashMap::new();
    for r in records {
        let entry = sums.entry(r.task_id.as_str()).or_insert((0.0, 0));
        entry.0 += r.tokens as f64;
        entry.1 += 1;
    }
    records
        .iter()
        .map(|r| {
            let (sum, count) = sums[r.task_id.as_str()];
            let mean = sum / count as f64;
            if mean <= 0.0 {
                1.0
            } else {
                (r.tokens as f64 / mean) as f32
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::topology::CoordinationShape;

    fn rec(task: &str, q: f32, shape: CoordinationShape, u: f32, tokens: u64) -> ExecutionRecord {
        ExecutionRecord::new(
            task,
            vec![q, 1.0 - q],
            shape.topology(4).unwrap(),
            u,
            tokens,
        )
    }

    #[test]
    fn normalization_cancels_task_difficulty() {
        // Two tasks, one 10x more expensive than the other, with the same
        // *within-task* ranking between two topologies.
        let set = RecordSet::new(vec![
            rec("easy", 0.1, CoordinationShape::Pipeline, 1.0, 100),
            rec("easy", 0.1, CoordinationShape::Debate, 1.0, 300),
            rec("hard", 0.9, CoordinationShape::Pipeline, 1.0, 1000),
            rec("hard", 0.9, CoordinationShape::Debate, 1.0, 3000),
        ])
        .unwrap();

        // Raw tokens differ 10x across tasks; normalized costs are identical.
        assert!((set.normalized_cost(0) - set.normalized_cost(2)).abs() < 1e-6);
        assert!((set.normalized_cost(1) - set.normalized_cost(3)).abs() < 1e-6);
        // And the within-task ordering survives.
        assert!(set.normalized_cost(0) < set.normalized_cost(1));
    }

    #[test]
    fn normalized_costs_average_to_one_within_a_task() {
        let set = RecordSet::new(vec![
            rec("t", 0.5, CoordinationShape::Pipeline, 1.0, 200),
            rec("t", 0.5, CoordinationShape::Debate, 1.0, 400),
            rec("t", 0.5, CoordinationShape::Supervisor, 1.0, 600),
        ])
        .unwrap();
        let mean: f32 = set.normalized_costs().iter().sum::<f32>() / 3.0;
        assert!((mean - 1.0).abs() < 1e-5, "got {mean}");
    }

    #[test]
    fn zero_token_task_normalizes_to_one_not_nan() {
        let set = RecordSet::new(vec![
            rec("t", 0.5, CoordinationShape::Pipeline, 1.0, 0),
            rec("t", 0.5, CoordinationShape::Debate, 1.0, 0),
        ])
        .unwrap();
        assert!(set.normalized_costs().iter().all(|c| c.is_finite()));
        assert_eq!(set.normalized_cost(0), 1.0);
    }

    #[test]
    fn reward_penalizes_the_expensive_topology() {
        let set = RecordSet::new(vec![
            rec("t", 0.5, CoordinationShape::Pipeline, 1.0, 100),
            rec("t", 0.5, CoordinationShape::Debate, 1.0, 900),
        ])
        .unwrap();
        assert!(set.reward(0, DEFAULT_COST_WEIGHT) > set.reward(1, DEFAULT_COST_WEIGHT));
    }

    #[test]
    fn survivors_filter_on_utility() {
        let set = RecordSet::new(vec![
            rec("t", 0.5, CoordinationShape::Pipeline, 1.0, 100),
            rec("t", 0.5, CoordinationShape::Debate, 0.0, 100),
            rec("t", 0.5, CoordinationShape::Supervisor, 0.5, 100),
        ])
        .unwrap();
        // Strictly greater than the threshold, per the paper's `u > 0.5`.
        assert_eq!(set.survivors(DEFAULT_SURVIVOR_THRESHOLD), vec![0]);
    }

    #[test]
    fn mixed_query_dimensions_are_rejected() {
        let mut a = rec("t", 0.5, CoordinationShape::Pipeline, 1.0, 100);
        let b = rec("t", 0.5, CoordinationShape::Debate, 1.0, 100);
        a.query = vec![0.1, 0.2, 0.3];
        assert!(matches!(
            RecordSet::new(vec![a, b]),
            Err(TopologyError::QueryDimMismatch { .. })
        ));
    }

    #[test]
    fn mixed_team_sizes_are_rejected() {
        let a = rec("t", 0.5, CoordinationShape::Pipeline, 1.0, 100);
        let b = ExecutionRecord::new(
            "t",
            vec![0.5, 0.5],
            CoordinationShape::Pipeline.topology(5).unwrap(),
            1.0,
            100,
        );
        assert!(matches!(
            RecordSet::new(vec![a, b]),
            Err(TopologyError::SizeMismatch { .. })
        ));
    }

    #[test]
    fn non_finite_inputs_are_rejected() {
        let mut a = rec("t", 0.5, CoordinationShape::Pipeline, 1.0, 100);
        a.query = vec![f32::NAN, 0.5];
        assert!(matches!(
            RecordSet::new(vec![a]),
            Err(TopologyError::NonFinite { .. })
        ));
    }

    #[test]
    fn empty_sets_are_rejected() {
        assert!(matches!(
            RecordSet::new(vec![]),
            Err(TopologyError::NoRecords { .. })
        ));
    }
}