Judge Citation Accuracy with Decisions API

httpx==0.28.1
pytest==9.1.1

"""Ask the Agent API questions with web search, and keep the text plus the search results it cited."""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any

import httpx

AGENT_URL = "https://api.perplexity.ai/v1/agent"
PRESET = "fast"
INSTRUCTIONS = (
    "Answer in plain prose paragraphs with no bullet points, headings, or bold text. "
    "Put a citation marker like [3] at the end of every sentence that states a fact, "
    "citing the search result that supports it. Write every sentence so it makes sense on its own: "
    "name the subject instead of starting with it, its, or they."
)


@dataclass(frozen=True)
class Source:
    id: int
    url: str
    title: str
    snippet: str


@dataclass(frozen=True)
class Answer:
    text: str
    sources: dict[int, Source]

def ask(question: str) -> Answer:
    """Answer a question with web search."""
    return _post(question)


def investigate(question: str, answer: str, sentence: str) -> Answer:
    """Search again for one sentence the sources did not back, and get a replacement for it."""
    prompt = (
        f'Question: "{question}"\n\nAnswer:\n{answer}\n\n'
        f'The sources did not clearly support this sentence from the answer:\n\n"{sentence}"\n\n'
        "Search for evidence about this sentence. Write the fewest sentences (one to three) that could replace it, "
        "stating only what the sources you find support and correcting anything that is wrong. "
        "Do not repeat facts the rest of the answer already states. Leave out any part you cannot support."
    )
    return _post(prompt)


def _post(prompt: str) -> Answer:
    body = {
        "input": prompt,
        "preset": PRESET,
        "instructions": INSTRUCTIONS,
        "tools": [{"type": "web_search"}],
    }
    headers = {"Authorization": f"Bearer {os.environ['PERPLEXITY_API_KEY']}"}
    response = httpx.post(AGENT_URL, headers=headers, json=body, timeout=120.0)
    response.raise_for_status()
    return parse(response.json())

def parse(raw: dict[str, Any]) -> Answer:
    """Pull the answer text and the numbered search results out of a response."""
    text = ""
    sources: dict[int, Source] = {}
    for item in raw["output"]:
        if item["type"] == "message":
            text = "\n\n".join(part["text"] for part in item["content"] if part["type"] == "output_text")
        elif item["type"] == "search_results":
            for result in item["results"]:
                sources[int(result["id"])] = Source(
                    id=int(result["id"]),
                    url=result["url"],
                    title=result.get("title", ""),
                    snippet=result.get("snippet", ""),
                )
    if not text:
        raise ValueError("The Agent API response has no message text.")
    if not sources:
        raise ValueError("The Agent API response has no search results; the model answered without searching.")
    return Answer(text=text, sources=sources)

"""Ask the Decisions API how well a set of passages supports one claim."""

from __future__ import annotations

import os
import time
from dataclasses import dataclass
from typing import Any

import httpx

DECISIONS_URL = "https://api.perplexity.ai/v1/decisions"
MODEL = "pplx-decider-v1.1-27b"
RETRIES = 2
RETRY_STATUSES = {429, 500, 502, 503}


@dataclass(frozen=True)
class Verdict:
    """Probabilities from one request. Nothing here is a decision yet."""

    checkable: float
    supports: dict[int, float]  # passage number -> P(passage supports the claim)
    contradicts: dict[int, float]  # passage number -> P(passage contradicts the claim)
    input_tokens: int
    output_tokens: int
    seconds: float
    request_id: str

def build_request(claim: str, passages: dict[int, str]) -> dict[str, Any]:
    """One claim, numbered passages, and 2N + 1 yes/no questions about them."""
    state = {
        "claim": claim,
        "passages": {str(n): text for n, text in passages.items()},
    }
    questions: dict[str, Any] = {
        "checkable": {
            "type": "noul",
            "instructions": (
                "Is the claim a factual statement that could be verified against a source? "
                "Opinions, transitions, and questions are not checkable."
            ),
        }
    }
    for n in passages:
        questions[f"supports_{n}"] = {
            "type": "noul",
            "instructions": (
                f"Does passage {n} state or directly imply the claim? "
                "Being about the same topic is not support; the passage has to back the specific claim."
            ),
        }
        questions[f"contradicts_{n}"] = {
            "type": "noul",
            "instructions": f"Does passage {n} state something that cannot be true if the claim is true?",
        }
    return {"model": MODEL, "state": state, "questions": questions}

def score(claim: str, passages: dict[int, str], client: httpx.Client | None = None) -> Verdict:
    """Send the request, retrying briefly on 429, 500, 502, and 503, and return the probabilities."""
    if client is None:
        with httpx.Client(timeout=90.0) as new_client:
            return score(claim, passages, new_client)
    body = build_request(claim, passages)
    headers = {"Authorization": f"Bearer {os.environ['PERPLEXITY_API_KEY']}"}
    started = time.perf_counter()
    for attempt in range(RETRIES + 1):
        response = client.post(DECISIONS_URL, headers=headers, json=body)
        if response.status_code in RETRY_STATUSES and attempt < RETRIES:
            time.sleep(float(response.headers.get("Retry-After", "1")))
            continue
        response.raise_for_status()
        break
    data = response.json()
    answers = data["answers"]
    return Verdict(
        checkable=answers["checkable"]["noul"],
        supports={n: answers[f"supports_{n}"]["noul"] for n in passages},
        contradicts={n: answers[f"contradicts_{n}"]["noul"] for n in passages},
        input_tokens=data["usage"]["input_tokens"],
        output_tokens=data["usage"]["output_tokens"],
        seconds=round(time.perf_counter() - started, 3),
        request_id=response.headers.get("x-request-id", ""),
    )

