Aigenthix/Graph_RAG

1"""RAG Mode Implementations"""2 3from typing import List, Dict, Any, Tuple4from abc import ABC, abstractmethod5import time6from groq import Groq7import logging8 9logger = logging.getLogger(__name__)10 11 12class RAGModeBase(ABC):13    """Abstract base class for RAG modes"""14 15    def __init__(self, groq_client: Groq):16        self.groq_client = groq_client17 18    @abstractmethod19    def execute(20        self,21        query: str,22        search_results: List[Dict[str, Any]],23        model_id: str,24        temperature: float = 0.7,25        max_tokens: int = 1024,26    ) -> Tuple[str, List[Dict[str, Any]]]:27        """Execute RAG mode and return answer and sources"""28        pass29 30 31class SimpleRAG(RAGModeBase):32    """Simple RAG: Direct retrieval + generation"""33 34    def execute(35        self,36        query: str,37        search_results: List[Dict[str, Any]],38        model_id: str,39        temperature: float = 0.7,40        max_tokens: int = 1024,41    ) -> Tuple[str, List[Dict[str, Any]]]:42        """Execute simple RAG pipeline"""43 44        step_start = time.time()45 46        # Assemble context from search results47        context_parts = []48        used_results = []49 50        for result in search_results:51            content = result.get("content", "")52            if not content:53                continue54            context_parts.append(55                f"[Source: {result['metadata'].get('doc_id', 'Unknown')}]\n{content}"56            )57            used_results.append(result)58 59        context = "\n\n".join(context_parts)60 61        # Generate answer62        system_prompt = """You are a helpful assistant answering questions based on provided context.63        - Answer only using the provided context64        - If the context doesn't contain relevant information, say so clearly65        - Be concise and accurate"""66 67        user_prompt = f"""Context:68{context}69 70Question: {query}71 72Please provide a detailed answer based on the context above."""73 74        try:75            response = self.groq_client.create_chat_completion(76                model=model_id,77                messages=[78                    {"role": "system", "content": system_prompt},79                    {"role": "user", "content": user_prompt},80                ],81                temperature=temperature,82                max_tokens=max_tokens,83            )84 85            answer = response["choices"][0]["message"]["content"]86            duration = (time.time() - step_start) * 100087 88            logger.info(f"Simple RAG completed in {duration:.2f}ms")89            return answer, used_results90 91        except Exception as e:92            logger.error(f"Simple RAG generation failed: {e}")93            raise94 95 96class AgenticRAG(RAGModeBase):97    """Agentic RAG: Multi-step reasoning with tool use"""98 99    def execute(100        self,101        query: str,102        search_results: List[Dict[str, Any]],103        model_id: str,104        temperature: float = 0.7,105        max_tokens: int = 1024,106    ) -> Tuple[str, List[Dict[str, Any]]]:107        """Execute agentic RAG pipeline with multi-step reasoning"""108 109        step_start = time.time()110 111        # Step 1: Analyze query and plan112        analysis_prompt = f"""Analyze the user's query and create a step-by-step plan to answer it.113Query: {query}114 115Available documents: {len(search_results)} retrieved documents116 117Create a brief plan (2-3 steps) for answering this query."""118 119        try:120            analysis_response = self.groq_client.create_chat_completion(121                model=model_id,122                messages=[123                    {"role": "system", "content": "You are an analytical assistant that plans query resolution."},124                    {"role": "user", "content": analysis_prompt},125                ],126                temperature=temperature,127                max_tokens=256,128            )129 130            plan = analysis_response["choices"][0]["message"]["content"]131 132            # Step 2: Gather context with analysis133            context_parts = []134            used_results = []135 136            for i, result in enumerate(search_results[:3]):  # Use top 3 for agentic137                content = result.get("content", "")138                if not content:139                    continue140                context_parts.append(141                    f"[Document {i+1}: {result['metadata'].get('doc_id', 'Unknown')}]\n{content}"142                )143                used_results.append(result)144 145            context = "\n\n".join(context_parts)146 147            # Step 3: Generate comprehensive answer with reasoning148            reasoning_prompt = f"""You are solving a user query through multi-step reasoning.149 150User Query: {query}151 152Your Analysis Plan:153{plan}154 155Available Context:156{context}157 158Now, follow your plan and provide a comprehensive answer with clear reasoning at each step.159Think step-by-step and cite sources."""160 161            final_response = self.groq_client.create_chat_completion(162                model=model_id,163                messages=[164                    {"role": "system", "content": "You are an expert reasoning assistant."},165                    {"role": "user", "content": reasoning_prompt},166                ],167                temperature=temperature,168                max_tokens=max_tokens,169            )170 171            answer = final_response["choices"][0]["message"]["content"]172            duration = (time.time() - step_start) * 1000173 174            logger.info(f"Agentic RAG completed in {duration:.2f}ms")175            return answer, used_results176 177        except Exception as e:178            logger.error(f"Agentic RAG generation failed: {e}")179            raise180 181 182class GraphRAG(RAGModeBase):183    """Graph RAG: Knowledge graph-based retrieval"""184 185    def execute(186        self,187        query: str,188        search_results: List[Dict[str, Any]],189        model_id: str,190        temperature: float = 0.7,191        max_tokens: int = 1024,192    ) -> Tuple[str, List[Dict[str, Any]]]:193        """Execute graph-based RAG pipeline"""194 195        step_start = time.time()196 197        # Step 1: Extract key concepts from query198        concept_prompt = f"""Extract key entities and concepts from this query:199Query: {query}200 201List 3-5 key concepts or entities mentioned or implied in the query."""202 203        try:204            concept_response = self.groq_client.create_chat_completion(205                model=model_id,206                messages=[207                    {"role": "system", "content": "You are an NLP expert extracting entities."},208                    {"role": "user", "content": concept_prompt},209                ],210                temperature=0.3,  # Lower temperature for extraction211                max_tokens=128,212            )213 214            concepts = concept_response["choices"][0]["message"]["content"]215 216            # Step 2: Find related documents based on concepts217            context_parts = []218            used_results = []219 220            # Group results by relevance221            high_relevance = []222            medium_relevance = []223            low_relevance = []224 225            for result in search_results:226                score = result.get("similarity_score", 0)227                if score > 0.8:228                    high_relevance.append(result)229                elif score > 0.6:230                    medium_relevance.append(result)231                else:232                    low_relevance.append(result)233 234            # Use high relevance results first235            for i, result in enumerate(high_relevance[:2]):236                content = result.get("content", "")237                context_parts.append(238                    f"[High Relevance - {result['metadata'].get('doc_id', 'Unknown')}]\n{content}"239                )240                used_results.append(result)241 242            # Add medium relevance if needed243            for i, result in enumerate(medium_relevance[:1]):244                content = result.get("content", "")245                context_parts.append(246                    f"[Medium Relevance - {result['metadata'].get('doc_id', 'Unknown')}]\n{content}"247                )248                used_results.append(result)249 250            context = "\n\n".join(context_parts)251 252            # Step 3: Generate graph-aware answer253            graph_prompt = f"""You are analyzing information from a knowledge graph perspective.254 255Key Concepts: {concepts}256 257Query: {query}258 259Relevant Context (organized by relevance):260{context}261 262Generate an answer that:2631. Shows relationships between concepts2642. Uses the most relevant sources first2653. Builds a coherent understanding2664. Cites all sources used"""267 268            final_response = self.groq_client.create_chat_completion(269                model=model_id,270                messages=[271                    {"role": "system", "content": "You are a graph-aware knowledge synthesis expert."},272                    {"role": "user", "content": graph_prompt},273                ],274                temperature=temperature,275                max_tokens=max_tokens,276            )277 278            answer = final_response["choices"][0]["message"]["content"]279            duration = (time.time() - step_start) * 1000280 281            logger.info(f"Graph RAG completed in {duration:.2f}ms")282            return answer, used_results283 284        except Exception as e:285            logger.error(f"Graph RAG generation failed: {e}")286            raise287 288 289class RAGModeFactory:290    """Factory for creating RAG mode instances"""291 292    _modes = {293        "simple": SimpleRAG,294        "agentic": AgenticRAG,295        "graph": GraphRAG,296    }297 298    @staticmethod299    def create(mode: str, groq_client: Groq) -> RAGModeBase:300        """Create a RAG mode instance"""301        if mode not in RAGModeFactory._modes:302            available = ", ".join(RAGModeFactory._modes.keys())303            raise ValueError(f"Unknown RAG mode '{mode}'. Available: {available}")304 305        mode_class = RAGModeFactory._modes[mode]306        return mode_class(groq_client)307 308    @staticmethod309    def available_modes() -> List[str]:310        """Get list of available modes"""311        return list(RAGModeFactory._modes.keys())312