| """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 |
| 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) |
| 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}""" |
|
|
| |
| 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 |
| 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: |
| 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 |
| |
| 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) |
|
|
