From 025abbc3fbd1a748540095a1980ba8681caaa800 Mon Sep 17 00:00:00 2001 From: mo Date: Sat, 27 Jun 2026 03:28:08 +0200 Subject: [PATCH] feat(rag): data-catalog sync + agentic tool-loop with approval-gated executor - /catalog/sync ingests OpenMetadata tables/columns/PII tags + Command Center PII catalog, data-flow lineage and movements into a 'catalog' Chroma collection. - /agent runs a model-agnostic JSON-action tool loop (search_catalog, get_pii, get_cdc, list_movements, propose_movement) and streams the final answer over SSE. - propose_movement files an approval in the Command Center (human gate) and never executes directly. --- rag-api/main.py | 335 +++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 332 insertions(+), 3 deletions(-) diff --git a/rag-api/main.py b/rag-api/main.py index 2490962..91ebb38 100644 --- a/rag-api/main.py +++ b/rag-api/main.py @@ -14,11 +14,11 @@ from typing import Any import httpx from fastapi import FastAPI, File, Form, UploadFile from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import FileResponse, JSONResponse +from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.vectorstores import Chroma from langchain_core.documents import Document -from langchain_core.messages import HumanMessage, SystemMessage +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage from langchain_openai import ChatOpenAI from langchain_text_splitters import RecursiveCharacterTextSplitter import chromadb @@ -36,6 +36,12 @@ UPLOADS_DIR = DATA_DIR / "uploads" REGISTRY_PATH = DATA_DIR / "document_registry.json" EMBED_MODEL = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2") +# Governance / agent integration +OPENMETADATA_URL = os.getenv("OPENMETADATA_URL", "http://10.0.21.47:8585").rstrip("/") +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") + # Use placeholder images in markdown — embedded base64 destroys RAG quality. DOCLING_IMAGE_MODE = os.getenv("DOCLING_IMAGE_MODE", "placeholder") @@ -539,10 +545,333 @@ async def chat(body: ChatRequest): return {"ok": True, "answer": answer, "sources": sources, "collection": col, "context_chunks": len(docs)} +# ───────────────────────────────────────────────────────────────────────────── +# Data-catalog sync — pull OpenMetadata + Command Center governance into ChromaDB +# so the assistant can answer questions about tables, columns, PII and lineage. +# ───────────────────────────────────────────────────────────────────────────── + +async def _om_get(client: httpx.AsyncClient, path: str, params: dict | None = None) -> dict[str, Any]: + headers = {"Authorization": f"Bearer {OPENMETADATA_TOKEN}"} if OPENMETADATA_TOKEN else {} + r = await client.get(f"{OPENMETADATA_URL}{path}", headers=headers, params=params) + r.raise_for_status() + return r.json() + + +async def _cc_get(client: httpx.AsyncClient, path: str) -> dict[str, Any]: + r = await client.get(f"{COMMAND_CENTER_URL}{path}") + r.raise_for_status() + return r.json() + + +async def _cc_post(client: httpx.AsyncClient, path: str, body: dict) -> dict[str, Any]: + r = await client.post(f"{COMMAND_CENTER_URL}{path}", json=body) + r.raise_for_status() + return r.json() + + +def _table_to_doc(t: dict[str, Any]) -> tuple[str, str]: + """Render an OpenMetadata table entity as a catalog document.""" + fqn = t.get("fullyQualifiedName") or t.get("name", "?") + lines = [f"# Table: {fqn}"] + if t.get("description"): + lines.append(t["description"].strip()) + cols = t.get("columns") or [] + if cols: + lines.append("\nColumns:") + for c in cols: + tags = [tg.get("tagFQN", "") for tg in (c.get("tags") or [])] + pii = " [PII]" if any("PII" in tg for tg in tags) else "" + tag_str = f" — tags: {', '.join(t for t in tags if t)}" if tags else "" + lines.append(f" - {c.get('name')} ({c.get('dataType', '?')}){pii}{tag_str}") + return fqn, "\n".join(lines) + + +async def _sync_catalog() -> dict[str, Any]: + col = _safe_collection(CATALOG_COLLECTION) + docs: list[Document] = [] + counts = {"tables": 0, "pii_datasets": 0, "flows": 0, "movements": 0} + now = datetime.now(timezone.utc).isoformat() + + async with httpx.AsyncClient(timeout=60.0, verify=False) as client: + # 1) OpenMetadata tables (name, columns, types, PII/classification tags) + try: + after: str | None = None + seen = 0 + while seen < 2000: + params = {"fields": "columns,tags,description", "limit": 100} + if after: + params["after"] = after + page = await _om_get(client, "/api/v1/tables", params) + for t in page.get("data", []): + fqn, text = _table_to_doc(t) + docs.append(Document(page_content=text, metadata={ + "source": f"openmetadata:{fqn}", "kind": "table", "fqn": fqn, + "doc_id": f"om_{_safe_collection(fqn)}", "ingested_at": now, + })) + seen += 1 + counts["tables"] = seen + after = (page.get("paging") or {}).get("after") + if not after: + break + except Exception as exc: # noqa: BLE001 + counts["openmetadata_error"] = str(exc)[:200] + + # 2) Command Center PII catalog (masked/unmasked, source per dataset) + try: + pii = await _cc_get(client, "/api/pii") + for ds in pii.get("datasets", []): + cols = ds.get("pii_columns", []) + key = ds.get("key") or ds.get("node_id") or "?" + lines = [f"# PII dataset: {ds.get('label') or key} (table: {ds.get('table', '?')})"] + for c in cols: + status = "masked" if c.get("masked") else "unmasked" + via = c.get("source", "heuristic") + lines.append(f" - {c.get('name')}: {c.get('category', 'PII')} [{status}, via {via}]") + docs.append(Document(page_content="\n".join(lines), metadata={ + "source": f"pii:{key}", "kind": "pii", "ingested_at": now, + "doc_id": f"pii_{_safe_collection(str(key))}", + })) + counts["pii_datasets"] = len(pii.get("datasets", [])) + except Exception as exc: # noqa: BLE001 + counts["pii_error"] = str(exc)[:200] + + # 3) Data-flow graph → lineage / movement edges + try: + flow = await _cc_get(client, "/api/dataflow") + nodes = {n["id"]: n for n in flow.get("nodes", [])} + edge_lines = ["# Data flow / lineage (Command Center)"] + for e in flow.get("edges", []): + src = nodes.get(e.get("from"), {}).get("label", e.get("from")) + dst = nodes.get(e.get("to"), {}).get("label", e.get("to")) + mid = f", movement_id={e['movement_id']}" if e.get("movement_id") else "" + edge_lines.append(f" - {src} → {dst} ({e.get('kind', '')}{mid})") + docs.append(Document(page_content="\n".join(edge_lines), metadata={ + "source": "dataflow:lineage", "kind": "lineage", "ingested_at": now, + "doc_id": "cc_lineage", + })) + counts["flows"] = len(flow.get("edges", [])) + except Exception as exc: # noqa: BLE001 + counts["dataflow_error"] = str(exc)[:200] + + # 4) Movements / pipelines + try: + mv = await _cc_get(client, "/api/movements") + mv_lines = ["# Data movements / pipelines (Command Center)"] + for m in mv.get("movements", []): + mv_lines.append( + f" - id={m.get('id')} · {m.get('label')} ({m.get('kind')}): " + f"{m.get('from', '?')} → {m.get('to', '?')} · agent={m.get('agent', '—')}" + ) + docs.append(Document(page_content="\n".join(mv_lines), metadata={ + "source": "movements", "kind": "movements", "ingested_at": now, + "doc_id": "cc_movements", + })) + counts["movements"] = len(mv.get("movements", [])) + except Exception as exc: # noqa: BLE001 + counts["movements_error"] = str(exc)[:200] + + if not docs: + return {"ok": False, "error": "No catalog data could be fetched", "counts": counts} + + # Rebuild the catalog collection cleanly each sync. + try: + get_chroma_client().delete_collection(col) + except Exception: + pass + get_vectorstore(col).add_documents(docs) + return {"ok": True, "collection": col, "documents": len(docs), "counts": counts, "synced_at": now} + + +@app.post("/catalog/sync") +async def catalog_sync(): + try: + return await _sync_catalog() + except Exception as exc: # noqa: BLE001 + return JSONResponse({"ok": False, "error": str(exc)}, status_code=502) + + +@app.get("/catalog/status") +async def catalog_status(): + col = _safe_collection(CATALOG_COLLECTION) + try: + count = get_chroma_client().get_collection(col).count() + return {"ok": True, "collection": col, "documents": count} + except Exception: + return {"ok": True, "collection": col, "documents": 0} + + +# ───────────────────────────────────────────────────────────────────────────── +# Agentic chat — a model-agnostic JSON-action tool loop over the live platform, +# with an approval-gated executor for write actions and SSE token streaming. +# ───────────────────────────────────────────────────────────────────────────── + +def _catalog_search(query: str, k: int = 6) -> str: + col = _safe_collection(CATALOG_COLLECTION) + try: + vs = get_vectorstore(col) + if get_chroma_client().get_collection(col).count() == 0: + return "Catalog is empty — run /catalog/sync first." + hits = vs.as_retriever(search_kwargs={"k": k}).invoke(query) + return "\n\n---\n\n".join(h.page_content for h in hits) or "No matches." + except Exception as exc: # noqa: BLE001 + return f"catalog search failed: {exc}" + + +async def _tool_search_catalog(client: httpx.AsyncClient, args: dict) -> str: + return _catalog_search(str(args.get("query", ""))) + + +async def _tool_get_pii(client: httpx.AsyncClient, args: dict) -> str: + return json.dumps((await _cc_get(client, "/api/pii")).get("datasets", []))[:3000] + + +async def _tool_get_cdc(client: httpx.AsyncClient, args: dict) -> str: + return json.dumps(await _cc_get(client, "/api/changes/stats"))[:2000] + + +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_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() + reason = str(args.get("reason", "Requested via assistant"))[:300] + if not mid: + return "error: movement_id is required" + body = { + "agent_id": "etl-guardian", + "action": f"Trigger data movement '{mid}'", + "reason": reason, + "action_type": "etl.restart", + "target": mid, + "payload": {"executor": "movement", "movement_id": mid}, + "priority": "normal", + } + res = await _cc_post(client, "/api/approvals", body) + appr = res.get("approval", {}) + return ( + f"Approval filed (id={appr.get('id', '?')}, status={appr.get('status', 'pending')}). " + "The movement runs automatically once Mo or Bart approves it in the Approvals queue." + ) + + +_TOOLS: dict[str, Any] = { + "search_catalog": _tool_search_catalog, + "get_pii": _tool_get_pii, + "get_cdc": _tool_get_cdc, + "list_movements": _tool_list_movements, + "propose_movement": _tool_propose_movement, +} + +_AGENT_SYSTEM = ( + "You are the Dell ATC data-platform assistant. You can inspect a live lakehouse " + "(PostgreSQL, MySQL, MongoDB sources → Debezium CDC → Kafka → Iceberg/Trino, governed " + "by OpenMetadata) and propose data movements.\n\n" + "You work in a strict loop. On every turn reply with EXACTLY ONE single-line JSON object, nothing else:\n" + ' To use a tool: {"action": "", "input": {}}\n' + ' 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_cdc{} → live CDC change volume per source\n' + ' - list_movements{} → available pipelines/movements with their ids\n' + ' - propose_movement{"movement_id": str, "reason": str} → WRITE: files an approval (human-gated)\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.' +) + +_JSON_RE = re.compile(r"\{.*\}", re.DOTALL) + + +def _parse_action(text: str) -> dict[str, Any] | None: + m = _JSON_RE.search(text or "") + if not m: + return None + try: + obj = json.loads(m.group(0)) + return obj if isinstance(obj, dict) else None + except Exception: + return None + + +class AgentRequest(BaseModel): + message: str + max_steps: int = 5 + + +def _sse(event: str, data: dict) -> str: + return f"event: {event}\ndata: {json.dumps(data, default=str)}\n\n" + + +async def _agent_stream(message: str, max_steps: int): + llm = get_llm(temperature=0.1) + transcript: list[Any] = [SystemMessage(content=_AGENT_SYSTEM), HumanMessage(content=message)] + steps: list[dict[str, Any]] = [] + + async with httpx.AsyncClient(timeout=60.0, verify=False) as client: + for _ in range(max(1, min(max_steps, 8))): + try: + resp = llm.invoke(transcript + [HumanMessage(content="Next action as a single JSON object:")]) + raw = resp.content if hasattr(resp, "content") else str(resp) + except Exception as exc: # noqa: BLE001 + yield _sse("error", {"error": f"LLM error: {exc}"}) + return + + action = _parse_action(raw) + if not action or action.get("final"): + break + + tool = action.get("action") + args = action.get("input") or {} + fn = _TOOLS.get(tool) + if not fn: + transcript.append(HumanMessage(content=f"Observation: unknown tool '{tool}'.")) + continue + + yield _sse("step", {"tool": tool, "input": args}) + try: + observation = await fn(client, args) + except Exception as exc: # noqa: BLE001 + observation = f"tool error: {exc}" + steps.append({"tool": tool, "input": args, "observation": observation[:500]}) + yield _sse("observation", {"tool": tool, "observation": observation[:1200]}) + transcript.append(AIMessage(content=raw)) + transcript.append(HumanMessage(content=f"Observation from {tool}: {observation[:4000]}")) + + # Final answer — stream tokens to the client. + final_prompt = transcript + [HumanMessage(content=( + "Using the observations above, write the final answer for the user. " + "Be concise and technical, cite concrete table/column/movement names, and if you filed an " + "approval say so explicitly. Plain prose only — no JSON." + ))] + yield _sse("answer_start", {"steps": len(steps)}) + try: + for chunk in llm.stream(final_prompt): + tok = chunk.content if hasattr(chunk, "content") else str(chunk) + if tok: + yield _sse("token", {"t": tok}) + except Exception as exc: # noqa: BLE001 + yield _sse("error", {"error": f"LLM stream error: {exc}"}) + return + yield _sse("done", {"steps": steps}) + + +@app.post("/agent") +async def agent(body: AgentRequest): + return StreamingResponse( + _agent_stream(body.message, body.max_steps), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + ) + + @app.get("/") async def root(): return { "service": "ATC RAG Knowledge API", "persistent_storage": "ChromaDB + document registry on disk", - "endpoints": ["/health", "/documents", "/collections", "/ingest", "/chat", "/summarize", "/docs"], + "endpoints": ["/health", "/documents", "/collections", "/ingest", "/chat", "/summarize", + "/catalog/sync", "/catalog/status", "/agent", "/docs"], }