car-topology 0.55.0

//! One-pass selection — predictor, decode, one batched proxy call, done.
//!
//! ```text
//! C_M(c) = top-M codes under p(k | c)
//! A(c)   = { decode(k) : k ∈ C_M(c) },  deduplicated
//! A*(c)  = argmax_{A ∈ A(c)}  û(A, c) − λ·ĉ(A, c)              (Eq. 11)
//! ```
//!
//! No sampling loop, no iterative refinement, no message passing on the
//! test-time path. The published iterative designers evaluate `T·K` candidates
//! through a graph network per query — the paper measures 301–396 ms — against
//! one predictor pass, at most `M` decodes, and one batched dot-product call
//! here.
//!
//! Latency is not the reason this matters for CAR. CAR's coordination patterns
//! are picked today by `car_agents::coordinator::Coordinator`, which spends a
//! full LLM round trip on a hand-written prompt to choose among four patterns,
//! and learns nothing from how the choice turned out. A selector fitted on
//! execution records replaces that round trip with a fold over data CAR already
//! has, and — unlike the prompt — gets better as records accumulate.

use serde::{Deserialize, Serialize};

use crate::codebook::{Codebook, CodebookConfig};
use crate::error::TopologyError;
use crate::predictor::{CodePredictor, PredictorConfig};
use crate::proxy::{ExecutionProxy, ProxyConfig, ProxyScore};
use crate::record::{RecordSet, DEFAULT_COST_WEIGHT};
use crate::topology::{shape_of, CoordinationShape, Topology};

/// Configuration for the whole three-stage fit plus the selection rule.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SelectorConfig {
    pub codebook: CodebookConfig,
    pub predictor: PredictorConfig,
    pub proxy: ProxyConfig,
    /// `M` — how many codes the predictor proposes before reranking. The
    /// paper's 5.
    pub top_m: usize,
    /// `λ` in Eq. (11). Kept separate from [`PredictorConfig::cost_weight`]
    /// because they weight different things — one shapes the training targets,
    /// one the test-time decision — even though the paper uses 0.1 for both.
    pub cost_weight: f32,
}

impl Default for SelectorConfig {
    fn default() -> Self {
        Self {
            codebook: CodebookConfig::default(),
            predictor: PredictorConfig::default(),
            proxy: ProxyConfig::default(),
            top_m: 5,
            cost_weight: DEFAULT_COST_WEIGHT,
        }
    }
}

/// One scored candidate, kept on the [`Selection`] so a caller can log why the
/// winner won rather than inferring it.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Candidate {
    /// Index into the codebook.
    pub code: usize,
    pub topology: Topology,
    /// The coordination pattern this topology corresponds to, when it
    /// corresponds to one. `None` for a learned topology that matches no
    /// named family — which is a real outcome, not a failure.
    pub shape: Option<CoordinationShape>,
    /// `p(k | c)` for this code.
    pub prior: f32,
    pub score: ProxyScore,
    /// `û − λ·ĉ`.
    pub objective: f32,
}

/// The outcome of one selection.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Selection {
    /// The chosen topology.
    pub topology: Topology,
    /// Its codebook index.
    pub code: usize,
    /// Its coordination pattern, when it has one.
    pub shape: Option<CoordinationShape>,
    pub score: ProxyScore,
    pub objective: f32,
    /// Every candidate that was scored, in the order the predictor proposed
    /// them.
    pub considered: Vec<Candidate>,
}

/// A fitted amortized topology selector: codebook + conditional prior +
/// execution-grounded proxy.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "TopologySelectorWire")]
pub struct TopologySelector {
    codebook: Codebook,
    predictor: CodePredictor,
    proxy: ExecutionProxy,
    top_m: usize,
    cost_weight: f32,
    /// The encoder the training queries came from, carried through from
    /// [`RecordSet::embedder`] so [`TopologySelector::select_with`] can refuse
    /// a query from a different one.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    embedder: Option<String>,
}

