llama-rs 0.16.1

//! Accumulated state of a council deliberation.
//!
//! The orchestrator builds up a `Transcript` as rounds complete; the
//! final transcript is what the synthesizer summarizes.

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::council::event::ExpertId;

/// One expert's answer in one round.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Answer {
    pub expert_id: ExpertId,
    pub text: String,
    pub tokens: u32,
}

/// Result of a single round's dispatch.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Round {
    pub index: u32,
    /// Successful answers, keyed by expert id (sorted by id for determinism).
    pub answers: BTreeMap<ExpertId, Answer>,
    /// Experts that failed this round, with their error messages.
    pub failures: BTreeMap<ExpertId, String>,
}

impl Round {
    pub fn new(index: u32) -> Self {
        Self {
            index,
            answers: BTreeMap::new(),
            failures: BTreeMap::new(),
        }
    }

    pub fn record_answer(&mut self, answer: Answer) {
        self.failures.remove(&answer.expert_id);
        self.answers.insert(answer.expert_id.clone(), answer);
    }

    pub fn record_failure(&mut self, expert_id: ExpertId, message: String) {
        self.answers.remove(&expert_id);
        self.failures.insert(expert_id, message);
    }

    pub fn responded_ids(&self) -> Vec<ExpertId> {
        self.answers.keys().cloned().collect()
    }

    pub fn failed_ids(&self) -> Vec<ExpertId> {
        self.failures.keys().cloned().collect()
    }

    pub fn responded_count(&self) -> u32 {
        self.answers.len() as u32
    }
}

/// Full record of a council deliberation.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct Transcript {
    pub rounds: Vec<Round>,
    pub final_answer: Option<String>,
}

impl Transcript {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push_round(&mut self, round: Round) {
        self.rounds.push(round);
    }

    pub fn last_round(&self) -> Option<&Round> {
        self.rounds.last()
    }

    pub fn round(&self, index: u32) -> Option<&Round> {
        self.rounds.iter().find(|r| r.index == index)
    }

    pub fn set_final_answer(&mut self, text: String) {
        self.final_answer = Some(text);
    }
}

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

    fn ans(id: &str, text: &str) -> Answer {
        Answer {
            expert_id: id.into(),
            text: text.into(),
            tokens: text.len() as u32,
        }
    }

    #[test]
    fn new_transcript_is_empty() {
        let t = Transcript::new();
        assert!(t.rounds.is_empty());
        assert!(t.final_answer.is_none());
        assert!(t.last_round().is_none());
    }

    #[test]
    fn round_records_answers_and_failures_distinctly() {
        let mut r = Round::new(0);
        r.record_answer(ans("A", "alpha"));
        r.record_failure("B".into(), "timeout".into());
        r.record_answer(ans("C", "gamma"));

        assert_eq!(r.responded_count(), 2);
        assert_eq!(r.responded_ids(), vec!["A", "C"]);
        assert_eq!(r.failed_ids(), vec!["B"]);
    }

    #[test]
    fn recording_an_answer_clears_a_prior_failure() {
        // If the orchestrator marks an expert as failed but a late response
        // arrives within the timeout window, the answer should win.
        let mut r = Round::new(0);
        r.record_failure("A".into(), "timeout".into());
        r.record_answer(ans("A", "ok"));
        assert_eq!(r.responded_count(), 1);
        assert!(r.failed_ids().is_empty());
        assert_eq!(r.answers["A"].text, "ok");
    }

    #[test]
    fn recording_a_failure_clears_a_prior_answer() {
        let mut r = Round::new(0);
        r.record_answer(ans("A", "ok"));
        r.record_failure("A".into(), "stream broke".into());
        assert_eq!(r.responded_count(), 0);
        assert_eq!(r.failed_ids(), vec!["A"]);
    }

    #[test]
    fn transcript_pushes_rounds_in_order_and_lookup_works() {
        let mut t = Transcript::new();
        let mut r0 = Round::new(0);
        r0.record_answer(ans("A", "round0"));
        let mut r1 = Round::new(1);
        r1.record_answer(ans("A", "round1"));

        t.push_round(r0);
        t.push_round(r1);

        assert_eq!(t.rounds.len(), 2);
        assert_eq!(t.last_round().unwrap().index, 1);
        assert_eq!(t.round(0).unwrap().answers["A"].text, "round0");
        assert_eq!(t.round(1).unwrap().answers["A"].text, "round1");
        assert!(t.round(2).is_none());
    }

    #[test]
    fn set_final_answer_records_it() {
        let mut t = Transcript::new();
        t.set_final_answer("the answer".into());
        assert_eq!(t.final_answer.as_deref(), Some("the answer"));
    }

    #[test]
    fn transcript_round_trips_through_json() {
        let mut t = Transcript::new();
        let mut r = Round::new(0);
        r.record_answer(ans("A", "hello"));
        r.record_failure("B".into(), "timeout".into());
        t.push_round(r);
        t.set_final_answer("synthesis".into());

        let json = serde_json::to_string(&t).unwrap();
        let back: Transcript = serde_json::from_str(&json).unwrap();
        assert_eq!(t, back);
    }
}