"""Check every sentence of a web-grounded answer, and send the weak ones back for another search.

The Agent API writes and rewrites. The Decisions API returns probabilities.
Only this file compares probabilities to cutoffs and chooses what to do with each sentence.
"""

from __future__ import annotations

import argparse
import html
import json
import os
import re
import sys
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path

import agent_client
import decisions_client
from agent_client import Answer

# Cutoffs. Starting points, not validated production values.
CHECKABLE_CUTOFF = 0.5
CONTRADICTED_CUTOFF = 0.6
SUPPORTED_CUTOFF = 0.7
WEAK_CUTOFF = 0.4

NEEDS_REVIEW = {"weak", "unsupported", "contradicted", "uncited"}
MAX_INVESTIGATIONS = 4  # caps the extra Agent API calls per answer

COLORS = {
    "supported": "#d4edda",
    "weak": "#fff3cd",
    "unsupported": "#f8d7da",
    "contradicted": "#f1948a",
    "uncited": "#e9ecef",
    "not_checkable": "transparent",
}


@dataclass
class Sentence:
    index: int
    paragraph: int
    text: str  # claim without citation markers
    source_ids: list[int]
    missing_ids: list[int] = field(default_factory=list)  # cited, but not in the search results
    passages: dict[int, str] = field(default_factory=dict)  # the snippets sent to the Decisions API
    urls: dict[int, str] = field(default_factory=dict)
    label: str = ""
    rule: str = ""
    checkable: float = 0.0
    supports: dict[int, float] = field(default_factory=dict)
    contradicts: dict[int, float] = field(default_factory=dict)
    input_tokens: int = 0
    request_id: str = ""  # include this when you contact support
    outcome: str = ""  # "replaced", "not confirmed", or "not reviewed" (over the limit)
    revision: list[Sentence] = field(default_factory=list)

# ---- Sentences --------------------------------------------------------------

# Move markers that sit before the final punctuation ("fact [1].") to after it ("fact. [1]").
MARKERS_BEFORE_END = re.compile(r"((?:\s*\[\d+\])+)\s*([.!?])")
# A sentence ends at . ! or ? plus any closing quotes, unless the period follows an initial ("John J.")
# or "e.g." / "i.e.". Trailing [n] markers belong to it. Text with no final punctuation is one sentence.
SENTENCE = re.compile(
    r"(.+?(?:(?<!\s[A-Z])(?<!e\.g)(?<!i\.e)[.!?][\"'”’)]*|$))((?:\s*\[\d+\])*)(?:\s+|$)",
    re.DOTALL,
)
MARKER = re.compile(r"\[(\d+)\]")


def split_sentences(text: str) -> list[Sentence]:
    """Cut the answer at sentence ends and attach each trailing [n] marker to its sentence."""
    sentences: list[Sentence] = []
    paragraphs = [p.strip() for p in text.replace("**", "").split("\n\n") if p.strip()]
    for number, paragraph in enumerate(paragraphs, start=1):
        paragraph = MARKERS_BEFORE_END.sub(r"\2\1", paragraph)
        matched = 0
        for match in SENTENCE.finditer(paragraph):
            body, markers = match.group(1), match.group(2)
            claim = MARKER.sub("", body).strip()
            ids = [int(m) for m in MARKER.findall(body + markers)]
            if claim:
                sentences.append(Sentence(len(sentences) + 1, number, claim, ids))
            matched += len(match.group(0))
        if matched != len(paragraph):
            raise ValueError(f"Could not split paragraph {number}: {paragraph[matched:][:80]!r}")
    return sentences

# ---- Policy -----------------------------------------------------------------


def label(sentence: Sentence, verdict: decisions_client.Verdict) -> None:
    """Compare probabilities to cutoffs. First matching rule wins."""
    sentence.checkable, sentence.supports, sentence.contradicts = (
        verdict.checkable,
        verdict.supports,
        verdict.contradicts,
    )
    sentence.input_tokens, sentence.request_id = verdict.input_tokens, verdict.request_id
    best_support = max(verdict.supports.values(), default=0.0)
    worst_contradiction = max(verdict.contradicts.values(), default=0.0)
    if verdict.checkable < CHECKABLE_CUTOFF:
        sentence.label, sentence.rule = "not_checkable", f"checkable {verdict.checkable:.3f}"
    elif worst_contradiction >= CONTRADICTED_CUTOFF:
        sentence.label, sentence.rule = "contradicted", f"contradicts {worst_contradiction:.3f}"
    elif not verdict.supports:
        sentence.label, sentence.rule = "uncited", f"checkable {verdict.checkable:.3f}, no citation"
    elif best_support >= SUPPORTED_CUTOFF:
        sentence.label, sentence.rule = "supported", f"supports {best_support:.3f}"
    elif best_support >= WEAK_CUTOFF:
        sentence.label, sentence.rule = "weak", f"supports {best_support:.3f}"
    else:
        sentence.label, sentence.rule = "unsupported", f"supports {best_support:.3f}"
    if sentence.missing_ids:
        sentence.rule += f"; cited {sentence.missing_ids} but no such search result"

# ---- The loop ---------------------------------------------------------------