/// On-disk shape, validated on the way in.
///
/// [`TopologySelector::from_parts`] refuses a codebook and predictor that
/// disagree, but serde bypasses constructors, so a persisted selector could
/// come back with exactly the mismatch that check exists to prevent — and it
/// would not fail, it would rank against a prior indexed into a different
/// codebook. Persisting a fitted selector is on the roadmap; validating here
/// while the crate is unreleased is far cheaper than discovering it later.
#[derive(Deserialize)]
struct TopologySelectorWire {
    codebook: Codebook,
    predictor: CodePredictor,
    proxy: ExecutionProxy,
    top_m: usize,
    cost_weight: f32,
    #[serde(default)]
    embedder: Option<String>,
}

impl TryFrom<TopologySelectorWire> for TopologySelector {
    type Error = TopologyError;

    fn try_from(wire: TopologySelectorWire) -> Result<Self, Self::Error> {
        // The parts validate themselves on their own way in; this is the
        // cross-part agreement `from_parts` checks.
        if wire.codebook.len() != wire.predictor.codes() {
            return Err(TopologyError::BadConfig {
                field: "predictor",
                expected: "one output per codebook entry",
                found: format!(
                    "{} codes vs {} outputs",
                    wire.codebook.len(),
                    wire.predictor.codes()
                ),
            });
        }
        if wire.codebook.team_size() != wire.proxy.team_size() {
            return Err(TopologyError::SizeMismatch {
                expected: wire.codebook.team_size(),
                found: wire.proxy.team_size(),
            });
        }
        if wire.predictor.query_dim() != wire.proxy.query_dim() {
            return Err(TopologyError::QueryDimMismatch {
                expected: wire.predictor.query_dim(),
                found: wire.proxy.query_dim(),
            });
        }
        if wire.top_m == 0 {
            return Err(TopologyError::BadConfig {
                field: "top_m",
                expected: "at least 1",
                found: "0".into(),
            });
        }
        if !wire.cost_weight.is_finite() {
            return Err(TopologyError::BadConfig {
                field: "cost_weight",
                expected: "finite",
                found: format!("{}", wire.cost_weight),
            });
        }
        Ok(TopologySelector {
            codebook: wire.codebook,
            predictor: wire.predictor,
            proxy: wire.proxy,
            top_m: wire.top_m,
            cost_weight: wire.cost_weight,
            embedder: wire.embedder,
        })
    }
}

impl TopologySelector {
    /// Fit all three stages on one set of execution records.
    ///
    /// No LLM calls happen here or anywhere downstream — the records are the
    /// only place a model was ever involved.
    pub fn fit(records: &RecordSet, config: &SelectorConfig) -> Result<Self, TopologyError> {
        if config.top_m == 0 {
            return Err(TopologyError::BadConfig {
                field: "top_m",
                expected: "at least 1",
                found: "0".into(),
            });
        }
        if !config.cost_weight.is_finite() {
            return Err(TopologyError::BadConfig {
                field: "cost_weight",
                expected: "finite",
                found: format!("{}", config.cost_weight),
            });
        }

        let codebook = Codebook::fit(records, &config.codebook)?;
        let predictor = CodePredictor::fit(records, &codebook, &config.predictor)?;
        let proxy = ExecutionProxy::fit(records, &config.proxy)?;

        Ok(Self {
            codebook,
            predictor,
            proxy,
            top_m: config.top_m,
            cost_weight: config.cost_weight,
            embedder: records.embedder().map(str::to_owned),
        })
    }

