anda_brain 0.10.2

//! LLM-as-judge for checkpoint scoring and semantic graph probes.
//!
//! The lexical scorer stays as a deterministic smoke gate; this module layers
//! rubric-driven judgment on top so paraphrases are not false negatives and
//! correct meta-references to superseded facts are not false positives.

use anda_core::{BoxError, CompletionRequest, Json, ModelEffort};
use serde::{Deserialize, Serialize};

use super::{EvalFinding, EvalFindingKind, MemoryExpectationMode, MemoryProbeReport};
use crate::assess::{
    AssessContext, JudgeCall, MAX_EVIDENCE_CHARS, parse_json_payload, truncate_chars,
};

/// Judge scores for one checkpoint answer sample. All values are clamped to
/// `0..=1` after parsing.
#[derive(Debug, Clone, Default, Serialize)]
pub struct JudgeVerdict {
    /// Did the answer actually use the relevant memories to help the user?
    pub memory_utility: f64,

    /// Are superseded/expired facts absent as *assertions*? Meta-references
    /// ("unlike your old BBQ preference…") are correct and must not lower this.
    pub forgetting_quality: f64,

    /// Confidence hygiene: no stale facts asserted as current, uncertainty
    /// admitted when evidence is thin.
    pub uncertainty_calibration: f64,

    /// How satisfied a real user with the hidden profile would be (0..1).
    pub satisfaction: f64,

    pub reasoning: String,

    pub findings: Vec<EvalFinding>,
}

/// Wire shape for the judge output: findings are parsed leniently so a single
/// unknown kind does not discard an otherwise valid verdict. The core score
/// fields are `Option` so a shape-mismatched payload (e.g. scores nested in a
/// wrapper object) parses to all-`None` and is rejected as a judge error —
/// falling back to lexical scoring — instead of silently scoring the sample
/// all-zero and poisoning the trajectory.
#[derive(Debug, Default, Deserialize)]
struct RawVerdict {
    memory_utility: Option<f64>,

    forgetting_quality: Option<f64>,

    uncertainty_calibration: Option<f64>,

    satisfaction: Option<f64>,

    #[serde(default)]
    reasoning: String,

    #[serde(default)]
    findings: Vec<Json>,
}

impl RawVerdict {
    /// True when at least one core score field was present in the payload.
    fn has_any_score(&self) -> bool {
        self.memory_utility.is_some()
            || self.forgetting_quality.is_some()
            || self.uncertainty_calibration.is_some()
            || self.satisfaction.is_some()
    }
}

impl From<RawVerdict> for JudgeVerdict {
    fn from(raw: RawVerdict) -> Self {
        let findings = raw
            .findings
            .into_iter()
            .filter_map(|value| serde_json::from_value::<EvalFinding>(value).ok())
            .filter(|finding| is_answer_finding(finding.kind))
            .collect();
        Self {
            memory_utility: raw.memory_utility.unwrap_or_default().clamp(0.0, 1.0),
            forgetting_quality: raw.forgetting_quality.unwrap_or_default().clamp(0.0, 1.0),
            uncertainty_calibration: raw
                .uncertainty_calibration
                .unwrap_or_default()
                .clamp(0.0, 1.0),
            satisfaction: raw.satisfaction.unwrap_or_default().clamp(0.0, 1.0),
            reasoning: raw.reasoning,
            findings,
        }
    }
}

/// Everything the judge sees for one checkpoint sample.
pub(crate) struct JudgeCheckpointInput<'a> {
    pub query: &'a str,
    pub answer: &'a str,
    pub scoring_rubric: Option<&'a str>,
    pub hidden_profile: &'a Json,
    pub required_terms: &'a [String],
    pub forbidden_terms: &'a [String],
    pub probes: &'a [MemoryProbeReport],
    pub expectations: Vec<JudgeExpectation>,
    pub trace_summary: Option<String>,
}

#[derive(Debug, Clone, Serialize)]
pub(crate) struct JudgeExpectation {
    pub id: String,
    pub mode: MemoryExpectationMode,
    pub description: Option<String>,
    /// Whether the pre-recall graph probe found this memory state satisfied.
    pub probe_satisfied: Option<bool>,
}

const JUDGE_INSTRUCTIONS: &str = r#"You are a strict evaluator of an AI memory system. You will receive a user query, the memory system's answer, a scoring rubric, the hidden ground-truth user profile, memory expectations with pre-answer graph probe outcomes, and a summary of the retrieval trace.

