teleo-infrastructure/scripts/prepare_kb_context_from_canonical_dump.py
2026-07-15 23:03:24 +02:00

320 lines
11 KiB
Python

#!/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())