    /// Assemble a selector from separately fitted parts.
    pub fn from_parts(
        codebook: Codebook,
        predictor: CodePredictor,
        proxy: ExecutionProxy,
        top_m: usize,
        cost_weight: f32,
    ) -> Result<Self, TopologyError> {
        if codebook.len() != predictor.codes() {
            return Err(TopologyError::BadConfig {
                field: "predictor",
                expected: "one output per codebook entry",
                found: format!("{} codes vs {} outputs", codebook.len(), predictor.codes()),
            });
        }
        if codebook.team_size() != proxy.team_size() {
            return Err(TopologyError::SizeMismatch {
                expected: codebook.team_size(),
                found: proxy.team_size(),
            });
        }
        // The same checks the serde `try_from` path enforces. Leaving them out
        // here let a mismatched pair construct fine and then fail on every
        // `select()`, and silently rounded `top_m: 0` up to 1 — two ways for
        // the same object to be valid or not depending on how it was built.
        if predictor.query_dim() != proxy.query_dim() {
            return Err(TopologyError::QueryDimMismatch {
                expected: predictor.query_dim(),
                found: proxy.query_dim(),
            });
        }
        if top_m == 0 {
            return Err(TopologyError::BadConfig {
                field: "top_m",
                expected: "at least 1",
                found: "0".into(),
            });
        }
        if !cost_weight.is_finite() {
            return Err(TopologyError::BadConfig {
                field: "cost_weight",
                expected: "finite",
                found: format!("{cost_weight}"),
            });
        }
        Ok(Self {
            codebook,
            predictor,
            proxy,
            top_m,
            cost_weight,
            embedder: None,
        })
    }

    /// Label this selector with the encoder its training queries came from.
    ///
    /// [`TopologySelector::fit`] carries the label over from the record set
    /// automatically; this is for a selector assembled by
    /// [`TopologySelector::from_parts`].
    pub fn with_embedder(mut self, embedder: impl Into<String>) -> Self {
        self.embedder = Some(embedder.into());
        self
    }

    /// The encoder this selector was fitted under, when known.
    pub fn embedder(&self) -> Option<&str> {
        self.embedder.as_deref()
    }

    /// The fitted codebook.
    pub fn codebook(&self) -> &Codebook {
        &self.codebook
    }

    /// The fitted conditional prior.
    pub fn predictor(&self) -> &CodePredictor {
        &self.predictor
    }

    /// The fitted proxy.
    pub fn proxy(&self) -> &ExecutionProxy {
        &self.proxy
    }

    /// Team size every selection is over.
    pub fn team_size(&self) -> usize {
        self.codebook.team_size()
    }

    /// Select a topology for a query, refusing one from the wrong encoder.
    ///
    /// The embedding space is the selector's whole coordinate system — the
    /// prior is a kernel over cosine similarity to the training conditions —
    /// so a query from another encoder is not merely noisier, it is meaningless
    /// here, and nothing downstream would notice. Prefer this over
    /// [`TopologySelector::select`] wherever the caller knows which embedder it
    /// used.
    ///
    /// A selector with no label cannot be checked, so the query passes: an
    /// unlabeled selector means "nobody recorded it", not "any encoder fits".
    pub fn select_with(&self, query: &[f32], embedder: &str) -> Result<Selection, TopologyError> {
        if let Some(fitted) = &self.embedder {
            if fitted != embedder {
                return Err(TopologyError::EmbedderMismatch {
                    fitted: fitted.clone(),
                    query: embedder.to_string(),
                });
            }
        }
        self.select(query)
    }

