Spaces:
Runtime error
Runtime error
| 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() | |