def score_answer(answer: Answer, indent: str = "  ") -> list[Sentence]:
    """One Decisions API request per sentence. Uncited sentences still get the checkable question."""
    sentences = split_sentences(answer.text)
    for sentence in sentences:
        found = [n for n in sentence.source_ids if n in answer.sources]
        sentence.missing_ids = [n for n in sentence.source_ids if n not in answer.sources]
        sentence.passages = {n: answer.sources[n].snippet for n in found}
        sentence.urls = {n: answer.sources[n].url for n in found}
        label(sentence, decisions_client.score(sentence.text, sentence.passages))
        print(f"{indent}{sentence.index:>2}  {sentence.label:<13} {sentence.rule:<28} {sentence.text[:60]}")
    return sentences


def review(question: str, answer: Answer, sentences: list[Sentence]) -> None:
    """Send low-scoring sentences back to the Agent API, then score what it returns."""
    flagged = [s for s in sentences if s.label in NEEDS_REVIEW]
    for sentence in flagged[MAX_INVESTIGATIONS:]:
        sentence.outcome = "not reviewed"
    for sentence in flagged[:MAX_INVESTIGATIONS]:
        print(f"\nInvestigating sentence {sentence.index} ({sentence.label}). Scoring the rewrite:")
        rewrite = agent_client.investigate(question, answer.text, sentence.text)
        sentence.revision = score_answer(rewrite, indent="    -> ")
        confirmed = any(r.label == "supported" for r in sentence.revision)
        sentence.outcome = "replaced" if confirmed else "not confirmed"
        print(f"    {sentence.outcome}")


def final_sentences(sentences: list[Sentence]) -> list[Sentence]:
    """The answer after review: a replaced sentence becomes the supported sentences of its rewrite."""
    return [r for s in sentences for r in (kept(s.revision) if s.outcome == "replaced" else [s])]


def kept(revision: list[Sentence]) -> list[Sentence]:
    return [r for r in revision if r.label == "supported"]

# ---- Output -----------------------------------------------------------------


def evidence_rows(s: Sentence, note: str = "") -> str:
    """One table row per sentence, with every passage it was scored against."""
    passages = "".join(
        f"<details><summary>supports {s.supports[n]:.3f}, contradicts {s.contradicts[n]:.3f} "
        f'<a href="{html.escape(s.urls[n])}">source</a></summary><blockquote>{html.escape(text)}</blockquote></details>'
        for n, text in s.passages.items()
    )
    return (
        f'<tr><td style="background:{COLORS[s.label]}">{s.label}</td>'
        f"<td>{note}{plain(s.text)}<br><small>{html.escape(s.rule)}; request {s.request_id}</small>{passages}</td></tr>"
    )


def plain(text: str) -> str:
    """Answer text for display: escaped, without Markdown backticks. Passages stay exactly as sent."""
    return html.escape(text.replace("`", ""))


def write_report(question: str, sentences: list[Sentence], path: Path) -> None:
    """The final answer colored by label, then a table of the evidence behind every score."""
    paragraphs: dict[int, list[str]] = {}
    numbers: dict[str, int] = {}  # rewrites come from separate searches, so number citations by URL
    for s in sentences:
        replaced = s.outcome == "replaced"
        for r in kept(s.revision) if replaced else [s]:
            cites = "".join(
                f' <a href="{html.escape(url)}">[{numbers.setdefault(url, len(numbers) + 1)}]</a>'
                for url in r.urls.values()
            )
            underline = "border-bottom:2px solid #555" if replaced else ""
            badge = f" <small>({s.outcome})</small>" if s.outcome and not replaced else ""
            paragraphs.setdefault(s.paragraph, []).append(
                f'<span style="background:{COLORS[r.label]};{underline}">{plain(r.text)}{cites}{badge}</span>'
            )
    answer = "".join(f"<p>{' '.join(spans)}</p>" for spans in paragraphs.values())
    rows = ""
    for s in sentences:
        rows += evidence_rows(s)
        for r in s.revision:
            used = "used" if s.outcome == "replaced" and r.label == "supported" else "not used"
            rows += evidence_rows(r, note=f"<b>Rewrite, {used}:</b> ")
    legend = " ".join(f'<span style="background:{c};padding:2px 6px">{k}</span>' for k, c in COLORS.items())
    path.write_text(
        "<!doctype html><meta charset='utf-8'><style>body{font:16px/1.6 sans-serif;max-width:48em;margin:3em auto}"
        "td{vertical-align:top;padding:6px;border-top:1px solid #ddd}blockquote{font-size:13px;color:#444}"
        f"a{{color:inherit}}</style><h2>{html.escape(question)}</h2><p>{legend} "
        f"<span style='border-bottom:2px solid #555'>replaced after review</span></p>{answer}"
        f"<h3>Evidence</h3><table>{rows}</table>"
    )


def summarize(sentences: list[Sentence], seconds: float) -> dict[str, object]:
    everything = sentences + [r for s in sentences for r in s.revision]
    final = final_sentences(sentences)
    return {
        "sentences": len(sentences),
        "investigated": sum(1 for s in sentences if s.revision),
        "replaced": sum(1 for s in sentences if s.outcome == "replaced"),
        "decisions_requests": len(everything),
        "decisions_input_tokens": sum(s.input_tokens for s in everything),
        "seconds": round(seconds, 1),
        "before": {k: n for k in COLORS if (n := sum(1 for s in sentences if s.label == k))},
        "after": {k: n for k in COLORS if (n := sum(1 for s in final if s.label == k))},
    }

