From fc4087ca4bc5d961e59c2ddf96c13c57943ee9e2 Mon Sep 17 00:00:00 2001 From: twentyOne2x Date: Wed, 15 Jul 2026 23:03:00 +0200 Subject: [PATCH] Support private local source reconstruction --- scripts/compile_kb_source_packet.py | 22 +- .../prepare_kb_context_from_canonical_dump.py | 320 ++++++++++++++++++ scripts/prepare_kb_source_manifest.py | 144 +++++++- tests/test_compile_kb_source_packet.py | 11 +- ..._prepare_kb_context_from_canonical_dump.py | 128 +++++++ tests/test_prepare_kb_source_manifest.py | 135 ++++++++ 6 files changed, 741 insertions(+), 19 deletions(-) create mode 100644 scripts/prepare_kb_context_from_canonical_dump.py create mode 100644 tests/test_prepare_kb_context_from_canonical_dump.py diff --git a/scripts/compile_kb_source_packet.py b/scripts/compile_kb_source_packet.py index e7b3efe..88d09b4 100644 --- a/scripts/compile_kb_source_packet.py +++ b/scripts/compile_kb_source_packet.py @@ -13,8 +13,10 @@ from __future__ import annotations import argparse import hashlib import json +import os import re import sys +import tempfile import uuid from pathlib import Path from typing import Any @@ -574,6 +576,23 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: return parser.parse_args(argv) +def _write_private_output(path: Path, rendered: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + fd, raw_temp = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent) + temp = Path(raw_temp) + try: + os.fchmod(fd, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(rendered) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temp, path) + os.chmod(path, 0o600) + finally: + if temp.exists(): + temp.unlink() + + def main(argv: list[str] | None = None) -> int: args = parse_args(argv) try: @@ -600,8 +619,7 @@ def main(argv: list[str] | None = None) -> int: rendered = json.dumps(bundle, indent=2, sort_keys=True, ensure_ascii=True) + "\n" if args.output: try: - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(rendered, encoding="utf-8") + _write_private_output(args.output, rendered) except OSError as exc: print( json.dumps( diff --git a/scripts/prepare_kb_context_from_canonical_dump.py b/scripts/prepare_kb_context_from_canonical_dump.py new file mode 100644 index 0000000..d5e26ce --- /dev/null +++ b/scripts/prepare_kb_context_from_canonical_dump.py @@ -0,0 +1,320 @@ +#!/usr/bin/env python3 +"""Build source-specific canonical claim context from a Postgres dump without a database write.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import os +import re +import shutil +import subprocess +import tempfile +from collections import Counter +from decimal import Decimal, InvalidOperation +from pathlib import Path +from types import ModuleType +from typing import Any + +SCHEMA = "livingip.kbSourceDumpContextReceipt.v1" +DEFAULT_DOCKER_IMAGE = "postgres:16-alpine" +DEFAULT_LIMIT = 12 +MAX_PREVIEW_BYTES = 16_000 +TABLES = ("claims", "claim_evidence", "claim_edges") + + +class ContextError(RuntimeError): + """Raised when a dump cannot safely produce canonical retrieval context.""" + + +def sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _read_file(path: Path, label: str) -> bytes: + try: + if not path.is_file(): + raise ContextError(f"{label} is not a regular file: {path}") + return path.read_bytes() + except OSError as exc: + raise ContextError(f"could not read {label}: {exc}") from exc + + +def _load_kb_tool(repo_root: Path) -> ModuleType: + path = repo_root / "hermes-agent" / "leoclean-bin" / "kb_tool.py" + spec = importlib.util.spec_from_file_location("teleo_kb_tool_context", path) + if spec is None or spec.loader is None: + raise ContextError(f"could not load retrieval semantics from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _decode_copy_field(value: str) -> str | None: + if value == r"\N": + return None + result: list[str] = [] + index = 0 + escapes = {"b": "\b", "f": "\f", "n": "\n", "r": "\r", "t": "\t", "v": "\v", "\\": "\\"} + while index < len(value): + char = value[index] + if char != "\\" or index + 1 >= len(value): + result.append(char) + index += 1 + continue + index += 1 + escaped = value[index] + if escaped in escapes: + result.append(escapes[escaped]) + index += 1 + continue + if escaped == "x": + match = re.match(r"[0-9A-Fa-f]{1,2}", value[index + 1 :]) + if match: + result.append(chr(int(match.group(0), 16))) + index += 1 + len(match.group(0)) + continue + if escaped in "01234567": + match = re.match(r"[0-7]{1,3}", value[index:]) + if match: + result.append(chr(int(match.group(0), 8))) + index += len(match.group(0)) + continue + result.append(escaped) + index += 1 + return "".join(result) + + +def parse_copy_table(sql: str, table: str) -> list[dict[str, str | None]]: + prefix = f"COPY public.{table} (" + lines = iter(sql.splitlines()) + for line in lines: + if line.startswith(prefix) and line.endswith(") FROM stdin;"): + columns = [item.strip() for item in line[len(prefix) : -len(") FROM stdin;")].split(",")] + break + else: + raise ContextError(f"pg_restore output has no COPY block for public.{table}") + rows: list[dict[str, str | None]] = [] + for line in lines: + if line == r"\.": + return rows + fields = line.split("\t") + if len(fields) != len(columns): + raise ContextError(f"public.{table} COPY row has {len(fields)} fields; expected {len(columns)}") + rows.append(dict(zip(columns, (_decode_copy_field(field) for field in fields), strict=True))) + raise ContextError(f"public.{table} COPY block is unterminated") + + +def restore_table(dump_path: Path, table: str, docker_image: str) -> str: + if table not in TABLES: + raise ContextError(f"unsupported canonical table: {table}") + if shutil.which("docker") is None: + raise ContextError("docker is not installed or not available on PATH") + command = [ + "docker", + "run", + "--rm", + "--network=none", + "--mount", + f"type=bind,src={dump_path.resolve()},dst=/canonical.dump,readonly", + docker_image, + "pg_restore", + "--data-only", + "-t", + table, + "-f", + "-", + "/canonical.dump", + ] + try: + completed = subprocess.run(command, capture_output=True, text=True, timeout=180, check=False) + except (OSError, subprocess.TimeoutExpired) as exc: + raise ContextError(f"pg_restore failed ({type(exc).__name__}); diagnostic detail redacted") from exc + if completed.returncode != 0: + raise ContextError(f"pg_restore returned exit {completed.returncode}; diagnostic detail redacted") + return completed.stdout + + +def _decimal(value: str | None) -> Decimal: + if value is None: + return Decimal(0) + try: + return Decimal(value) + except InvalidOperation as exc: + raise ContextError("canonical claim confidence is not numeric") from exc + + +def build_context( + *, dump_path: Path, artifact_path: Path, title: str, limit: int, docker_image: str, repo_root: Path +) -> tuple[dict[str, Any], dict[str, Any]]: + if limit < 1 or limit > 20: + raise ContextError("limit must be between 1 and 20") + dump_bytes = _read_file(dump_path, "canonical dump") + artifact_bytes = _read_file(artifact_path, "source artifact") + try: + preview = artifact_bytes[:MAX_PREVIEW_BYTES].decode("utf-8", errors="strict") + except UnicodeDecodeError as exc: + raise ContextError(f"source artifact preview must be strict UTF-8: {exc}") from exc + kb_tool = _load_kb_tool(repo_root) + query = " ".join(kb_tool.query_terms(f"{title} {preview}")) + terms = kb_tool.query_terms(query) + min_score = 2 if len(terms) >= 3 else 1 + + tables = {table: parse_copy_table(restore_table(dump_path, table, docker_image), table) for table in TABLES} + evidence_counts = Counter(row["claim_id"] for row in tables["claim_evidence"] if row["claim_id"]) + edge_counts: Counter[str] = Counter() + for row in tables["claim_edges"]: + if row["from_claim"]: + edge_counts[row["from_claim"]] += 1 + if row["to_claim"]: + edge_counts[row["to_claim"]] += 1 + + ranked: list[dict[str, Any]] = [] + for row in tables["claims"]: + if row["status"] != "open" or not row["id"] or not row["text"]: + continue + text = row["text"] + searchable = f"{text}\n{row['tags'] or ''}".lower() + score = sum(term in searchable for term in terms) + if score < min_score: + continue + ranked.append( + { + "id": row["id"], + "text": text, + "score": score, + "evidence_count": evidence_counts[row["id"]], + "edge_count": edge_counts[row["id"]], + "_confidence": _decimal(row["confidence"]), + } + ) + ranked.sort( + key=lambda row: ( + -row["score"], + -row["evidence_count"], + -row["edge_count"], + -row["_confidence"], + len(row["text"]), + row["text"], + row["id"], + ) + ) + candidates = [{key: value for key, value in row.items() if key != "_confidence"} for row in ranked[:limit]] + context = {"database_search_query": query, "candidate_claims": candidates} + receipt = { + "schema": SCHEMA, + "status": "prepared", + "database_write_performed": False, + "model_call_performed": False, + "canonical_dump": { + "path": str(dump_path.resolve()), + "sha256": sha256_bytes(dump_bytes), + }, + "source": { + "path": str(artifact_path.resolve()), + "sha256": sha256_bytes(artifact_bytes), + }, + "retrieval": { + "query_sha256": sha256_bytes(query.encode("utf-8")), + "term_count": len(terms), + "minimum_score": min_score, + "canonical_open_claim_count": sum(row["status"] == "open" for row in tables["claims"]), + "candidate_count": len(candidates), + "candidate_ids": [row["id"] for row in candidates], + }, + "runtime": { + "docker_image": docker_image, + "network_disabled": True, + "pg_restore_output_retained": False, + }, + } + return context, receipt + + +def _private_json(path: Path, value: Any) -> None: + if path.exists(): + raise ContextError(f"refusing to overwrite existing output: {path}") + path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + os.chmod(path.parent, 0o700) + fd, raw_temp = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent) + temp = Path(raw_temp) + try: + os.fchmod(fd, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(value, handle, indent=2, sort_keys=True, ensure_ascii=True) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temp, path) + os.chmod(path, 0o600) + finally: + if temp.exists(): + temp.unlink() + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dump", required=True, type=Path) + parser.add_argument("--artifact", required=True, type=Path) + parser.add_argument("--title", required=True) + parser.add_argument("--context-output", required=True, type=Path) + parser.add_argument("--receipt-output", required=True, type=Path) + parser.add_argument("--limit", type=int, default=DEFAULT_LIMIT) + parser.add_argument("--docker-image", default=DEFAULT_DOCKER_IMAGE) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + repo_root = Path(__file__).resolve().parents[1] + try: + if args.context_output.resolve() == args.receipt_output.resolve(): + raise ContextError("context and receipt outputs must be different files") + context, receipt = build_context( + dump_path=args.dump, + artifact_path=args.artifact, + title=args.title, + limit=args.limit, + docker_image=args.docker_image, + repo_root=repo_root, + ) + _private_json(args.context_output, context) + receipt["outputs"] = { + "context": str(args.context_output.resolve()), + "receipt": str(args.receipt_output.resolve()), + } + _private_json(args.receipt_output, receipt) + except (ContextError, OSError, ValueError) as exc: + print( + json.dumps( + { + "schema": SCHEMA, + "status": "rejected", + "database_write_performed": False, + "error": str(exc), + }, + sort_keys=True, + ) + ) + return 2 + print( + json.dumps( + { + "schema": SCHEMA, + "status": "prepared", + "database_write_performed": False, + "candidate_count": receipt["retrieval"]["candidate_count"], + "context_sha256": sha256_bytes(json.dumps(context, sort_keys=True).encode("utf-8")), + "outputs": receipt["outputs"], + }, + indent=2, + sort_keys=True, + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/prepare_kb_source_manifest.py b/scripts/prepare_kb_source_manifest.py index 6823e1d..e9b8ade 100644 --- a/scripts/prepare_kb_source_manifest.py +++ b/scripts/prepare_kb_source_manifest.py @@ -15,6 +15,8 @@ import hashlib import json import os import re +import shutil +import subprocess import sys import tempfile import urllib.error @@ -34,6 +36,8 @@ DEFAULT_MODEL = "anthropic/claude-sonnet-4.5" EXTRACTOR_VERSION = "2" DEFAULT_OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" DEFAULT_KEY_FILE = Path("/opt/teleo-eval/secrets/openrouter-key") +DEFAULT_CLAUDE_MAX_BUDGET_USD = 0.10 +MODEL_PROVIDERS = {"openrouter", "claude-cli"} TEXT_SUFFIXES = {".csv", ".htm", ".html", ".json", ".jsonl", ".md", ".rst", ".txt", ".xml"} MAX_SOURCE_CHARS = 120_000 MAX_CANDIDATES = 20 @@ -248,10 +252,10 @@ def call_openrouter(prompt: str, *, model: str, key_file: Path, url: str, timeou with urllib.request.urlopen(request, timeout=timeout_seconds) as response: body = response.read() except urllib.error.HTTPError as exc: - detail = exc.read(600).decode("utf-8", errors="replace") - raise PreparationError(f"OpenRouter returned HTTP {exc.code}: {detail}") from exc + exc.read(600) + raise PreparationError(f"OpenRouter returned HTTP {exc.code}; provider detail redacted") from exc except (urllib.error.URLError, TimeoutError, OSError) as exc: - raise PreparationError(f"OpenRouter request failed: {exc}") from exc + raise PreparationError(f"OpenRouter request failed ({type(exc).__name__}); provider detail redacted") from exc try: envelope = json.loads(body.decode("utf-8", errors="strict")) content = envelope["choices"][0]["message"]["content"] @@ -263,6 +267,100 @@ def call_openrouter(prompt: str, *, model: str, key_file: Path, url: str, timeou return content, usage if isinstance(usage, dict) else {} +def call_claude_cli( + prompt: str, + *, + model: str, + timeout_seconds: int, + max_budget_usd: float, +) -> tuple[str, dict[str, Any]]: + """Run a non-persistent, tool-free Claude extraction with a cost cap. + + The private prompt is supplied on stdin rather than in the process argv. + Provider diagnostics are never copied into exceptions because they may + contain source fragments or credential-adjacent values. + """ + executable = shutil.which("claude") + if executable is None: + raise PreparationError("claude CLI is not installed or not available on PATH") + if isinstance(max_budget_usd, bool) or not isinstance(max_budget_usd, (int, float)) or max_budget_usd <= 0: + raise PreparationError("Claude CLI max budget must be a positive number") + command = [ + executable, + "-p", + "--no-session-persistence", + "--tools", + "", + "--model", + model, + "--output-format", + "json", + "--max-budget-usd", + str(max_budget_usd), + ] + try: + completed = subprocess.run( + command, + input=prompt, + capture_output=True, + text=True, + timeout=timeout_seconds, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise PreparationError(f"Claude CLI request failed ({type(exc).__name__}); provider detail redacted") from exc + if completed.returncode != 0: + raise PreparationError(f"Claude CLI returned exit {completed.returncode}; provider detail redacted") + try: + envelope = json.loads(completed.stdout) + except json.JSONDecodeError as exc: + raise PreparationError("Claude CLI returned an invalid JSON envelope; provider detail redacted") from exc + if not isinstance(envelope, dict) or envelope.get("is_error") is True or envelope.get("subtype") != "success": + raise PreparationError("Claude CLI returned an unsuccessful envelope; provider detail redacted") + content = envelope.get("result") + if not isinstance(content, str) or not content.strip(): + raise PreparationError("Claude CLI returned no extraction content") + usage = { + "total_cost_usd": envelope.get("total_cost_usd"), + "usage": envelope.get("usage") if isinstance(envelope.get("usage"), dict) else {}, + "model_usage": envelope.get("modelUsage") if isinstance(envelope.get("modelUsage"), dict) else {}, + "max_budget_usd": max_budget_usd, + "session_persistence": False, + "tools_enabled": False, + } + return content, usage + + +def call_model( + prompt: str, args: argparse.Namespace, *, max_budget_usd: float | None = None +) -> tuple[str, dict[str, Any]]: + provider = getattr(args, "provider", "openrouter") + if provider == "openrouter": + return call_openrouter( + prompt, + model=args.model, + key_file=args.key_file, + url=args.openrouter_url, + timeout_seconds=args.timeout_seconds, + ) + if provider == "claude-cli": + return call_claude_cli( + prompt, + model=args.model, + timeout_seconds=args.timeout_seconds, + max_budget_usd=( + max_budget_usd + if max_budget_usd is not None + else getattr(args, "max_budget_usd", DEFAULT_CLAUDE_MAX_BUDGET_USD) + ), + ) + raise PreparationError(f"model provider must be one of {sorted(MODEL_PROVIDERS)}") + + +def extractor_name(args: argparse.Namespace) -> str: + return f"{getattr(args, 'provider', 'openrouter')}:{args.model}" + + def parse_model_output(content: str) -> dict[str, Any]: stripped = content.strip() if stripped.startswith("```"): @@ -447,7 +545,7 @@ def build_manifest( "schema": compiler.MANIFEST_SCHEMA_V2, "artifact_sha256": artifact_sha256, "extracted_text_sha256": text_sha256, - "extractor": {"name": f"openrouter:{args.model}", "version": EXTRACTOR_VERSION}, + "extractor": {"name": extractor_name(args), "version": EXTRACTOR_VERSION}, "source": { "identity": args.identity, "source_key": args.source_key, @@ -526,14 +624,23 @@ def prepare(args: argparse.Namespace) -> dict[str, Any]: extraction: dict[str, Any] | None = None bundle: dict[str, Any] | None = None manifest: dict[str, Any] | None = None + reported_total_cost_usd = 0.0 for attempt in range(1, MAX_ATTEMPTS + 1): - raw, usage = call_openrouter( + provider = getattr(args, "provider", "openrouter") + max_total_budget_usd = getattr(args, "max_budget_usd", DEFAULT_CLAUDE_MAX_BUDGET_USD) + remaining_budget_usd = max_total_budget_usd - reported_total_cost_usd if provider == "claude-cli" else None + raw, usage = call_model( build_prompt(fragments, context, repair_error), - model=args.model, - key_file=args.key_file, - url=args.openrouter_url, - timeout_seconds=args.timeout_seconds, + args, + max_budget_usd=remaining_budget_usd, ) + if provider == "claude-cli": + reported_cost = usage.get("total_cost_usd") + if isinstance(reported_cost, bool) or not isinstance(reported_cost, (int, float)) or reported_cost < 0: + raise PreparationError("Claude CLI did not report a valid non-negative model cost") + reported_total_cost_usd += float(reported_cost) + if reported_total_cost_usd > max_total_budget_usd + 1e-9: + raise PreparationError("Claude CLI reported model cost above the configured total budget") attempts.append({"attempt": attempt, "response_sha256": sha256_bytes(raw.encode("utf-8")), "usage": usage}) try: selection = parse_model_output(raw) @@ -616,9 +723,17 @@ def prepare(args: argparse.Namespace) -> dict[str, Any]: "duplicate_judgments": extraction["duplicates"], }, "extraction": { - "provider": "openrouter", + "provider": getattr(args, "provider", "openrouter"), "model": args.model, "attempts": attempts, + "cost_control": { + "max_total_budget_usd": ( + getattr(args, "max_budget_usd", DEFAULT_CLAUDE_MAX_BUDGET_USD) + if getattr(args, "provider", "openrouter") == "claude-cli" + else None + ), + "reported_total_cost_usd": reported_total_cost_usd, + }, "claim_count": len(extraction["claims"]), "evidence_count": sum(len(claim["evidence"]) for claim in extraction["claims"]), "duplicate_judgment_count": len(extraction["duplicates"]), @@ -631,7 +746,7 @@ def prepare(args: argparse.Namespace) -> dict[str, Any]: "extracted_text_sha256": text_sha256, "kb_context_sha256": sha256_bytes(_json_bytes(context)), "fragment_index_sha256": sha256_bytes(_json_bytes(fragment_locations)), - "extractor": {"name": f"openrouter:{args.model}", "version": EXTRACTOR_VERSION}, + "extractor": {"name": extractor_name(args), "version": EXTRACTOR_VERSION}, "exact_selected_locations_validated": bundle is not None, }, "outputs": { @@ -668,6 +783,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--locator", required=True) parser.add_argument("--output-dir", required=True, type=Path) parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--provider", choices=sorted(MODEL_PROVIDERS), default="openrouter") + parser.add_argument( + "--max-budget-usd", + type=float, + default=DEFAULT_CLAUDE_MAX_BUDGET_USD, + help="hard total Claude CLI budget cap across validation retries; ignored by the OpenRouter provider", + ) parser.add_argument("--openrouter-url", default=DEFAULT_OPENROUTER_URL) parser.add_argument("--timeout-seconds", type=int, default=180) parser.set_defaults(key_file=DEFAULT_KEY_FILE) diff --git a/tests/test_compile_kb_source_packet.py b/tests/test_compile_kb_source_packet.py index 49adf02..a220b3d 100644 --- a/tests/test_compile_kb_source_packet.py +++ b/tests/test_compile_kb_source_packet.py @@ -3,6 +3,7 @@ from __future__ import annotations import copy import hashlib import json +import stat import sys from pathlib import Path @@ -228,6 +229,7 @@ def test_cli_emits_structured_json_to_stdout_and_optional_output(tmp_path: Path, assert status == 0 assert output.read_text(encoding="utf-8") == stdout assert json.loads(stdout)["status"] == "pending_review" + assert stat.S_IMODE(output.stat().st_mode) == 0o600 @pytest.mark.parametrize( @@ -363,14 +365,11 @@ def test_cli_reports_output_write_failure_as_structured_rejection( ) -> None: artifact, text, manifest_path = _write_inputs(tmp_path) output = tmp_path / "packet.json" - original_write_text = Path.write_text - def guarded_write_text(path: Path, data: str, *args, **kwargs): - if path == output: - raise PermissionError("fixture output denied") - return original_write_text(path, data, *args, **kwargs) + def fail_write(_path: Path, _rendered: str) -> None: + raise PermissionError("fixture output denied") - monkeypatch.setattr(Path, "write_text", guarded_write_text) + monkeypatch.setattr(compiler, "_write_private_output", fail_write) status = compiler.main( [ "--artifact", diff --git a/tests/test_prepare_kb_context_from_canonical_dump.py b/tests/test_prepare_kb_context_from_canonical_dump.py new file mode 100644 index 0000000..5272afc --- /dev/null +++ b/tests/test_prepare_kb_context_from_canonical_dump.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import importlib.util +import json +import stat +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +SCRIPT_PATH = REPO_ROOT / "scripts" / "prepare_kb_context_from_canonical_dump.py" +SPEC = importlib.util.spec_from_file_location("prepare_kb_context_from_canonical_dump", SCRIPT_PATH) +assert SPEC is not None and SPEC.loader is not None +context_builder = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = context_builder +SPEC.loader.exec_module(context_builder) + +CLAIM_A = "11111111-1111-4111-8111-111111111111" +CLAIM_B = "22222222-2222-4222-8222-222222222222" +SOURCE_ID = "33333333-3333-4333-8333-333333333333" + + +def _copy(table: str, columns: list[str], rows: list[list[str]]) -> str: + rendered = [f"COPY public.{table} ({', '.join(columns)}) FROM stdin;"] + rendered.extend("\t".join(row) for row in rows) + rendered.append(r"\.") + return "\n".join(rendered) + "\n" + + +def test_build_context_matches_readonly_claim_ranking_and_counts( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + dump = tmp_path / "canonical.dump" + dump.write_bytes(b"fixture dump") + artifact = tmp_path / "segment.txt" + artifact.write_text("Exact source quotes keep staged proposals reviewable.\n", encoding="utf-8") + sql = { + "claims": _copy( + "claims", + [ + "id", + "type", + "text", + "status", + "confidence", + "tags", + "created_by", + "superseded_by", + "created_at", + "updated_at", + ], + [ + [ + CLAIM_A, + "structural", + "Exact source quotes keep proposals reviewable.", + "open", + "0.8", + "{provenance}", + r"\N", + r"\N", + "2026-01-01", + "2026-01-01", + ], + [ + CLAIM_B, + "structural", + "Staged proposals need guarded apply.", + "open", + "0.9", + "{review}", + r"\N", + r"\N", + "2026-01-01", + "2026-01-01", + ], + ], + ), + "claim_evidence": _copy( + "claim_evidence", + ["claim_id", "source_id", "role", "weight", "created_by", "created_at"], + [[CLAIM_A, SOURCE_ID, "grounds", "1", r"\N", "2026-01-01"]], + ), + "claim_edges": _copy( + "claim_edges", + ["id", "from_claim", "to_claim", "edge_type", "weight", "created_by", "created_at"], + [["44444444-4444-4444-8444-444444444444", CLAIM_A, CLAIM_B, "supports", "1", r"\N", "2026-01-01"]], + ), + } + monkeypatch.setattr(context_builder, "restore_table", lambda _dump, table, _image: sql[table]) + + context, receipt = context_builder.build_context( + dump_path=dump, + artifact_path=artifact, + title="Telegram segment", + limit=12, + docker_image="fixture-postgres", + repo_root=REPO_ROOT, + ) + + assert context["candidate_claims"][0]["id"] == CLAIM_A + assert context["candidate_claims"][0]["evidence_count"] == 1 + assert context["candidate_claims"][0]["edge_count"] == 1 + assert receipt["database_write_performed"] is False + assert receipt["model_call_performed"] is False + assert receipt["runtime"]["network_disabled"] is True + assert receipt["runtime"]["pg_restore_output_retained"] is False + + +def test_copy_parser_preserves_escaped_source_text() -> None: + sql = _copy("claims", ["id", "text"], [[CLAIM_A, r"line one\nline two\tvalue\\tail"]]) + + rows = context_builder.parse_copy_table(sql, "claims") + + assert rows == [{"id": CLAIM_A, "text": "line one\nline two\tvalue\\tail"}] + + +def test_private_json_is_owner_only_and_refuses_overwrite(tmp_path: Path) -> None: + output = tmp_path / "private" / "context.json" + + context_builder._private_json(output, {"safe": True}) + + assert json.loads(output.read_text(encoding="utf-8")) == {"safe": True} + assert stat.S_IMODE(output.parent.stat().st_mode) == 0o700 + assert stat.S_IMODE(output.stat().st_mode) == 0o600 + with pytest.raises(context_builder.ContextError, match="refusing to overwrite"): + context_builder._private_json(output, {"safe": False}) diff --git a/tests/test_prepare_kb_source_manifest.py b/tests/test_prepare_kb_source_manifest.py index d85a0c2..f24335a 100644 --- a/tests/test_prepare_kb_source_manifest.py +++ b/tests/test_prepare_kb_source_manifest.py @@ -1,8 +1,10 @@ from __future__ import annotations +import io import json import stat import sys +import urllib.error from pathlib import Path from types import SimpleNamespace @@ -35,9 +37,11 @@ def _args(tmp_path: Path, artifact: Path, context: Path) -> SimpleNamespace: locator="artifact://source-preparation/test-v1", output_dir=tmp_path / "private-output", key_file=tmp_path / "openrouter-key", + provider="openrouter", model="fixture/model", openrouter_url="https://openrouter.invalid/v1/chat/completions", timeout_seconds=30, + max_budget_usd=0.10, ) @@ -274,3 +278,134 @@ def test_openrouter_key_cannot_be_redirected_to_an_unapproved_endpoint(tmp_path: url="https://example.invalid/collect", timeout_seconds=1, ) + + +def test_claude_cli_uses_stdin_without_tools_or_session_persistence_and_retains_cost( + monkeypatch: pytest.MonkeyPatch, +) -> None: + prompt = "private source prompt" + observed: dict[str, object] = {} + + def fake_run(command, **kwargs): + observed["command"] = command + observed.update(kwargs) + return SimpleNamespace( + returncode=0, + stderr="", + stdout=json.dumps( + { + "subtype": "success", + "is_error": False, + "result": _model_output(), + "total_cost_usd": 0.0123, + "usage": {"input_tokens": 321, "output_tokens": 123}, + "modelUsage": {"claude-sonnet-test": {"costUSD": 0.0123}}, + } + ), + ) + + monkeypatch.setattr(preparer.shutil, "which", lambda _name: "/usr/local/bin/claude") + monkeypatch.setattr(preparer.subprocess, "run", fake_run) + + content, usage = preparer.call_claude_cli( + prompt, + model="sonnet", + timeout_seconds=30, + max_budget_usd=0.10, + ) + + command = observed["command"] + assert prompt not in command + assert observed["input"] == prompt + assert "--no-session-persistence" in command + assert command[command.index("--tools") + 1] == "" + assert command[command.index("--max-budget-usd") + 1] == "0.1" + assert json.loads(content)["schema"] == preparer.MODEL_OUTPUT_SCHEMA + assert usage["total_cost_usd"] == 0.0123 + assert usage["session_persistence"] is False + assert usage["tools_enabled"] is False + + +def test_prepare_routes_claude_cli_and_records_provider_cost(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + artifact, context = _inputs(tmp_path) + args = _args(tmp_path, artifact, context) + args.provider = "claude-cli" + args.model = "sonnet" + monkeypatch.setattr( + preparer, + "call_openrouter", + lambda *_args, **_kwargs: pytest.fail("OpenRouter must not be called for the Claude CLI provider"), + ) + monkeypatch.setattr( + preparer, + "call_claude_cli", + lambda *_args, **_kwargs: (_model_output(), {"total_cost_usd": 0.0123}), + ) + + receipt = preparer.prepare(args) + + assert receipt["extraction"]["provider"] == "claude-cli" + assert receipt["extraction"]["attempts"][0]["usage"]["total_cost_usd"] == 0.0123 + assert receipt["extraction"]["cost_control"] == { + "max_total_budget_usd": 0.10, + "reported_total_cost_usd": 0.0123, + } + assert receipt["replay"]["extractor"]["name"] == "claude-cli:sonnet" + manifest = json.loads(Path(receipt["outputs"]["manifest"]).read_text(encoding="utf-8")) + assert manifest["extractor"]["name"] == "claude-cli:sonnet" + + +def test_claude_retry_budget_is_reduced_by_reported_first_attempt_cost( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifact, context = _inputs(tmp_path) + args = _args(tmp_path, artifact, context) + args.provider = "claude-cli" + args.model = "sonnet" + outputs = iter(((_model_output(bad_reference=True), 0.04), (_model_output(), 0.03))) + budgets: list[float] = [] + + def fake_claude(_prompt: str, *, model: str, timeout_seconds: int, max_budget_usd: float): + assert model == "sonnet" + assert timeout_seconds == 30 + budgets.append(max_budget_usd) + output, cost = next(outputs) + return output, {"total_cost_usd": cost} + + monkeypatch.setattr(preparer, "call_claude_cli", fake_claude) + + receipt = preparer.prepare(args) + + assert budgets == pytest.approx([0.10, 0.06]) + assert receipt["extraction"]["cost_control"]["reported_total_cost_usd"] == pytest.approx(0.07) + + +def test_provider_errors_redact_secret_and_private_detail(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + key_file = tmp_path / "key" + key_file.write_text("fixture-secret", encoding="utf-8") + error = urllib.error.HTTPError( + preparer.DEFAULT_OPENROUTER_URL, + 401, + "unauthorized", + {}, + io.BytesIO(b"fixture-secret and private source quote"), + ) + + def fail_urlopen(*_args, **_kwargs): + raise error + + monkeypatch.setattr(preparer.urllib.request, "urlopen", fail_urlopen) + + with pytest.raises(preparer.PreparationError) as caught: + preparer.call_openrouter( + "private source quote", + model="fixture/model", + key_file=key_file, + url=preparer.DEFAULT_OPENROUTER_URL, + timeout_seconds=1, + ) + + rendered = str(caught.value) + assert "fixture-secret" not in rendered + assert "private source quote" not in rendered + assert "provider detail redacted" in rendered