Score the answer on these dimensions, each 0.0..1.0:
- memory_utility: did the answer use the relevant remembered facts to genuinely help this user? Paraphrases count fully; literal keyword overlap is NOT required. Dumping remembered facts irrelevant to the query lowers this score.
- forgetting_quality: superseded or expired facts must not be asserted as currently true. IMPORTANT: meta-references are correct behavior — e.g. "unlike your old BBQ preference, here are vegetarian options" correctly acknowledges the superseded fact and must NOT be penalized. Only penalize stale facts presented as still valid.
- uncertainty_calibration: confident claims must be supported by the probes/trace; when evidence is thin the answer should hedge or ask.
- satisfaction: how satisfied would the real user (see hidden profile) be with this answer, 0..1.

Also report findings, attributing each failure to the responsible stage:
- "formation_miss": the probes show the graph never formed a needed memory.
- "bad_consolidation": the graph still holds a stale/conflicting memory as active.
- "bad_grounding": the graph has the memory (probe satisfied) but the retrieval trace shows it was never retrieved.
- "bad_synthesis": the memory was retrieved (visible in trace) but the answer failed to use it.
- "overconfidence": the answer asserts stale or unsupported facts as current truth.

A probe with "error": true (its "satisfied" is null, and the matching expectation's "probe_satisfied" is null) could not observe the graph at all — an infrastructure failure, not a memory failure. Treat that expectation's graph state as UNKNOWN: never report formation_miss or bad_consolidation from an errored probe; judge only from the remaining evidence.

Each finding must set "expectation_id" to the id of the memory expectation it is about (copy the id from the expectations/probes). Use null ONLY when the finding does not correspond to any listed expectation. Findings are merged with the harness's own probe findings by exact (kind, expectation_id), so omitting the id on an expectation-linked finding double-counts it.

Respond with ONLY a JSON object:
{"memory_utility": 0.0, "forgetting_quality": 0.0, "uncertainty_calibration": 0.0, "satisfaction": 0.0, "reasoning": "...", "findings": [{"kind": "bad_synthesis", "message": "...", "expectation_id": "expectation-id-or-null"}]}"#;

pub(crate) async fn judge_checkpoint<C>(
    driver: &C,
    input: JudgeCheckpointInput<'_>,
) -> Result<JudgeCall<JudgeVerdict>, BoxError>
where
    C: AssessContext + ?Sized,
{
    let probes: Vec<Json> = input
        .probes
        .iter()
        .map(|probe| {
            // An errored probe never observed the graph: its `satisfied`
            // value is meaningless, so it is reported as null plus an
            // explicit error flag the instructions tell the judge not to
            // attribute memory failures from.
            let errored = probe.errored();
            serde_json::json!({
                "expectation_id": probe.expectation_id,
                "mode": probe.mode,
                "satisfied": if errored { Json::Null } else { Json::Bool(probe.satisfied) },
                "hit_count": probe.hit_count,
                "assertion": probe.assertion,
                "error": errored,
            })
        })
        .collect();

    let prompt = format!(
        "# User query\n{}\n\n# Memory system answer\n{}\n\n# Scoring rubric\n{}\n\n# Hidden ground-truth user profile\n{}\n\n# Memory expectations\n{}\n\n# Pre-answer graph probes\n{}\n\n# Retrieval trace summary\n{}\n\n# Term hints (advisory only; paraphrase counts)\nrequired: {:?}\nforbidden-as-current-assertion: {:?}",
        input.query,
        input.answer,
        input.scoring_rubric.unwrap_or("(none provided)"),
        truncate_chars(
            &serde_json::to_string(input.hidden_profile).unwrap_or_default(),
            MAX_EVIDENCE_CHARS
        ),
        truncate_chars(
            &serde_json::to_string(&input.expectations).unwrap_or_default(),
            MAX_EVIDENCE_CHARS
        ),
        truncate_chars(
            &serde_json::to_string(&probes).unwrap_or_default(),
            MAX_EVIDENCE_CHARS
        ),
        input
            .trace_summary
            .as_deref()
            .unwrap_or("(no trace available)"),
        input.required_terms,
        input.forbidden_terms,
    );

    let output = driver
        .judge_complete(CompletionRequest {
            instructions: JUDGE_INSTRUCTIONS.to_string(),
            prompt,
            effort: Some(ModelEffort::Low),
            ..Default::default()
        })
        .await?;

    let raw: RawVerdict = parse_json_payload(&output.content)?;
    if !raw.has_any_score() {
        // Any legal JSON object parses into `RawVerdict`; without at least
        // one core score field it is a shape mismatch, and treating it as an
        // all-zero verdict would tank the sample. Erroring here routes the
        // sample through the existing JudgeError + lexical fallback.
        return Err(format!(
            "judge output has none of the core score fields (memory_utility, forgetting_quality, uncertainty_calibration, satisfaction): {}",
            truncate_chars(&output.content, 200)
        )
        .into());
    }
    Ok(JudgeCall {
        verdict: raw.into(),
        usage: output.usage,
    })
}

/// Maps a judge finding to the attribution kind used by the harness; unknown
/// kinds are dropped by the deserializer, so this only sanity-checks bounds.
pub(crate) fn is_answer_finding(kind: EvalFindingKind) -> bool {
    matches!(
        kind,
        EvalFindingKind::FormationMiss
            | EvalFindingKind::BadConsolidation
            | EvalFindingKind::BadGrounding
            | EvalFindingKind::BadSynthesis
            | EvalFindingKind::Overconfidence
    )
}

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

    #[test]
    fn parse_json_payload_tolerates_fences_and_prose() {
        let text = "Here is my verdict:\n```json\n{\"memory_utility\": 0.9, \"forgetting_quality\": 1.2, \"uncertainty_calibration\": -0.5, \"satisfaction\": 0.8, \"reasoning\": \"good\", \"findings\": [{\"kind\": \"bad_synthesis\", \"message\": \"m\"}, {\"kind\": \"weird_unknown\", \"message\": \"skip me\"}, {\"kind\": \"latency_cost\", \"message\": \"not a judge finding\"}]}\n```\nDone.";
        let raw: RawVerdict = parse_json_payload(text).unwrap();
        let verdict = JudgeVerdict::from(raw);
        assert_eq!(verdict.memory_utility, 0.9);
        assert_eq!(verdict.forgetting_quality, 1.0);
        assert_eq!(verdict.uncertainty_calibration, 0.0);
        assert_eq!(verdict.reasoning, "good");
        assert_eq!(verdict.findings.len(), 1);
        assert_eq!(verdict.findings[0].kind, EvalFindingKind::BadSynthesis);
    }

    /// Judge driver returning one canned completion.
    struct FakeJudge {
        response: String,
    }

    #[async_trait::async_trait]
    impl AssessContext for FakeJudge {
        async fn complete(
            &self,
            _req: CompletionRequest,
        ) -> Result<AgentOutput, anda_core::BoxError> {
            Ok(AgentOutput {
                content: self.response.clone(),
                ..Default::default()
            })
        }

        async fn execute_kip_readonly(
            &self,
            _request: anda_kip::Request,
        ) -> Result<anda_kip::Response, anda_core::BoxError> {
            Err("not used".into())
        }
    }

    fn checkpoint_input<'a>(profile: &'a Json) -> JudgeCheckpointInput<'a> {
        JudgeCheckpointInput {
            query: "what do I like?",
            answer: "green tea",
            scoring_rubric: None,
            hidden_profile: profile,
            required_terms: &[],
            forbidden_terms: &[],
            probes: &[],
            expectations: Vec::new(),
            trace_summary: None,
        }
    }

    #[tokio::test]
    async fn judge_checkpoint_rejects_shape_mismatched_output_instead_of_scoring_zero() {
        // A nested wrapper is legal JSON that used to deserialize into an
        // all-zero verdict; it must error so the caller can fall back to
        // lexical scoring.
        let driver = FakeJudge {
            response: serde_json::json!({
                "verdict": {
                    "memory_utility": 0.9,
                    "forgetting_quality": 0.9,
                    "uncertainty_calibration": 0.9,
                    "satisfaction": 0.9
                }
            })
            .to_string(),
        };
        let profile = Json::Null;
        let err = judge_checkpoint(&driver, checkpoint_input(&profile))
            .await
            .unwrap_err();
        assert!(
            err.to_string().contains("core score fields"),
            "unexpected error: {err}"
        );
    }

    #[tokio::test]
    async fn judge_checkpoint_accepts_partial_scores() {
        // A verdict carrying at least one core score field is not a shape
        // mismatch; missing fields default to zero as before.
        let driver = FakeJudge {
            response: serde_json::json!({
                "memory_utility": 0.8,
                "reasoning": "partial"
            })
            .to_string(),
        };
        let profile = Json::Null;
        let call = judge_checkpoint(&driver, checkpoint_input(&profile))
            .await
            .unwrap();
        assert_eq!(call.verdict.memory_utility, 0.8);
        assert_eq!(call.verdict.satisfaction, 0.0);
        assert_eq!(call.verdict.reasoning, "partial");
    }
}