# ---- Command line -----------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description="Check every sentence of a web-grounded answer.")
    parser.add_argument("question", help="The question to ask the Agent API.")
    parser.add_argument("--out", default="out", help="Output directory.")
    args = parser.parse_args()
    if not os.environ.get("PERPLEXITY_API_KEY"):
        sys.exit("Set PERPLEXITY_API_KEY first: export PERPLEXITY_API_KEY=pplx-...")
    out = Path(args.out)
    out.mkdir(exist_ok=True)

    started = time.perf_counter()
    print("Asking the Agent API...")
    answer = agent_client.ask(args.question)
    print("Scoring each sentence with the Decisions API:")
    sentences = score_answer(answer)
    try:
        review(args.question, answer, sentences)
    finally:  # if a review request fails, save the first pass and the completed investigations
        summary = summarize(sentences, time.perf_counter() - started)
        print(f"\n{json.dumps(summary)}")
        write_report(args.question, sentences, out / "report.html")
        run = {"summary": summary, "sentences": [asdict(s) for s in sentences]}
        (out / "run.json").write_text(json.dumps(run, indent=1))
        print(f"Report: {out / 'report.html'}")


if __name__ == "__main__":
    main()

"""Offline tests. No API key, no network. Run with: python -m pytest -q"""

from __future__ import annotations

from pathlib import Path

import httpx
import pytest

import agent_client
import answer_gate
import decisions_client
from agent_client import Answer, Source
from decisions_client import Verdict


def verdict(
    supports: dict[int, float], contradicts: dict[int, float] | None = None, checkable: float = 0.95
) -> Verdict:
    """Made-up probabilities in the shape the Decisions API returns."""
    contradicts = contradicts or {n: 0.0 for n in supports}
    return Verdict(checkable, supports, contradicts, input_tokens=100, output_tokens=0, seconds=0.1, request_id="t")


def label_of(v: Verdict) -> str:
    sentence = answer_gate.Sentence(1, 1, "claim", [1])
    answer_gate.label(sentence, v)
    return sentence.label


def test_split_keeps_every_sentence_and_its_markers() -> None:
    text = "John J. Hopfield won, e.g. for nets.[1] He said “it works.” [2][3] Known fact [4].\n\nNo end mark [5]"
    sentences = answer_gate.split_sentences(text)
    assert [s.source_ids for s in sentences] == [[1], [2, 3], [4], [5]]
    assert [s.paragraph for s in sentences] == [1, 1, 1, 2]
    assert sentences[2].text == "Known fact."


def test_policy_order_and_bands() -> None:
    assert label_of(verdict({1: 0.9}, {1: 0.7})) == "contradicted"  # contradiction beats support
    assert label_of(verdict({1: 0.9})) == "supported"
    assert label_of(verdict({1: 0.5})) == "weak"
    assert label_of(verdict({1: 0.1})) == "unsupported"
    assert label_of(verdict({}, checkable=0.9)) == "uncited"
    assert label_of(verdict({1: 0.9}, checkable=0.2)) == "not_checkable"


def test_request_has_one_question_per_passage_plus_checkable() -> None:
    body = decisions_client.build_request("claim", {4: "a", 7: "b"})
    assert set(body["questions"]) == {"checkable", "supports_4", "contradicts_4", "supports_7", "contradicts_7"}
    assert all(q["type"] == "noul" for q in body["questions"].values())


