llama-rs 0.16.0

//! Convergence detection — pairwise cosine similarity over expert
//! answer embeddings, aggregated by `min` (most conservative).
//!
//! The orchestrator embeds each responding-expert's answer, computes
//! the pairwise cosine across the resulting vectors, takes the minimum,
//! and compares it to the configured threshold.

use crate::council::embedder::{Embedder, EmbedderError};

/// Outcome of one convergence check.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ConvergenceOutcome {
    /// Minimum pairwise cosine across all answer pairs. `1.0` when
    /// there are fewer than 2 answers (no disagreement is possible).
    pub min_cosine: f32,
    /// `true` iff `min_cosine >= threshold`.
    pub converged: bool,
}

/// Cosine similarity of two vectors. Caller is responsible for
/// dimension match. Vectors are *not* required to be normalized — this
/// function divides by both norms.
pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if na == 0.0 || nb == 0.0 {
        0.0
    } else {
        dot / (na * nb)
    }
}

/// Minimum pairwise cosine across `vectors`. Returns `1.0` when fewer
/// than two vectors are supplied (degenerate "consensus of one").
pub fn min_pairwise_cosine(vectors: &[Vec<f32>]) -> f32 {
    if vectors.len() < 2 {
        return 1.0;
    }
    let mut min = f32::INFINITY;
    for i in 0..vectors.len() {
        for j in (i + 1)..vectors.len() {
            let c = cosine(&vectors[i], &vectors[j]);
            if c < min {
                min = c;
            }
        }
    }
    min
}

/// Embed each text and run the convergence check against `threshold`.
///
/// On embedder failure, returns `EmbedderError`. The orchestrator
/// treats this as "round did not converge" and proceeds to the next
/// round, logging a warning.
pub fn check<E: Embedder + ?Sized>(
    embedder: &E,
    texts: &[&str],
    threshold: f32,
) -> Result<ConvergenceOutcome, EmbedderError> {
    let vectors: Vec<Vec<f32>> = texts
        .iter()
        .map(|t| embedder.embed(t))
        .collect::<Result<_, _>>()?;
    let min_cosine = min_pairwise_cosine(&vectors);
    Ok(ConvergenceOutcome {
        min_cosine,
        converged: min_cosine >= threshold,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::council::embedder::testing::{AlwaysEmbedder, MockEmbedder};

    fn close(a: f32, b: f32) -> bool {
        (a - b).abs() < 1e-5
    }

    #[test]
    fn cosine_of_identical_unit_vectors_is_one() {
        assert!(close(cosine(&[1.0, 0.0], &[1.0, 0.0]), 1.0));
    }

    #[test]
    fn cosine_of_orthogonal_vectors_is_zero() {
        assert!(close(cosine(&[1.0, 0.0], &[0.0, 1.0]), 0.0));
    }

    #[test]
    fn cosine_of_opposite_vectors_is_minus_one() {
        assert!(close(cosine(&[1.0, 0.0], &[-1.0, 0.0]), -1.0));
    }

    #[test]
    fn cosine_handles_zero_vector_safely() {
        assert_eq!(cosine(&[0.0, 0.0], &[1.0, 0.0]), 0.0);
        assert_eq!(cosine(&[1.0, 0.0], &[0.0, 0.0]), 0.0);
    }

    #[test]
    fn cosine_normalizes_unequal_magnitudes() {
        // (3,0) vs (5,0) — same direction, cos = 1.0 regardless of magnitude.
        assert!(close(cosine(&[3.0, 0.0], &[5.0, 0.0]), 1.0));
    }

    #[test]
    fn min_pairwise_cosine_is_one_for_fewer_than_two() {
        assert_eq!(min_pairwise_cosine(&[]), 1.0);
        assert_eq!(min_pairwise_cosine(&[vec![1.0, 0.0]]), 1.0);
    }

    #[test]
    fn min_pairwise_cosine_picks_worst_pair() {
        // Three vectors: A and B agree perfectly (cos=1), C disagrees with both
        // (cos=0 with A, cos=0 with B). Min should be 0.
        let v = vec![
            vec![1.0, 0.0],
            vec![1.0, 0.0],
            vec![0.0, 1.0],
        ];
        assert!(close(min_pairwise_cosine(&v), 0.0));
    }

    #[test]
    fn check_converged_when_all_identical() {
        let e = AlwaysEmbedder {
            vector: vec![1.0, 0.0],
            dim: 2,
        };
        let out = check(&e, &["a", "b", "c"], 0.99).unwrap();
        assert!(out.converged);
        assert!(close(out.min_cosine, 1.0));
    }

    #[test]
    fn check_not_converged_when_below_threshold() {
        let e = MockEmbedder::new(
            2,
            vec![vec![1.0, 0.0], vec![0.0, 1.0]], // orthogonal
        );
        let out = check(&e, &["a", "b"], 0.92).unwrap();
        assert!(!out.converged);
        assert!(close(out.min_cosine, 0.0));
    }

    #[test]
    fn check_returns_embedder_error() {
        let e = MockEmbedder::new(2, vec![vec![1.0, 0.0]]); // only one vector
        let err = check(&e, &["a", "b"], 0.5).unwrap_err();
        assert!(matches!(err, EmbedderError::Failed(_)));
    }

    #[test]
    fn check_handles_zero_or_one_text_as_converged() {
        // An "always-fail" embedder won't be called when there are <2 texts.
        let e = MockEmbedder::new(2, vec![vec![1.0, 0.0]]);
        let out = check(&e, &["only-one"], 0.5).unwrap();
        assert!(out.converged);
        assert_eq!(out.min_cosine, 1.0);
    }
}