""" RAG Engine with advanced retrieval strategies. Supports vector search, hybrid search, reranking, and agentic modes. """ import logging from typing import Dict, List, Optional from llama_index.core import Settings, StorageContext, VectorStoreIndex, get_response_synthesizer from llama_index.core.query_engine import SubQuestionQueryEngine from llama_index.core.retrievers import VectorIndexRetriever from llama_index.core.tools import QueryEngineTool from llama_index.embeddings.ollama import OllamaEmbedding from llama_index.llms.ollama import Ollama from llama_index.core.postprocessor import SentenceTransformerRerank from llama_index.vector_stores.chroma import ChromaVectorStore # BM25 retriever - needs to be installed separately try: from llama_index.retrievers.bm25 import BM25Retriever HAS_BM25 = True except ImportError: HAS_BM25 = False from app.core.chroma_manager import get_chroma_manager from app.core.config import Config logger = logging.getLogger(__name__) class RAGEngine: """Advanced RAG engine with multiple retrieval strategies.""" def __init__(self, collection_name: str): """ Initialize RAG engine for a specific knowledge base. Args: collection_name: Name of the ChromaDB collection """ self.collection_name = collection_name # Initialize Ollama LLM with optimized context window # Context needs to fit: retrieved chunks (5 × 1024 = 5120) + question + answer # Using 4096 is a good balance between memory usage and functionality Settings.llm = Ollama( model=Config.OLLAMA_MODEL, base_url=Config.OLLAMA_HOST, request_timeout=120.0, context_window=4096, # Balanced: fits most RAG queries without excessive memory num_ctx=4096 # Ollama-specific parameter for context size ) # Initialize Ollama embeddings Settings.embed_model = OllamaEmbedding( model_name=Config.OLLAMA_EMBED_MODEL, base_url=Config.OLLAMA_HOST ) # Get ChromaDB collection chroma_manager = get_chroma_manager() collection = chroma_manager.get_collection(collection_name) if not collection: raise ValueError(f"Collection {collection_name} not found") # Create vector store and index vector_store = ChromaVectorStore(chroma_collection=collection) storage_context = StorageContext.from_defaults(vector_store=vector_store) self.index = VectorStoreIndex.from_vector_store( vector_store, storage_context=storage_context ) # Create base query engine self.query_engine = self.index.as_query_engine( similarity_top_k=Config.TOP_K_RETRIEVAL ) # Initialize retrievers self.vector_retriever = VectorIndexRetriever( index=self.index, similarity_top_k=Config.TOP_K_RETRIEVAL ) # Initialize BM25 retriever (for hybrid search) if HAS_BM25: try: all_nodes = self.index.docstore.docs.values() self.bm25_retriever = BM25Retriever.from_defaults( nodes=list(all_nodes), similarity_top_k=Config.TOP_K_RETRIEVAL ) except Exception as e: logger.warning(f"Failed to initialize BM25 retriever: {e}") self.bm25_retriever = None else: logger.warning("BM25 retriever not available - llama-index-retrievers-bm25 not installed") self.bm25_retriever = None # Initialize reranker try: self.reranker = SentenceTransformerRerank( model="cross-encoder/ms-marco-MiniLM-L-6-v2", top_n=Config.RERANK_TOP_N ) except Exception as e: logger.warning(f"Failed to initialize reranker: {e}") self.reranker = None # Initialize agentic query engine try: query_tool = QueryEngineTool.from_defaults( query_engine=self.query_engine, name="knowledge_base", description=f"Searches the {collection_name} knowledge base for relevant information" ) self.sub_question_engine = SubQuestionQueryEngine.from_defaults( query_engine_tools=[query_tool] ) except Exception as e: logger.warning(f"Failed to initialize agentic engine: {e}") self.sub_question_engine = None def _reciprocal_rank_fusion( self, results_list: List[List], k: int = 60 ) -> List: """ Fuse multiple retrieval results using reciprocal rank fusion. Args: results_list: List of retrieval result lists k: RRF parameter (default 60) Returns: Fused and sorted list of nodes """ # Aggregate scores by node ID node_scores = {} for results in results_list: for rank, node in enumerate(results, start=1): node_id = node.node_id score = 1.0 / (k + rank) if node_id in node_scores: node_scores[node_id]["score"] += score else: node_scores[node_id] = { "node": node, "score": score } # Sort by score sorted_nodes = sorted( node_scores.values(), key=lambda x: x["score"], reverse=True ) return [item["node"] for item in sorted_nodes[:Config.TOP_K_RETRIEVAL]] def _query_hybrid(self, question: str) -> Dict[str, any]: """ Execute hybrid search (vector + BM25). Args: question: Query string Returns: dict: Answer and sources """ if not self.bm25_retriever: logger.warning("BM25 retriever not available, falling back to vector search") return self._query_base(question) # Retrieve from both methods vector_results = self.vector_retriever.retrieve(question) bm25_results = self.bm25_retriever.retrieve(question) # Fuse results fused_nodes = self._reciprocal_rank_fusion([vector_results, bm25_results]) # Synthesize answer synthesizer = get_response_synthesizer() response = synthesizer.synthesize(question, nodes=fused_nodes) return self._format_response(response, fused_nodes) def _apply_reranking(self, nodes: List, query: str) -> List: """ Apply reranking to retrieved nodes. Args: nodes: Retrieved nodes query: Query string Returns: Reranked nodes """ if not self.reranker: logger.warning("Reranker not available") return nodes reranked = self.reranker.postprocess_nodes(nodes, query_str=query) return reranked def _query_agentic(self, question: str) -> Dict[str, any]: """ Execute agentic RAG with sub-question decomposition. Args: question: Query string Returns: dict: Answer, sub-questions, and sources """ if not self.sub_question_engine: logger.warning("Agentic engine not available, falling back to base query") return self._query_base(question) response = self.sub_question_engine.query(question) # Extract sub-questions if available sub_questions = [] if hasattr(response, "metadata") and response.metadata: sub_qa = response.metadata.get("sub_qa", []) sub_questions = [ {"question": item[0], "answer": item[1]} for item in sub_qa ] result = self._format_response(response, response.source_nodes) result["sub_questions"] = sub_questions return result def _query_base(self, question: str) -> Dict[str, any]: """ Execute base vector search query. Args: question: Query string Returns: dict: Answer and sources """ response = self.query_engine.query(question) return self._format_response(response, response.source_nodes) def _format_response(self, response, source_nodes) -> Dict[str, any]: """Format query response for return.""" sources = [] for node in source_nodes: sources.append({ "text": node.text, "metadata": node.metadata if hasattr(node, "metadata") else {}, "score": node.score if hasattr(node, "score") else None }) return { "answer": str(response), "sources": sources } def query( self, question: str, use_hybrid: bool = False, use_rerank: bool = False, use_agentic: bool = False ) -> Dict[str, any]: """ Execute query with specified strategies. Args: question: Query string use_hybrid: Enable hybrid search use_rerank: Enable reranking use_agentic: Enable agentic RAG Returns: dict: Answer and sources """ try: # Route to appropriate query method if use_agentic: return self._query_agentic(question) elif use_hybrid: result = self._query_hybrid(question) if use_rerank and result.get("sources"): # Note: Reranking after hybrid fusion pass # Already applied in hybrid return result else: result = self._query_base(question) if use_rerank and result.get("sources"): # Apply reranking to base results nodes = self.vector_retriever.retrieve(question) reranked_nodes = self._apply_reranking(nodes, question) synthesizer = get_response_synthesizer() response = synthesizer.synthesize(question, nodes=reranked_nodes) return self._format_response(response, reranked_nodes) return result except Exception as e: logger.error(f"Query failed: {e}") return { "answer": f"Error executing query: {str(e)}", "sources": [] }