mo
549 lines
18 KiB
Python
549 lines
18 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
|
|
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"],
|
|
}
|