sovereign_bench/engine.py

from __future__ import annotations

import json
import re
import time
from collections import Counter
from collections.abc import Callable, Iterable

from pydantic import ValidationError

from .cases import get_case
from .llm import ModelCall, ModelResult, call_small_model
from .models import AgentTurn, CasePacket, JurorVote, TrialEvent, TrialRequest, Verdict
from .retrieval import build_live_case

GPT_OSS_MODEL = "openai/gpt-oss-20b"
OPENBMB_MODEL = "openbmb/AgentCPM-Explore"
NEMOTRON_MODEL = "nvidia/Nemotron-Orchestrator-8B"
OPENAI_PROVIDER = "auto"
OPENBMB_PROVIDER = "featherless-ai"
NEMOTRON_PROVIDER = "featherless-ai"

MODEL_BUDGET = [
    ("Presiding Advocate", GPT_OSS_MODEL, 20.0),
    ("Clerk of Style", OPENBMB_MODEL, 4.0),
    ("Juror/Auditor Ring", NEMOTRON_MODEL, 8.0),
]
TOTAL_PARAMS_B = sum(item[2] for item in MODEL_BUDGET)

JUDGE_NAME = "Marcus Aurelius"
JUDGE_PERSONA = "Stoic duty, restraint, public reason, and disciplined judgment"

JUROR_PERSONAS = {
    "Karl Marx": "class power, material conditions, exploitation, institutional incentives",
    "John Stuart Mill": "liberty, harm principle, utility, individual rights",
    "Confucius": "social harmony, role duty, ritual order, moral cultivation",
    "Cleopatra VII": "sovereign pragmatism, diplomacy, survival, legitimacy under pressure",
    "Niccolo Machiavelli": "political realism, stability, power, consequences over ideals",
    "Jensen Huang": "technological optimism, operator mindset, systems thinking, innovation tradeoffs",
}
JUROR_NAMES = list(JUROR_PERSONAS)


class RequiredModelError(RuntimeError):
    """Raised when a required courtroom model call cannot produce usable output."""


ModelRunner = Callable[..., ModelResult]


def _turn(agent: str, role: str, result: ModelResult, model: str, confidence: float) -> AgentTurn:
    return AgentTurn(
        agent=agent,
        role=role,
        content=result.text,
        model=model,
        confidence=confidence,
        input=getattr(result, "input_text", ""),
    )


def _case_summary(packet: CasePacket) -> str:
    return (
        f"{packet.title}. Charge: {packet.charge}\n"
        f"Claimant: {packet.claimant_claim}\n"
        f"Respondent: {packet.respondent_claim}"
    )


def _evidence_summary(packet: CasePacket) -> str:
    return "\n".join(
        f"{item.id}: {item.title}; direction={item.supports}; reliability={item.reliability:.2f}; note={item.note}"
        for item in packet.evidence
    )


def _call_trace(calls: list[ModelCall]) -> list[dict]:
    return [call.__dict__ for call in calls]


def resolve_case(request: TrialRequest) -> tuple[CasePacket, dict]:
    if request.case_id == "live":
        packet = build_live_case(request.search_query, request.hypothetical)
        if packet:
            return packet, {"mode": "live"}
        raise RuntimeError("Live retrieval produced too little usable evidence; no fallback case will be substituted.")
    return get_case(request.case_id), {"mode": "cached"}


def _generate_role(model_runner: ModelRunner | None = None, **kwargs) -> ModelResult:
    if model_runner is not None:
        return model_runner(**kwargs)
    return call_small_model(**kwargs)


def _required_role(model_runner: ModelRunner | None, model_calls: list[ModelCall], **kwargs) -> ModelResult:
    try:
        result = _generate_role(model_runner, **kwargs)
    except Exception as exc:
        raise RequiredModelError(f"{kwargs.get('agent', 'Model')} unavailable: {exc}") from exc
    model_calls.append(result.call)
    if not result.call.ok:
        error = result.call.error or "model call did not complete"
        raise RequiredModelError(f"{kwargs.get('agent', 'Model')} unavailable: {error}")
    if not result.text.strip():
        raise RequiredModelError(f"{kwargs.get('agent', 'Model')} returned an empty response.")
    return result


def _trace(packet: CasePacket, source_trace: dict, model_calls: list[ModelCall]) -> dict:
    return {
        "case_id": packet.id,
        "model_budget_b": TOTAL_PARAMS_B,
        "models": [{"role": role, "model": model, "params_b": params} for role, model, params in MODEL_BUDGET],
        "model_calls": _call_trace(model_calls),
        "live_model_call_count": sum(1 for call in model_calls if call.ok),
        "attempted_model_call_count": len(model_calls),
        **source_trace,
    }


