import threading from typing import Optional # Vector search backend — prefer ChromaDB, fall back to FAISS, then keyword try: import chromadb from chromadb.utils import embedding_functions VECTOR_BACKEND = "chroma" except ImportError: try: import faiss import numpy as np from sentence_transformers import SentenceTransformer VECTOR_BACKEND = "faiss" except ImportError: VECTOR_BACKEND = "none" from src.chatbot.document_hub import get_document # Configuration TOP_K = 5 EMBED_MODEL = "all-MiniLM-L6-v2" # Backend state (lazy-initialised) _lock = threading.Lock() # ChromaDB _chroma_client = None _chroma_collection = None # FAISS _faiss_index = None _faiss_metadata: list[dict] = [] _embed_model = None def _get_chroma_collection(): global _chroma_client, _chroma_collection if _chroma_collection is None: _chroma_client = chromadb.Client() ef = embedding_functions.SentenceTransformerEmbeddingFunction( model_name=EMBED_MODEL ) _chroma_collection = _chroma_client.get_or_create_collection( name="rag_chunks", embedding_function=ef, metadata={"hnsw:space": "cosine"}, ) return _chroma_collection def _get_faiss_model(): global _embed_model if _embed_model is None: _embed_model = SentenceTransformer(EMBED_MODEL) return _embed_model # Public API def index_document(doc_id: str) -> dict: """ Embed all chunks from a processed document and add them to the vector index. Chunks are now dicts {"text", "page"} — page number is stored as metadata so it survives retrieval. Returns ------- dict: { status, doc_id, chunks_indexed } """ doc = get_document(doc_id) if not doc: return {"status": "error", "message": f"Document '{doc_id}' not found in DocumentHub."} chunks = doc["chunks"] # list[dict] {"text": str, "page": int} filename = doc["filename"] with _lock: if VECTOR_BACKEND == "chroma": _index_chroma(doc_id, filename, chunks) elif VECTOR_BACKEND == "faiss": _index_faiss(doc_id, filename, chunks) else: print("[RAGPipeline] No vector backend — keyword search will be used.") print(f"[RAGPipeline] Indexed {len(chunks)} chunks for '{filename}' (backend={VECTOR_BACKEND})") return {"status": "success", "doc_id": doc_id, "chunks_indexed": len(chunks)} def retrieve_context(query: str, doc_id: Optional[str] = None, top_k: int = TOP_K) -> dict: """ Find the most relevant chunks for *query*. Returns ------- dict: context_str — formatted text block ready for LLM injection source — human-readable source reference, e.g. "page 4" (taken from the top-ranked chunk) chunks — raw list of chunk dicts for callers that need more detail """ with _lock: if VECTOR_BACKEND == "chroma": chunks = _query_chroma(query, doc_id, top_k) elif VECTOR_BACKEND == "faiss": chunks = _query_faiss(query, doc_id, top_k) else: chunks = _keyword_search(query, doc_id, top_k) if not chunks: return { "context_str": "No relevant context found in the uploaded documents.", "source": None, "chunks": [], } lines = ["--- RELEVANT DOCUMENT CONTEXT ---"] for i, chunk in enumerate(chunks, 1): filename = chunk.get("filename", "Unknown") page = chunk.get("page") page_label = f"Page {page}" if page else "unknown page" lines.append(f"\n[Excerpt {i} — {filename}, {page_label}]\n{chunk.get('text', '').strip()}") lines.append("\n--- END OF CONTEXT ---") # Source = top chunk's page (most relevant result) top = chunks[0] source = f"page {top['page']}" if top.get("page") else top.get("filename", "unknown") return { "context_str": "\n".join(lines), "source": source, "chunks": chunks, } def build_rag_prompt(user_query: str, context_str: str, base_system_prompt: str) -> str: """ Inject *context_str* into *base_system_prompt* so the LLM answers are grounded in document content. """ return ( f"{base_system_prompt.strip()}\n\n" "You have access to the following excerpts retrieved from the user's documents. " "Use them to answer accurately. If the answer is not in the excerpts, say so.\n\n" f"{context_str}" ) # ChromaDB helpers def _index_chroma(doc_id: str, filename: str, chunks: list[dict]) -> None: col = _get_chroma_collection() texts = [c["text"] for c in chunks] ids = [f"{doc_id}_chunk_{i}" for i in range(len(chunks))] metadatas = [ {"doc_id": doc_id, "filename": filename, "page": c.get("page", 0), "chunk_index": i} for i, c in enumerate(chunks) ] col.upsert(documents=texts, ids=ids, metadatas=metadatas) def _query_chroma(query: str, doc_id: Optional[str], top_k: int) -> list[dict]: col = _get_chroma_collection() where = {"doc_id": doc_id} if doc_id else None results = col.query( query_texts=[query], n_results=top_k, where=where, include=["documents", "metadatas", "distances"], ) chunks = [] for text, meta in zip(results["documents"][0], results["metadatas"][0]): chunks.append({ "text": text, "filename": meta.get("filename", ""), "doc_id": meta.get("doc_id", ""), "page": meta.get("page"), }) return chunks # FAISS helpers def _index_faiss(doc_id: str, filename: str, chunks: list[dict]) -> None: global _faiss_index, _faiss_metadata model = _get_faiss_model() texts = [c["text"] for c in chunks] embeddings = model.encode(texts, convert_to_numpy=True) if _faiss_index is None: _faiss_index = faiss.IndexFlatL2(embeddings.shape[1]) _faiss_index.add(embeddings) for c in chunks: _faiss_metadata.append({ "text": c["text"], "doc_id": doc_id, "filename": filename, "page": c.get("page"), }) def _query_faiss(query: str, doc_id: Optional[str], top_k: int) -> list[dict]: if _faiss_index is None or _faiss_index.ntotal == 0: return [] model = _get_faiss_model() q_vec = model.encode([query], convert_to_numpy=True) _, indices = _faiss_index.search(q_vec, min(top_k * 3, _faiss_index.ntotal)) results = [] for idx in indices[0]: if idx == -1: continue meta = _faiss_metadata[idx] if doc_id and meta["doc_id"] != doc_id: continue results.append(meta) if len(results) >= top_k: break return results # Keyword fallback def _keyword_search(query: str, doc_id: Optional[str], top_k: int) -> list[dict]: """Bag-of-words overlap search — used when no vector library is installed.""" from src.chatbot.document_hub import documents_db, doc_lock query_tokens = set(query.lower().split()) scored: list[tuple[int, dict]] = [] with doc_lock: targets = ( {doc_id: documents_db[doc_id]}.items() if doc_id and doc_id in documents_db else documents_db.items() ) for did, doc in targets: for chunk in doc["chunks"]: # chunk is now a dict chunk_tokens = set(chunk["text"].lower().split()) overlap = len(query_tokens & chunk_tokens) if overlap > 0: scored.append((overlap, { "text": chunk["text"], "page": chunk.get("page"), "doc_id": did, "filename": doc["filename"], })) scored.sort(key=lambda x: x[0], reverse=True) return [item for _, item in scored[:top_k]]