Blame for interfaces.py

"""One interface for the two places an LLM is allowed: search planning and claim extraction.
Everything downstream (guards, validation, bounds comparison, conflicts, staleness) stays deterministic.

Extractor.extract(text, meta) -> (claims, notes)      # same claim dict shape as tools.extract_claims
Planner.plan(prompt) -> [query, ...]                   # search queries, capped by the agent budget

Defaults are rule-based (no network, no spend). LLM versions take any OpenAI-compatible chat client;
MockLLMClient lets you test the plumbing offline. It never produces numbers for reports."""
from __future__ import annotations
import json, os, re
from tools import extract_claims, display_uncertainty, SKIP_LINE

DEFAULT_QUERY = "squares in squares packing"


class RuleExtractor:
    name = "rule"
    def extract(self, text, meta):
        return extract_claims(text, meta)


class FixedPlanner:
    name = "fixed"
    def plan(self, prompt):
        return [prompt.get("query", DEFAULT_QUERY)]


class OpenAICompatClient:
    """Minimal chat client. Env: SQ_LLM_BASE_URL, SQ_LLM_MODEL, SQ_LLM_API_KEY.
    json_mode=False omits response_format (needed for providers that reject json_object)."""
    def __init__(self, model=None, base_url=None, api_key=None, timeout=120, sleep_s=0.0, max_retries=2, json_mode=True):
        self.sleep_s, self.max_retries = sleep_s, max_retries
        self.model = model or os.environ.get("SQ_LLM_MODEL")
        if not self.model:
            raise ValueError("SQ_LLM_MODEL (or model=) required for live LLM calls")
        self.base_url = (base_url or os.environ.get("SQ_LLM_BASE_URL") or "").rstrip("/")
        if not self.base_url:
            raise ValueError("SQ_LLM_BASE_URL (or base_url=) required for live LLM calls")
        self.api_key = api_key or os.environ.get("SQ_LLM_API_KEY")
        self.timeout = timeout
        self.json_mode = json_mode
        self.usage = {"calls": 0, "prompt_tokens": 0, "completion_tokens": 0}
        if not self.api_key:
            raise RuntimeError("no LLM key: set SQ_LLM_API_KEY")

    def complete_json(self, prompt):
        import httpx, time
        body = {"model": self.model, "temperature": 0,
                "messages": [{"role": "user", "content": prompt}]}
        if self.json_mode:
            body["response_format"] = {"type": "json_object"}
        last = None
        for k in range(self.max_retries + 1):
            if self.sleep_s:
                time.sleep(self.sleep_s)
            self.usage["calls"] += 1  # count every attempt
            try:
                r = httpx.post(f"{self.base_url}/chat/completions", timeout=self.timeout,
                               headers={"Authorization": f"Bearer {self.api_key}"},
                               json=body)
                if r.status_code in (429, 500, 502, 503, 504):
                    last = RuntimeError(f"status {r.status_code}")
                    if k < self.max_retries:
                        ra = float(r.headers.get("retry-after", 0) or 0)
                        time.sleep(ra or min(60, 5 * 2 ** k)); continue
                    raise last
                r.raise_for_status()
                j = r.json()
                if isinstance(j, dict) and j.get("error"):
                    last = RuntimeError(f"api error: {j['error']}")
                    if k < self.max_retries:
                        time.sleep(min(60, 5 * 2 ** k)); continue
                    raise last
                u = j.get("usage", {}) or {}
                for kk in ("prompt_tokens", "completion_tokens"):
                    self.usage[kk] += u.get(kk, 0)
                c = ((j.get("choices") or [{}])[0].get("message") or {}).get("content") or ""
                if not str(c).strip():
                    last = ValueError("empty content")
                    if k < self.max_retries:
                        time.sleep(min(60, 5 * 2 ** k)); continue
                    raise last
                m = re.search(r"\{.*\}", c, re.S)  # outermost {...}; tolerate prose/fences
                return json.loads(m.group(0) if m else c)
            except (httpx.HTTPError, ValueError, KeyError, json.JSONDecodeError) as e:
                last = e
                if k < self.max_retries:
                    time.sleep(min(60, 5 * 2 ** k)); continue
                raise
        raise last if last else RuntimeError("complete_json failed")


class MockLLMClient:
    """Offline stand-in: returns canned JSON (list or callable). For plumbing tests only."""
    def __init__(self, responses):
        self.responses = responses; self.usage = {"calls": 0}
    def complete_json(self, prompt):
        self.usage["calls"] += 1
        r = self.responses(prompt) if callable(self.responses) else self.responses[min(self.usage["calls"] - 1, len(self.responses) - 1)]
        if isinstance(r, Exception):
            raise r
        return r


