aibridze/document_intelligence
0
1# import chromadb2# from chromadb.config import Settings as ChromaSettings3# from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings4# from langchain_text_splitters import RecursiveCharacterTextSplitter5# from langchain_core.documents import Document6# from typing import List, Dict, Optional7# import hashlib8 9# from app.config import get_settings10 11# settings = get_settings()12 13 14# class RAGEngine:15 16# def __init__(self):17# if not settings.GOOGLE_API_KEY:18# raise ValueError("GOOGLE_API_KEY not configured")19 20# # ChromaDB21# self.chroma_client = chromadb.PersistentClient(22# path=str(settings.get_vectorstore_path()),23# settings=ChromaSettings(anonymized_telemetry=False)24# )25# self.collection = self.chroma_client.get_or_create_collection(26# name="contract_documents",27# metadata={"hnsw:space": "cosine"}28# )29 30# # Embeddings31# self.embeddings = GoogleGenerativeAIEmbeddings(32# model=settings.EMBEDDING_MODEL,33# google_api_key=settings.GOOGLE_API_KEY34# )35 36# # LLM for answer generation37# self.llm = ChatGoogleGenerativeAI(38# model=settings.VISION_MODEL,39# google_api_key=settings.GOOGLE_API_KEY,40# temperature=041# )42 43# # Text splitter44# self.text_splitter = RecursiveCharacterTextSplitter(45# chunk_size=1000,46# chunk_overlap=200,47# separators=["\n\n", "\n", ". ", " ", ""]48# )49 50# def _generate_chunk_id(self, doc_id: str, chunk_index: int) -> str:51# return hashlib.md5(f"{doc_id}_{chunk_index}".encode()).hexdigest()52 53# async def index_documents(self, doc_id: str, documents: List[Document]) -> int:54# if not documents:55# return 056 57# all_chunks: List[str] = []58# all_metadatas: List[dict] = []59 60# for doc in documents:61# page_chunks = self.text_splitter.split_text(doc.page_content)62# for chunk in page_chunks:63# all_chunks.append(chunk)64# all_metadatas.append({65# "doc_id": doc_id,66# "page": doc.metadata.get("page", 0),67# "source": doc.metadata.get("filename", ""),68# })69 70# if not all_chunks:71# return 072 73# # Generate embeddings74# chunk_embeddings = self.embeddings.embed_documents(all_chunks)75 76# # Generate IDs77# ids = [self._generate_chunk_id(doc_id, i) for i in range(len(all_chunks))]78 79# # Upsert to ChromaDB80# self.collection.upsert(81# ids=ids,82# embeddings=chunk_embeddings,83# documents=all_chunks,84# metadatas=all_metadatas85# )86 87# return len(all_chunks)88 89# async def index_document(self, doc_id: str, text: str) -> int:90# doc = Document(page_content=text, metadata={"source": doc_id})91# return await self.index_documents(doc_id, [doc])92 93# async def search(94# self,95# query: str,96# doc_id: Optional[str] = None,97# top_k: int = 598# ) -> List[Dict]:99# """Search for relevant chunks. Returns results with page numbers."""100# query_embedding = self.embeddings.embed_query(query)101 102# where_filter = {"doc_id": doc_id} if doc_id else None103 104# results = self.collection.query(105# query_embeddings=[query_embedding],106# n_results=top_k,107# where=where_filter,108# include=["documents", "metadatas", "distances"]109# )110 111# formatted = []112# if results and results["documents"] and results["documents"][0]:113# for i in range(len(results["documents"][0])):114# meta = results["metadatas"][0][i] if results["metadatas"] else {}115# formatted.append({116# "text": results["documents"][0][i],117# "page": meta.get("page", None),118# "source": meta.get("source", ""),119# "score": 1 - results["distances"][0][i] if results["distances"] else 0120# })121 122# return formatted123 124# async def answer_query(125# self,126# query: str,127# doc_id: Optional[str] = None,128# top_k: int = 5129# ) -> Dict:130# """Answer a query using RAG with source citations."""131# sources = await self.search(query, doc_id, top_k)132 133# if not sources:134# return {135# "answer": "No relevant information found in the indexed documents.",136# "sources": []137# }138 139# context = "\n\n---\n\n".join([140# f"[Source {i+1} | Page {s.get('page', '?')}]\n{s['text']}"141# for i, s in enumerate(sources)142# ])143 144# prompt = f"""Based on the following contract excerpts, answer the question.145 146# IMPORTANT RULES:147# 1. Only use information from the provided excerpts148# 2. Quote the exact text when relevant149# 3. Cite which source and page contains the information150# 4. If the answer isn't in the excerpts, say so clearly151 152# EXCERPTS:153# {context}154 155# QUESTION: {query}156 157# ANSWER:"""158 159# response = await self.llm.ainvoke(prompt)160 161# return {162# "answer": response.content,163# "sources": sources164# }165 166# def delete_document(self, doc_id: str) -> bool:167# self.collection.delete(where={"doc_id": doc_id})168# return True169 170# def get_indexed_documents(self) -> List[str]:171# results = self.collection.get(include=["metadatas"])172# if not results or not results["metadatas"]:173# return []174# doc_ids = set()175# for metadata in results["metadatas"]:176# if metadata and "doc_id" in metadata:177# doc_ids.add(metadata["doc_id"])178# return list(doc_ids)179 180 181"""182RAG Engine — Production-grade contract Q&A.183 184Features:185 1. Page-aware chunking (preserves page boundaries from OCR text)186 2. Deduplication (deletes old chunks before re-indexing)187 3. Contract-type metadata for filtered search188 4. Legal-optimized chunk sizing (1500 chars, 300 overlap)189 5. Contract-domain answer generation with structured citations190 6. Relevance threshold filtering (drops low-score chunks)191 7. Comprehensive error handling192"""193 194import chromadb195from chromadb.config import Settings as ChromaSettings196from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings197from langchain_text_splitters import RecursiveCharacterTextSplitter198from langchain_core.documents import Document199from typing import List, Dict, Optional200import hashlib201import logging202import re203 204from app.config import get_settings205 206settings = get_settings()207logger = logging.getLogger("trident_poc")208 209# Relevance threshold — chunks below this cosine similarity are discarded210RELEVANCE_THRESHOLD = 0.25211# Maximum chunks to send to LLM for answer generation212MAX_CONTEXT_CHUNKS = 8213 214 215class RAGEngine:216 217 def __init__(self):218 if not settings.GOOGLE_API_KEY:219 raise ValueError("GOOGLE_API_KEY not configured")220 221 # ChromaDB persistent store222 self.chroma_client = chromadb.PersistentClient(223 path=str(settings.get_vectorstore_path()),224 settings=ChromaSettings(anonymized_telemetry=False),225 )226 self.collection = self.chroma_client.get_or_create_collection(227 name="contract_documents",228 metadata={"hnsw:space": "cosine"},229 )230 231 # Embeddings232 self.embeddings = GoogleGenerativeAIEmbeddings(233 model=settings.EMBEDDING_MODEL,234 google_api_key=settings.GOOGLE_API_KEY,235 )236 237 # LLM for answer generation238 self.llm = ChatGoogleGenerativeAI(239 model=settings.VISION_MODEL,240 google_api_key=settings.GOOGLE_API_KEY,241 temperature=0,242 )243 244 # Legal-optimized text splitter:245 # - 1500 chars per chunk (legal clauses are long)246 # - 300 char overlap (clauses often reference prior text)247 # - Split on paragraph → sentence → word boundaries248 self.text_splitter = RecursiveCharacterTextSplitter(249 chunk_size=1500,250 chunk_overlap=300,251 separators=["\n\n\n", "\n\n", "\n", ". ", "; ", ", ", " ", ""],252 )253 254 # ─── Chunk ID Generation ─────────────────────────────────────255 256 @staticmethod257 def _generate_chunk_id(doc_id: str, chunk_index: int) -> str:258 return hashlib.md5(f"{doc_id}_chunk_{chunk_index}".encode()).hexdigest()259 260 # ─── Index: Page-Based Documents ──────────────────────────────261 262 async def index_documents(263 self,264 doc_id: str,265 documents: List[Document],266 contract_type: str = None,267 ) -> int:268 """269 Index a list of per-page LangChain Documents.270 Deletes any existing chunks for this doc_id first (dedup).271 Each chunk preserves the source page number in metadata.272 """273 if not documents:274 return 0275 276 # Step 1: Delete old chunks for this doc (prevents duplicates)277 try:278 self._delete_doc_chunks(doc_id)279 except Exception as e:280 logger.warning(f"Could not delete old chunks for {doc_id}: {e}")281 282 # Step 2: Chunk each page separately (preserves page boundaries)283 all_chunks: List[str] = []284 all_metadatas: List[dict] = []285 286 for doc in documents:287 page_num = doc.metadata.get("page", 0)288 filename = doc.metadata.get("filename", "")289 290 page_chunks = self.text_splitter.split_text(doc.page_content)291 292 for chunk in page_chunks:293 chunk_clean = chunk.strip()294 if len(chunk_clean) < 20:295 continue # Skip tiny fragments296 297 all_chunks.append(chunk_clean)298 all_metadatas.append({299 "doc_id": doc_id,300 "page": page_num,301 "source": filename,302 "contract_type": contract_type or "",303 })304 305 if not all_chunks:306 return 0307 308 # Step 3: Embed and upsert309 try:310 chunk_embeddings = self.embeddings.embed_documents(all_chunks)311 ids = [312 self._generate_chunk_id(doc_id, i)313 for i in range(len(all_chunks))314 ]315 316 # ChromaDB has a batch limit — upsert in batches of 100317 batch_size = 100318 for start in range(0, len(all_chunks), batch_size):319 end = min(start + batch_size, len(all_chunks))320 self.collection.upsert(321 ids=ids[start:end],322 embeddings=chunk_embeddings[start:end],323 documents=all_chunks[start:end],324 metadatas=all_metadatas[start:end],325 )326 327 logger.info(f"Indexed {len(all_chunks)} chunks for doc {doc_id}")328 return len(all_chunks)329 330 except Exception as e:331 logger.error(f"Failed to index documents for {doc_id}: {e}", exc_info=True)332 return 0333 334 # ─── Index: Raw Text (from Vision OCR) ────────────────────────335 336 async def index_document(337 self,338 doc_id: str,339 text: str,340 contract_type: str = None,341 ) -> int:342 """343 Index raw OCR text. Splits on page markers first,344 then chunks each page. Preserves page numbers from OCR output.345 """346 if not text or not text.strip():347 return 0348 349 # Parse page markers from OCR text: "--- Page N ---" or "--- Page N (OCR) ---"350 page_pattern = re.compile(351 r'---\s*Page\s+(\d+)\s*(?:\(OCR\))?\s*---'352 )353 354 # Split text into per-page documents355 documents = []356 parts = page_pattern.split(text)357 358 if len(parts) > 1:359 # parts = [pre-text, page_num, page_text, page_num, page_text, ...]360 for i in range(1, len(parts), 2):361 page_num = int(parts[i])362 page_text = parts[i + 1].strip() if i + 1 < len(parts) else ""363 if page_text and len(page_text) > 20:364 documents.append(Document(365 page_content=page_text,366 metadata={367 "page": page_num,368 "source": doc_id,369 "filename": doc_id,370 },371 ))372 else:373 # No page markers — treat as single document374 documents.append(Document(375 page_content=text,376 metadata={"page": 0, "source": doc_id, "filename": doc_id},377 ))378 379 return await self.index_documents(doc_id, documents, contract_type)380 381 # ─── Search ───────────────────────────────────────────────────382 383 async def search(384 self,385 query: str,386 doc_id: Optional[str] = None,387 top_k: int = 5,388 ) -> List[Dict]:389 """390 Semantic search for relevant chunks.391 Returns results sorted by relevance score with page numbers.392 Filters out results below RELEVANCE_THRESHOLD.393 """394 try:395 query_embedding = self.embeddings.embed_query(query)396 397 where_filter = {"doc_id": doc_id} if doc_id else None398 399 # Fetch more than needed so we can filter low-relevance results400 fetch_k = min(top_k * 2, 20)401 402 results = self.collection.query(403 query_embeddings=[query_embedding],404 n_results=fetch_k,405 where=where_filter,406 include=["documents", "metadatas", "distances"],407 )408 409 if not results or not results["documents"] or not results["documents"][0]:410 return []411 412 formatted = []413 for i in range(len(results["documents"][0])):414 meta = results["metadatas"][0][i] if results["metadatas"] else {}415 distance = results["distances"][0][i] if results["distances"] else 1.0416 score = 1 - distance # Convert distance to similarity417 418 # Filter low-relevance results419 if score < RELEVANCE_THRESHOLD:420 continue421 422 formatted.append({423 "text": results["documents"][0][i],424 "page": meta.get("page", None),425 "source": meta.get("source", ""),426 "contract_type": meta.get("contract_type", ""),427 "score": round(score, 4),428 })429 430 # Sort by score descending and limit431 formatted.sort(key=lambda x: x["score"], reverse=True)432 return formatted[:top_k]433 434 except Exception as e:435 logger.error(f"Search failed: {e}", exc_info=True)436 return []437 438 # ─── Answer Generation ────────────────────────────────────────439 440 async def answer_query(441 self,442 query: str,443 doc_id: Optional[str] = None,444 top_k: int = 5,445 ) -> Dict:446 """447 Answer a question using RAG with contract-domain-specific prompting.448 Returns answer with source citations including page numbers.449 """450 sources = await self.search(query, doc_id, top_k)451 452 if not sources:453 return {454 "answer": (455 "I could not find relevant information in the indexed documents. "456 "This might be because the document hasn't been indexed yet, "457 "or the question is about content not present in the document."458 ),459 "sources": [],460 }461 462 # Build context with clear source attribution463 context_parts = []464 for i, s in enumerate(sources[:MAX_CONTEXT_CHUNKS]):465 page_info = f"Page {s['page']}" if s.get("page") else "Unknown page"466 source_info = s.get("source", "")467 context_parts.append(468 f"[Source {i+1} | {source_info} | {page_info} | "469 f"Relevance: {s['score']:.0%}]\n{s['text']}"470 )471 472 context = "\n\n---\n\n".join(context_parts)473 474 # Contract-domain-specific prompt475 prompt = f"""You are a legal contract analyst answering questions about contract documents.476 477CONTEXT (relevant excerpts from the contract):478{context}479 480RULES:4811. Answer ONLY based on the information in the excerpts above.4822. Be specific and precise — include exact names, dates, amounts, and clauses.4833. When quoting contract language, use quotation marks.4844. Cite the source number and page for each piece of information, e.g., (Source 1, Page 3).4855. If the answer involves multiple clauses or sections, organize your response clearly.4866. If the excerpts don't contain enough information to answer fully, say what you found and what's missing.4877. For financial questions, always mention currency, amounts, and payment conditions.4888. For party-related questions, provide full legal names and roles.4899. Never make up or infer information that isn't in the excerpts.490 491QUESTION: {query}492 493ANSWER:"""494 495 try:496 response = await self.llm.ainvoke(prompt)497 answer = response.content.strip()498 except Exception as e:499 logger.error(f"Answer generation failed: {e}", exc_info=True)500 answer = (501 "I found relevant passages but encountered an error generating the answer. "502 "Please check the source excerpts below for the information you need."503 )504 505 return {506 "answer": answer,507 "sources": sources,508 }509 510 # ─── Document Management ──────────────────────────────────────511 512 def delete_document(self, doc_id: str) -> bool:513 """Delete all indexed chunks for a document."""514 try:515 self._delete_doc_chunks(doc_id)516 logger.info(f"Deleted all chunks for doc {doc_id}")517 return True518 except Exception as e:519 logger.error(f"Failed to delete chunks for {doc_id}: {e}")520 return False521 522 def _delete_doc_chunks(self, doc_id: str):523 """Internal: delete chunks by doc_id filter."""524 try:525 self.collection.delete(where={"doc_id": doc_id})526 except Exception:527 # ChromaDB may throw if no matching docs — that's fine528 pass529 530 def get_indexed_documents(self) -> List[Dict]:531 """List all indexed documents with chunk counts."""532 try:533 results = self.collection.get(include=["metadatas"])534 if not results or not results["metadatas"]:535 return []536 537 doc_info = {}538 for metadata in results["metadatas"]:539 if metadata and "doc_id" in metadata:540 did = metadata["doc_id"]541 if did not in doc_info:542 doc_info[did] = {543 "doc_id": did,544 "source": metadata.get("source", ""),545 "contract_type": metadata.get("contract_type", ""),546 "chunk_count": 0,547 "pages": set(),548 }549 doc_info[did]["chunk_count"] += 1550 page = metadata.get("page")551 if page:552 doc_info[did]["pages"].add(page)553 554 # Convert sets to sorted lists for JSON serialization555 result_list = []556 for info in doc_info.values():557 info["pages"] = sorted(info["pages"])558 info["page_count"] = len(info["pages"])559 result_list.append(info)560 561 return result_list562 563 except Exception as e:564 logger.error(f"Failed to list indexed documents: {e}")565 return []566 567 def get_document_chunks(self, doc_id: str, limit: int = 10) -> List[Dict]:568 """Get sample chunks for a document (for debugging)."""569 try:570 results = self.collection.get(571 where={"doc_id": doc_id},572 include=["documents", "metadatas"],573 limit=limit,574 )575 576 if not results or not results["documents"]:577 return []578 579 chunks = []580 for i in range(len(results["documents"])):581 meta = results["metadatas"][i] if results["metadatas"] else {}582 chunks.append({583 "text": results["documents"][i][:200] + "..."584 if len(results["documents"][i]) > 200585 else results["documents"][i],586 "page": meta.get("page"),587 "source": meta.get("source", ""),588 })589 590 return chunks591 592 except Exception as e:593 logger.error(f"Failed to get chunks for {doc_id}: {e}")594 return []