def test_retry_on_429_honors_retry_after(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("PERPLEXITY_API_KEY", "test")
    sleeps: list[float] = []
    monkeypatch.setattr("time.sleep", sleeps.append)
    ok = {
        "answers": {"checkable": {"noul": 0.9}, "supports_1": {"noul": 0.8}, "contradicts_1": {"noul": 0.1}},
        "usage": {"input_tokens": 50, "output_tokens": 0},
    }
    replies = iter([httpx.Response(429, headers={"Retry-After": "2"}), httpx.Response(200, json=ok)])
    client = httpx.Client(transport=httpx.MockTransport(lambda request: next(replies)))
    assert decisions_client.score("claim", {1: "passage"}, client=client).supports == {1: 0.8}
    assert sleeps == [2.0]


def test_review_keeps_only_supported_rewrite_sentences(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    source = {1: Source(1, "https://example.com", "t", "snippet with `code`")}
    answer = Answer("Good claim. [1] Weak claim. [1] Bad claim. [1] Missing claim. [9]", source)
    rewrites = {"Weak claim.": "Better claim. [1] Extra guess. [1]", "Bad claim.": "Still bad. [1]"}
    probabilities = {
        "Good claim.": 0.9,
        "Weak claim.": 0.5,
        "Bad claim.": 0.1,
        "Better claim.": 0.95,
        "Still bad.": 0.2,
        "Extra guess.": 0.3,
    }
    monkeypatch.setattr(agent_client, "investigate", lambda q, a, s: Answer(rewrites.get(s, "Nothing. [1]"), source))
    monkeypatch.setattr(
        decisions_client, "score", lambda claim, p: verdict({n: probabilities.get(claim, 0.9) for n in p})
    )

    sentences = answer_gate.score_answer(answer)
    answer_gate.review("q", answer, sentences)

    assert [s.outcome for s in sentences] == ["", "replaced", "not confirmed", "replaced"]
    assert sentences[3].missing_ids == [9] and "no such search result" in sentences[3].rule
    final = [s.text for s in answer_gate.final_sentences(sentences)]
    assert final == ["Good claim.", "Better claim.", "Bad claim.", "Nothing."]
    answer_gate.write_report("q", sentences, tmp_path / "report.html")
    report = (tmp_path / "report.html").read_text()
    assert "(not confirmed)" in report and "snippet with `code`" in report

httpx==0.28.1
pytest==9.1.1

"""Ask the Agent API questions with web search, and keep the text plus the search results it cited."""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any

import httpx

AGENT_URL = "https://api.perplexity.ai/v1/agent"
PRESET = "fast"
INSTRUCTIONS = (
    "Answer in plain prose paragraphs with no bullet points, headings, or bold text. "
    "Put a citation marker like [3] at the end of every sentence that states a fact, "
    "citing the search result that supports it. Write every sentence so it makes sense on its own: "
    "name the subject instead of starting with it, its, or they."
)


@dataclass(frozen=True)
class Source:
    id: int
    url: str
    title: str
    snippet: str


@dataclass(frozen=True)
class Answer:
    text: str
    sources: dict[int, Source]


def ask(question: str) -> Answer:
    """Answer a question with web search."""
    return _post(question)


def investigate(question: str, answer: str, sentence: str) -> Answer:
    """Search again for one sentence the sources did not back, and get a replacement for it."""
    prompt = (
        f'Question: "{question}"\n\nAnswer:\n{answer}\n\n'
        f'The sources did not clearly support this sentence from the answer:\n\n"{sentence}"\n\n'
        "Search for evidence about this sentence. Write the fewest sentences (one to three) that could replace it, "
        "stating only what the sources you find support and correcting anything that is wrong. "
        "Do not repeat facts the rest of the answer already states. Leave out any part you cannot support."
    )
    return _post(prompt)


def _post(prompt: str) -> Answer:
    body = {
        "input": prompt,
        "preset": PRESET,
        "instructions": INSTRUCTIONS,
        "tools": [{"type": "web_search"}],
    }
    headers = {"Authorization": f"Bearer {os.environ['PERPLEXITY_API_KEY']}"}
    response = httpx.post(AGENT_URL, headers=headers, json=body, timeout=120.0)
    response.raise_for_status()
    return parse(response.json())


def parse(raw: dict[str, Any]) -> Answer:
    """Pull the answer text and the numbered search results out of a response."""
    text = ""
    sources: dict[int, Source] = {}
    for item in raw["output"]:
        if item["type"] == "message":
            text = "\n\n".join(part["text"] for part in item["content"] if part["type"] == "output_text")
        elif item["type"] == "search_results":
            for result in item["results"]:
                sources[int(result["id"])] = Source(
                    id=int(result["id"]),
                    url=result["url"],
                    title=result.get("title", ""),
                    snippet=result.get("snippet", ""),
                )
    if not text:
        raise ValueError("The Agent API response has no message text.")
    if not sources:
        raise ValueError("The Agent API response has no search results; the model answered without searching.")
    return Answer(text=text, sources=sources)

"""Ask the Decisions API how well a set of passages supports one claim."""

from __future__ import annotations

import os
import time
from dataclasses import dataclass
from typing import Any

import httpx

DECISIONS_URL = "https://api.perplexity.ai/v1/decisions"
MODEL = "pplx-decider-v1.1-27b"
RETRIES = 2
RETRY_STATUSES = {429, 500, 502, 503}


@dataclass(frozen=True)
class Verdict:
    """Probabilities from one request. Nothing here is a decision yet."""

    checkable: float
    supports: dict[int, float]  # passage number -> P(passage supports the claim)
    contradicts: dict[int, float]  # passage number -> P(passage contradicts the claim)
    input_tokens: int
    output_tokens: int
    seconds: float
    request_id: str


def build_request(claim: str, passages: dict[int, str]) -> dict[str, Any]:
    """One claim, numbered passages, and 2N + 1 yes/no questions about them."""
    state = {
        "claim": claim,
        "passages": {str(n): text for n, text in passages.items()},
    }
    questions: dict[str, Any] = {
        "checkable": {
            "type": "noul",
            "instructions": (
                "Is the claim a factual statement that could be verified against a source? "
                "Opinions, transitions, and questions are not checkable."
            ),
        }
    }
    for n in passages:
        questions[f"supports_{n}"] = {
            "type": "noul",
            "instructions": (
                f"Does passage {n} state or directly imply the claim? "
                "Being about the same topic is not support; the passage has to back the specific claim."
            ),
        }
        questions[f"contradicts_{n}"] = {
            "type": "noul",
            "instructions": f"Does passage {n} state something that cannot be true if the claim is true?",
        }
    return {"model": MODEL, "state": state, "questions": questions}


def score(claim: str, passages: dict[int, str], client: httpx.Client | None = None) -> Verdict:
    """Send the request, retrying briefly on 429, 500, 502, and 503, and return the probabilities."""
    if client is None:
        with httpx.Client(timeout=90.0) as new_client:
            return score(claim, passages, new_client)
    body = build_request(claim, passages)
    headers = {"Authorization": f"Bearer {os.environ['PERPLEXITY_API_KEY']}"}
    started = time.perf_counter()
    for attempt in range(RETRIES + 1):
        response = client.post(DECISIONS_URL, headers=headers, json=body)
        if response.status_code in RETRY_STATUSES and attempt < RETRIES:
            time.sleep(float(response.headers.get("Retry-After", "1")))
            continue
        response.raise_for_status()
        break
    data = response.json()
    answers = data["answers"]
    return Verdict(
        checkable=answers["checkable"]["noul"],
        supports={n: answers[f"supports_{n}"]["noul"] for n in passages},
        contradicts={n: answers[f"contradicts_{n}"]["noul"] for n in passages},
        input_tokens=data["usage"]["input_tokens"],
        output_tokens=data["usage"]["output_tokens"],
        seconds=round(time.perf_counter() - started, 3),
        request_id=response.headers.get("x-request-id", ""),
    )

"""Check every sentence of a web-grounded answer, and send the weak ones back for another search.

The Agent API writes and rewrites. The Decisions API returns probabilities.
Only this file compares probabilities to cutoffs and chooses what to do with each sentence.
"""

from __future__ import annotations

import argparse
import html
import json
import os
import re
import sys
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path

import agent_client
import decisions_client
from agent_client import Answer

# Cutoffs. Starting points, not validated production values.
CHECKABLE_CUTOFF = 0.5
CONTRADICTED_CUTOFF = 0.6
SUPPORTED_CUTOFF = 0.7
WEAK_CUTOFF = 0.4

NEEDS_REVIEW = {"weak", "unsupported", "contradicted", "uncited"}
MAX_INVESTIGATIONS = 4  # caps the extra Agent API calls per answer

COLORS = {
    "supported": "#d4edda",
    "weak": "#fff3cd",
    "unsupported": "#f8d7da",
    "contradicted": "#f1948a",
    "uncited": "#e9ecef",
    "not_checkable": "transparent",
}


@dataclass
class Sentence:
    index: int
    paragraph: int
    text: str  # claim without citation markers
    source_ids: list[int]
    missing_ids: list[int] = field(default_factory=list)  # cited, but not in the search results
    passages: dict[int, str] = field(default_factory=dict)  # the snippets sent to the Decisions API
    urls: dict[int, str] = field(default_factory=dict)
    label: str = ""
    rule: str = ""
    checkable: float = 0.0
    supports: dict[int, float] = field(default_factory=dict)
    contradicts: dict[int, float] = field(default_factory=dict)
    input_tokens: int = 0
    request_id: str = ""  # include this when you contact support
    outcome: str = ""  # "replaced", "not confirmed", or "not reviewed" (over the limit)
    revision: list[Sentence] = field(default_factory=list)


# ---- Sentences --------------------------------------------------------------

# Move markers that sit before the final punctuation ("fact [1].") to after it ("fact. [1]").
MARKERS_BEFORE_END = re.compile(r"((?:\s*\[\d+\])+)\s*([.!?])")
# A sentence ends at . ! or ? plus any closing quotes, unless the period follows an initial ("John J.")
# or "e.g." / "i.e.". Trailing [n] markers belong to it. Text with no final punctuation is one sentence.
SENTENCE = re.compile(
    r"(.+?(?:(?<!\s[A-Z])(?<!e\.g)(?<!i\.e)[.!?][\"'”’)]*|$))((?:\s*\[\d+\])*)(?:\s+|$)",
    re.DOTALL,
)
MARKER = re.compile(r"\[(\d+)\]")


def split_sentences(text: str) -> list[Sentence]:
    """Cut the answer at sentence ends and attach each trailing [n] marker to its sentence."""
    sentences: list[Sentence] = []
    paragraphs = [p.strip() for p in text.replace("**", "").split("\n\n") if p.strip()]
    for number, paragraph in enumerate(paragraphs, start=1):
        paragraph = MARKERS_BEFORE_END.sub(r"\2\1", paragraph)
        matched = 0
        for match in SENTENCE.finditer(paragraph):
            body, markers = match.group(1), match.group(2)
            claim = MARKER.sub("", body).strip()
            ids = [int(m) for m in MARKER.findall(body + markers)]
            if claim:
                sentences.append(Sentence(len(sentences) + 1, number, claim, ids))
            matched += len(match.group(0))
        if matched != len(paragraph):
            raise ValueError(f"Could not split paragraph {number}: {paragraph[matched:][:80]!r}")
    return sentences


# ---- Policy -----------------------------------------------------------------


def label(sentence: Sentence, verdict: decisions_client.Verdict) -> None:
    """Compare probabilities to cutoffs. First matching rule wins."""
    sentence.checkable, sentence.supports, sentence.contradicts = (
        verdict.checkable,
        verdict.supports,
        verdict.contradicts,
    )
    sentence.input_tokens, sentence.request_id = verdict.input_tokens, verdict.request_id
    best_support = max(verdict.supports.values(), default=0.0)
    worst_contradiction = max(verdict.contradicts.values(), default=0.0)
    if verdict.checkable < CHECKABLE_CUTOFF:
        sentence.label, sentence.rule = "not_checkable", f"checkable {verdict.checkable:.3f}"
    elif worst_contradiction >= CONTRADICTED_CUTOFF:
        sentence.label, sentence.rule = "contradicted", f"contradicts {worst_contradiction:.3f}"
    elif not verdict.supports:
        sentence.label, sentence.rule = "uncited", f"checkable {verdict.checkable:.3f}, no citation"
    elif best_support >= SUPPORTED_CUTOFF:
        sentence.label, sentence.rule = "supported", f"supports {best_support:.3f}"
    elif best_support >= WEAK_CUTOFF:
        sentence.label, sentence.rule = "weak", f"supports {best_support:.3f}"
    else:
        sentence.label, sentence.rule = "unsupported", f"supports {best_support:.3f}"
    if sentence.missing_ids:
        sentence.rule += f"; cited {sentence.missing_ids} but no such search result"


# ---- The loop ---------------------------------------------------------------


def score_answer(answer: Answer, indent: str = "  ") -> list[Sentence]:
    """One Decisions API request per sentence. Uncited sentences still get the checkable question."""
    sentences = split_sentences(answer.text)
    for sentence in sentences:
        found = [n for n in sentence.source_ids if n in answer.sources]
        sentence.missing_ids = [n for n in sentence.source_ids if n not in answer.sources]
        sentence.passages = {n: answer.sources[n].snippet for n in found}
        sentence.urls = {n: answer.sources[n].url for n in found}
        label(sentence, decisions_client.score(sentence.text, sentence.passages))
        print(f"{indent}{sentence.index:>2}  {sentence.label:<13} {sentence.rule:<28} {sentence.text[:60]}")
    return sentences


def review(question: str, answer: Answer, sentences: list[Sentence]) -> None:
    """Send low-scoring sentences back to the Agent API, then score what it returns."""
    flagged = [s for s in sentences if s.label in NEEDS_REVIEW]
    for sentence in flagged[MAX_INVESTIGATIONS:]:
        sentence.outcome = "not reviewed"
    for sentence in flagged[:MAX_INVESTIGATIONS]:
        print(f"\nInvestigating sentence {sentence.index} ({sentence.label}). Scoring the rewrite:")
        rewrite = agent_client.investigate(question, answer.text, sentence.text)
        sentence.revision = score_answer(rewrite, indent="    -> ")
        confirmed = any(r.label == "supported" for r in sentence.revision)
        sentence.outcome = "replaced" if confirmed else "not confirmed"
        print(f"    {sentence.outcome}")


def final_sentences(sentences: list[Sentence]) -> list[Sentence]:
    """The answer after review: a replaced sentence becomes the supported sentences of its rewrite."""
    return [r for s in sentences for r in (kept(s.revision) if s.outcome == "replaced" else [s])]


def kept(revision: list[Sentence]) -> list[Sentence]:
    return [r for r in revision if r.label == "supported"]


# ---- Output -----------------------------------------------------------------


def evidence_rows(s: Sentence, note: str = "") -> str:
    """One table row per sentence, with every passage it was scored against."""
    passages = "".join(
        f"<details><summary>supports {s.supports[n]:.3f}, contradicts {s.contradicts[n]:.3f} "
        f'<a href="{html.escape(s.urls[n])}">source</a></summary><blockquote>{html.escape(text)}</blockquote></details>'
        for n, text in s.passages.items()
    )
    return (
        f'<tr><td style="background:{COLORS[s.label]}">{s.label}</td>'
        f"<td>{note}{plain(s.text)}<br><small>{html.escape(s.rule)}; request {s.request_id}</small>{passages}</td></tr>"
    )


def plain(text: str) -> str:
    """Answer text for display: escaped, without Markdown backticks. Passages stay exactly as sent."""
    return html.escape(text.replace("`", ""))


def write_report(question: str, sentences: list[Sentence], path: Path) -> None:
    """The final answer colored by label, then a table of the evidence behind every score."""
    paragraphs: dict[int, list[str]] = {}
    numbers: dict[str, int] = {}  # rewrites come from separate searches, so number citations by URL
    for s in sentences:
        replaced = s.outcome == "replaced"
        for r in kept(s.revision) if replaced else [s]:
            cites = "".join(
                f' <a href="{html.escape(url)}">[{numbers.setdefault(url, len(numbers) + 1)}]</a>'
                for url in r.urls.values()
            )
            underline = "border-bottom:2px solid #555" if replaced else ""
            badge = f" <small>({s.outcome})</small>" if s.outcome and not replaced else ""
            paragraphs.setdefault(s.paragraph, []).append(
                f'<span style="background:{COLORS[r.label]};{underline}">{plain(r.text)}{cites}{badge}</span>'
            )
    answer = "".join(f"<p>{' '.join(spans)}</p>" for spans in paragraphs.values())
    rows = ""
    for s in sentences:
        rows += evidence_rows(s)
        for r in s.revision:
            used = "used" if s.outcome == "replaced" and r.label == "supported" else "not used"
            rows += evidence_rows(r, note=f"<b>Rewrite, {used}:</b> ")
    legend = " ".join(f'<span style="background:{c};padding:2px 6px">{k}</span>' for k, c in COLORS.items())
    path.write_text(
        "<!doctype html><meta charset='utf-8'><style>body{font:16px/1.6 sans-serif;max-width:48em;margin:3em auto}"
        "td{vertical-align:top;padding:6px;border-top:1px solid #ddd}blockquote{font-size:13px;color:#444}"
        f"a{{color:inherit}}</style><h2>{html.escape(question)}</h2><p>{legend} "
        f"<span style='border-bottom:2px solid #555'>replaced after review</span></p>{answer}"
        f"<h3>Evidence</h3><table>{rows}</table>"
    )


def summarize(sentences: list[Sentence], seconds: float) -> dict[str, object]:
    everything = sentences + [r for s in sentences for r in s.revision]
    final = final_sentences(sentences)
    return {
        "sentences": len(sentences),
        "investigated": sum(1 for s in sentences if s.revision),
        "replaced": sum(1 for s in sentences if s.outcome == "replaced"),
        "decisions_requests": len(everything),
        "decisions_input_tokens": sum(s.input_tokens for s in everything),
        "seconds": round(seconds, 1),
        "before": {k: n for k in COLORS if (n := sum(1 for s in sentences if s.label == k))},
        "after": {k: n for k in COLORS if (n := sum(1 for s in final if s.label == k))},
    }


# ---- Command line -----------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description="Check every sentence of a web-grounded answer.")
    parser.add_argument("question", help="The question to ask the Agent API.")
    parser.add_argument("--out", default="out", help="Output directory.")
    args = parser.parse_args()
    if not os.environ.get("PERPLEXITY_API_KEY"):
        sys.exit("Set PERPLEXITY_API_KEY first: export PERPLEXITY_API_KEY=pplx-...")
    out = Path(args.out)
    out.mkdir(exist_ok=True)

    started = time.perf_counter()
    print("Asking the Agent API...")
    answer = agent_client.ask(args.question)
    print("Scoring each sentence with the Decisions API:")
    sentences = score_answer(answer)
    try:
        review(args.question, answer, sentences)
    finally:  # if a review request fails, save the first pass and the completed investigations
        summary = summarize(sentences, time.perf_counter() - started)
        print(f"\n{json.dumps(summary)}")
        write_report(args.question, sentences, out / "report.html")
        run = {"summary": summary, "sentences": [asdict(s) for s in sentences]}
        (out / "run.json").write_text(json.dumps(run, indent=1))
        print(f"Report: {out / 'report.html'}")


if __name__ == "__main__":
    main()

"""Offline tests. No API key, no network. Run with: python -m pytest -q"""

from __future__ import annotations

from pathlib import Path

import httpx
import pytest

import agent_client
import answer_gate
import decisions_client
from agent_client import Answer, Source
from decisions_client import Verdict


def verdict(
    supports: dict[int, float], contradicts: dict[int, float] | None = None, checkable: float = 0.95
) -> Verdict:
    """Made-up probabilities in the shape the Decisions API returns."""
    contradicts = contradicts or {n: 0.0 for n in supports}
    return Verdict(checkable, supports, contradicts, input_tokens=100, output_tokens=0, seconds=0.1, request_id="t")


def label_of(v: Verdict) -> str:
    sentence = answer_gate.Sentence(1, 1, "claim", [1])
    answer_gate.label(sentence, v)
    return sentence.label


def test_split_keeps_every_sentence_and_its_markers() -> None:
    text = "John J. Hopfield won, e.g. for nets.[1] He said “it works.” [2][3] Known fact [4].\n\nNo end mark [5]"
    sentences = answer_gate.split_sentences(text)
    assert [s.source_ids for s in sentences] == [[1], [2, 3], [4], [5]]
    assert [s.paragraph for s in sentences] == [1, 1, 1, 2]
    assert sentences[2].text == "Known fact."


def test_policy_order_and_bands() -> None:
    assert label_of(verdict({1: 0.9}, {1: 0.7})) == "contradicted"  # contradiction beats support
    assert label_of(verdict({1: 0.9})) == "supported"
    assert label_of(verdict({1: 0.5})) == "weak"
    assert label_of(verdict({1: 0.1})) == "unsupported"
    assert label_of(verdict({}, checkable=0.9)) == "uncited"
    assert label_of(verdict({1: 0.9}, checkable=0.2)) == "not_checkable"


def test_request_has_one_question_per_passage_plus_checkable() -> None:
    body = decisions_client.build_request("claim", {4: "a", 7: "b"})
    assert set(body["questions"]) == {"checkable", "supports_4", "contradicts_4", "supports_7", "contradicts_7"}
    assert all(q["type"] == "noul" for q in body["questions"].values())


def test_retry_on_429_honors_retry_after(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("PERPLEXITY_API_KEY", "test")
    sleeps: list[float] = []
    monkeypatch.setattr("time.sleep", sleeps.append)
    ok = {
        "answers": {"checkable": {"noul": 0.9}, "supports_1": {"noul": 0.8}, "contradicts_1": {"noul": 0.1}},
        "usage": {"input_tokens": 50, "output_tokens": 0},
    }
    replies = iter([httpx.Response(429, headers={"Retry-After": "2"}), httpx.Response(200, json=ok)])
    client = httpx.Client(transport=httpx.MockTransport(lambda request: next(replies)))
    assert decisions_client.score("claim", {1: "passage"}, client=client).supports == {1: 0.8}
    assert sleeps == [2.0]


def test_review_keeps_only_supported_rewrite_sentences(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    source = {1: Source(1, "https://example.com", "t", "snippet with `code`")}
    answer = Answer("Good claim. [1] Weak claim. [1] Bad claim. [1] Missing claim. [9]", source)
    rewrites = {"Weak claim.": "Better claim. [1] Extra guess. [1]", "Bad claim.": "Still bad. [1]"}
    probabilities = {
        "Good claim.": 0.9,
        "Weak claim.": 0.5,
        "Bad claim.": 0.1,
        "Better claim.": 0.95,
        "Still bad.": 0.2,
        "Extra guess.": 0.3,
    }
    monkeypatch.setattr(agent_client, "investigate", lambda q, a, s: Answer(rewrites.get(s, "Nothing. [1]"), source))
    monkeypatch.setattr(
        decisions_client, "score", lambda claim, p: verdict({n: probabilities.get(claim, 0.9) for n in p})
    )

    sentences = answer_gate.score_answer(answer)
    answer_gate.review("q", answer, sentences)

    assert [s.outcome for s in sentences] == ["", "replaced", "not confirmed", "replaced"]
    assert sentences[3].missing_ids == [9] and "no such search result" in sentences[3].rule
    final = [s.text for s in answer_gate.final_sentences(sentences)]
    assert final == ["Good claim.", "Better claim.", "Bad claim.", "Nothing."]
    answer_gate.write_report("q", sentences, tmp_path / "report.html")
    report = (tmp_path / "report.html").read_text()
    assert "(not confirmed)" in report and "snippet with `code`" in report