EXTRACT_PROMPT = """You find claims about packing n unit squares into a square container of side s.
Return JSON {{"claims":[{{"n":int,"kind":"upper_bound_packing"|"lower_bound"|"other","quote":str}}]}}.
Rules: quote must be copied VERBATIM from the text (one line or sentence) and must contain the side value.
Do NOT return the side value separately; it is parsed from your quote. Use kind=lower_bound for statements like
"cannot be packed in less than". Skip hypotheticals, targets, and retracted values. Only n in {watch}.

TEXT:
{text}"""

# a numeric token: not glued to letters/digits; allows comma decimals (3,70711)
NUM_TOKEN = re.compile(r"(?<![\w.,])(\d+(?:[.,]\d+)?)(?![\w]|[.,]\d)")


def _tokval(tok):
    return float(tok.replace(",", "."))


def n_in_quote(q, n):
    return re.search(rf"(?<![\w.,]){n}(?![\w]|[.,]\d)", q) is not None


def chunk_lines(text, max_chars=12000, overlap_lines=3):
    """Split on line boundaries (never mid-number). Consecutive chunks share `overlap_lines` lines.
    A single line longer than max_chars is split at whitespace."""
    lines = []
    for ln in text.splitlines():
        while len(ln) > max_chars:
            cut = ln.rfind(" ", 0, max_chars)
            cut = cut if cut > 0 else max_chars  # no whitespace at all: hard cut is the only option
            lines.append(ln[:cut]); ln = ln[cut:].lstrip()
        lines.append(ln)
    chunks, cur, size = [], [], 0
    for ln in lines:
        if cur and size + len(ln) + 1 > max_chars:
            chunks.append("\n".join(cur))
            cur = cur[-overlap_lines:] if overlap_lines else []
            while cur and sum(len(x) + 1 for x in cur) + len(ln) + 1 > max_chars:
                cur.pop(0)
            size = sum(len(x) + 1 for x in cur)
        cur.append(ln); size += len(ln) + 1
    if cur or not chunks:
        chunks.append("\n".join(cur))
    return chunks


def side_from_quote(q, n):
    """Parse the container side from the verbatim quote. Returns (raw_token, start, end) or (None, reason).
    The rule extractor's patterns are tried first (they know 's(n) = x' etc.); otherwise exactly one
    numeric token other than n itself must remain, else the quote is ambiguous and rejected.
    Comma decimals ('3,70711') are accepted via _tokval."""
    rule, _ = extract_claims(q, {"url": "quote"})
    vals = {c["s"] for c in rule if c["n"] == n and c["s"] is not None and not c.get("rejected")}
    toks = list(NUM_TOKEN.finditer(q))
    if len(vals) == 1:
        v = vals.pop()
        hits = [m for m in toks if _tokval(m.group(1)) == v]
        if hits:
            m = hits[-1]; return m.group(1), m.start(1), m.end(1)
    if len(vals) > 1:
        return None, "ambiguous (several sides in quote)"
    rest = [m for m in toks if not (re.fullmatch(r"\d+", m.group(1)) and int(m.group(1)) == n)]
    if len(rest) != 1:
        return None, f"ambiguous ({len(rest)} numeric tokens besides n)"
    m = rest[0]; return m.group(1), m.start(1), m.end(1)


