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.
This commit is contained in:
mo
+332
-3
@@ -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": "<tool>", "input": {<args>}}\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"],
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user