add multi-agent system

multi_agent_chatbot.py

@@ -0,0 +1,1167 @@
+"""
+Multi-Agent RAG Chatbot using LangGraph
+
+This system implements a 3-agent architecture:
+1. Main Agent: Handles conversation flow, follow-ups, and determines when to call RAG
+2. RAG Agent: Rewrites queries and applies filters for document retrieval
+3. Response Agent: Generates final answers from retrieved documents
+
+Each agent has specialized prompts and responsibilities.
+"""
+
+import os
+import json
+import time
+import logging
+from pathlib import Path
+from datetime import datetime
+from dataclasses import dataclass
+from typing import Dict, List, Any, Optional, TypedDict
+
+
+import re
+from langchain_core.tools import tool
+from langgraph.graph import StateGraph, END
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+from langchain_core.prompts import ChatPromptTemplate
+
+
+from src.pipeline import PipelineManager
+from src.config.loader import load_config
+from src.llm.adapters import get_llm_client
+
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class QueryContext:
+    """Context extracted from conversation"""
+    has_district: bool = False
+    has_source: bool = False
+    has_year: bool = False
+    extracted_district: Optional[str] = None
+    extracted_source: Optional[str] = None
+    extracted_year: Optional[str] = None
+    ui_filters: Dict[str, List[str]] = None
+    confidence_score: float = 0.0
+    needs_follow_up: bool = False
+    follow_up_question: Optional[str] = None
+
+class MultiAgentState(TypedDict):
+    """State for the multi-agent conversation flow"""
+    conversation_id: str
+    messages: List[Any]
+    current_query: str
+    query_context: Optional[QueryContext]
+    rag_query: Optional[str]
+    rag_filters: Optional[Dict[str, Any]]
+    retrieved_documents: Optional[List[Any]]
+    final_response: Optional[str]
+    agent_logs: List[str]
+    conversation_context: Dict[str, Any]
+    session_start_time: float
+    last_ai_message_time: float
+
+class MultiAgentRAGChatbot:
+    """Multi-agent RAG chatbot with specialized agents"""
+    
+    def __init__(self, config_path: str = "auditqa/config/settings.yaml"):
+        """Initialize the multi-agent chatbot"""
+        self.config = load_config(config_path)
+        
+        # Get LLM provider from config
+        reader_config = self.config.get("reader", {})
+        default_type = reader_config.get("default_type", "INF_PROVIDERS")
+        provider_name = default_type.lower()
+        
+        self.llm_adapter = get_llm_client(provider_name, self.config)
+        
+        # Create a simple wrapper for LangChain compatibility
+        class LLMWrapper:
+            def __init__(self, adapter):
+                self.adapter = adapter
+            
+            def invoke(self, messages):
+                # Convert LangChain messages to the format expected by the adapter
+                if isinstance(messages, list):
+                    formatted_messages = []
+                    for msg in messages:
+                        if hasattr(msg, 'content'):
+                            role = "user" if msg.__class__.__name__ == "HumanMessage" else "assistant"
+                            formatted_messages.append({"role": role, "content": msg.content})
+                        else:
+                            formatted_messages.append({"role": "user", "content": str(msg)})
+                else:
+                    formatted_messages = [{"role": "user", "content": str(messages)}]
+                
+                # Use the adapter to get response
+                response = self.adapter.generate(formatted_messages)
+                
+                # Return a mock response object
+                class MockResponse:
+                    def __init__(self, content):
+                        self.content = content
+                
+                return MockResponse(response.content)
+        
+        self.llm = LLMWrapper(self.llm_adapter)
+        
+        # Initialize pipeline manager early to load models
+        logger.info("🔄 Initializing pipeline manager and loading models...")
+        self.pipeline_manager = PipelineManager(self.config)
+        logger.info("✅ Pipeline manager initialized and models loaded")
+        
+        # Connect to vector store
+        logger.info("🔄 Connecting to vector store...")
+        if not self.pipeline_manager.connect_vectorstore():
+            logger.error("❌ Failed to connect to vector store")
+            raise RuntimeError("Vector store connection failed")
+        logger.info("✅ Vector store connected successfully")
+        
+        # Load dynamic data
+        self._load_dynamic_data()
+        
+        # Build the multi-agent graph
+        self.graph = self._build_graph()
+        
+        # Conversations directory
+        self.conversations_dir = Path("conversations")
+        self.conversations_dir.mkdir(exist_ok=True)
+        
+        logger.info("🤖 Multi-Agent RAG Chatbot initialized")
+    
+    def _load_dynamic_data(self):
+        """Load dynamic data from filter_options.json and add_district_metadata.py"""
+        # Load filter options
+        try:
+            fo = Path("filter_options.json")
+            if fo.exists():
+                with open(fo) as f:
+                    data = json.load(f)
+                    self.year_whitelist = [str(y).strip() for y in data.get("years", [])]
+                    self.source_whitelist = [str(s).strip() for s in data.get("sources", [])]
+                    self.district_whitelist = [str(d).strip() for d in data.get("districts", [])]
+            else:
+                # Fallback to default values
+                self.year_whitelist = ['2018', '2019', '2020', '2021', '2022', '2023', '2024']
+                self.source_whitelist = ['Consolidated', 'Local Government', 'Ministry, Department and Agency']
+                self.district_whitelist = ['Kampala', 'Gulu', 'Kalangala']
+        except Exception as e:
+            logger.warning(f"Could not load filter options: {e}")
+            self.year_whitelist = ['2018', '2019', '2020', '2021', '2022', '2023', '2024']
+            self.source_whitelist = ['Consolidated', 'Local Government', 'Ministry, Department and Agency']
+            self.district_whitelist = ['Kampala', 'Gulu', 'Kalangala']
+        
+        # Enrich district list from add_district_metadata.py
+        try:
+            from add_district_metadata import DistrictMetadataProcessor
+            proc = DistrictMetadataProcessor()
+            names = set()
+            for key, mapping in proc.district_mappings.items():
+                if getattr(mapping, 'is_district', True):
+                    names.add(mapping.name)
+            if names:
+                merged = list(self.district_whitelist)
+                for n in sorted(names):
+                    if n not in merged:
+                        merged.append(n)
+                self.district_whitelist = merged
+                logger.info(f"🧭 District whitelist enriched: {len(self.district_whitelist)} entries")
+        except Exception as e:
+            logger.info(f"ℹ️ Could not enrich districts: {e}")
+        
+        # Calculate current year dynamically
+        self.current_year = str(datetime.now().year)
+        self.previous_year = str(datetime.now().year - 1)
+        
+        # Log the actual filter values for debugging
+        logger.info(f"📊 ACTUAL FILTER VALUES:")
+        logger.info(f"   Years: {self.year_whitelist}")
+        logger.info(f"   Sources: {self.source_whitelist}")
+        logger.info(f"   Districts: {len(self.district_whitelist)} districts (first 10: {self.district_whitelist[:10]})")
+    
+    def _build_graph(self) -> StateGraph:
+        """Build the multi-agent LangGraph"""
+        graph = StateGraph(MultiAgentState)
+        
+        # Add nodes for each agent
+        graph.add_node("main_agent", self._main_agent)
+        graph.add_node("rag_agent", self._rag_agent)
+        graph.add_node("response_agent", self._response_agent)
+        
+        # Define the flow
+        graph.set_entry_point("main_agent")
+        
+        # Main agent decides next step
+        graph.add_conditional_edges(
+            "main_agent",
+            self._should_call_rag,
+            {
+                "follow_up": END,
+                "call_rag": "rag_agent"
+            }
+        )
+        
+        # RAG agent calls response agent
+        graph.add_edge("rag_agent", "response_agent")
+        
+        # Response agent returns to main agent for potential follow-ups
+        graph.add_edge("response_agent", "main_agent")
+        
+        return graph.compile()
+    
+    def _should_call_rag(self, state: MultiAgentState) -> str:
+        """Determine if we should call RAG or ask follow-up"""
+        # If we already have a final response (from response agent), end
+        if state.get("final_response"):
+            return "follow_up"
+        
+        context = state["query_context"]
+        if context and context.needs_follow_up:
+            return "follow_up"
+        return "call_rag"
+    
+    def _main_agent(self, state: MultiAgentState) -> MultiAgentState:
+        """Main Agent: Handles conversation flow and follow-ups"""
+        logger.info("🎯 MAIN AGENT: Starting analysis")
+        
+        # If we already have a final response from response agent, end gracefully
+        if state.get("final_response"):
+            logger.info("🎯 MAIN AGENT: Final response already exists, ending conversation flow")
+            return state
+        
+        query = state["current_query"]
+        messages = state["messages"]
+        
+        logger.info(f"🎯 MAIN AGENT: Extracting UI filters from query")
+        ui_filters = self._extract_ui_filters(query)
+        logger.info(f"🎯 MAIN AGENT: UI filters extracted: {ui_filters}")
+        
+        # Analyze query context
+        logger.info(f"🎯 MAIN AGENT: Analyzing query context")
+        context = self._analyze_query_context(query, messages, ui_filters)
+        
+        # Log agent decision
+        state["agent_logs"].append(f"MAIN AGENT: Context analyzed - district={context.has_district}, source={context.has_source}, year={context.has_year}")
+        logger.info(f"🎯 MAIN AGENT: Context analysis complete - district={context.has_district}, source={context.has_source}, year={context.has_year}")
+        
+        # Store context
+        state["query_context"] = context
+        
+        # If follow-up needed, generate response
+        if context.needs_follow_up:
+            logger.info(f"🎯 MAIN AGENT: Follow-up needed, generating question")
+            response = context.follow_up_question
+            state["final_response"] = response
+            state["last_ai_message_time"] = time.time()
+            logger.info(f"🎯 MAIN AGENT: Follow-up question generated: {response[:100]}...")
+        else:
+            logger.info("🎯 MAIN AGENT: No follow-up needed, proceeding to RAG")
+        
+        return state
+    
+    def _rag_agent(self, state: MultiAgentState) -> MultiAgentState:
+        """RAG Agent: Rewrites queries and applies filters"""
+        logger.info("🔍 RAG AGENT: Starting query rewriting and filter preparation")
+        
+        context = state["query_context"]
+        messages = state["messages"]
+        
+        logger.info(f"🔍 RAG AGENT: Context received - district={context.has_district}, source={context.has_source}, year={context.has_year}")
+        
+        # Rewrite query for RAG
+        logger.info(f"🔍 RAG AGENT: Rewriting query for optimal retrieval")
+        rag_query = self._rewrite_query_for_rag(messages, context)
+        logger.info(f"🔍 RAG AGENT: Query rewritten: '{rag_query}'")
+        
+        # Build filters
+        logger.info(f"🔍 RAG AGENT: Building filters from context")
+        filters = self._build_filters(context)
+        logger.info(f"🔍 RAG AGENT: Filters built: {filters}")
+        
+        # Log RAG preparation
+        state["agent_logs"].append(f"RAG AGENT: Query='{rag_query}', Filters={filters}")
+        
+        # Store for response agent
+        state["rag_query"] = rag_query
+        state["rag_filters"] = filters
+        
+        logger.info(f"🔍 RAG AGENT: Preparation complete, ready for retrieval")
+        
+        return state
+    
+    def _response_agent(self, state: MultiAgentState) -> MultiAgentState:
+        """Response Agent: Generates final answer from retrieved documents"""
+        logger.info("📝 RESPONSE AGENT: Starting document retrieval and answer generation")
+        
+        rag_query = state["rag_query"]
+        filters = state["rag_filters"]
+        
+        logger.info(f"📝 RESPONSE AGENT: Starting RAG retrieval with query: '{rag_query}'")
+        logger.info(f"📝 RESPONSE AGENT: Using filters: {filters}")
+        
+        # Perform RAG retrieval
+        logger.info(f"📝 RESPONSE AGENT: Calling pipeline manager for retrieval")
+        logger.info(f"🔍 ACTUAL RAG QUERY: '{rag_query}'")
+        logger.info(f"🔍 ACTUAL FILTERS: {filters}")
+        try:
+            # Extract filenames from filters if present
+            filenames = filters.get("filenames") if filters else None
+            
+            result = self.pipeline_manager.run(
+                query=rag_query,
+                sources=filters.get("sources") if filters else None,
+                auto_infer_filters=False,
+                filters=filters if filters else None
+            )
+            
+            logger.info(f"📝 RESPONSE AGENT: RAG retrieval completed - {len(result.sources)} documents retrieved")
+            logger.info(f"🔍 RETRIEVAL DEBUG: Result type: {type(result)}")
+            logger.info(f"🔍 RETRIEVAL DEBUG: Result sources type: {type(result.sources)}")
+            # logger.info(f"🔍 RETRIEVAL DEBUG: Result metadata: {getattr(result, 'metadata', 'No metadata')}")
+            
+            if len(result.sources) == 0:
+                logger.warning(f"⚠️ NO DOCUMENTS RETRIEVED: Query='{rag_query}', Filters={filters}")
+                logger.warning(f"⚠️ RETRIEVAL DEBUG: This could be due to:")
+                logger.warning(f"   - Query too specific for available documents")
+                logger.warning(f"   - Filters too restrictive")
+                logger.warning(f"   - Vector store connection issues")
+                logger.warning(f"   - Embedding model issues")
+            else:
+                logger.info(f"✅ DOCUMENTS RETRIEVED: {len(result.sources)} documents found")
+                for i, doc in enumerate(result.sources[:3]):  # Log first 3 docs
+                    logger.info(f"   Doc {i+1}: {getattr(doc, 'metadata', {}).get('filename', 'Unknown')[:50]}...")
+            
+            state["retrieved_documents"] = result.sources
+            state["agent_logs"].append(f"RESPONSE AGENT: Retrieved {len(result.sources)} documents")
+            
+            # Check highest similarity score
+            highest_score = 0.0
+            if result.sources:
+                # Check reranked_score first (more accurate), fallback to original_score
+                for doc in result.sources:
+                    score = doc.metadata.get('reranked_score') or doc.metadata.get('original_score', 0.0)
+                    if score > highest_score:
+                        highest_score = score
+            
+            logger.info(f"📝 RESPONSE AGENT: Highest similarity score: {highest_score:.4f}")
+            
+            # If highest score is too low, don't use retrieved documents
+            if highest_score <= 0.15:
+                logger.warning(f"⚠️ RESPONSE AGENT: Low similarity score ({highest_score:.4f} <= 0.15), using LLM knowledge only")
+                response = self._generate_conversational_response_without_docs(
+                    state["current_query"],
+                    state["messages"]
+                )
+            else:
+                # Generate conversational response with documents
+                logger.info(f"📝 RESPONSE AGENT: Generating conversational response from {len(result.sources)} documents")
+                response = self._generate_conversational_response(
+                    state["current_query"], 
+                    result.sources, 
+                    result.answer,
+                    state["messages"]
+                )
+            
+            logger.info(f"📝 RESPONSE AGENT: Response generated: {response[:100]}...")
+            
+            state["final_response"] = response
+            state["last_ai_message_time"] = time.time()
+            
+            logger.info(f"📝 RESPONSE AGENT: Answer generation complete")
+            
+        except Exception as e:
+            logger.error(f"❌ RESPONSE AGENT ERROR: {e}")
+            state["final_response"] = "I apologize, but I encountered an error while retrieving information. Please try again."
+            state["last_ai_message_time"] = time.time()
+        
+        return state
+    
+    def _extract_ui_filters(self, query: str) -> Dict[str, List[str]]:
+        """Extract UI filters from query"""
+        filters = {}
+        
+        # Look for FILTER CONTEXT in query
+        if "FILTER CONTEXT:" in query:
+            # Extract the entire filter section (until USER QUERY: or end of query)
+            filter_section = query.split("FILTER CONTEXT:")[1]
+            if "USER QUERY:" in filter_section:
+                filter_section = filter_section.split("USER QUERY:")[0]
+            filter_section = filter_section.strip()
+            
+            # Parse sources
+            if "Sources:" in filter_section:
+                sources_line = [line for line in filter_section.split('\n') if line.strip().startswith('Sources:')][0]
+                sources_str = sources_line.split("Sources:")[1].strip()
+                if sources_str and sources_str != "None":
+                    filters["sources"] = [s.strip() for s in sources_str.split(",")]
+            
+            # Parse years
+            if "Years:" in filter_section:
+                years_line = [line for line in filter_section.split('\n') if line.strip().startswith('Years:')][0]
+                years_str = years_line.split("Years:")[1].strip()
+                if years_str and years_str != "None":
+                    filters["years"] = [y.strip() for y in years_str.split(",")]
+            
+            # Parse districts
+            if "Districts:" in filter_section:
+                districts_line = [line for line in filter_section.split('\n') if line.strip().startswith('Districts:')][0]
+                districts_str = districts_line.split("Districts:")[1].strip()
+                if districts_str and districts_str != "None":
+                    filters["districts"] = [d.strip() for d in districts_str.split(",")]
+            
+            # Parse filenames
+            if "Filenames:" in filter_section:
+                filenames_line = [line for line in filter_section.split('\n') if line.strip().startswith('Filenames:')][0]
+                filenames_str = filenames_line.split("Filenames:")[1].strip()
+                if filenames_str and filenames_str != "None":
+                    filters["filenames"] = [f.strip() for f in filenames_str.split(",")]
+        
+        return filters
+    
+    def _analyze_query_context(self, query: str, messages: List[Any], ui_filters: Dict[str, List[str]]) -> QueryContext:
+        """Analyze query context using LLM"""
+        logger.info(f"🔍 QUERY ANALYSIS: '{query[:50]}...' | UI filters: {ui_filters} | Messages: {len(messages)}")
+        
+        # Build conversation context
+        conversation_context = ""
+        for i, msg in enumerate(messages[-6:]):  # Last 6 messages
+            if isinstance(msg, HumanMessage):
+                conversation_context += f"User: {msg.content}\n"
+            elif isinstance(msg, AIMessage):
+                conversation_context += f"Assistant: {msg.content}\n"
+        
+        # Create analysis prompt
+        analysis_prompt = ChatPromptTemplate.from_messages([
+            SystemMessage(content=f"""You are the Main Agent in an advanced multi-agent RAG system for audit report analysis.
+
+🎯 PRIMARY GOAL: Intelligently analyze user queries and determine the optimal conversation flow, whether that's answering directly, asking follow-ups, or proceeding to RAG retrieval.
+
+🧠 INTELLIGENCE LEVEL: You are a sophisticated conversational AI that can handle any type of user interaction - from greetings to complex audit queries.
+
+📊 YOUR EXPERTISE: You specialize in analyzing audit reports from various sources (Local Government, Ministry, Hospital, etc.) across different years and districts in Uganda.
+
+🔍 AVAILABLE FILTERS:
+- Years: {', '.join(self.year_whitelist)}
+- Current year: {self.current_year}, Previous year: {self.previous_year}
+- Sources: {', '.join(self.source_whitelist)}
+- Districts: {', '.join(self.district_whitelist[:50])}... (and {len(self.district_whitelist)-50} more)
+
+🎛️ UI FILTERS PROVIDED: {ui_filters}
+
+📋 UI FILTER HANDLING:
+- If UI filters contain multiple values (e.g., districts: ['Lwengo', 'Kiboga']), extract ALL values
+- For multiple districts: extract each district separately and validate each one
+- For multiple years: extract each year separately and validate each one  
+- For multiple sources: extract each source separately and validate each one
+- UI filters take PRIORITY over conversation context - use them first
+
+🧭 CONVERSATION FLOW INTELLIGENCE:
+
+1. **GREETINGS & GENERAL CHAT**: 
+   - If user greets you ("Hi", "Hello", "How are you"), respond warmly and guide them to audit-related questions
+   - Example: "Hello! I'm here to help you analyze audit reports. What would you like to know about budget allocations, expenditures, or audit findings?"
+
+2. **EDGE CASES**:
+   - Handle "What can you do?", "Help", "I don't know what to ask" with helpful guidance
+   - Example: "I can help you analyze audit reports! Try asking about budget allocations, salary management, PDM implementation, or any specific audit findings."
+
+3. **AUDIT QUERIES**:
+   - Extract ONLY values that EXACTLY match the available lists above
+   - DO NOT hallucinate or infer values not in the lists
+   - If user mentions "salary payroll management" - this is NOT a valid source filter
+   
+   **YEAR EXTRACTION**:
+   - If user mentions "2023" and it's in the years list - extract "2023"
+   - If user mentions "2022 / 23" - extract ["2022", "2023"] (as a JSON array)
+   - If user mentions "2022-2023" - extract ["2022", "2023"] (as a JSON array)
+   - If user mentions "latest couple of years" - extract the 2 most recent years from available data as JSON array
+   - Always return years as JSON arrays when multiple years are mentioned
+   
+   **DISTRICT EXTRACTION**:
+   - If user mentions "Kampala" and it's in the districts list - extract "Kampala"
+   - If user mentions "Pader District" - extract "Pader" (remove "District" suffix)
+   - If user mentions "Lwengo, Kiboga and Namutumba" - extract ["Lwengo", "Kiboga", "Namutumba"] (as JSON array)
+   - If user mentions "Lwengo District and Kiboga District" - extract ["Lwengo", "Kiboga"] (as JSON array, remove "District" suffix)
+   - Always return districts as JSON arrays when multiple districts are mentioned
+   - If no exact matches found, set extracted values to null
+
+4. **FILENAME FILTERING (MUTUALLY EXCLUSIVE)**:
+   - If UI provides filenames filter - ONLY use that, ignore all other filters (year, district, source)
+   - With filenames filter, no follow-ups needed - proceed directly to RAG
+   - When filenames are specified, skip filter inference entirely
+
+5. **HALLUCINATION PREVENTION**:
+   - If user asks about a specific report but NO filename is selected in UI and NONE is extracted from conversation - DO NOT hallucinate
+   - Clearly state: "I don't have any specific report selected. Could you please select a report from the list or tell me which report you'd like to analyze?"
+   - DO NOT pretend to know which report they mean
+   - DO NOT infer reports from context alone - only use explicitly mentioned reports
+
+6. **CONVERSATION CONTEXT AWARENESS**:
+   - ALWAYS consider the full conversation context when extracting filters
+   - If district was mentioned in previous messages, include it in current analysis
+   - If year was mentioned in previous messages, include it in current analysis
+   - If source was mentioned in previous messages, include it in current analysis
+   - Example: If conversation shows "User: Tell me about Pader District" then "User: 2023", extract both: district="Pader" and year="2023"
+
+5. **SMART FOLLOW-UP STRATEGY**:
+   - NEVER ask the same question twice in a row
+   - If user provides source info, ask for year or district next
+   - If user provides year info, ask for source or district next  
+   - If user provides district info, ask for year or source next
+   - If user provides 2+ pieces of info, proceed to RAG instead of asking more
+   - Make follow-ups conversational and contextual, not robotic
+
+5. **DYNAMIC FOLLOW-UP EXAMPLES**:
+   - Budget queries: "What year are you interested in?" or "Which department - Local Government or Ministry?"
+   - PDM queries: "Which district are you interested in?" or "What year?"
+   - General queries: "Could you be more specific about what you'd like to know?"
+
+🎯 DECISION LOGIC:
+- If query is a greeting/general chat → needs_follow_up: true, provide helpful guidance
+- If query has 2+ pieces of info → needs_follow_up: false, proceed to RAG
+- If query has 1 piece of info → needs_follow_up: true, ask for missing piece
+- If query has 0 pieces of info → needs_follow_up: true, ask for clarification
+
+RESPOND WITH JSON ONLY:
+{{
+  "has_district": boolean,
+  "has_source": boolean,
+  "has_year": boolean,
+  "extracted_district": "single district name or JSON array of districts or null",
+  "extracted_source": "single source name or JSON array of sources or null", 
+  "extracted_year": "single year or JSON array of years or null",
+  "confidence_score": 0.0-1.0,
+  "needs_follow_up": boolean,
+  "follow_up_question": "conversational question or helpful guidance or null"
+}}"""),
+            HumanMessage(content=f"""Query: {query}
+
+Conversation Context:
+{conversation_context}
+
+CRITICAL: You MUST analyze the FULL conversation context above, not just the current query. 
+- If ANY district was mentioned in previous messages, extract it
+- If ANY year was mentioned in previous messages, extract it  
+- If ANY source was mentioned in previous messages, extract it
+- Combine information from ALL messages in the conversation
+
+Analyze this query using ONLY the exact values provided above:""")
+        ])
+        
+        try:
+            response = self.llm.invoke(analysis_prompt.format_messages())
+            
+            # Clean the response to extract JSON
+            content = response.content.strip()
+            if content.startswith("```json"):
+                # Remove markdown formatting
+                content = content.replace("```json", "").replace("```", "").strip()
+            elif content.startswith("```"):
+                # Remove generic markdown formatting
+                content = content.replace("```", "").strip()
+            
+            # Clean and parse JSON with better error handling
+            try:
+                # Remove comments (// and /* */) from JSON
+                import re
+                # Remove single-line comments
+                content = re.sub(r'//.*?$', '', content, flags=re.MULTILINE)
+                # Remove multi-line comments
+                content = re.sub(r'/\*.*?\*/', '', content, flags=re.DOTALL)
+                
+                analysis = json.loads(content)
+                logger.info(f"🔍 QUERY ANALYSIS: ✅ Parsed successfully")
+            except json.JSONDecodeError as e:
+                logger.error(f"❌ JSON parsing failed: {e}")
+                logger.error(f"❌ Raw content: {content[:200]}...")
+                
+                # Try to extract JSON from text if embedded
+                import re
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    try:
+                        # Clean the extracted JSON
+                        cleaned_json = json_match.group()
+                        cleaned_json = re.sub(r'//.*?$', '', cleaned_json, flags=re.MULTILINE)
+                        cleaned_json = re.sub(r'/\*.*?\*/', '', cleaned_json, flags=re.DOTALL)
+                        analysis = json.loads(cleaned_json)
+                        logger.info(f"🔍 QUERY ANALYSIS: ✅ Extracted and cleaned JSON from text")
+                    except json.JSONDecodeError as e2:
+                        logger.error(f"❌ Failed to extract JSON from text: {e2}")
+                        # Return fallback context
+                        context = QueryContext(
+                            has_district=False,
+                            has_source=False,
+                            has_year=False,
+                            extracted_district=None,
+                            extracted_source=None,
+                            extracted_year=None,
+                            confidence_score=0.0,
+                            needs_follow_up=True,
+                            follow_up_question="I apologize, but I'm having trouble processing your request. Could you please rephrase it or ask for help?"
+                        )
+                        return context
+                else:
+                    # Return fallback context
+                    context = QueryContext(
+                        has_district=False,
+                        has_source=False,
+                        has_year=False,
+                        extracted_district=None,
+                        extracted_source=None,
+                        extracted_year=None,
+                        confidence_score=0.0,
+                        needs_follow_up=True,
+                        follow_up_question="I apologize, but I'm having trouble processing your request. Could you please rephrase it or ask for help?"
+                    )
+                    return context
+            
+            # Validate extracted values against whitelists
+            extracted_district = analysis.get("extracted_district")
+            extracted_source = analysis.get("extracted_source")
+            extracted_year = analysis.get("extracted_year")
+            
+            logger.info(f"🔍 QUERY ANALYSIS: Raw extracted values - district: {extracted_district}, source: {extracted_source}, year: {extracted_year}")
+            
+            # Validate district (handle both single values and arrays)
+            if extracted_district:
+                if isinstance(extracted_district, list):
+                    # Validate each district in the array
+                    valid_districts = []
+                    for district in extracted_district:
+                        if district in self.district_whitelist:
+                            valid_districts.append(district)
+                        else:
+                            # Try removing "District" suffix
+                            district_name = district.replace(" District", "").replace(" district", "")
+                            if district_name in self.district_whitelist:
+                                valid_districts.append(district_name)
+                    
+                    if valid_districts:
+                        extracted_district = valid_districts[0] if len(valid_districts) == 1 else valid_districts
+                        logger.info(f"🔍 QUERY ANALYSIS: Extracted districts: {extracted_district}")
+                    else:
+                        logger.warning(f"⚠️ No valid districts found in: '{extracted_district}'")
+                        extracted_district = None
+                else:
+                    # Single district validation
+                    if extracted_district not in self.district_whitelist:
+                        # Try removing "District" suffix
+                        district_name = extracted_district.replace(" District", "").replace(" district", "")
+                        if district_name in self.district_whitelist:
+                            logger.info(f"🔍 QUERY ANALYSIS: Normalized district '{extracted_district}' to '{district_name}'")
+                            extracted_district = district_name
+                        else:
+                            logger.warning(f"⚠️ Invalid district extracted: '{extracted_district}' not in whitelist")
+                            extracted_district = None
+            
+            # Validate source (handle both single values and arrays)
+            if extracted_source:
+                if isinstance(extracted_source, list):
+                    # Validate each source in the array
+                    valid_sources = []
+                    for source in extracted_source:
+                        if source in self.source_whitelist:
+                            valid_sources.append(source)
+                        else:
+                            logger.warning(f"⚠️ Invalid source in array: '{source}' not in whitelist")
+                    
+                    if valid_sources:
+                        extracted_source = valid_sources[0] if len(valid_sources) == 1 else valid_sources
+                        logger.info(f"🔍 QUERY ANALYSIS: Extracted sources: {extracted_source}")
+                    else:
+                        logger.warning(f"⚠️ No valid sources found in: '{extracted_source}'")
+                        extracted_source = None
+                else:
+                    # Single source validation
+                    if extracted_source not in self.source_whitelist:
+                        logger.warning(f"⚠️ Invalid source extracted: '{extracted_source}' not in whitelist")
+                        extracted_source = None
+            
+            # Validate year (handle both single values and arrays)
+            if extracted_year:
+                if isinstance(extracted_year, list):
+                    # Validate each year in the array
+                    valid_years = []
+                    for year in extracted_year:
+                        year_str = str(year)
+                        if year_str in self.year_whitelist:
+                            valid_years.append(year_str)
+                    
+                    if valid_years:
+                        extracted_year = valid_years[0] if len(valid_years) == 1 else valid_years
+                        logger.info(f"🔍 QUERY ANALYSIS: Extracted years: {extracted_year}")
+                    else:
+                        logger.warning(f"⚠️ No valid years found in: '{extracted_year}'")
+                        extracted_year = None
+                else:
+                    # Single year validation
+                    year_str = str(extracted_year)
+                    if year_str not in self.year_whitelist:
+                        logger.warning(f"⚠️ Invalid year extracted: '{extracted_year}' not in whitelist")
+                        extracted_year = None
+                    else:
+                        extracted_year = year_str
+            
+            logger.info(f"🔍 QUERY ANALYSIS: Validated values - district: {extracted_district}, source: {extracted_source}, year: {extracted_year}")
+            
+            # Create QueryContext object
+            context = QueryContext(
+                has_district=bool(extracted_district),
+                has_source=bool(extracted_source),
+                has_year=bool(extracted_year),
+                extracted_district=extracted_district,
+                extracted_source=extracted_source,
+                extracted_year=extracted_year,
+                ui_filters=ui_filters,
+                confidence_score=analysis.get("confidence_score", 0.0),
+                needs_follow_up=analysis.get("needs_follow_up", False),
+                follow_up_question=analysis.get("follow_up_question")
+            )
+            
+            logger.info(f"🔍 QUERY ANALYSIS: Analysis complete - needs_follow_up: {context.needs_follow_up}, confidence: {context.confidence_score}")
+            
+            # If filenames are provided in UI, skip follow-ups and proceed to RAG
+            if ui_filters and ui_filters.get("filenames"):
+                logger.info(f"🔍 QUERY ANALYSIS: Filenames provided, skipping follow-ups, proceeding to RAG")
+                context.needs_follow_up = False
+                context.follow_up_question = None
+            
+            # Additional smart decision logic
+            if context.needs_follow_up:
+                # Check if we have enough information to proceed
+                info_count = sum([
+                    bool(context.extracted_district),
+                    bool(context.extracted_source), 
+                    bool(context.extracted_year)
+                ])
+                
+                # Check if user is asking for more info vs providing it
+                query_lower = query.lower()
+                is_requesting_info = any(phrase in query_lower for phrase in [
+                    "please provide", "could you provide", "can you provide", 
+                    "what is", "what are", "how much", "which", "what year",
+                    "what district", "what source", "tell me about"
+                ])
+                
+                # If we have 2+ pieces of info AND user is not requesting more info, proceed to RAG
+                if info_count >= 2 and not is_requesting_info:
+                    logger.info(f"🔍 QUERY ANALYSIS: Smart override - have {info_count} pieces of info and user not requesting more, proceeding to RAG")
+                    context.needs_follow_up = False
+                    context.follow_up_question = None
+                elif info_count >= 2 and is_requesting_info:
+                    logger.info(f"🔍 QUERY ANALYSIS: User requesting more info despite having {info_count} pieces, proceeding to RAG with comprehensive answer")
+                    context.needs_follow_up = False
+                    context.follow_up_question = None
+            
+            return context
+            
+        except Exception as e:
+            logger.error(f"❌ Query analysis failed: {e}")
+            # Fallback: proceed with RAG
+            return QueryContext(
+                has_district=bool(ui_filters.get("districts")),
+                has_source=bool(ui_filters.get("sources")),
+                has_year=bool(ui_filters.get("years")),
+                ui_filters=ui_filters,
+                confidence_score=0.5,
+                needs_follow_up=False
+            )
+    
+    def _rewrite_query_for_rag(self, messages: List[Any], context: QueryContext) -> str:
+        """Rewrite query for optimal RAG retrieval"""
+        logger.info("🔄 QUERY REWRITING: Starting query rewrite for RAG")
+        logger.info(f"🔄 QUERY REWRITING: Processing {len(messages)} messages")
+        
+        # Build conversation context
+        logger.info(f"🔄 QUERY REWRITING: Building conversation context from last 6 messages")
+        conversation_lines = []
+        for i, msg in enumerate(messages[-6:]):
+            if isinstance(msg, HumanMessage):
+                conversation_lines.append(f"User: {msg.content}")
+                logger.info(f"🔄 QUERY REWRITING: Message {i+1}: User - {msg.content[:50]}...")
+            elif isinstance(msg, AIMessage):
+                conversation_lines.append(f"Assistant: {msg.content}")
+                logger.info(f"🔄 QUERY REWRITING: Message {i+1}: Assistant - {msg.content[:50]}...")
+        
+        convo_text = "\n".join(conversation_lines)
+        logger.info(f"🔄 QUERY REWRITING: Conversation context built ({len(convo_text)} chars)")
+        
+        # Create rewrite prompt
+        rewrite_prompt = ChatPromptTemplate.from_messages([
+            SystemMessage(content=f"""You are a query rewriter for RAG retrieval.
+
+GOAL: Create the best possible search query for document retrieval.
+
+CRITICAL RULES:
+1. Focus on the core information need from the conversation
+2. Remove meta-verbs like "summarize", "list", "compare", "how much", "what" - keep the content focus
+3. DO NOT include filter details (years, districts, sources) - these are applied separately as filters
+4. DO NOT include specific years, district names, or source types in the query
+5. Output ONE clear sentence suitable for vector search
+6. Keep it generic and focused on the topic/subject matter
+
+EXAMPLES:
+- "What are the top challenges in budget allocation?" → "budget allocation challenges"
+- "How were PDM administrative costs utilized in 2023?" → "PDM administrative costs utilization"
+- "Compare salary management across districts" → "salary management"
+- "How much was budget allocation for Local Government in 2023?" → "budget allocation"
+
+OUTPUT FORMAT:
+Provide your response in this exact format:
+
+EXPLANATION: [Your reasoning here]
+QUERY: [One clean sentence for retrieval]
+
+The QUERY line will be extracted and used directly for RAG retrieval."""),
+            HumanMessage(content=f"""Conversation:
+{convo_text}
+
+Rewrite the best retrieval query:""")
+        ])
+        
+        try:
+            logger.info(f"🔄 QUERY REWRITING: Calling LLM for query rewrite")
+            response = self.llm.invoke(rewrite_prompt.format_messages())
+            logger.info(f"🔄 QUERY REWRITING: LLM response received: {response.content[:100]}...")
+            
+            rewritten = response.content.strip()
+            
+            # Extract only the QUERY line from the structured response
+            lines = rewritten.split('\n')
+            query_line = None
+            for line in lines:
+                if line.strip().startswith('QUERY:'):
+                    query_line = line.replace('QUERY:', '').strip()
+                    break
+            
+            if query_line and len(query_line) > 5:
+                logger.info(f"🔄 QUERY REWRITING: Query rewritten successfully: '{query_line[:50]}...'")
+                return query_line
+            else:
+                logger.info(f"🔄 QUERY REWRITING: No QUERY line found or too short, using fallback")
+                # Fallback to last user message
+                for msg in reversed(messages):
+                    if isinstance(msg, HumanMessage):
+                        logger.info(f"🔄 QUERY REWRITING: Using fallback message: '{msg.content[:50]}...'")
+                        return msg.content
+                logger.info(f"🔄 QUERY REWRITING: Using default fallback")
+                return "audit report information"
+                
+        except Exception as e:
+            logger.error(f"❌ QUERY REWRITING: Error during rewrite: {e}")
+            # Fallback
+            for msg in reversed(messages):
+                if isinstance(msg, HumanMessage):
+                    logger.info(f"🔄 QUERY REWRITING: Using error fallback message: '{msg.content[:50]}...'")
+                    return msg.content
+            logger.info(f"🔄 QUERY REWRITING: Using default error fallback")
+            return "audit report information"
+    
+    def _build_filters(self, context: QueryContext) -> Dict[str, Any]:
+        """Build filters for RAG retrieval"""
+        logger.info("🔧 FILTER BUILDING: Starting filter construction")
+        filters = {}
+        
+        # Check for filename filtering first (mutually exclusive)
+        if context.ui_filters and context.ui_filters.get("filenames"):
+            logger.info(f"🔧 FILTER BUILDING: Filename filtering requested (mutually exclusive mode)")
+            filters["filenames"] = context.ui_filters["filenames"]
+            logger.info(f"🔧 FILTER BUILDING: Added filenames filter: {context.ui_filters['filenames']}")
+            logger.info(f"🔧 FILTER BUILDING: Final filters: {filters}")
+            return filters  # Return early, skip all other filters
+        
+        # UI filters take priority, but merge with extracted context if UI filters are incomplete
+        if context.ui_filters:
+            logger.info(f"🔧 FILTER BUILDING: UI filters present: {context.ui_filters}")
+            
+            # Add UI filters first
+            if context.ui_filters.get("sources"):
+                filters["sources"] = context.ui_filters["sources"]
+                logger.info(f"🔧 FILTER BUILDING: Added sources filter from UI: {context.ui_filters['sources']}")
+            
+            if context.ui_filters.get("years"):
+                filters["year"] = context.ui_filters["years"]
+                logger.info(f"🔧 FILTER BUILDING: Added years filter from UI: {context.ui_filters['years']}")
+            
+            if context.ui_filters.get("districts"):
+                # Normalize district names to title case (match Qdrant metadata format)
+                normalized_districts = [d.title() for d in context.ui_filters['districts']]
+                filters["district"] = normalized_districts
+                logger.info(f"🔧 FILTER BUILDING: Added districts filter from UI: {context.ui_filters['districts']} → normalized: {normalized_districts}")
+            
+            # Merge with extracted context for missing filters
+            if not filters.get("year") and context.extracted_year:
+                # Handle both single values and arrays
+                if isinstance(context.extracted_year, list):
+                    filters["year"] = context.extracted_year
+                else:
+                    filters["year"] = [context.extracted_year]
+                logger.info(f"🔧 FILTER BUILDING: Added extracted year filter (UI missing): {context.extracted_year}")
+            
+            if not filters.get("district") and context.extracted_district:
+                # Handle both single values and arrays
+                if isinstance(context.extracted_district, list):
+                    # Normalize district names to title case (match Qdrant metadata format)
+                    normalized = [d.title() for d in context.extracted_district]
+                    filters["district"] = normalized
+                else:
+                    filters["district"] = [context.extracted_district.title()]
+                logger.info(f"🔧 FILTER BUILDING: Added extracted district filter (UI missing): {context.extracted_district}")
+            
+            if not filters.get("sources") and context.extracted_source:
+                # Handle both single values and arrays
+                if isinstance(context.extracted_source, list):
+                    filters["sources"] = context.extracted_source
+                else:
+                    filters["sources"] = [context.extracted_source]
+                logger.info(f"🔧 FILTER BUILDING: Added extracted source filter (UI missing): {context.extracted_source}")
+        else:
+            logger.info(f"🔧 FILTER BUILDING: No UI filters, using extracted context")
+            # Use extracted context
+            if context.extracted_source:
+                # Handle both single values and arrays
+                if isinstance(context.extracted_source, list):
+                    filters["sources"] = context.extracted_source
+                else:
+                    filters["sources"] = [context.extracted_source]
+                logger.info(f"🔧 FILTER BUILDING: Added extracted source filter: {context.extracted_source}")
+            
+            if context.extracted_year:
+                # Handle both single values and arrays
+                if isinstance(context.extracted_year, list):
+                    filters["year"] = context.extracted_year
+                else:
+                    filters["year"] = [context.extracted_year]
+                logger.info(f"🔧 FILTER BUILDING: Added extracted year filter: {context.extracted_year}")
+            
+            if context.extracted_district:
+                # Handle both single values and arrays
+                if isinstance(context.extracted_district, list):
+                    filters["district"] = context.extracted_district
+                else:
+                    filters["district"] = [context.extracted_district]
+                logger.info(f"🔧 FILTER BUILDING: Added extracted district filter: {context.extracted_district}")
+        
+        logger.info(f"🔧 FILTER BUILDING: Final filters: {filters}")
+        return filters
+    
+    def _generate_conversational_response(self, query: str, documents: List[Any], rag_answer: str, messages: List[Any]) -> str:
+        """Generate conversational response from RAG results"""
+        logger.info("💬 RESPONSE GENERATION: Starting conversational response generation")
+        logger.info(f"💬 RESPONSE GENERATION: Processing {len(documents)} documents")
+        logger.info(f"💬 RESPONSE GENERATION: Query: '{query[:50]}...'")
+        
+        # Create response prompt
+        logger.info(f"💬 RESPONSE GENERATION: Building response prompt")
+        response_prompt = ChatPromptTemplate.from_messages([
+            SystemMessage(content="""You are a helpful audit report assistant. Generate a natural, conversational response.
+
+RULES:
+1. Answer the user's question directly and clearly
+2. Use the retrieved documents as evidence
+3. Be conversational, not technical
+4. Don't mention scores, retrieval details, or technical implementation
+5. If relevant documents were found, reference them naturally
+6. If no relevant documents, explain based on your knowledge (if you have it) or just say you do not have enough information.
+7. If the passages have useful facts or numbers, use them in your answer.
+8. When you use information from a passage, mention where it came from by using [Doc i] at the end of the sentence. i stands for the number of the document.
+9. Do not use the sentence 'Doc i says ...' to say where information came from.
+10. If the same thing is said in more than one document, you can mention all of them like this: [Doc i, Doc j, Doc k]
+11. Do not just summarize each passage one by one. Group your summaries to highlight the key parts in the explanation.
+12. If it makes sense, use bullet points and lists to make your answers easier to understand.
+13. You do not need to use every passage. Only use the ones that help answer the question.
+14. If the documents do not have the information needed to answer the question, just say you do not have enough information.
+
+
+TONE: Professional but friendly, like talking to a colleague."""),
+            HumanMessage(content=f"""User Question: {query}
+
+Retrieved Documents: {len(documents)} documents found
+
+RAG Answer: {rag_answer}
+
+Generate a conversational response:""")
+        ])
+        
+        try:
+            logger.info(f"💬 RESPONSE GENERATION: Calling LLM for final response")
+            response = self.llm.invoke(response_prompt.format_messages())
+            logger.info(f"💬 RESPONSE GENERATION: LLM response received: {response.content[:100]}...")
+            return response.content.strip()
+        except Exception as e:
+            logger.error(f"❌ RESPONSE GENERATION: Error during generation: {e}")
+            logger.info(f"💬 RESPONSE GENERATION: Using RAG answer as fallback")
+            return rag_answer  # Fallback to RAG answer
+    
+    def _generate_conversational_response_without_docs(self, query: str, messages: List[Any]) -> str:
+        """Generate conversational response using only LLM knowledge and conversation history"""
+        logger.info("💬 RESPONSE GENERATION (NO DOCS): Starting response generation without documents")
+        logger.info(f"💬 RESPONSE GENERATION (NO DOCS): Query: '{query[:50]}...'")
+        
+        # Build conversation context
+        conversation_context = ""
+        for i, msg in enumerate(messages[-6:]):  # Last 6 messages for context
+            if isinstance(msg, HumanMessage):
+                conversation_context += f"User: {msg.content}\n"
+            elif isinstance(msg, AIMessage):
+                conversation_context += f"Assistant: {msg.content}\n"
+        
+        # Create response prompt
+        logger.info(f"💬 RESPONSE GENERATION (NO DOCS): Building response prompt")
+        response_prompt = ChatPromptTemplate.from_messages([
+            SystemMessage(content="""You are a helpful audit report assistant. Generate a natural, conversational response.
+
+RULES:
+1. Answer the user's question directly and clearly based on your knowledge
+2. Use conversation history for context
+3. Be conversational, not technical
+4. Acknowledge if the answer is based on general knowledge rather than specific documents
+5. Stay professional but friendly
+
+TONE: Professional but friendly, like talking to a colleague."""),
+            HumanMessage(content=f"""Current Question: {query}
+
+Conversation History:
+{conversation_context}
+
+Generate a conversational response based on your knowledge:""")
+        ])
+        
+        try:
+            logger.info(f"💬 RESPONSE GENERATION (NO DOCS): Calling LLM")
+            response = self.llm.invoke(response_prompt.format_messages())
+            logger.info(f"💬 RESPONSE GENERATION (NO DOCS): LLM response received: {response.content[:100]}...")
+            return response.content.strip()
+        except Exception as e:
+            logger.error(f"❌ RESPONSE GENERATION (NO DOCS): Error during generation: {e}")
+            return "I apologize, but I encountered an error. Please try asking your question differently."
+    
+    def chat(self, user_input: str, conversation_id: str = "default") -> Dict[str, Any]:
+        """Main chat interface"""
+        logger.info(f"💬 MULTI-AGENT CHAT: Processing '{user_input[:50]}...'")
+        
+        # Load conversation
+        logger.info(f"💬 MULTI-AGENT CHAT: Loading conversation {conversation_id}")
+        conversation_file = self.conversations_dir / f"{conversation_id}.json"
+        conversation = self._load_conversation(conversation_file)
+        logger.info(f"💬 MULTI-AGENT CHAT: Loaded {len(conversation['messages'])} previous messages")
+        
+        # Add user message
+        conversation["messages"].append(HumanMessage(content=user_input))
+        logger.info(f"💬 MULTI-AGENT CHAT: Added user message to conversation")
+        
+        # Prepare state
+        logger.info(f"💬 MULTI-AGENT CHAT: Preparing state for graph execution")
+        state = MultiAgentState(
+            conversation_id=conversation_id,
+            messages=conversation["messages"],
+            current_query=user_input,
+            query_context=None,
+            rag_query=None,
+            rag_filters=None,
+            retrieved_documents=None,
+            final_response=None,
+            agent_logs=[],
+            conversation_context=conversation.get("context", {}),
+            session_start_time=conversation["session_start_time"],
+            last_ai_message_time=conversation["last_ai_message_time"]
+        )
+        
+        # Run multi-agent graph
+        logger.info(f"💬 MULTI-AGENT CHAT: Executing multi-agent graph")
+        final_state = self.graph.invoke(state)
+        logger.info(f"💬 MULTI-AGENT CHAT: Graph execution completed")
+        
+        # Add AI response to conversation
+        if final_state["final_response"]:
+            conversation["messages"].append(AIMessage(content=final_state["final_response"]))
+            logger.info(f"💬 MULTI-AGENT CHAT: Added AI response to conversation")
+        
+        # Update conversation
+        conversation["last_ai_message_time"] = final_state["last_ai_message_time"]
+        conversation["context"] = final_state["conversation_context"]
+        
+        # Save conversation
+        logger.info(f"💬 MULTI-AGENT CHAT: Saving conversation")
+        self._save_conversation(conversation_file, conversation)
+        
+        logger.info("✅ MULTI-AGENT CHAT: Completed")
+        
+        # Return response and RAG results
+        return {
+            'response': final_state["final_response"],
+            'rag_result': {
+                'sources': final_state["retrieved_documents"] or [],
+                'answer': final_state["final_response"]
+            },
+            'agent_logs': final_state["agent_logs"],
+            'actual_rag_query': final_state.get("rag_query", "")
+        }
+    
+    def _load_conversation(self, conversation_file: Path) -> Dict[str, Any]:
+        """Load conversation from file"""
+        if conversation_file.exists():
+            try:
+                with open(conversation_file) as f:
+                    data = json.load(f)
+                    # Convert message dicts back to LangChain messages
+                    messages = []
+                    for msg_data in data.get("messages", []):
+                        if msg_data["type"] == "human":
+                            messages.append(HumanMessage(content=msg_data["content"]))
+                        elif msg_data["type"] == "ai":
+                            messages.append(AIMessage(content=msg_data["content"]))
+                    data["messages"] = messages
+                    return data
+            except Exception as e:
+                logger.warning(f"Could not load conversation: {e}")
+        
+        # Return default conversation
+        return {
+            "messages": [],
+            "session_start_time": time.time(),
+            "last_ai_message_time": time.time(),
+            "context": {}
+        }
+    
+    def _save_conversation(self, conversation_file: Path, conversation: Dict[str, Any]):
+        """Save conversation to file"""
+        try:
+            # Convert messages to serializable format
+            messages_data = []
+            for msg in conversation["messages"]:
+                if isinstance(msg, HumanMessage):
+                    messages_data.append({"type": "human", "content": msg.content})
+                elif isinstance(msg, AIMessage):
+                    messages_data.append({"type": "ai", "content": msg.content})
+            
+            conversation_data = {
+                "messages": messages_data,
+                "session_start_time": conversation["session_start_time"],
+                "last_ai_message_time": conversation["last_ai_message_time"],
+                "context": conversation.get("context", {})
+            }
+            
+            with open(conversation_file, 'w') as f:
+                json.dump(conversation_data, f, indent=2)
+                
+        except Exception as e:
+            logger.error(f"Could not save conversation: {e}")
+    
+
+def get_multi_agent_chatbot():
+    """Get multi-agent chatbot instance"""
+    return MultiAgentRAGChatbot()
+
+if __name__ == "__main__":
+    # Test the multi-agent system
+    chatbot = MultiAgentRAGChatbot()
+    
+    # Test conversation
+    result = chatbot.chat("List me top 10 challenges in budget allocation for the last 3 years")
+    print("Response:", result['response'])
+    print("Agent Logs:", result['agent_logs'])