teleo-infrastructure/scripts/prepare_kb_context_from_canonical_dump.py
2026-07-16 12:24:20 +02:00

388 lines
13 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 codecs
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"
STATUS_SCHEMA = "livingip.kbSourceDumpContextStatus.v1"
DEFAULT_DOCKER_IMAGE = "postgres:16-alpine"
DEFAULT_LIMIT = 12
DEFAULT_COPY_ENCODING = "utf-8"
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 _copy_encoding(value: str) -> str:
try:
return codecs.lookup(value).name
except LookupError as exc:
raise ContextError("COPY encoding is not supported") from exc
def _decode_copy_field(value: str, *, encoding: str = DEFAULT_COPY_ENCODING) -> str | None:
if value == r"\N":
return None
encoding = _copy_encoding(encoding)
result = bytearray()
index = 0
escapes = {"b": 8, "f": 12, "n": 10, "r": 13, "t": 9, "v": 11, "\\": 92}
def append_text(text: str) -> None:
try:
result.extend(text.encode(encoding, errors="strict"))
except UnicodeEncodeError as exc:
raise ContextError(f"COPY field cannot be encoded as {encoding}") from exc
while index < len(value):
char = value[index]
if char != "\\":
append_text(char)
index += 1
continue
if index + 1 >= len(value):
raise ContextError("COPY field contains a dangling backslash escape")
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 is None:
raise ContextError("COPY field contains a malformed hexadecimal byte escape")
result.append(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 is None:
raise ContextError("COPY field contains a malformed octal byte escape")
byte_value = int(match.group(0), 8)
if byte_value > 0xFF:
raise ContextError("COPY field contains an out-of-range octal byte escape")
result.append(byte_value)
index += len(match.group(0))
continue
append_text(escaped)
index += 1
if 0 in result:
raise ContextError("COPY field contains a zero byte")
try:
return result.decode(encoding, errors="strict")
except UnicodeDecodeError as exc:
raise ContextError(f"COPY byte escapes do not form valid {encoding} text") from exc
def parse_copy_table(
sql: str, table: str, *, encoding: str = DEFAULT_COPY_ENCODING
) -> list[dict[str, str | None]]:
encoding = _copy_encoding(encoding)
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, encoding=encoding) 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, *, encoding: str = DEFAULT_COPY_ENCODING
) -> 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",
]
encoding = _copy_encoding(encoding)
try:
completed = subprocess.run(
command,
capture_output=True,
text=True,
encoding=encoding,
errors="strict",
timeout=180,
check=False,
)
except (OSError, subprocess.TimeoutExpired, UnicodeError) 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,
copy_encoding: str = DEFAULT_COPY_ENCODING,
) -> 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
copy_encoding = _copy_encoding(copy_encoding)
tables = {
table: parse_copy_table(
restore_table(dump_path, table, docker_image, encoding=copy_encoding),
table,
encoding=copy_encoding,
)
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,
"copy_encoding": copy_encoding,
"copy_decode_errors": "strict",
},
}
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)
parser.add_argument("--copy-encoding", default=DEFAULT_COPY_ENCODING)
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,
copy_encoding=args.copy_encoding,
)
_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):
print(
json.dumps(
{
"schema": STATUS_SCHEMA,
"status": "rejected",
"database_write_performed": False,
"model_call_performed": False,
"private_context_written": False,
"private_receipt_written": False,
"reason": "context_preparation_failed",
},
sort_keys=True,
)
)
return 2
print(
json.dumps(
{
"schema": STATUS_SCHEMA,
"status": "prepared",
"database_write_performed": False,
"model_call_performed": False,
"private_context_written": True,
"private_receipt_written": True,
},
indent=2,
sort_keys=True,
)
)
return 0
if __name__ == "__main__":
raise SystemExit(main())