from typing import List, Dict, Optional from sqlalchemy.orm import Session import chromadb import os from app.utils.chunking import text_chunker from app.utils.embeddings import embedding_generator from app.utils.reranker import reranker class RAGService: """Service for RAG operations including document processing and retrieval.""" def __init__(self): """Initialize RAG service with ChromaDB.""" # Initialize ChromaDB client chroma_path = os.path.join(os.path.dirname(__file__), "..", "..", "data", "chromadb") os.makedirs(chroma_path, exist_ok=True) self.chroma_client = chromadb.PersistentClient(path=chroma_path) # Get or create collection self.collection = self.chroma_client.get_or_create_collection( name="construction_documents", metadata={"description": "Construction documents and manuals"} ) print(f"[RAG Service] Initialized with embedding model: {embedding_generator.model_name}") def chunk_text(self, text: str, chunk_size: int = 800, overlap: int = 200) -> List[str]: """ Chunk text using the modular text chunker utility. Args: text: Text to chunk chunk_size: Target size of each chunk in characters overlap: Overlap between chunks in characters Returns: List of text chunks """ return text_chunker.chunk_by_sentences(text, chunk_size, overlap) def process_document( self, document_id: str, filename: str, content: str, user_id: str, db: Optional[Session] = None ) -> int: """ Process document by chunking and storing in vector database. Args: document_id: Unique document ID filename: Document filename content: Document text content user_id: User ID who uploaded the document Returns: Number of chunks created """ # Chunk the document chunks = self.chunk_text(content) if not chunks: return 0 # Generate embeddings using modular utility embeddings = embedding_generator.generate_embeddings(chunks) # Prepare metadata metadatas = [ { "document_id": document_id, "filename": filename, "user_id": user_id, "chunk_index": i } for i in range(len(chunks)) ] # Generate IDs for chunks ids = [f"{document_id}_chunk_{i}" for i in range(len(chunks))] # Add to ChromaDB self.collection.add( embeddings=embeddings, documents=chunks, metadatas=metadatas, ids=ids ) # Also store in structured database if it's an official policy and DB session is provided if user_id == "official_policies" and db: from app.database.models import PolicyChunk # Delete existing chunks for this policy first to avoid duplicates db.query(PolicyChunk).filter(PolicyChunk.policy_id == document_id).delete() # Create new chunks chunk_objects = [] for i, chunk_text_content in enumerate(chunks): chunk_objects.append(PolicyChunk( id=f"{document_id}_{i}", policy_id=document_id, chunk_index=i, content=chunk_text_content )) db.bulk_save_objects(chunk_objects) db.commit() print(f"[RAG Service] {len(chunks)} chunks bulk-stored in structured database for policy {document_id}") return len(chunks) def semantic_search( self, query: str, user_id: Optional[str] = None, top_k: int = 5 ) -> List[Dict]: """ Perform hybrid search (semantic + keyword BM25) with reranking. Args: query: Search query user_id: Optional user ID to filter documents top_k: Number of results to return after reranking Returns: List of relevant chunks with metadata """ # Build where filter where_filter = {"user_id": user_id} if user_id else None # STEP 1: Get all documents for BM25 indexing all_docs = self.collection.get(where=where_filter) if not all_docs or not all_docs['documents']: return [] # STEP 2: Semantic search (ChromaDB embedding-based) query_embedding = embedding_generator.generate_embedding(query) semantic_results = self.collection.query( query_embeddings=[query_embedding], n_results=min(30, top_k * 3), where=where_filter ) semantic_chunks = [] if semantic_results and semantic_results['documents']: for i in range(len(semantic_results['documents'][0])): semantic_chunks.append({ "content": semantic_results['documents'][0][i], "metadata": semantic_results['metadatas'][0][i], "distance": semantic_results['distances'][0][i] if 'distances' in semantic_results else None }) # STEP 3: Keyword search (BM25) from app.utils.bm25_search import BM25Search, HybridSearch bm25 = BM25Search() bm25_chunks = [ { "content": all_docs['documents'][i], "metadata": all_docs['metadatas'][i] } for i in range(len(all_docs['documents'])) ] bm25.index_documents(bm25_chunks) keyword_chunks = bm25.search(query, top_k=30) # STEP 4: Combine with hybrid scoring (70% semantic, 30% keyword) hybrid = HybridSearch(semantic_weight=0.7, keyword_weight=0.3) combined_chunks = hybrid.combine_results(semantic_chunks, keyword_chunks, top_k=30) if not combined_chunks: return [] # STEP 5: Final reranking with cross-encoder reranked = reranker.rerank(query, combined_chunks, top_k=top_k) print(f"[RAG Service] Hybrid: {len(semantic_chunks)} semantic + {len(keyword_chunks)} keyword → {len(reranked)} final") return reranked def search_policies( self, query: str, policy_ids: List[str], top_k: int = 10 ) -> List[Dict]: """ Search within specific official policy documents with hybrid search and reranking. Args: query: Search query policy_ids: List of policy document IDs to search within top_k: Number of results to return after reranking Returns: List of relevant chunks with metadata from selected policies """ # Build filter for official policies where_filter = { "$and": [ {"user_id": {"$eq": "official_policies"}}, {"document_id": {"$in": policy_ids}} ] } print(f"[RAG Service] search_policies called with {len(policy_ids)} policies") # STEP 1: Semantic search query_embedding = embedding_generator.generate_embedding(query) initial_results = self.collection.query( query_embeddings=[query_embedding], n_results=min(30, top_k * 3), where=where_filter ) initial_chunks = [] if initial_results and initial_results['documents']: for i in range(len(initial_results['documents'][0])): initial_chunks.append({ "content": initial_results['documents'][0][i], "metadata": initial_results['metadatas'][0][i], "distance": initial_results['distances'][0][i] if 'distances' in initial_results else None }) if not initial_chunks: print(f"[RAG Service] No chunks found") return [] # STEP 2: Rerank reranked_chunks = reranker.rerank(query, initial_chunks, top_k=top_k) print(f"[RAG Service] Returning {len(reranked_chunks)} reranked chunks") return reranked_chunks def delete_document_chunks(self, document_id: str): """ Delete all chunks for a document. Args: document_id: Document ID """ results = self.collection.get( where={"document_id": document_id} ) if results and results['ids']: self.collection.delete(ids=results['ids']) # Global RAG service instance rag_service = RAGService()