class LLMExtractor:
    """LLM proposes only (n, kind, quote). The side s is NEVER taken from the model: it is parsed from
    the verbatim quote with word-boundary tokens, and the quote must occur verbatim in the chunk. The
    canonical line then goes through extract_claims so the deterministic guards apply. Coordinate
    blocks are always parsed by the rule extractor (the LLM never supplies coordinates)."""
    name = "llm"
    def __init__(self, client, watch=None, max_chars=12000, chunk=True, overlap_lines=3):
        self.client, self.watch, self.max_chars, self.chunk, self.overlap = client, watch, max_chars, chunk, overlap_lines

    def extract(self, text, meta):
        notes, claims, seen = [], [], set()
        chunks = chunk_lines(text, self.max_chars, self.overlap) if self.chunk else [text[:self.max_chars]]
        watch = "all" if not self.watch else ",".join(map(str, sorted(self.watch)))
        for ch in chunks:
            try:
                items = self.client.complete_json(EXTRACT_PROMPT.format(watch=watch, text=ch)).get("claims", [])
                notes.append("llm chunk ok")
            except Exception as e:  # LLM failure degrades to the rule extractor for this chunk
                notes.append(f"llm error ({type(e).__name__}); rule fallback")
                c, n = extract_claims(ch, meta)
                for x in c:
                    k = (x["n"], x["s"], x["extracted_text"])
                    if k not in seen: seen.add(k); claims.append(x)
                notes += n; continue
            for it in items:
                try:
                    n, q = int(it["n"]), str(it.get("quote", ""))
                except Exception:
                    notes.append("llm item malformed; dropped"); continue
                if it.get("kind") != "upper_bound_packing":
                    notes.append(f"llm kind={it.get('kind')} n={n}; dropped"); continue
                pos = ch.find(q) if q.strip() else -1
                if pos < 0:
                    notes.append(f"llm quote not verbatim for n={n}; dropped"); continue
                line_start = ch.rfind("\n", 0, pos) + 1
                line_end = ch.find("\n", pos + len(q)); line_end = len(ch) if line_end < 0 else line_end
                if SKIP_LINE.search(ch[line_start:line_end]):
                    notes.append(f"llm quote on hedged/retracted line n={n}; dropped"); continue
                r = side_from_quote(q, n)
                if r[0] is None:
                    notes.append(f"llm quote for n={n}: {r[1]}; dropped"); continue
                if not n_in_quote(q, n):
                    notes.append(f"llm quote does not contain n={n}; dropped"); continue
                raw, a, b = r
                # the quote may stop mid-number ("5.9" of "5.93383346"): re-read the full token from the chunk
                a, b = pos + a, pos + b
                while a > 0 and (ch[a - 1].isdigit() or (ch[a - 1] in ".," and a > 1 and ch[a - 2].isdigit())):
                    a -= 1
                while b < len(ch) and (ch[b].isdigit() or (ch[b] in ".," and b + 1 < len(ch) and ch[b + 1].isdigit())):
                    b += 1
                raw = ch[a:b].replace(",", ".")
                if raw.count(".") > 1:
                    notes.append(f"llm quote for n={n}: malformed number {raw!r}; dropped"); continue
                c, n2 = extract_claims(f"s({n}) = {raw}", meta); notes += n2
                for x in c:
                    x["extracted_text"] = q[:160]
                    x["s_uncertainty"] = display_uncertainty(ch, a, b, raw)
                    k = (x["n"], x["s"], x["extracted_text"])
                    if k not in seen: seen.add(k); claims.append(x)
        blocks, _ = extract_claims("\n".join(m.group(0) for m in re.finditer(r"```packing\s*\n.*?(?:```|\Z)", text, re.S)), meta)
        claims += [b for b in blocks if b["coordinates_available"]]
        notes.append("llm_extract used")
        return claims, notes


class RegexGuardExtractor:
    """Middle arm: the naive regex, plus the same claim-level guards the rule extractor applies
    (hedge/retraction skip, area lower bound, trivial-grid upper bound). Report-level guards
    (validated coordinates, staleness, printed precision) come from the agent, as for every arm."""
    name = "regex_guards"
    def extract(self, text, meta):
        from baseline_regex import extract_naive
        from tools import area_lower_bound, trivial_upper_bound
        kept_lines = "\n".join(l for l in text.splitlines() if not SKIP_LINE.search(l))
        raw, notes = extract_naive(kept_lines, meta)
        out = []
        for c in raw:
            if c["s"] < area_lower_bound(c["n"]) - 1e-9:
                out.append({**c, "rejected": "impossible (area bound)"}); continue
            if not 1 <= c["n"] <= 1000 or c["s"] > trivial_upper_bound(c["n"]) + 1e-9:
                notes.append(f"ignored n={c['n']} s={c['s']}"); continue
            out.append({**c, "s_uncertainty": 0.0})
        return out, notes


PLAN_PROMPT = """You plan web searches for a monitor of public progress on packing n unit squares in a square.
Standing request: {req}. Return JSON {{"queries":[str, ...]}} with at most {k} short queries. No other text."""


class LLMPlanner:
    name = "llm"
    def __init__(self, client, k=3):
        self.client, self.k = client, k
    def plan(self, prompt):
        try:
            q = [str(x) for x in self.client.complete_json(PLAN_PROMPT.format(req=json.dumps(prompt), k=self.k)).get("queries", [])]
            q = [x.strip()[:120] for x in q if x.strip()][: self.k]
            return q or FixedPlanner().plan(prompt)
        except Exception:
            return FixedPlanner().plan(prompt)