    /// Select a topology for one query (Eq. 11).
    ///
    /// Deterministic, including under ties: candidates are compared on the
    /// objective first and on codebook index second, so the same query always
    /// yields the same topology.
    ///
    /// Does no embedder check — see [`TopologySelector::select_with`].
    pub fn select(&self, query: &[f32]) -> Result<Selection, TopologyError> {
        let prior = self.predictor.predict(query)?;
        let codes = self.predictor.top_codes(query, self.top_m)?;

        // Solve the proxy's heads ONCE. They depend only on the query, so
        // calling `proxy.score` per candidate re-solves the locally-weighted
        // ridge for every one — measured ~6x slower on a 12k-record fit, and
        // flatly contrary to the "one batched pass" this crate documents.
        let conditioned = self.proxy.condition(query)?;

        let mut considered: Vec<Candidate> = Vec::with_capacity(codes.len());
        for code in codes {
            let topology = self
                .codebook
                .decode(code)
                .ok_or(TopologyError::EmptyCodebook)?
                .clone();
            // Deduplicate decoded topologies: two codes can decode to the same
            // graph after a quantizer merge, and scoring it twice would let an
            // arbitrary index break the tie.
            if considered.iter().any(|c| c.topology == topology) {
                continue;
            }
            let score = conditioned.score(&topology, query)?;
            considered.push(Candidate {
                code,
                shape: shape_of(&topology),
                topology,
                prior: prior.get(code).copied().unwrap_or(0.0),
                score,
                objective: score.objective(self.cost_weight),
            });
        }

        let winner = considered
            .iter()
            .enumerate()
            .max_by(|(ai, a), (bi, b)| {
                a.objective
                    .total_cmp(&b.objective)
                    // Lower codebook index wins a tie; then lower position, so
                    // the comparison is total even for equal codes.
                    .then(b.code.cmp(&a.code))
                    .then(bi.cmp(ai))
            })
            .map(|(_, c)| c.clone())
            .ok_or(TopologyError::EmptyCodebook)?;

        Ok(Selection {
            topology: winner.topology,
            code: winner.code,
            shape: winner.shape,
            score: winner.score,
            objective: winner.objective,
            considered,
        })
    }
}

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

    /// Math-flavoured queries are solved cheaply by the complete graph; code
    /// queries by the chain. Both solve everything, so only measured cost
    /// separates them — which is the regime the paper says the incumbent
    /// scorer gets backwards.
    fn records() -> RecordSet {
        let n = 4;
        let complete = Topology::complete(n).unwrap();
        let chain = Topology::chain(n).unwrap();
        let star = Topology::star(n, 0).unwrap();
        let mut out = Vec::new();
        for i in 0..6 {
            let drift = i as f32 * 0.01;
            let math_q = vec![1.0, 0.0, drift];
            let math = format!("math{i}");
            out.push(ExecutionRecord::new(
                &math,
                math_q.clone(),
                complete.clone(),
                1.0,
                600,
            ));
            out.push(ExecutionRecord::new(
                &math,
                math_q.clone(),
                star.clone(),
                1.0,
                1200,
            ));
            out.push(ExecutionRecord::new(
                &math,
                math_q,
                chain.clone(),
                1.0,
                2400,
            ));

            let code_q = vec![0.0, 1.0, drift];
            let code = format!("code{i}");
            out.push(ExecutionRecord::new(
                &code,
                code_q.clone(),
                chain.clone(),
                1.0,
                600,
            ));
            out.push(ExecutionRecord::new(
                &code,
                code_q.clone(),
                star.clone(),
                1.0,
                1200,
            ));
            out.push(ExecutionRecord::new(
                &code,
                code_q,
                complete.clone(),
                1.0,
                2400,
            ));
        }
        RecordSet::new(out).unwrap()
    }

    fn selector() -> TopologySelector {
        TopologySelector::fit(&records(), &SelectorConfig::default()).unwrap()
    }

    #[test]
    fn selection_adapts_to_the_query() {
        let s = selector();
        let math = s.select(&[1.0, 0.0, 0.0]).unwrap();
        let code = s.select(&[0.0, 1.0, 0.0]).unwrap();
        assert_eq!(math.topology, Topology::complete(4).unwrap());
        assert_eq!(code.topology, Topology::chain(4).unwrap());
        assert_ne!(math.code, code.code);
    }

    #[test]
    fn the_winner_maximizes_the_objective_among_the_candidates() {
        let s = selector();
        let selection = s.select(&[1.0, 0.0, 0.0]).unwrap();
        let best = selection
            .considered
            .iter()
            .map(|c| c.objective)
            .fold(f32::NEG_INFINITY, f32::max);
        assert!((selection.objective - best).abs() < 1e-6);
    }

    #[test]
    fn selection_is_deterministic() {
        let s = selector();
        let a = s.select(&[0.3, 0.7, 0.1]).unwrap();
        let b = s.select(&[0.3, 0.7, 0.1]).unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn candidates_are_deduplicated_and_bounded_by_top_m() {
        let s = TopologySelector::fit(
            &records(),
            &SelectorConfig {
                top_m: 2,
                ..Default::default()
            },
        )
        .unwrap();
        let selection = s.select(&[1.0, 0.0, 0.0]).unwrap();
        assert!(selection.considered.len() <= 2);
        for pair in 0..selection.considered.len() {
            for other in (pair + 1)..selection.considered.len() {
                assert_ne!(
                    selection.considered[pair].topology,
                    selection.considered[other].topology
                );
            }
        }
    }

    #[test]
    fn the_selection_reports_the_shape_a_caller_can_execute() {
        let s = selector();
        let selection = s.select(&[0.0, 1.0, 0.0]).unwrap();
        assert_eq!(selection.shape, Some(CoordinationShape::Pipeline));
    }

    #[test]
    fn reranking_beats_taking_the_prior_top_1_on_cost() {
        // top_m = 1 is the paper's "drop rerank" ablation.
        let full = TopologySelector::fit(&records(), &SelectorConfig::default()).unwrap();
        let top1 = TopologySelector::fit(
            &records(),
            &SelectorConfig {
                top_m: 1,
                ..Default::default()
            },
        )
        .unwrap();
        let q = [1.0, 0.0, 0.0];
        assert!(full.select(&q).unwrap().objective >= top1.select(&q).unwrap().objective - 1e-6);
    }

    #[test]
    fn a_cold_start_selector_still_selects() {
        let n = 4;
        let codebook = Codebook::from_topologies(
            CoordinationShape::ALL
                .iter()
                .map(|s| s.topology(n).unwrap())
                .collect(),
        )
        .unwrap();
        let predictor = CodePredictor::uniform(codebook.len(), 3).unwrap();
        let proxy = ExecutionProxy::fit(&records(), &ProxyConfig::default()).unwrap();
        let s = TopologySelector::from_parts(codebook, predictor, proxy, 5, DEFAULT_COST_WEIGHT)
            .unwrap();
        let selection = s.select(&[1.0, 0.0, 0.0]).unwrap();
        assert!(selection.objective.is_finite());
        assert!(!selection.considered.is_empty());
    }

    #[test]
    fn mismatched_parts_are_rejected() {
        let n = 4;
        let codebook = Codebook::from_topologies(vec![Topology::complete(n).unwrap()]).unwrap();
        let predictor = CodePredictor::uniform(3, 3).unwrap();
        let proxy = ExecutionProxy::fit(&records(), &ProxyConfig::default()).unwrap();
        assert!(matches!(
            TopologySelector::from_parts(codebook, predictor, proxy, 5, 0.1),
            Err(TopologyError::BadConfig {
                field: "predictor",
                ..
            })
        ));
    }

    #[test]
    fn zero_top_m_is_rejected() {
        assert!(matches!(
            TopologySelector::fit(
                &records(),
                &SelectorConfig {
                    top_m: 0,
                    ..Default::default()
                }
            ),
            Err(TopologyError::BadConfig { field: "top_m", .. })
        ));
    }

    #[test]
    fn a_wrong_dimension_query_is_rejected_at_selection_time() {
        let s = selector();
        assert!(matches!(
            s.select(&[1.0, 0.0]),
            Err(TopologyError::QueryDimMismatch { .. })
        ));
    }

    #[test]
    fn a_selector_round_trips_through_json() {
        let s = selector();
        let json = serde_json::to_string(&s).unwrap();
        let back: TopologySelector = serde_json::from_str(&json).unwrap();
        assert_eq!(s, back);
        assert_eq!(
            s.select(&[1.0, 0.0, 0.0]).unwrap(),
            back.select(&[1.0, 0.0, 0.0]).unwrap()
        );
    }
}