def _emit(
    packet: CasePacket,
    source_trace: dict,
    model_calls: list[ModelCall],
    event: TrialEvent,
    delay: float,
) -> TrialEvent:
    event.trace = _trace(packet, source_trace, model_calls)
    if delay > 0:
        time.sleep(delay)
    return event


def _extract_json(text: str) -> object:
    stripped = text.strip()
    if stripped.startswith("```"):
        stripped = re.sub(r"^```(?:json)?\s*", "", stripped, flags=re.I)
        stripped = re.sub(r"\s*```$", "", stripped)
    try:
        return json.loads(stripped)
    except json.JSONDecodeError:
        match = re.search(r"(\{.*\}|\[.*\])", stripped, flags=re.S)
        if not match:
            raise
        return json.loads(match.group(1))


def _parse_jury_votes(result: ModelResult, packet: CasePacket) -> list[JurorVote]:
    try:
        data = _extract_json(result.text)
    except json.JSONDecodeError as exc:
        raise RequiredModelError(f"Nemotron Jury returned invalid JSON: {exc.msg}") from exc

    raw_votes = data.get("votes") if isinstance(data, dict) else data
    if not isinstance(raw_votes, list):
        raise RequiredModelError("Nemotron Jury output must contain a votes list.")
    if len(raw_votes) != len(JUROR_NAMES):
        raise RequiredModelError("Nemotron Jury must return exactly six juror votes.")

    known_evidence = {item.id for item in packet.evidence}
    votes: list[JurorVote] = []
    try:
        for item in raw_votes:
            vote = JurorVote.model_validate(item)
            votes.append(vote)
    except ValidationError as exc:
        raise RequiredModelError(f"Nemotron Jury vote schema is invalid: {exc.errors()[0]['msg']}") from exc

    if [vote.juror for vote in votes] != JUROR_NAMES:
        raise RequiredModelError("Nemotron Jury must return votes in the fixed juror order.")
    for vote in votes:
        expected_persona = JUROR_PERSONAS[vote.juror]
        if vote.persona.strip().lower() != expected_persona:
            raise RequiredModelError(f"{vote.juror} persona must be '{expected_persona}'.")
        if not vote.reason.strip():
            raise RequiredModelError(f"{vote.juror} must include a rationale.")
        if not vote.evidence_ids or any(evidence_id not in known_evidence for evidence_id in vote.evidence_ids):
            raise RequiredModelError(f"{vote.juror} must cite known evidence IDs.")
    return votes


def _majority_finding(votes: list[JurorVote]) -> str:
    counts = Counter(vote.vote for vote in votes)
    top = counts.most_common()
    if not top:
        return "uncertain"
    if len(top) > 1 and top[0][1] == top[1][1]:
        return "mixed"
    if top[0][0] == "uncertain":
        return "uncertain"
    return top[0][0]


def _verdict_from_votes(votes: list[JurorVote]) -> Verdict:
    finding = _majority_finding(votes)
    evidence_ids = []
    for vote in votes:
        for evidence_id in vote.evidence_ids:
            if evidence_id not in evidence_ids:
                evidence_ids.append(evidence_id)
    cited = evidence_ids[:4]
    counts = Counter(vote.vote for vote in votes)
    vote_line = ", ".join(f"{name}: {counts.get(name, 0)}" for name in ("liable", "not_liable", "uncertain"))
    decree_by_finding = {
        "liable": "The jury majority finds liability on the miniature record.",
        "not_liable": "The jury majority does not find liability on the miniature record.",
        "mixed": "The jury divides too closely for a clean finding.",
        "uncertain": "The jury leaves the court with unresolved uncertainty.",
    }
    remedy_by_finding = {
        "liable": "Enter symbolic censure and proportional repair.",
        "not_liable": "Dismiss without prejudice to stronger proof.",
        "mixed": "Record a divided result and preserve the exhibits for later review.",
        "uncertain": "Withhold sanction and identify the proof gaps before any retrial.",
    }
    return Verdict(
        finding=finding,  # type: ignore[arg-type]
        decree=decree_by_finding[finding],
        rationale=f"Jury vote: {vote_line}. Cited evidence IDs: {', '.join(cited)}.",
        evidence_ids=cited,
        uncertainty=(
            "Uncertainty remains visible: this is an AI-native miniature trial. Retrieved facts, cached "
            "packets, and model inferences are separated in the trace and should not be treated as legal advice."
        ),
        remedy=remedy_by_finding[finding],
    )


