diff --git a/rag-api/main.py b/rag-api/main.py index 91ebb38..3a51837 100644 --- a/rag-api/main.py +++ b/rag-api/main.py @@ -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.'