bastion agent 0.1.14

2026-09-26 18:31:40 +00:00
commit 3f1d48bf9b
65 changed files with 15548 additions and 0 deletions
@@ -0,0 +1,112 @@
+"""LLM span extraction — the recall pass of the sterilizer.
+
+The model is a sensor, never an actor: it proposes verbatim substrings
+plus a type from the closed vocabulary, grammar-constrained via
+llama.cpp's json_schema response_format (off-schema output is
+impossible). It never rewrites text — pipeline.py grounds every proposed
+span against the input (verbatim occurrence or it is dropped and
+counted) and performs all replacement itself, so the model cannot rewrite
+the text whatever it outputs.
+
+Talks OpenAI-style chat completions to a llama.cpp-compatible server at
+LLM_URL, with thinking disabled per request via chat_template_kwargs.
+Stdlib urllib only.
+"""
+
+import json
+import urllib.request
+from typing import List, Tuple
+
+from .detect import PII_TYPES
+
+Findings = List[Tuple[str, str]]
+
+_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "findings": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "text": {"type": "string"},
+                    "type": {"type": "string", "enum": PII_TYPES},
+                },
+                "required": ["text", "type"],
+            },
+        }
+    },
+    "required": ["findings"],
+}
+
+# The exact vocabulary MUST appear in the prompts. The grammar constrains
+# type to this enum, but grammar is a backstop, not steering: a model that
+# free-texts "School" gets prefix-coerced to "SSN" (seen with a real model).
+# test_pii guards prompt/enum drift.
+_TYPE_LIST = ", ".join(PII_TYPES)
+
+EXTRACT_SYSTEM = f"""You are a PII detector for a privacy gateway. List every substring of the user's text that could identify a person, household, account, or device.
+
+Label each finding with exactly one of these types (use the label verbatim):
+{_TYPE_LIST}
+
+Find: personal names; email addresses; phone numbers; street or postal addresses; government/tax/license/passport ID numbers; credit card and bank account numbers; dates of birth; usernames and social handles; passwords, API keys, and other secrets; medical facts tied to a person; employers, schools, or organizations that identify someone; specific places that identify someone (their home, school, workplace); IP addresses; MAC addresses and device serial numbers.
+
+Rules:
+- Copy each span VERBATIM, character for character, exactly as it appears in the text. Never paraphrase, correct, or reformat.
+- Report the minimal identifying span, without surrounding punctuation or filler words.
+- If the same value appears multiple times, report it once.
+- Do not report placeholder tokens that look like [PII_TYPE_N] — those are already sterilized.
+- Generic words, product names, and public facts are not PII. When a span does not identify anyone, leave it out.
+- If there is no PII, return an empty findings list."""
+
+VERIFY_SYSTEM = f"""You are auditing a sterilized document for residual PII. Placeholder tokens like [PII_EMAIL_1] are already sterilized — ignore them. Report ONLY remaining substrings that could still identify a person, household, account, or device, copied verbatim, each labeled with exactly one of these types (use the label verbatim): {_TYPE_LIST}. Be aggressive: a partial address, a first name, a nickname, an unusual job title can all identify someone. If nothing identifying remains, return an empty findings list."""
+
+
+class RecallError(RuntimeError):
+    """The recall pass could not run. Every backend raises this (or a
+    subclass); the service fails closed on it. Defined here so the older
+    LlmError keeps its name while sharing the contract."""
+
+
+class LlmError(RecallError):
+    pass
+
+
+def _chat(llm_url: str, system: str, text: str, timeout: float) -> Findings:
+    body = json.dumps({
+        "messages": [
+            {"role": "system", "content": system},
+            {"role": "user", "content": text},
+        ],
+        "temperature": 0,
+        "chat_template_kwargs": {"enable_thinking": False},
+        "response_format": {
+            "type": "json_schema",
+            "json_schema": {"name": "pii_findings", "schema": _SCHEMA},
+        },
+    }).encode()
+    req = urllib.request.Request(
+        llm_url.rstrip("/") + "/v1/chat/completions",
+        data=body,
+        headers={"Content-Type": "application/json"},
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            payload = json.load(resp)
+    except OSError as e:
+        raise LlmError(f"llm unreachable: {e}") from e
+    try:
+        content = payload["choices"][0]["message"]["content"]
+        findings = json.loads(content)["findings"]
+        return [(f["text"], f["type"]) for f in findings]
+    except (KeyError, IndexError, TypeError, ValueError) as e:
+        raise LlmError(f"malformed llm response: {e}") from e
+
+
+def extract(text: str, llm_url: str, timeout: float = 300.0) -> Findings:
+    return _chat(llm_url, EXTRACT_SYSTEM, text, timeout)
+
+
+def verify(text: str, llm_url: str, timeout: float = 300.0) -> Findings:
+    return _chat(llm_url, VERIFY_SYSTEM, text, timeout)