mo
de702da759
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.
903 lines
35 KiB
Python
903 lines
35 KiB
Python
"""RAG Knowledge API — LangChain + ChromaDB + Docling + LLM with persistent document registry."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
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, StreamingResponse
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings
|
|
from langchain_community.vectorstores import Chroma
|
|
from langchain_core.documents import Document
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
import chromadb
|
|
from pydantic import BaseModel
|
|
|
|
CHROMA_HOST = os.getenv("CHROMA_HOST", "chromadb")
|
|
CHROMA_PORT = int(os.getenv("CHROMA_PORT", "8000"))
|
|
CHROMA_URL = f"http://{CHROMA_HOST}:{CHROMA_PORT}"
|
|
DOCLING_URL = os.getenv("DOCLING_URL", "http://docling-serve:5001").rstrip("/")
|
|
LLM_URL = os.getenv("LLM_URL", "http://10.0.20.106:8001/v1")
|
|
LLM_MODEL = os.getenv("LLM_MODEL", "gpt-4o")
|
|
LLM_API_KEY = os.getenv("LLM_API_KEY", "sk-local")
|
|
DATA_DIR = Path(os.getenv("RAG_DATA_DIR", "/data"))
|
|
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")
|
|
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")
|
|
|
|
app = FastAPI(title="ATC RAG Knowledge API", version="1.2.0")
|
|
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
|
|
|
|
_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=120)
|
|
_embeddings: HuggingFaceEmbeddings | None = None
|
|
|
|
_BASE64_BLOB = re.compile(r"[A-Za-z0-9+/]{120,}={0,2}")
|
|
_BASE64_IMG = re.compile(r"!\[[^\]]*\]\(data:image/[^)]+\)", re.IGNORECASE)
|
|
_IMAGE_REF = re.compile(r"!\[Image\]\([^)]+\)")
|
|
|
|
|
|
def _ensure_dirs() -> None:
|
|
UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
|
|
if not REGISTRY_PATH.exists():
|
|
REGISTRY_PATH.write_text(json.dumps({"documents": []}, indent=2))
|
|
|
|
|
|
def _load_registry() -> list[dict[str, Any]]:
|
|
_ensure_dirs()
|
|
try:
|
|
data = json.loads(REGISTRY_PATH.read_text())
|
|
return data.get("documents", [])
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def _save_registry(docs: list[dict[str, Any]]) -> None:
|
|
_ensure_dirs()
|
|
REGISTRY_PATH.write_text(
|
|
json.dumps({"documents": docs, "updated_at": datetime.now(timezone.utc).isoformat()}, indent=2, default=str)
|
|
)
|
|
|
|
|
|
def _file_hash(content: bytes) -> str:
|
|
return hashlib.sha256(content).hexdigest()
|
|
|
|
|
|
def _safe_collection(name: str) -> str:
|
|
return "".join(c if c.isalnum() or c in "-_" else "_" for c in name.strip())[:64] or "default"
|
|
|
|
|
|
def clean_text_for_rag(text: str) -> str:
|
|
"""Strip embedded images and base64 blobs that pollute vector search."""
|
|
text = _BASE64_IMG.sub("<!-- image -->", text)
|
|
text = _IMAGE_REF.sub("<!-- image -->", text)
|
|
text = _BASE64_BLOB.sub("", text)
|
|
text = re.sub(r"\n{3,}", "\n\n", text)
|
|
return text.strip()
|
|
|
|
|
|
def is_garbage_chunk(text: str) -> bool:
|
|
"""Detect chunks that are mostly binary/base64 noise."""
|
|
if not text or len(text) < 20:
|
|
return True
|
|
if "data:image" in text:
|
|
return True
|
|
if _BASE64_BLOB.search(text):
|
|
return True
|
|
alpha = sum(1 for c in text if c.isalpha() or c.isspace())
|
|
if alpha / max(len(text), 1) < 0.35:
|
|
return True
|
|
return False
|
|
|
|
|
|
def get_embeddings() -> HuggingFaceEmbeddings:
|
|
global _embeddings
|
|
if _embeddings is None:
|
|
_embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
|
|
return _embeddings
|
|
|
|
|
|
def get_chroma_client() -> chromadb.HttpClient:
|
|
return chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)
|
|
|
|
|
|
def get_vectorstore(collection: str) -> Chroma:
|
|
return Chroma(
|
|
client=get_chroma_client(),
|
|
collection_name=collection,
|
|
embedding_function=get_embeddings(),
|
|
)
|
|
|
|
|
|
def get_llm(temperature: float = 0.2) -> ChatOpenAI:
|
|
return ChatOpenAI(
|
|
base_url=LLM_URL,
|
|
api_key=LLM_API_KEY,
|
|
model=LLM_MODEL,
|
|
temperature=temperature,
|
|
)
|
|
|
|
|
|
def _delete_doc_vectors(collection: str, doc_id: str) -> None:
|
|
try:
|
|
col = get_chroma_client().get_collection(collection)
|
|
col.delete(where={"doc_id": doc_id})
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def extract_text(content: bytes, filename: str) -> str:
|
|
ext = Path(filename).suffix.lower()
|
|
if ext in {".txt", ".md", ".csv", ".json"}:
|
|
try:
|
|
raw = content.decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
raw = content.decode("latin-1", errors="replace")
|
|
return clean_text_for_rag(raw)
|
|
|
|
async with httpx.AsyncClient(timeout=300.0) as client:
|
|
r = await client.post(
|
|
f"{DOCLING_URL}/v1/convert/file",
|
|
files={"files": (filename, content, "application/octet-stream")},
|
|
data={
|
|
"to_formats": ["md"],
|
|
"image_export_mode": DOCLING_IMAGE_MODE,
|
|
"do_ocr": "true",
|
|
"table_mode": "accurate",
|
|
},
|
|
)
|
|
if r.status_code >= 400:
|
|
raise ValueError(f"Docling failed: {r.text[:300]}")
|
|
doc = r.json().get("document") or {}
|
|
md = doc.get("md_content") or doc.get("text_content") or ""
|
|
md = clean_text_for_rag(md)
|
|
if len(md) < 50:
|
|
raise ValueError("No readable text extracted from document")
|
|
return md
|
|
|
|
|
|
def _find_duplicate(content_hash: str, collection: str) -> dict[str, Any] | None:
|
|
for d in _load_registry():
|
|
if d.get("content_hash") == content_hash and d.get("collection") == collection:
|
|
return d
|
|
return None
|
|
|
|
|
|
def _find_by_id(doc_id: str) -> dict[str, Any] | None:
|
|
for d in _load_registry():
|
|
if d.get("id") == doc_id:
|
|
return d
|
|
return None
|
|
|
|
|
|
async def _ingest_bytes(
|
|
content: bytes,
|
|
filename: str,
|
|
collection: str,
|
|
source: str = "upload",
|
|
*,
|
|
force_reindex: bool = False,
|
|
) -> dict[str, Any]:
|
|
col = _safe_collection(collection)
|
|
content_hash = _file_hash(content)
|
|
existing = _find_duplicate(content_hash, col)
|
|
|
|
if existing and not force_reindex:
|
|
return {
|
|
"ok": True,
|
|
"duplicate": True,
|
|
"skipped": True,
|
|
"message": f"Document already indexed as '{existing['filename']}' — chat immediately, no re-upload needed.",
|
|
**{k: existing[k] for k in ("id", "filename", "collection", "chunks", "characters", "ingested_at") if k in existing},
|
|
}
|
|
|
|
if existing and force_reindex:
|
|
doc_id = existing["id"]
|
|
_delete_doc_vectors(col, doc_id)
|
|
stored_path = Path(existing.get("stored_path", ""))
|
|
if stored_path.exists():
|
|
stored_path.write_bytes(content)
|
|
else:
|
|
stored_name = f"{doc_id}_{Path(filename).name}"
|
|
stored_path = UPLOADS_DIR / stored_name
|
|
stored_path.write_bytes(content)
|
|
else:
|
|
doc_id = uuid.uuid4().hex[:12]
|
|
stored_name = f"{doc_id}_{Path(filename).name}"
|
|
stored_path = UPLOADS_DIR / stored_name
|
|
stored_path.write_bytes(content)
|
|
|
|
text = await extract_text(content, filename)
|
|
chunks = [c for c in _splitter.split_text(text) if not is_garbage_chunk(c)]
|
|
if not chunks:
|
|
raise ValueError("No usable text chunks after cleaning — document may be image-only")
|
|
|
|
docs = [
|
|
Document(
|
|
page_content=chunk,
|
|
metadata={
|
|
"source": filename,
|
|
"doc_id": doc_id,
|
|
"chunk": i,
|
|
"content_hash": content_hash,
|
|
"ingested_at": datetime.now(timezone.utc).isoformat(),
|
|
},
|
|
)
|
|
for i, chunk in enumerate(chunks)
|
|
]
|
|
vs = get_vectorstore(col)
|
|
vs.add_documents(docs)
|
|
|
|
record = {
|
|
"id": doc_id,
|
|
"filename": filename,
|
|
"collection": col,
|
|
"content_hash": content_hash,
|
|
"stored_path": str(stored_path),
|
|
"chunks": len(chunks),
|
|
"characters": len(text),
|
|
"bytes": len(content),
|
|
"source": source,
|
|
"ingested_at": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
registry = _load_registry()
|
|
registry = [d for d in registry if not (d.get("id") == doc_id and d.get("collection") == col)]
|
|
registry.insert(0, record)
|
|
_save_registry(registry)
|
|
return {"ok": True, "duplicate": False, "reindexed": force_reindex, **record}
|
|
|
|
|
|
async def _summarize_text(text: str, filename: str) -> str:
|
|
"""Summarize document text using LLM with map-reduce for long docs."""
|
|
llm = get_llm(temperature=0.1)
|
|
max_chunk = 12000
|
|
if len(text) <= max_chunk:
|
|
prompt = (
|
|
f"Summarize this document ({filename}) clearly in English. "
|
|
"Include: main topic, key sections, important technologies/products mentioned, and target audience. "
|
|
"Use bullet points and short paragraphs.\n\nDocument:\n{text}"
|
|
)
|
|
resp = llm.invoke([HumanMessage(content=prompt.format(text=text[:max_chunk]))])
|
|
return resp.content if hasattr(resp, "content") else str(resp)
|
|
|
|
# Map-reduce for long documents
|
|
parts = [text[i : i + max_chunk] for i in range(0, min(len(text), 60000), max_chunk)]
|
|
partials: list[str] = []
|
|
for i, part in enumerate(parts[:5]):
|
|
resp = llm.invoke([
|
|
HumanMessage(content=(
|
|
f"Summarize part {i + 1}/{min(len(parts), 5)} of '{filename}'. "
|
|
f"List key topics, products, and technical points:\n\n{part}"
|
|
))
|
|
])
|
|
partials.append(resp.content if hasattr(resp, "content") else str(resp))
|
|
|
|
combined = "\n\n".join(partials)
|
|
final = llm.invoke([
|
|
HumanMessage(content=(
|
|
f"Create a clear executive summary of '{filename}' from these section summaries. "
|
|
"Structure: Overview, Main Topics, Key Technologies, Audience. Use bullet points.\n\n"
|
|
f"{combined}"
|
|
))
|
|
])
|
|
return final.content if hasattr(final, "content") else str(final)
|
|
|
|
|
|
class ChatRequest(BaseModel):
|
|
message: str
|
|
collection: str = "default"
|
|
top_k: int = 5
|
|
|
|
|
|
class IngestTextRequest(BaseModel):
|
|
text: str
|
|
collection: str = "default"
|
|
source: str = "manual"
|
|
|
|
|
|
class SummarizeRequest(BaseModel):
|
|
collection: str = "default"
|
|
doc_id: str | None = None
|
|
filename: str | None = None
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
chroma_ok = docling_ok = llm_ok = False
|
|
doc_count = len(_load_registry())
|
|
try:
|
|
async with httpx.AsyncClient(timeout=5.0) as c:
|
|
cr = await c.get(f"{CHROMA_URL}/api/v1/heartbeat")
|
|
chroma_ok = cr.status_code < 400
|
|
dr = await c.get(f"{DOCLING_URL}/health")
|
|
docling_ok = dr.status_code < 400 and dr.json().get("status") == "ok"
|
|
lr = await c.get(f"{LLM_URL.rstrip('/')}/models")
|
|
llm_ok = lr.status_code < 400
|
|
except Exception:
|
|
pass
|
|
return {
|
|
"ok": chroma_ok,
|
|
"chroma": chroma_ok,
|
|
"docling": docling_ok,
|
|
"llm": llm_ok,
|
|
"embed_model": EMBED_MODEL,
|
|
"stored_documents": doc_count,
|
|
"persistent": True,
|
|
"docling_image_mode": DOCLING_IMAGE_MODE,
|
|
}
|
|
|
|
|
|
@app.get("/documents")
|
|
async def list_documents(collection: str | None = None):
|
|
docs = _load_registry()
|
|
if collection:
|
|
col = _safe_collection(collection)
|
|
docs = [d for d in docs if d.get("collection") == col]
|
|
return {"documents": docs, "total": len(docs)}
|
|
|
|
|
|
@app.get("/documents/{doc_id}")
|
|
async def get_document(doc_id: str):
|
|
doc = _find_by_id(doc_id)
|
|
if doc:
|
|
return doc
|
|
return JSONResponse({"error": "not found"}, status_code=404)
|
|
|
|
|
|
@app.get("/documents/{doc_id}/file")
|
|
async def download_document(doc_id: str):
|
|
doc = _find_by_id(doc_id)
|
|
if doc:
|
|
path = Path(doc.get("stored_path", ""))
|
|
if path.exists():
|
|
return FileResponse(path, filename=doc.get("filename", path.name))
|
|
return JSONResponse({"error": "not found"}, status_code=404)
|
|
|
|
|
|
@app.post("/documents/{doc_id}/reindex")
|
|
async def reindex_document(doc_id: str):
|
|
doc = _find_by_id(doc_id)
|
|
if not doc:
|
|
return JSONResponse({"error": "not found"}, status_code=404)
|
|
path = Path(doc.get("stored_path", ""))
|
|
if not path.exists():
|
|
return JSONResponse({"error": "stored file missing"}, status_code=404)
|
|
try:
|
|
content = path.read_bytes()
|
|
result = await _ingest_bytes(
|
|
content,
|
|
doc["filename"],
|
|
doc["collection"],
|
|
source=doc.get("source", "reindex"),
|
|
force_reindex=True,
|
|
)
|
|
return result
|
|
except ValueError as exc:
|
|
return JSONResponse({"ok": False, "error": str(exc)}, status_code=422)
|
|
|
|
|
|
@app.get("/collections")
|
|
async def list_collections():
|
|
try:
|
|
client = get_chroma_client()
|
|
cols = client.list_collections()
|
|
registry = _load_registry()
|
|
items = []
|
|
for col in cols:
|
|
files = {d["filename"] for d in registry if d.get("collection") == col.name}
|
|
items.append({
|
|
"name": col.name,
|
|
"documents": col.count(),
|
|
"files": len(files),
|
|
"filenames": sorted(files)[:20],
|
|
})
|
|
return {"collections": items}
|
|
except Exception as exc:
|
|
return JSONResponse({"error": str(exc), "collections": []}, status_code=502)
|
|
|
|
|
|
@app.post("/collections")
|
|
async def create_collection(name: str = Form(...)):
|
|
safe = _safe_collection(name)
|
|
get_vectorstore(safe)
|
|
return {"ok": True, "collection": safe}
|
|
|
|
|
|
@app.post("/ingest")
|
|
async def ingest_file(
|
|
file: UploadFile = File(...),
|
|
collection: str = Form("default"),
|
|
force_reindex: bool = Form(False),
|
|
):
|
|
_ensure_dirs()
|
|
safe_name = file.filename or "upload.txt"
|
|
content = await file.read()
|
|
try:
|
|
result = await _ingest_bytes(content, safe_name, collection, force_reindex=force_reindex)
|
|
if not result.get("ok"):
|
|
return JSONResponse(result, status_code=422)
|
|
return result
|
|
except ValueError as exc:
|
|
return JSONResponse({"ok": False, "error": str(exc)}, status_code=422)
|
|
|
|
|
|
@app.post("/ingest/text")
|
|
async def ingest_text(body: IngestTextRequest):
|
|
col = _safe_collection(body.collection)
|
|
content = body.text.encode("utf-8")
|
|
filename = f"{body.source}.txt"
|
|
return await _ingest_bytes(content, filename, col, source=body.source)
|
|
|
|
|
|
@app.post("/summarize")
|
|
async def summarize(body: SummarizeRequest):
|
|
col = _safe_collection(body.collection)
|
|
registry = _load_registry()
|
|
doc: dict[str, Any] | None = None
|
|
if body.doc_id:
|
|
doc = _find_by_id(body.doc_id)
|
|
elif body.filename:
|
|
for d in registry:
|
|
if d.get("filename") == body.filename and d.get("collection") == col:
|
|
doc = d
|
|
break
|
|
else:
|
|
docs_in_col = [d for d in registry if d.get("collection") == col]
|
|
if len(docs_in_col) == 1:
|
|
doc = docs_in_col[0]
|
|
|
|
if not doc:
|
|
return JSONResponse({"ok": False, "error": "Document not found — specify doc_id or filename"}, status_code=404)
|
|
|
|
path = Path(doc.get("stored_path", ""))
|
|
if not path.exists():
|
|
return JSONResponse({"ok": False, "error": "Stored file missing"}, status_code=404)
|
|
|
|
try:
|
|
content = path.read_bytes()
|
|
text = await extract_text(content, doc["filename"])
|
|
summary = await _summarize_text(text, doc["filename"])
|
|
return {
|
|
"ok": True,
|
|
"summary": summary,
|
|
"filename": doc["filename"],
|
|
"doc_id": doc["id"],
|
|
"collection": col,
|
|
"characters": len(text),
|
|
}
|
|
except ValueError as exc:
|
|
return JSONResponse({"ok": False, "error": str(exc)}, status_code=422)
|
|
except Exception as exc:
|
|
return JSONResponse({"ok": False, "error": f"Summarize failed: {exc}"}, status_code=502)
|
|
|
|
|
|
@app.post("/chat")
|
|
async def chat(body: ChatRequest):
|
|
col = _safe_collection(body.collection)
|
|
try:
|
|
vs = get_vectorstore(col)
|
|
count = get_chroma_client().get_collection(col).count()
|
|
except Exception as exc:
|
|
return JSONResponse({"ok": False, "error": f"Collection unavailable: {exc}"}, status_code=404)
|
|
|
|
if count == 0:
|
|
stored = [d for d in _load_registry() if d.get("collection") == col]
|
|
if stored:
|
|
return JSONResponse({
|
|
"ok": False,
|
|
"error": "Vectors missing but files exist — click Re-index on the document in the library.",
|
|
"stored_documents": len(stored),
|
|
}, status_code=400)
|
|
return JSONResponse({"ok": False, "error": "Collection is empty — upload documents first."}, status_code=400)
|
|
|
|
retriever = vs.as_retriever(search_kwargs={"k": min(body.top_k * 3, 20)})
|
|
raw_docs = retriever.invoke(body.message)
|
|
docs = [d for d in raw_docs if not is_garbage_chunk(d.page_content)][: body.top_k]
|
|
|
|
if not docs:
|
|
return JSONResponse({
|
|
"ok": False,
|
|
"error": "Retrieved chunks are corrupted (old base64 index). Click Re-index on the document.",
|
|
}, status_code=400)
|
|
|
|
context = "\n\n---\n\n".join(
|
|
f"[Source: {d.metadata.get('source', '?')} | chunk {d.metadata.get('chunk', '?')}]\n{d.page_content}"
|
|
for d in docs
|
|
)
|
|
|
|
system = (
|
|
"You are a helpful data assistant for the Dell ATC platform. "
|
|
"Answer ONLY based on the provided context. If the context does not contain the answer, say so clearly. "
|
|
"Cite sources by filename when relevant. Be concise and technical."
|
|
)
|
|
user = f"Context:\n{context}\n\nQuestion: {body.message}"
|
|
|
|
try:
|
|
llm = get_llm()
|
|
resp = llm.invoke([SystemMessage(content=system), HumanMessage(content=user)])
|
|
answer = resp.content if hasattr(resp, "content") else str(resp)
|
|
except Exception as exc:
|
|
return JSONResponse({"ok": False, "error": f"LLM error: {exc}"}, status_code=502)
|
|
|
|
sources = [
|
|
{"source": d.metadata.get("source"), "chunk": d.metadata.get("chunk"), "preview": d.page_content[:200]}
|
|
for d in docs
|
|
]
|
|
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_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()
|
|
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,
|
|
"lookup_records": _tool_lookup_records,
|
|
"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 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.'
|
|
)
|
|
|
|
_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",
|
|
"/catalog/sync", "/catalog/status", "/agent", "/docs"],
|
|
}
|