src/agent/frontier_baseline.py

"""Frontier model baseline — the "Flawed Titan" side of the demo.

Sends the raw text narrative + question to a frontier model API (Gemini or
Claude) WITHOUT any graph tools or ontology awareness. The model must rely
solely on text comprehension — which is where semantic priming trips it up.
"""

from __future__ import annotations

import time
from dataclasses import dataclass


@dataclass
class FrontierResult:
    """Result from a frontier model call."""
    answer: str
    full_response: str
    model_name: str
    elapsed_seconds: float


# ---------------------------------------------------------------------------
# Gemini baseline
# ---------------------------------------------------------------------------

def run_gemini_baseline(
    narrative: str,
    question: str,
    model_name: str = "gemini-2.5-flash",
) -> FrontierResult:
    """Send narrative + question to Gemini and get an ungrounded answer.

    Args:
        narrative: The text narrative (with semantic distractors).
        question: The inheritance question.
        model_name: Gemini model to use.

    Returns:
        A FrontierResult with the model's response.
    """
    import os
    from google import genai

    api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
    client = genai.Client(api_key=api_key)

    prompt = f"""\
Read the following narrative about a family and answer the inheritance question.
Provide your reasoning and then your final answer.

--- NARRATIVE ---
{narrative}

--- QUESTION ---
{question}

Respond with your reasoning, then state your final answer as:
FINAL ANSWER: [full name of the heir]
"""

    start = time.time()
    response = client.models.generate_content(
        model=model_name,
        contents=prompt,
    )
    elapsed = time.time() - start

    text = response.text

    # Try to extract the final answer line
    answer = ""
    for line in text.split("\n"):
        if "FINAL ANSWER:" in line.upper():
            answer = line.split(":", 1)[-1].strip()
            break

    if not answer:
        # Fallback: last non-empty line
        lines = [l.strip() for l in text.strip().split("\n") if l.strip()]
        answer = lines[-1] if lines else text[:100]

    return FrontierResult(
        answer=answer,
        full_response=text,
        model_name=model_name,
        elapsed_seconds=elapsed,
    )


# ---------------------------------------------------------------------------
# Claude baseline
# ---------------------------------------------------------------------------

def run_claude_baseline(
    narrative: str,
    question: str,
    model_name: str = "claude-sonnet-4-20250514",
) -> FrontierResult:
    """Send narrative + question to Claude and get an ungrounded answer.

    Requires the ``anthropic`` package and ``ANTHROPIC_API_KEY`` env var.
    """
    import anthropic

    client = anthropic.Anthropic()

    prompt = f"""\
Read the following narrative about a family and answer the inheritance question.
Provide your reasoning and then your final answer.

--- NARRATIVE ---
{narrative}

--- QUESTION ---
{question}

Respond with your reasoning, then state your final answer as:
FINAL ANSWER: [full name of the heir]
"""

    start = time.time()
    message = client.messages.create(
        model=model_name,
        max_tokens=1024,
        messages=[{"role": "user", "content": prompt}],
    )
    elapsed = time.time() - start

    text = message.content[0].text

    # Extract answer
    answer = ""
    for line in text.split("\n"):
        if "FINAL ANSWER:" in line.upper():
            answer = line.split(":", 1)[-1].strip()
            break

    if not answer:
        lines = [l.strip() for l in text.strip().split("\n") if l.strip()]
        answer = lines[-1] if lines else text[:100]

    return FrontierResult(
        answer=answer,
        full_response=text,
        model_name=model_name,
        elapsed_seconds=elapsed,
    )


# ---------------------------------------------------------------------------
# Generic runner
# ---------------------------------------------------------------------------

def run_frontier_baseline(
    narrative: str,
    question: str,
    provider: str = "gemini",
    model_name: str | None = None,
) -> FrontierResult:
    """Run a frontier model baseline.

    Args:
        narrative: Text narrative with semantic distractors.
        question: Inheritance question.
        provider: "gemini" or "claude".
        model_name: Override the default model name.

    Returns:
        FrontierResult with the model's answer.
    """
    if provider == "gemini":
        return run_gemini_baseline(
            narrative, question,
            model_name=model_name or "gemini-2.5-flash",
        )
    elif provider == "claude":
        return run_claude_baseline(
            narrative, question,
            model_name=model_name or "claude-sonnet-4-20250514",
        )
    else:
        raise ValueError(f"Unknown provider: {provider}")


# ---------------------------------------------------------------------------
# Quick test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    sys.path.insert(0, ".")

    from src.graph.graph_builder import get_arthur_scenario
    from src.graph.scenario_generator import generate_narrative, generate_question

    scenario = get_arthur_scenario()
    narrative = generate_narrative(scenario)
    question = generate_question(scenario)

    print("NARRATIVE:")
    print(narrative)
    print(f"\nQUESTION: {question}")
    print(f"GOLD ANSWER: Edward Bellini")

    try:
        result = run_frontier_baseline(narrative, question, provider="gemini")
        print(f"\n{'=' * 60}")
        print(f"MODEL: {result.model_name}")
        print(f"ANSWER: {result.answer}")
        print(f"CORRECT: {'Edward' in result.answer}")
        print(f"TIME: {result.elapsed_seconds:.2f}s")
        print(f"\nFULL RESPONSE:\n{result.full_response}")
    except Exception as e:
        print(f"\nCould not run baseline: {e}")