feat(rag): lookup_records tool + strict masking-policy enforcement in agent

Agent can look up actual records via the Command Center, but masked column values
are redacted server-side (never reach the model). System prompt forbids revealing,
guessing or reconstructing masked values and explains how to disable masking.
This commit is contained in:
mo
2026-06-27 11:36:26 +02:00
parent 025abbc3fb
commit de702da759
+26 -1
View File
@@ -41,6 +41,7 @@ OPENMETADATA_URL = os.getenv("OPENMETADATA_URL", "http://10.0.21.47:8585").rstri
OPENMETADATA_TOKEN = os.getenv("OPENMETADATA_TOKEN", "")
COMMAND_CENTER_URL = os.getenv("COMMAND_CENTER_URL", "http://api:3201").rstrip("/")
CATALOG_COLLECTION = os.getenv("CATALOG_COLLECTION", "catalog")
MASK_TOKEN_HINT = "🔒 MASKED (masking policy ON)"
# Use placeholder images in markdown — embedded base64 destroys RAG quality.
DOCLING_IMAGE_MODE = os.getenv("DOCLING_IMAGE_MODE", "placeholder")
@@ -733,6 +734,20 @@ async def _tool_list_movements(client: httpx.AsyncClient, args: dict) -> str:
return json.dumps((await _cc_get(client, "/api/movements")).get("movements", []))[:3000]
async def _tool_lookup_records(client: httpx.AsyncClient, args: dict) -> str:
"""Look up actual records in a source dataset. Masked columns come back redacted —
the real values are enforced server-side and never reach the model."""
key = str(args.get("key") or args.get("dataset") or "").strip()
if not key:
return "error: 'key' is required (one of: postgres, mysql, mongodb, curated)"
body = {"key": key, "search": args.get("search"), "limit": int(args.get("limit", 5) or 5)}
try:
res = await _cc_post(client, "/api/pii/lookup", body)
except Exception as exc: # noqa: BLE001
return f"lookup failed: {exc}"
return json.dumps(res)[:3500]
async def _tool_propose_movement(client: httpx.AsyncClient, args: dict) -> str:
"""WRITE action — never executes directly; files an approval request (Mo & Bart gate)."""
mid = str(args.get("movement_id", "")).strip()
@@ -761,6 +776,7 @@ _TOOLS: dict[str, Any] = {
"get_pii": _tool_get_pii,
"get_cdc": _tool_get_cdc,
"list_movements": _tool_list_movements,
"lookup_records": _tool_lookup_records,
"propose_movement": _tool_propose_movement,
}
@@ -773,10 +789,19 @@ _AGENT_SYSTEM = (
' To finish: {"final": true}\n\n'
"Available tools:\n"
' - search_catalog{"query": str} → tables, columns, PII tags, lineage from the catalog\n'
' - get_pii{} → live PII columns per dataset (masked/unmasked)\n'
' - get_pii{} → live PII columns per dataset (masked/unmasked status)\n'
' - get_cdc{} → live CDC change volume per source\n'
' - list_movements{} → available pipelines/movements with their ids\n'
' - lookup_records{"key": str, "search": str} → actual records in a source dataset\n'
" (key is one of: postgres, mysql, mongodb, curated; search matches the name column)\n"
' - propose_movement{"movement_id": str, "reason": str} → WRITE: files an approval (human-gated)\n\n'
"DATA MASKING POLICY — STRICT:\n"
"Some columns are masked by the operator's policy. In any tool result, a masked value appears "
f'literally as "{MASK_TOKEN_HINT}". You must NEVER reveal, guess, infer, reconstruct, or work '
"around a masked value. If the user asks for a field that comes back masked, reply that the value "
"is withheld for privacy/security because masking is enabled for that column, and tell them an "
"operator can disable masking per column in the Command Center: Data Flow → click the source node "
"→ PII panel. If a value is returned in clear text, you may share it normally.\n\n"
"Gather evidence with read tools before answering. Use propose_movement ONLY when the user "
"clearly asks to run/trigger a pipeline; it never executes directly — it requires human approval. "
'When you have enough information, reply {"final": true} and you will then be asked to write the answer.'