def _jury_task() -> str:
    personas = "\n".join(f"- {name}: {persona}" for name, persona in JUROR_PERSONAS.items())
    return (
        "Return JSON only with a top-level 'votes' array. Create exactly one vote for each juror, in this order: "
        f"{', '.join(JUROR_NAMES)}. Valid vote values are liable, not_liable, uncertain. Each item must contain "
        "juror, persona, vote, reason, and evidence_ids. The persona value must exactly match the profile below. "
        "Each reason should be one concise sentence and each evidence_ids list must cite evidence IDs from the record. "
        "Vote through the named public-history worldview, not a generic juror role.\n"
        f"{personas}"
    )


def run_trial(request: TrialRequest, model_runner: ModelRunner | None = None) -> list[TrialEvent]:
    return list(stream_trial(request, delay=0.0, model_runner=model_runner))


def stream_trial(
    request: TrialRequest,
    delay: float = 0.0,
    model_runner: ModelRunner | None = None,
) -> Iterable[TrialEvent]:
    packet, source_trace = resolve_case(request)
    case_summary = _case_summary(packet)
    evidence_summary = _evidence_summary(packet)
    model_calls: list[ModelCall] = []
    hypo = request.hypothetical.strip()
    hypo_line = f"\n\nUser hypothetical admitted as a blue-ribbon sidebar: {hypo}" if hypo else ""

    clerk = _required_role(
        model_runner,
        model_calls,
        agent="Clerk Meridian",
        role="clerk",
        model=OPENBMB_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task="Announce the case by name, identify the parties, and read the charge.",
        provider=OPENBMB_PROVIDER,
        max_tokens=110,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="intake",
            title="The Court Convenes",
            body=f"{packet.title}\n{packet.subtitle}\n\nCharge: {packet.charge}{hypo_line}",
            turns=[_turn("Clerk Meridian", "clerk", clerk, OPENBMB_MODEL, 0.88)],
            evidence=packet.evidence,
        ),
        delay,
    )

    judge_open = _required_role(
        model_runner,
        model_calls,
        agent=JUDGE_NAME,
        role="judge",
        model=GPT_OSS_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task=(
            f"As {JUDGE_NAME}, a Stoic courtroom judge guided by {JUDGE_PERSONA}, explain the proceeding "
            "and the burden of proof in one or two disciplined sentences."
        ),
        provider=OPENAI_PROVIDER,
        max_tokens=110,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="intake",
            title="The Burden Is Set",
            body="The bench defines how the miniature court will weigh the record.",
            turns=[_turn(JUDGE_NAME, "judge", judge_open, GPT_OSS_MODEL, 0.88)],
            evidence=packet.evidence,
        ),
        delay,
    )

    claimant_opening = _required_role(
        model_runner,
        model_calls,
        agent="Advocate Auric",
        role="claimant advocate",
        model=GPT_OSS_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task="Make the claimant's opening statement alone. Cite the strongest claimant-side exhibit.",
        provider=OPENAI_PROVIDER,
        max_tokens=130,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="claims",
            title="Claimant Opening",
            body=packet.claimant_claim,
            turns=[_turn("Advocate Auric", "claimant advocate", claimant_opening, GPT_OSS_MODEL, 0.88)],
            evidence=packet.evidence,
        ),
        delay,
    )

    respondent_opening = _required_role(
        model_runner,
        model_calls,
        agent="Counsel Sable",
        role="respondent advocate",
        model=GPT_OSS_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task="Make the respondent's opening statement alone. Emphasize uncertainty and cite a helpful exhibit.",
        provider=OPENAI_PROVIDER,
        max_tokens=130,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="opening",
            title="Respondent Opening",
            body=packet.respondent_claim,
            turns=[_turn("Counsel Sable", "respondent advocate", respondent_opening, GPT_OSS_MODEL, 0.88)],
            evidence=packet.evidence,
        ),
        delay,
    )

    auditor = _required_role(
        model_runner,
        model_calls,
        agent="Auditor Prism",
        role="evidence auditor",
        model=NEMOTRON_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task="Present the evidence record. Identify the strongest exhibit and the weakest inference.",
        provider=NEMOTRON_PROVIDER,
        max_tokens=150,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="evidence",
            title="The Record Is Audited",
            body="\n".join(f"{item.id}: {item.title} | reliability {item.reliability:.2f} | {item.note}" for item in packet.evidence),
            turns=[_turn("Auditor Prism", "evidence auditor", auditor, NEMOTRON_MODEL, 0.86)],
            evidence=packet.evidence,
        ),
        delay,
    )

    judge_question = _required_role(
        model_runner,
        model_calls,
        agent=JUDGE_NAME,
        role="judge",
        model=GPT_OSS_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task=(
            f"As {JUDGE_NAME}, ask one sharp hinge question that would change the outcome if answered. "
            "Use Stoic restraint and public reason."
        ),
        provider=OPENAI_PROVIDER,
        max_tokens=100,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="questions",
            title="The Hinge Question",
            body="The bench asks the single question that could turn the record.",
            turns=[_turn(JUDGE_NAME, "judge", judge_question, GPT_OSS_MODEL, 0.88)],
            evidence=packet.evidence,
        ),
        delay,
    )

    claimant_answer = _required_role(
        model_runner,
        model_calls,
        agent="Advocate Auric",
        role="claimant advocate",
        model=GPT_OSS_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task=f"Answer {JUDGE_NAME}'s hinge question for the claimant: {judge_question.text}",
        provider=OPENAI_PROVIDER,
        max_tokens=130,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="questions",
            title="Claimant Answers the Bench",
            body="The claimant answers the hinge question.",
            turns=[_turn("Advocate Auric", "claimant advocate", claimant_answer, GPT_OSS_MODEL, 0.88)],
            evidence=packet.evidence,
        ),
        delay,
    )

    respondent_answer = _required_role(
        model_runner,
        model_calls,
        agent="Counsel Sable",
        role="respondent advocate",
        model=GPT_OSS_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task=f"Answer {JUDGE_NAME}'s hinge question for the respondent: {judge_question.text}",
        provider=OPENAI_PROVIDER,
        max_tokens=130,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="questions",
            title="Respondent Answers the Bench",
            body="The respondent answers the hinge question.",
            turns=[_turn("Counsel Sable", "respondent advocate", respondent_answer, GPT_OSS_MODEL, 0.88)],
            evidence=packet.evidence,
        ),
        delay,
    )

    jury_panel = _required_role(
        model_runner,
        model_calls,
        agent="Nemotron Jury",
        role="juror panel",
        model=NEMOTRON_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task="Announce that the six named jurors retire to vote. Do not reveal the votes yet.",
        provider=NEMOTRON_PROVIDER,
        max_tokens=100,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="deliberation",
            title="The Jury Retires",
            body="Six fixed-perspective jurors leave the public floor to vote from the record.",
            turns=[_turn("Nemotron Jury", "juror panel", jury_panel, NEMOTRON_MODEL, 0.86)],
            evidence=packet.evidence,
        ),
        delay,
    )

    jury_votes_result = _required_role(
        model_runner,
        model_calls,
        agent="Nemotron Jury",
        role="juror vote generator",
        model=NEMOTRON_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task=_jury_task(),
        provider=NEMOTRON_PROVIDER,
        max_tokens=650,
    )
    votes = _parse_jury_votes(jury_votes_result, packet)
    for vote in votes:
        juror_result = ModelResult(
            text=f"{vote.vote.replace('_', ' ').title()}. {vote.reason}",
            call=jury_votes_result.call,
            input_text=jury_votes_result.input_text,
        )
        yield _emit(
            packet,
            source_trace,
            model_calls,
            TrialEvent(
                phase="deliberation",
                title=f"Juror {vote.juror} Votes",
                body=f"{vote.persona}. Evidence: {', '.join(vote.evidence_ids)}.",
                turns=[_turn(vote.juror, "juror", juror_result, NEMOTRON_MODEL, 0.86)],
                votes=[vote],
                evidence=packet.evidence,
            ),
            delay,
        )

    verdict = _verdict_from_votes(votes)
    verdict_voice = _required_role(
        model_runner,
        model_calls,
        agent=JUDGE_NAME,
        role="verdict writer",
        model=GPT_OSS_MODEL,
        case_summary=case_summary,
        evidence_summary=evidence_summary,
        task=(
            f"As {JUDGE_NAME}, announce the final legal finding after the jury vote with Stoic restraint. "
            f"Finding: {verdict.finding}. "
            f"Jury rationale: {verdict.rationale} Remedy: {verdict.remedy}. Include uncertainty without disclaiming the role."
        ),
        provider=OPENAI_PROVIDER,
        max_tokens=160,
    )
    yield _emit(
        packet,
        source_trace,
        model_calls,
        TrialEvent(
            phase="verdict",
            title="The Court Announces Judgment",
            body=f"{verdict_voice.text}\n\n{verdict.rationale}\n\nRemedy: {verdict.remedy}",
            verdict=verdict,
            votes=votes,
            evidence=packet.evidence,
            turns=[_turn(JUDGE_NAME, "verdict writer", verdict_voice, GPT_OSS_MODEL, 0.88)],
        ),
        delay,
    )


def stream_trial_jsonl(
    request: TrialRequest,
    delay: float = 0.0,
    model_runner: ModelRunner | None = None,
) -> Iterable[str]:
    for event in stream_trial(request, delay, model_runner=model_runner):
        yield json.dumps(event.model_dump(), ensure_ascii=True) + "\n"