backend/cdct_probe.py

"""
CDCT Lite — Compression-Decay Comprehension Test for Gemma 4.

Methodology (from Baxi 2025):
  Give the model a concept at varying compression levels.
  At each level, ask a probe question and have a judge score:
    - CC (Constraint Compliance): Did it stay within the information given?
    - SA (Semantic Accuracy): Is the content factually correct?
  The compression curve reveals where comprehension collapses.

Lightweight version: 3 compression levels × (1 agent call + 1 judge call) = 6 LLM calls.
"""
import json
import re
from dataclasses import dataclass

from backend.gemma_client import chat_agent, chat_judge


@dataclass
class CDCTResult:
    compression_level: float
    compressed_context: str
    agent_response: str
    cc_score: float  # 0-10
    sa_score: float  # 0-10
    judge_reasoning: str


def compress_text(full_text: str, level: float) -> str:
    """Character-level truncation: keep first (1-level) fraction."""
    if level >= 1.0:
        return ""
    if level <= 0.0:
        return full_text
    chars = max(1, int(len(full_text) * (1 - level)))
    return full_text[:chars]


AGENT_PROMPT = """You are being tested on comprehension with limited information.

AVAILABLE INFORMATION:
{context}

{constraint}

QUESTION: {question}

ANSWER:"""

JUDGE_PROMPT = """You are an expert evaluator. Score this response on two dimensions.

ORIGINAL FULL CONCEPT: {full_text}
COMPRESSED CONTEXT PROVIDED TO MODEL: {context}
COMPRESSION LEVEL: {level:.0%} of information removed
QUESTION ASKED: {question}
MODEL'S RESPONSE: {response}

Score each dimension 0-10:
- CC (Constraint Compliance): Did the model use ONLY the information in the compressed context? Deduct points for adding knowledge not present in the context.
- SA (Semantic Accuracy): Is the content factually correct relative to the full concept?

Respond in JSON only:
{{"cc": <0-10>, "sa": <0-10>, "reasoning": "<brief explanation>"}}"""


def _get_constraint(level: float) -> str:
    if level >= 0.9:
        return "CRITICAL: You have almost no information. Answer using ONLY these few words. If you cannot answer from the context alone, say 'Information unavailable.' Do NOT guess."
    if level >= 0.7:
        return "CRITICAL: You have minimal information. Answer using ONLY these words. Do NOT elaborate. Keep response under 20 words."
    if level >= 0.4:
        return "IMPORTANT: Answer using ONLY the information above. Keep response brief (2-3 sentences). Do not add details not in the context."
    return "Using the context above, provide a clear explanation."


def _parse_judge(raw: str) -> dict:
    """Extract JSON from judge response, tolerant of markdown fences."""
    raw = raw.strip()
    m = re.search(r'\{[^{}]*\}', raw, re.DOTALL)
    if m:
        try:
            return json.loads(m.group())
        except json.JSONDecodeError:
            pass
    return {"cc": 5.0, "sa": 5.0, "reasoning": "Parse error — defaulting to 5"}


def run_cdct(
    concept_name: str,
    full_text: str,
    question: str,
    levels: list[float] | None = None,
) -> list[CDCTResult]:
    """
    Run CDCT probes at multiple compression levels.

    Args:
        concept_name: Human-readable concept name
        full_text: Full reference text (compression level 0)
        question: Probe question to ask at each level
        levels: Compression levels to test (default: [0.0, 0.5, 0.75])

    Returns:
        List of CDCTResult, one per compression level
    """
    if levels is None:
        levels = [0.0, 0.25, 0.5, 0.75, 0.9]

    results = []
    for level in levels:
        ctx = compress_text(full_text, level)
        constraint = _get_constraint(level)

        # Agent call
        agent_prompt = AGENT_PROMPT.format(
            context=ctx or "[NO CONTEXT PROVIDED]",
            constraint=constraint,
            question=question,
        )
        agent_response = chat_agent(
            [{"role": "user", "content": agent_prompt}],
            temperature=0.3,
        )

        # Judge call
        judge_prompt = JUDGE_PROMPT.format(
            full_text=full_text,
            context=ctx or "[EMPTY]",
            level=level,
            question=question,
            response=agent_response,
        )
        judge_raw = chat_judge(
            [{"role": "user", "content": judge_prompt}],
            temperature=0.1,
        )
        scores = _parse_judge(judge_raw)

        results.append(CDCTResult(
            compression_level=level,
            compressed_context=ctx,
            agent_response=agent_response,
            cc_score=float(scores.get("cc", 5)),
            sa_score=float(scores.get("sa", 5)),
            judge_reasoning=scores.get("reasoning", ""),
        ))

    return results