privatae/bastion-agent
2
0
Fork 0
Code Issues Actions Packages Projects Releases Wiki Activity

bastion agent 0.1.14

Browse Source
This commit is contained in:
Privatae
2026-09-26 18:31:40 +00:00
commit 3f1d48bf9b
65 changed files with 15548 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pii/pipeline.py
+217
View File
@@ -0,0 +1,217 @@
"""Sterilization pipeline — deterministic control flow, LLM as sensor.
Stages (docs/pii-sterilization.md):
1. rules pass detect.detect() — precision anchors, model-free
2. LLM pass llmclient.extract() — recall for contextual PII
3. grounding every LLM span must occur verbatim in the input;
ungrounded spans are dropped and counted, never trusted
4. replacement script-side, longest value first, every occurrence
(plain substring, inside longer words too); overlapping
matches merge into one token; never inside an existing
token; placeholder-shaped text in the input is escaped
first (all of it for raw documents, only what the vault
does not hold otherwise); tokens minted by the vault
(stable per scope)
5. residual loop rules re-run on the output until clean (bounded);
any residual is tokenized, not just reported
6. LLM verify optional second model pass over the sterilized text;
grounded hits go through replacement again
The model proposes; this module disposes. Every mutation of the text
happens in replace_values(), which guarantees non-PII bytes are
byte-identical to the input.
"""
import re
from . import detect
TOKEN_RE = re.compile(r"\[PII_([A-Z_]+)_(\d+)\]")
# a token's name without its brackets; as a value it could only ever come
# from the model reading an already-sterilized text
BARE_TOKEN_RE = re.compile(r"PII_[A-Z_]+_\d+")
_MAX_RESIDUAL_ROUNDS = 3
def _usable(value: str, typ: str) -> bool:
if typ not in detect.PII_TYPES:
return False
if len(value) < 2:
return False
if TOKEN_RE.search(value) or BARE_TOKEN_RE.search(value):
return False
return True
def ground(proposed, text):
"""Split LLM proposals into (grounded, dropped_count)."""
grounded, dropped = [], 0
for value, typ in proposed:
if value and value in text and _usable(value, typ):
grounded.append((value, typ))
else:
dropped += 1
return grounded, dropped
def _occurrences(value, text, blocked):
"""Every occurrence of value, overlapping ones included ("aba" twice in
"ababa"), minus any that overlap a blocked span. Overlapping hits are
merged by the caller, so no fragment of a self-overlapping value is
left between two tokens."""
spans = []
i = text.find(value)
while i >= 0:
end = i + len(value)
if not any(i < be and end > bs for bs, be in blocked):
spans.append((i, end))
i = text.find(value, i + 1)
return spans
def replace_values(text, values, scope, vault, block_tokens=True):
"""Replace every occurrence of each value with its vault token.
Plain substring matching: a value is replaced inside longer words too
("Garcia" in "Garcias", "Marisol" in "Marisol_resume.pdf"), because a
missed inflection is a leak and an over-eager match only costs
readability. Longest value first, so "John Smith" is consumed before
"John" can split it. No match may overlap a token already in the text,
so a value like "PERSON" cannot break a placeholder apart.
Two matches that partly overlap (neither contains the other) become one
token over their union: tokenizing one would leave the other's
remainder in the clear. It takes the longer value's type.
Every byte outside a replaced span is copied unchanged.
block_tokens=False is only for escaping placeholder-shaped text in the
input (see sterilize), where the literal itself is the value to replace.
"""
types, order = {}, {}
for value, typ in values:
if value and value not in types:
types[value] = typ
order[value] = len(order)
blocked = [m.span() for m in TOKEN_RE.finditer(text)] if block_tokens else []
matches = []
for value in sorted(types, key=lambda v: (-len(v), order[v])):
matches.extend((s, e, value) for s, e in _occurrences(value, text, blocked))
# union overlapping matches; the longest (then earliest-given) value names
# the group, and a group wider than it is a merged value of its own
groups = []
for s, e, value in sorted(matches):
if groups and s < groups[-1][1]:
g = groups[-1]
g[1] = max(g[1], e)
g[2].append(value)
else:
groups.append([s, e, [value]])
spans = []
for s, e, members in groups:
lead = min(members, key=lambda v: (-len(v), order[v]))
value = text[s:e]
if value not in types:
types[value] = types[lead]
order[value] = order[lead]
spans.append((s, e, value))
tokens, applied = {}, []
for value in sorted({v for _, _, v in spans}, key=lambda v: (-len(v), order[v])):
tokens[value] = vault.token_for(scope, types[value], value)
applied.append({"token": tokens[value], "type": types[value]})
parts, pos = [], 0
for s, e, value in spans:
parts.append(text[pos:s])
parts.append(tokens[value])
pos = e
parts.append(text[pos:])
return "".join(parts), applied
def sterilize(text, scope, vault, llm_extract, llm_verify=None, escape_held_literals=False):
"""Returns (sterilized_text, result_dict).
escape_held_literals: escape every placeholder-shaped string in the
input, even one the vault holds for this scope. For raw documents,
which never contain this service's own output: a document that quotes
"[PII_EMAIL_1]" must not rehydrate to another document's email. Off for
callers that send already-sterilized text back (a stored conversation on
every save), where a held placeholder is that token and must stay it.
"""
rule_values = [(f.value, f.type) for f in detect.detect(text)]
llm_values, ungrounded = ground(llm_extract(text), text)
# Placeholder-shaped text that is not one of this scope's tokens is
# escaped as a value of its own, so rehydrate gives the literal back
# instead of leaving it to collide with a token minted later. Unless the
# caller asks for every literal to be escaped, one the vault holds passes
# through untouched: already-sterilized text is sterilized again (a
# stored conversation on every save), and its tokens must stay the same
# tokens. Each distinct literal is looked up once.
shapes = {m.group(0) for m in TOKEN_RE.finditer(text)}
if not escape_held_literals:
shapes = {lit for lit in shapes if vault.value_for(scope, lit) is None}
literals = [(lit, "OTHER") for lit in sorted(shapes)]
out, escaped = replace_values(text, literals, scope, vault, block_tokens=False)
# Rules first: replace_values keeps the first type seen for a value,
# and on a duplicate the rule's type is the higher-precision one.
out, applied = replace_values(out, rule_values + llm_values, scope, vault)
# every finding a later round tokenizes is counted, not only the last
# round's: the round that comes back clean must not reset the total
residual = 0
for _ in range(_MAX_RESIDUAL_ROUNDS):
leftover = [(f.value, f.type) for f in detect.detect(out)]
if not leftover:
break
residual += len(leftover)
out, more = replace_values(out, leftover, scope, vault)
applied.extend(more)
verify_hits = 0
if llm_verify is not None:
vgrounded, _ = ground(llm_verify(out), out)
if vgrounded:
# count what was replaced, not what was proposed: a grounded value
# that only occurs inside an existing token replaces nothing
out, more = replace_values(out, vgrounded, scope, vault)
verify_hits = len(more)
applied.extend(more)
return out, {
"findings": applied,
"stats": {
"rules": len(rule_values),
"llm": len(llm_values),
"llm_ungrounded": ungrounded,
"residual_tokenized": residual,
"verify_tokenized": verify_hits,
"literal_tokens_escaped": len(escaped),
},
}
def rehydrate(text, scope, vault):
"""Replace vault tokens with their original values. Unknown tokens are
left in place and reported — never guessed."""
unknown = []
restored = 0
def sub(m):
nonlocal restored
value = vault.value_for(scope, m.group(0))
if value is None:
unknown.append(m.group(0))
return m.group(0)
restored += 1
return value
out = TOKEN_RE.sub(sub, text)
return out, {"restored": restored, "unknown_tokens": sorted(set(unknown))}