445 lines
17 KiB
Python
445 lines
17 KiB
Python
"""bastion agent `catalog` stage — the curated dataset.
|
|
|
|
The catalog is what a host reasons over and what `seal` ships. One card
|
|
per file, one per directory, small and structured, built from:
|
|
|
|
walk.jsonl metadata (path, size, mtime, mime, hash)
|
|
extract.jsonl whether text came out, and how
|
|
sanitize.jsonl entity counts (token types minted per file)
|
|
sanitize/<sha>.txt the SANITIZED text — the only text this stage reads
|
|
chunks.jsonl chunk count per document (optional)
|
|
|
|
Outputs:
|
|
<out>/catalog.jsonl file cards, sorted by (root, path)
|
|
<out>/catalog.dirs.jsonl directory cards, sorted by (root, path)
|
|
<out>/catalog.summary.json
|
|
|
|
Every text-derived field (title, head, keywords, summary) comes from the
|
|
sanitized store, so the catalog carries PII tokens, never values. File
|
|
and directory NAMES are plaintext by nature; they go through the rules
|
|
pass and any precision anchor in a name (an email, a card number) is
|
|
masked as [PII_TYPE] in the card. Contextual PII in a filename
|
|
("Marisol_medical.pdf") is not caught by rules — the honest limit of the
|
|
sanitize stage applies here too, and the doc says so.
|
|
|
|
The optional summarizer (manifest `summarize`) is a local chat model — a
|
|
sensor: it fills the `summary` string and nothing else, is reached via
|
|
SUMMARIZE_URL (never the manifest), sees only sanitized text, and is
|
|
cached per (document, model, version). The first failure stops
|
|
summarizing (cards still get written, `summary_error` set, exit 1).
|
|
"""
|
|
|
|
import datetime as _dt
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import re
|
|
import tempfile
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
from .manifest import signing_digest
|
|
from .sanitize import store_path as sanitized_store_path
|
|
|
|
|
|
def _rules_detect():
|
|
"""pii.detect.detect — a sibling package here and in the .pyz (see
|
|
sanitize.pii_modules for why the old `tools.pii` branch is gone)."""
|
|
from pii.detect import detect
|
|
return detect
|
|
|
|
CATALOG_VERSION = 1
|
|
HEAD_CHARS = 300
|
|
TITLE_CHARS = 120
|
|
KEYWORDS_PER_CARD = 8
|
|
KEYWORDS_PER_DIR = 12
|
|
SUMMARY_INPUT_CHARS = 12000
|
|
SUMMARY_MAX_CHARS = 600
|
|
DEFAULT_SUMMARIZE_URL = "http://127.0.0.1:11434"
|
|
SUMMARIZE_TIMEOUT = 300.0
|
|
|
|
WORD_RE = re.compile(r"[^\W\d_]{3,}", re.UNICODE)
|
|
TOKEN_MARK_RE = re.compile(r"\[PII_[A-Z_]+_\d+\]")
|
|
STOPWORDS = set("""
|
|
the and for that with this from are was were will have has had not but you your our their they them
|
|
its his her she him who what which when where why how all any can could would should may might must
|
|
into onto over under than then there here also just more most much many some such very only own same
|
|
other another each both few between about after before during while above below off out too again
|
|
further once here per via use used using does did done being been each same so if or as at by on in
|
|
of to is it be an a we us me my i no yes do go get got let new old one two three first last next see
|
|
set run make made take took like well also back down up still even ever never always often sometimes
|
|
""".split())
|
|
|
|
|
|
class CatalogError(Exception):
|
|
pass
|
|
|
|
|
|
def _atomic_write_text(path, text):
|
|
d = os.path.dirname(path)
|
|
os.makedirs(d, exist_ok=True)
|
|
fd, tmp = tempfile.mkstemp(prefix=".tmp-", dir=d)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as fh:
|
|
fh.write(text)
|
|
os.replace(tmp, path)
|
|
except BaseException:
|
|
try:
|
|
os.unlink(tmp)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def _read_jsonl(out_dir, name, required=True):
|
|
path = os.path.join(out_dir, name)
|
|
if not os.path.exists(path):
|
|
if required:
|
|
raise FileNotFoundError("no %s at %s — run the earlier stages first" % (name, path))
|
|
return []
|
|
with open(path, encoding="utf-8") as fh:
|
|
return [json.loads(l) for l in fh]
|
|
|
|
|
|
# ── text facts ──────────────────────────────────────────────────────────────
|
|
|
|
def title_of(text):
|
|
for line in text.splitlines():
|
|
s = line.strip().lstrip("#*-= \t").strip()
|
|
if len(s) >= 3:
|
|
return s[:TITLE_CHARS]
|
|
return ""
|
|
|
|
|
|
def head_of(text):
|
|
s = re.sub(r"\s+", " ", text).strip()
|
|
return s[:HEAD_CHARS]
|
|
|
|
|
|
def term_counts(text):
|
|
counts = {}
|
|
text = TOKEN_MARK_RE.sub(" ", text) # [PII_*_n] placeholders are not keywords
|
|
for w in WORD_RE.findall(text.lower()):
|
|
if w in STOPWORDS:
|
|
continue
|
|
counts[w] = counts.get(w, 0) + 1
|
|
return counts
|
|
|
|
|
|
def keywords(counts, df, ndocs, k):
|
|
"""Top-k terms by TF-IDF; ties broken alphabetically for determinism."""
|
|
scored = []
|
|
for w, tf in counts.items():
|
|
idf = math.log((1 + ndocs) / (1 + df.get(w, 0))) + 1.0
|
|
scored.append((-(tf * idf), w))
|
|
scored.sort()
|
|
return [w for _, w in scored[:k]]
|
|
|
|
|
|
def mask_name(name, detect):
|
|
"""Rules pass over a file/dir name: precision anchors -> [PII_TYPE]."""
|
|
findings = detect(name)
|
|
if not findings:
|
|
return name
|
|
out = name
|
|
for f in sorted(findings, key=lambda f: -len(f.value)):
|
|
out = out.replace(f.value, "[PII_%s]" % f.type)
|
|
return out
|
|
|
|
|
|
def mask_path(path, detect):
|
|
"""Every component of a relative path through mask_name — a parent
|
|
directory can carry an address as easily as a filename, and the file
|
|
card's path is what seal ships. Returns (masked
|
|
path, number of components changed)."""
|
|
if not path:
|
|
return path, 0
|
|
out, changed = [], 0
|
|
for seg in path.split("/"):
|
|
m = mask_name(seg, detect) if seg else seg
|
|
if m != seg:
|
|
changed += 1
|
|
out.append(m)
|
|
return "/".join(out), changed
|
|
|
|
|
|
# ── the summarizer (a sensor) ───────────────────────────────────────────────
|
|
|
|
SUMMARY_PROMPT = (
|
|
"Summarize the following document in at most three plain sentences for a catalog card: "
|
|
"what it is, what it covers, and who or what it concerns. Placeholders like [PII_PERSON_1] "
|
|
"stand for redacted personal data — keep them as they are, never guess what they hide. "
|
|
"Reply with the summary only.\n\n"
|
|
)
|
|
DIR_SUMMARY_PROMPT = (
|
|
"These are catalog entries for the files in one folder. In at most three plain sentences "
|
|
"say what the folder is for and what kinds of material it holds. Placeholders like "
|
|
"[PII_PERSON_1] are redacted personal data — keep them as they are. Reply with the "
|
|
"summary only.\n\n"
|
|
)
|
|
|
|
|
|
class Summarizer:
|
|
def __init__(self, url, model, api=None, timeout=SUMMARIZE_TIMEOUT):
|
|
self.url = url.rstrip("/")
|
|
self.model = model
|
|
self.api = api
|
|
self.timeout = timeout
|
|
|
|
def _post(self, path, body):
|
|
req = urllib.request.Request(
|
|
self.url + path, data=json.dumps(body).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"})
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
|
return json.load(resp), None
|
|
except urllib.error.HTTPError as e:
|
|
return None, "HTTP %d" % e.code
|
|
except (OSError, ValueError) as e:
|
|
return None, str(e)
|
|
|
|
def _ollama(self, prompt):
|
|
data, err = self._post("/api/chat", {
|
|
"model": self.model, "stream": False, "options": {"temperature": 0},
|
|
"messages": [{"role": "user", "content": prompt}]})
|
|
if err:
|
|
return None, err
|
|
try:
|
|
return data["message"]["content"], None
|
|
except (KeyError, TypeError) as e:
|
|
return None, "malformed ollama response: %s" % e
|
|
|
|
def _openai(self, prompt):
|
|
data, err = self._post("/v1/chat/completions", {
|
|
"model": self.model, "temperature": 0,
|
|
"messages": [{"role": "user", "content": prompt}]})
|
|
if err:
|
|
return None, err
|
|
try:
|
|
return data["choices"][0]["message"]["content"], None
|
|
except (KeyError, IndexError, TypeError) as e:
|
|
return None, "malformed chat response: %s" % e
|
|
|
|
def complete(self, prompt):
|
|
if self.api == "ollama":
|
|
text, err = self._ollama(prompt)
|
|
elif self.api == "openai":
|
|
text, err = self._openai(prompt)
|
|
else:
|
|
text, err = self._ollama(prompt)
|
|
if text is not None:
|
|
self.api = "ollama"
|
|
else:
|
|
text, err2 = self._openai(prompt)
|
|
if text is not None:
|
|
self.api = "openai"
|
|
else:
|
|
err = "ollama: %s; openai: %s" % (err, err2)
|
|
if text is None:
|
|
raise CatalogError("summarizer %s (model %s): %s" % (self.url, self.model, err))
|
|
return clean_summary(text)
|
|
|
|
def describe(self):
|
|
return "%s@%s(%s)" % (self.model, self.url, self.api or "auto")
|
|
|
|
|
|
def clean_summary(text):
|
|
text = re.sub(r"<think>.*?</think>", "", text, flags=re.S) # thinking models
|
|
text = re.sub(r"\s+", " ", text).strip().strip('"')
|
|
return text[:SUMMARY_MAX_CHARS]
|
|
|
|
|
|
def build_summarizer(manifest, env):
|
|
if manifest.summarize is None:
|
|
return None
|
|
url = env.get("SUMMARIZE_URL", DEFAULT_SUMMARIZE_URL).strip() or DEFAULT_SUMMARIZE_URL
|
|
api = env.get("SUMMARIZE_API", "").strip().lower() or None
|
|
if api not in (None, "ollama", "openai"):
|
|
raise CatalogError("SUMMARIZE_API=%r is not ollama|openai" % api)
|
|
return Summarizer(url, manifest.summarize.model, api=api,
|
|
timeout=float(env.get("SUMMARIZE_TIMEOUT", SUMMARIZE_TIMEOUT)))
|
|
|
|
|
|
def _cache_path(out_dir):
|
|
return os.path.join(out_dir, "catalog.cache.json")
|
|
|
|
|
|
def _load_cache(out_dir):
|
|
try:
|
|
with open(_cache_path(out_dir), encoding="utf-8") as fh:
|
|
return json.load(fh)
|
|
except (OSError, ValueError):
|
|
return {}
|
|
|
|
|
|
# ── the stage ───────────────────────────────────────────────────────────────
|
|
|
|
def catalog(manifest, out_dir, env=None, summarizer=None, now=None):
|
|
env = dict(os.environ) if env is None else env
|
|
summ = summarizer if summarizer is not None else build_summarizer(manifest, env)
|
|
per = set(manifest.summarize.per) if manifest.summarize else set()
|
|
walk_rows = _read_jsonl(out_dir, "walk.jsonl")
|
|
extract_rows = {(r["root"], r["path"]): r for r in _read_jsonl(out_dir, "extract.jsonl")}
|
|
sanitize_rows = {(r["root"], r["path"]): r for r in _read_jsonl(out_dir, "sanitize.jsonl")}
|
|
chunk_counts = {}
|
|
for c in _read_jsonl(out_dir, "chunks.jsonl", required=False):
|
|
chunk_counts[c["doc"]] = chunk_counts.get(c["doc"], 0) + 1
|
|
detect = _rules_detect()
|
|
cache = _load_cache(out_dir)
|
|
stamp = now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
# pass 1: text facts per distinct sanitized document
|
|
texts = {}
|
|
doc_terms = {}
|
|
for key, s in sanitize_rows.items():
|
|
if s.get("status") != "ok":
|
|
continue
|
|
sha = s["sha256"]
|
|
if sha in texts:
|
|
continue
|
|
with open(sanitized_store_path(out_dir, sha), encoding="utf-8") as fh:
|
|
texts[sha] = fh.read()
|
|
doc_terms[sha] = term_counts(texts[sha])
|
|
df = {}
|
|
for counts in doc_terms.values():
|
|
for w in counts:
|
|
df[w] = df.get(w, 0) + 1
|
|
ndocs = len(doc_terms)
|
|
|
|
counts = {"files": 0, "with_text": 0, "dirs": 0, "summaries": 0, "summaries_reused": 0,
|
|
"summary_errors": 0, "names_masked": 0}
|
|
stopped = None
|
|
summary_key = lambda sha, kind: "%s|%s|%s|v%d" % (
|
|
kind, sha, manifest.summarize.model if manifest.summarize else "", CATALOG_VERSION)
|
|
|
|
def summarize(kind, sha, prompt_body):
|
|
nonlocal stopped
|
|
if summ is None or stopped is not None:
|
|
return None, None
|
|
k = summary_key(sha, kind)
|
|
if k in cache:
|
|
counts["summaries_reused"] += 1
|
|
return cache[k], None
|
|
try:
|
|
text = summ.complete((SUMMARY_PROMPT if kind == "file" else DIR_SUMMARY_PROMPT)
|
|
+ prompt_body[:SUMMARY_INPUT_CHARS])
|
|
except CatalogError as e:
|
|
stopped = str(e)
|
|
counts["summary_errors"] += 1
|
|
return None, stopped
|
|
cache[k] = text
|
|
counts["summaries"] += 1
|
|
return text, None
|
|
|
|
# pass 2: file cards
|
|
cards = []
|
|
dir_acc = {}
|
|
for w in walk_rows:
|
|
if w.get("kind") != "file":
|
|
continue
|
|
key = (w["root"], w["path"])
|
|
counts["files"] += 1
|
|
name = w["path"].rsplit("/", 1)[-1]
|
|
shown_path, changed = mask_path(w["path"], detect)
|
|
masked = shown_path.rsplit("/", 1)[-1]
|
|
counts["names_masked"] += changed
|
|
e = extract_rows.get(key, {})
|
|
s = sanitize_rows.get(key, {})
|
|
card = {
|
|
"root": w["root"], "path": shown_path, "name": masked, "kind": "file",
|
|
"size": w.get("size"), "mtime": w.get("mtime"), "mime": w.get("mime"),
|
|
"sha256": w.get("sha256"), "v": CATALOG_VERSION,
|
|
"text": e.get("status") if e else "not_extracted",
|
|
"extractor": e.get("extractor"),
|
|
"sanitized": s.get("status") if s else None,
|
|
}
|
|
sha = s.get("sha256") if s.get("status") == "ok" else None
|
|
if sha and sha in texts:
|
|
text = texts[sha]
|
|
counts["with_text"] += 1
|
|
card.update({
|
|
"title": title_of(text),
|
|
"head": head_of(text),
|
|
"chars": len(text),
|
|
"keywords": keywords(doc_terms[sha], df, ndocs, KEYWORDS_PER_CARD),
|
|
"entities": s.get("types", {}),
|
|
"chunks": chunk_counts.get(sha, 0),
|
|
})
|
|
if "file" in per:
|
|
summary, err = summarize("file", sha, text)
|
|
card["summary"] = summary
|
|
if err:
|
|
card["summary_error"] = err
|
|
elif "file" in per:
|
|
card["summary"] = None # nothing to summarize
|
|
cards.append(card)
|
|
# directory aggregation: every ancestor, root included as ""
|
|
parts = w["path"].split("/")[:-1]
|
|
for depth in range(len(parts) + 1):
|
|
dpath = "/".join(parts[:depth])
|
|
acc = dir_acc.setdefault((w["root"], dpath), {
|
|
"files": 0, "bytes": 0, "mimes": {}, "latest": "", "terms": {}, "entities": {},
|
|
"children": [], "with_text": 0})
|
|
acc["files"] += 1
|
|
acc["bytes"] += w.get("size") or 0
|
|
m = w.get("mime") or "?"
|
|
acc["mimes"][m] = acc["mimes"].get(m, 0) + 1
|
|
if (w.get("mtime") or "") > acc["latest"]:
|
|
acc["latest"] = w["mtime"]
|
|
if sha and sha in doc_terms:
|
|
acc["with_text"] += 1
|
|
for t, n in doc_terms[sha].items():
|
|
acc["terms"][t] = acc["terms"].get(t, 0) + n
|
|
for t, n in (s.get("types") or {}).items():
|
|
acc["entities"][t] = acc["entities"].get(t, 0) + n
|
|
if depth == len(parts):
|
|
acc["children"].append(card)
|
|
cards.sort(key=lambda c: (c["root"], c["path"]))
|
|
|
|
# pass 3: directory cards
|
|
dcards = []
|
|
for (root, dpath), acc in dir_acc.items():
|
|
shown, _ = mask_path(dpath, detect)
|
|
masked = shown.rsplit("/", 1)[-1] if shown else ""
|
|
mimes = sorted(acc["mimes"].items(), key=lambda kv: (-kv[1], kv[0]))[:6]
|
|
card = {
|
|
"root": root, "path": shown, "name": masked, "kind": "dir", "v": CATALOG_VERSION,
|
|
"files": acc["files"], "with_text": acc["with_text"], "bytes": acc["bytes"],
|
|
"latest": acc["latest"], "types": dict(mimes),
|
|
"keywords": keywords(acc["terms"], df, ndocs, KEYWORDS_PER_DIR),
|
|
"entities": acc["entities"],
|
|
}
|
|
if "directory" in per:
|
|
body = "\n".join(
|
|
"- %s: %s" % (c["name"], (c.get("summary") or c.get("head") or "")[:200])
|
|
for c in sorted(acc["children"], key=lambda c: c["path"])[:60])
|
|
sha = hashlib.sha256(
|
|
("%s|%s|" % (root, dpath)).encode("utf-8") + body.encode("utf-8")).hexdigest()
|
|
summary, err = summarize("dir", sha, body)
|
|
card["summary"] = summary
|
|
if err:
|
|
card["summary_error"] = err
|
|
dcards.append(card)
|
|
dcards.sort(key=lambda c: (c["root"], c["path"]))
|
|
counts["dirs"] = len(dcards)
|
|
|
|
_atomic_write_text(os.path.join(out_dir, "catalog.jsonl"), "".join(
|
|
json.dumps(c, ensure_ascii=False, sort_keys=True) + "\n" for c in cards))
|
|
_atomic_write_text(os.path.join(out_dir, "catalog.dirs.jsonl"), "".join(
|
|
json.dumps(c, ensure_ascii=False, sort_keys=True) + "\n" for c in dcards))
|
|
_atomic_write_text(_cache_path(out_dir), json.dumps(cache, ensure_ascii=False, sort_keys=True))
|
|
summary = {
|
|
"freeagent": 1, "stage": "catalog", "catalog_version": CATALOG_VERSION,
|
|
"manifest_id": manifest.id,
|
|
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
|
"cataloged_at": stamp,
|
|
"summarizer": summ.describe() if summ else None,
|
|
"summarize_per": sorted(per),
|
|
"counts": counts, "stopped": stopped,
|
|
"bytes_catalog": sum(len(json.dumps(c, ensure_ascii=False)) + 1 for c in cards + dcards),
|
|
}
|
|
_atomic_write_text(os.path.join(out_dir, "catalog.summary.json"),
|
|
json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n")
|
|
return summary
|