113 lines
4.9 KiB
Python
113 lines
4.9 KiB
Python
"""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)
|