Add DQ + RAG APIs with Docling, ChromaDB, persistent ingest
This commit is contained in:
mo
+548
@@ -0,0 +1,548 @@
|
||||
"""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
|
||||
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_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")
|
||||
|
||||
# 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)}
|
||||
|
||||
|
||||
@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"],
|
||||
}
|
||||
Reference in New Issue
Block a user