agent/mathcheck.py

"""Deterministic checks for trivial arithmetic claims.

Orange/uncertain extract labels for pure arithmetic often stay orange after web
search (search is a poor verifier for `2^5 = 32`). Resolve those locally so
correct maths go green and wrong maths go red before the search agent runs.
"""

from __future__ import annotations

import ast
import math
import operator
import re
from typing import Literal

from agent.state import Claim

_EQ_RE = re.compile(
    r"^\s*(?P<lhs>.+?)\s*(?:=|equals|is)\s*(?P<rhs>.+?)\s*$",
    re.IGNORECASE,
)

_ALLOWED_BINOPS: dict[type, object] = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.FloorDiv: operator.floordiv,
    ast.Mod: operator.mod,
    ast.Pow: operator.pow,
}

_ALLOWED_UNARY: dict[type, object] = {
    ast.UAdd: operator.pos,
    ast.USub: operator.neg,
}


def _normalize_expr(expr: str) -> str:
    text = (expr or "").strip()
    # Strip trailing punctuation common in prose bullets.
    text = text.rstrip(".;,")
    replacements = {
        "×": "*",
        "⋅": "*",
        "·": "*",
        "÷": "/",
        "^": "**",
        "−": "-",
        "–": "-",
        "—": "-",
    }
    for src, dst in replacements.items():
        text = text.replace(src, dst)
    # Factorial: turn 5! / (5)! into factorial(5)
    text = re.sub(r"(\d+)!", r"factorial(\1)", text)
    text = re.sub(r"\(([^()]+)\)!", r"factorial(\1)", text)
    return text


def _eval_node(node: ast.AST) -> float | int:
    if isinstance(node, ast.Expression):
        return _eval_node(node.body)
    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
        return node.value
    if isinstance(node, ast.UnaryOp) and type(node.op) in _ALLOWED_UNARY:
        fn = _ALLOWED_UNARY[type(node.op)]
        return fn(_eval_node(node.operand))  # type: ignore[operator]
    if isinstance(node, ast.BinOp) and type(node.op) in _ALLOWED_BINOPS:
        fn = _ALLOWED_BINOPS[type(node.op)]
        return fn(_eval_node(node.left), _eval_node(node.right))  # type: ignore[operator]
    if isinstance(node, ast.Call):
        if isinstance(node.func, ast.Name) and node.func.id == "factorial":
            if len(node.args) != 1 or node.keywords:
                raise ValueError("bad factorial")
            n = _eval_node(node.args[0])
            if not isinstance(n, int) or isinstance(n, bool) or n < 0 or n > 20:
                raise ValueError("factorial out of range")
            return math.factorial(n)
    raise ValueError("disallowed expression")


def _safe_eval(expr: str) -> float | int:
    normalized = _normalize_expr(expr)
    if not normalized or len(normalized) > 80:
        raise ValueError("empty/too long")
    tree = ast.parse(normalized, mode="eval")
    return _eval_node(tree)


def _values_equal(left: float | int, right: float | int) -> bool:
    if isinstance(left, int) and isinstance(right, int):
        return left == right
    return math.isclose(float(left), float(right), rel_tol=0, abs_tol=1e-9)


def verify_basic_math(
    text: str,
) -> tuple[Literal["supported", "contradicted"], str] | None:
    """If `text` is a simple arithmetic equality, return a local verdict."""
    raw = (text or "").strip()
    # Avoid prose-heavy claims; stick to short equation-like bullets.
    if not raw or len(raw) > 100 or "\n" in raw:
        return None
    # Skip if it looks like narrative with too many letters.
    letters = sum(ch.isalpha() for ch in raw)
    if letters > 12 and not re.search(r"\d", raw):
        return None

    match = _EQ_RE.match(raw)
    if not match:
        return None
    lhs = match.group("lhs")
    rhs = match.group("rhs")
    # RHS should be a plain number for these trivial checks.
    try:
        left_val = _safe_eval(lhs)
        right_val = _safe_eval(rhs)
    except Exception:  # noqa: BLE001
        return None

    if _values_equal(left_val, right_val):
        return "supported", "verified basic arithmetic"
    return (
        "contradicted",
        f"basic arithmetic: left={left_val}, right={right_val}",
    )


def apply_basic_math_verdicts(claims: list[Claim]) -> list[Claim]:
    """Upgrade/downgrade clear arithmetic claims before web search."""
    out: list[Claim] = []
    for claim in claims:
        current = claim.get("verdict", "uncertain")
        if current not in {"uncertain", "supported"}:
            out.append(claim)
            continue
        result = verify_basic_math(claim.get("text") or "")
        if result is None:
            out.append(claim)
            continue
        verdict, reason = result
        updated = {
            **claim,
            "verdict": verdict,
            "reason": reason,
            "comment": claim.get("comment") or reason,
        }
        if verdict == "supported":
            updated["citations"] = []
        out.append(updated)  # type: ignore[arg-type]
    return out