inherit from base class

src/agents/multi_agent_chatbot.py

@@ -1,115 +1,45 @@
 """
 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.
+This is the TEXT-BASED RAG chatbot that inherits from BaseMultiAgentChatbot.
+It implements the retrieval using the PipelineManager (Qdrant + text embeddings).
 """
-import re
-import json
-import time
 import logging
 import traceback
-from pathlib import Path
-from datetime import datetime
-from dataclasses import dataclass
-from typing import Dict, List, Any, Optional, TypedDict
+from typing import Dict, List, Any
 
-from langchain_core.tools import tool
-from langgraph.graph import StateGraph, END
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 
-
+from src.agents.base_multi_agent_chatbot import BaseMultiAgentChatbot
 from src.pipeline import PipelineManager
-from src.llm.adapters import get_llm_client
-from src.config.paths import PROJECT_DIR, CONVERSATIONS_DIR
-from src.config.loader import load_config, get_embedding_model_for_collection
-
 
 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"""
+class MultiAgentRAGChatbot(BaseMultiAgentChatbot):
+    """
+    Text-based Multi-agent RAG chatbot.
+    
+    Inherits all the sophisticated logic from BaseMultiAgentChatbot:
+    - LLM-based query analysis
+    - Filter extraction and validation
+    - Query rewriting
+    - Main agent, RAG agent, Response agent
+    
+    Implements:
+    - _perform_retrieval(): Uses PipelineManager for text-based RAG
+    - _generate_conversational_response(): Text-focused response generation
+    - _generate_conversational_response_without_docs(): Fallback response
+    """
     
     def __init__(self, config_path: str = "src/config/settings.yaml"):
-        """Initialize the multi-agent chatbot"""
-        self.config = load_config(config_path)
+        """Initialize the text-based multi-agent chatbot"""
+        # Initialize base class first (loads config, LLM, filters, builds graph)
+        super().__init__(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
+        # Initialize pipeline manager for text-based retrieval
         logger.info("🔄 Initializing pipeline manager and loading models...")
         try:
             self.pipeline_manager = PipelineManager(self.config)
@@ -124,929 +54,44 @@ class MultiAgentRAGChatbot:
         try:
             if not self.pipeline_manager.connect_vectorstore():
                 logger.error("❌ Failed to connect to vector store")
-                logger.error("💡 Check that QDRANT_API_KEY environment variable is set")
-                logger.error("💡 Check that Qdrant URL and collection name are correct in config")
                 raise RuntimeError("Vector store connection failed")
             logger.info("✅ Vector store connected successfully")
         except RuntimeError:
-            raise  # Re-raise RuntimeError as-is
+            raise
         except Exception as e:
             logger.error(f"❌ Error during vector store connection: {e}")
             traceback.print_exc()
             raise RuntimeError(f"Vector store connection failed: {e}")
         
-        # Load dynamic data
-        self._load_dynamic_data()
-        
-        # Build the multi-agent graph
-        self.graph = self._build_graph()
-        
-        # Conversations directory - use PROJECT_DIR for local vs deployed compatibility
-        self.conversations_dir = CONVERSATIONS_DIR
-        try:
-            # Use 777 permissions for maximum compatibility (HF Spaces runs as different user)
-            self.conversations_dir.mkdir(parents=True, mode=0o777, exist_ok=True)
-        except (PermissionError, OSError) as e:
-            logger.warning(f"Could not create conversations directory at {self.conversations_dir}: {e}")
-            # Fallback to a relative path (current directory)
-            self.conversations_dir = Path("conversations")
-            try:
-                self.conversations_dir.mkdir(parents=True, mode=0o777, exist_ok=True)
-            except (PermissionError, OSError) as e2:
-                logger.error(f"Could not create conversations directory at {self.conversations_dir}: {e2}")
-                raise RuntimeError(f"Failed to create conversations directory: {e2}")
-        
-        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 - use PROJECT_DIR relative path
-        try:
-            fo = PROJECT_DIR / "src" / "config" / "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 (if available)
-        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 _normalize_district_name(self, district: str) -> Optional[str]:
-        """Normalize district name with fuzzy matching for common misspellings."""
-        if not district:
-            return None
-        
-        district = district.strip()
-        
-        # Direct match
-        if district in self.district_whitelist:
-            return district
-        
-        # Remove "District" suffix
-        district_name = district.replace(" District", "").replace(" district", "").strip()
-        if district_name in self.district_whitelist:
-            return district_name
-        
-        # Common misspellings mapping
-        misspelling_map = {
-            "kalagala": "Kalangala",
-            "Kalagala": "Kalangala",
-            "KALAGALA": "Kalangala",
-            "kalangala": "Kalangala",
-            "gulu": "Gulu",
-            "GULU": "Gulu",
-            "kampala": "Kampala",
-            "KAMPALA": "Kampala",
-        }
-        
-        # Check misspelling map (case-insensitive)
-        district_lower = district_name.lower()
-        if district_lower in misspelling_map:
-            corrected = misspelling_map[district_lower]
-            if corrected in self.district_whitelist:
-                return corrected
-        
-        # Fuzzy matching for similar names (simple Levenshtein-like check)
-        # Check if the district name is very similar to any whitelist entry
-        for whitelist_district in self.district_whitelist:
-            # Case-insensitive comparison
-            if district_name.lower() == whitelist_district.lower():
-                return whitelist_district
-            
-            # Check if one is a substring of the other (for partial matches)
-            if len(district_name) >= 4 and len(whitelist_district) >= 4:
-                if district_name.lower() in whitelist_district.lower() or whitelist_district.lower() in district_name.lower():
-                    # Only return if it's a strong match (at least 80% of characters match)
-                    min_len = min(len(district_name), len(whitelist_district))
-                    max_len = max(len(district_name), len(whitelist_district))
-                    if min_len / max_len >= 0.8:
-                        return whitelist_district
-        
-        return None
+        logger.info("🤖 Text-based Multi-Agent RAG Chatbot initialized")
     
-    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"
-            }
+    def _perform_retrieval(self, query: str, filters: Dict[str, Any]) -> Any:
+        """
+        Perform text-based retrieval using PipelineManager.
+        
+        Args:
+            query: The rewritten query
+            filters: The filters to apply
+        
+        Returns:
+            PipelineResult with .sources and .answer
+        """
+        logger.info(f"🔍 TEXT RETRIEVAL: Query='{query}', Filters={filters}")
+        
+        result = self.pipeline_manager.run(
+            query=query,
+            sources=filters.get("sources") if filters else None,
+            auto_infer_filters=False,
+            filters=filters if filters else None
         )
         
-        # 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"
+        logger.info(f"🔍 TEXT RETRIEVAL: Retrieved {len(result.sources)} documents")
+        return result
     
-    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
-   - **COMMON MISSPELLINGS**: Handle common misspellings intelligently:
-     * "Kalagala" (missing 'n') should be extracted as "Kalangala"
-     * "kalagala", "Kalagala", "KALAGALA" should all be normalized to "Kalangala"
-     * Similar case-insensitive variations should be normalized to the correct district name
-   - 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
-                # 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
-                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:
-                        normalized = self._normalize_district_name(district)
-                        if normalized:
-                            valid_districts.append(normalized)
-                    
-                    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 with fuzzy matching
-                    normalized = self._normalize_district_name(extracted_district)
-                    if normalized:
-                        if normalized != extracted_district:
-                            logger.info(f"🔍 QUERY ANALYSIS: Normalized district '{extracted_district}' to '{normalized}'")
-                        extracted_district = normalized
-                    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("district") and context.extracted_district:
-                # Normalize district names using the normalization function
-                if isinstance(context.extracted_district, list):
-                    normalized_districts = []
-                    for d in context.extracted_district:
-                        normalized = self._normalize_district_name(d)
-                        if normalized:
-                            normalized_districts.append(normalized)
-                    if normalized_districts:
-                        filters["district"] = normalized_districts
-                        logger.info(f"🔧 FILTER BUILDING: Added districts filter from context: {context.extracted_district} → normalized: {normalized_districts}")
-                else:
-                    normalized = self._normalize_district_name(context.extracted_district)
-                    if normalized:
-                        filters["district"] = [normalized]
-                        logger.info(f"🔧 FILTER BUILDING: Added district filter from context: {context.extracted_district} → normalized: {normalized}")
-            
-            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("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:
-                # Normalize district names using the normalization function
-                if isinstance(context.extracted_district, list):
-                    normalized_districts = []
-                    for d in context.extracted_district:
-                        normalized = self._normalize_district_name(d)
-                        if normalized:
-                            normalized_districts.append(normalized)
-                    if normalized_districts:
-                        filters["district"] = normalized_districts
-                        logger.info(f"🔧 FILTER BUILDING: Added districts filter from context: {context.extracted_district} → normalized: {normalized_districts}")
-                else:
-                    normalized = self._normalize_district_name(context.extracted_district)
-                    if normalized:
-                        filters["district"] = [normalized]
-                        logger.info(f"🔧 FILTER BUILDING: Added district filter from context: {context.extracted_district} → normalized: {normalized}")
-        
-        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:
+    def _generate_conversational_response(self, query: str, documents: List[Any], rag_answer: str, messages: List[Any], filters: Dict[str, Any] = None) -> 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]}...'")
-        logger.info(f"💬 RESPONSE GENERATION: Conversation history: {len(messages)} messages")
+        logger.info(f"💬 RESPONSE GENERATION: Filters for validation: {filters}")
         
         # Build conversation history context
         conversation_context = self._build_conversation_context(messages)
@@ -1054,11 +99,10 @@ Rewrite the best retrieval query:""")
         # Build detailed document information
         document_details = self._build_document_details(documents)
         
-        # Extract correct district/source/year names from documents (to correct misspellings)
+        # Extract correct district/source/year names from documents
         correct_names = self._extract_correct_names_from_documents(documents)
         
         # 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.
 
@@ -1068,7 +112,7 @@ CRITICAL RULES - NO HALLUCINATION:
 3. **If a document doesn't contain the information, DO NOT make it up**
 4. **If the user asks about a year/district that's NOT in the retrieved documents, explicitly state that**
 5. **Check the document years/districts before making any claims about them**
-6. **USE CORRECT NAMES**: If the conversation mentions a misspelled district/source name (e.g., "Kalagala"), use the CORRECT spelling from the document metadata (e.g., "Kalangala"). Always use the exact names from document metadata, not misspellings from conversation.
+6. **USE CORRECT NAMES**: Use the CORRECT spelling from the document metadata
 
 RULES:
 1. Answer the user's question directly and clearly
@@ -1086,9 +130,8 @@ RULES:
 13. You do not need to use every passage. Only use the ones that help answer the question.
 14. **VERIFY**: Before mentioning any year, district, or number, check that it exists in the retrieved documents. If it doesn't, say "I don't have information about [year/district] in the retrieved documents."
 15. **NO HALLUCINATION**: If documents show years 2021, 2022, 2023 but user asks about 2020, DO NOT provide 2020 data. Instead say "The retrieved documents cover 2021-2023, but I don't have information for 2020."
-16. **USE CORRECT SPELLING**: Always use the district/source names exactly as they appear in the document metadata below, even if the conversation history has misspellings.
 
-TONE: Professional but friendly, like talking to a colleague."""),
+TONE: Professional but friendly."""),
             HumanMessage(content=f"""Conversation History:
 {conversation_context}
 
@@ -1096,7 +139,7 @@ Current User Question: {query}
 
 Retrieved Documents: {len(documents)} documents found
 
-CORRECT NAMES TO USE (from document metadata - use these exact spellings):
+CORRECT NAMES TO USE (from document metadata):
 {correct_names}
 
 Full Document Details:
@@ -1109,28 +152,65 @@ CRITICAL:
 - If user asks about a specific year but documents show other years, or districts or sources then explicitly state "can't provide response on ... because ..."
 - Every factual claim MUST have [Doc i] reference
 - If information is not in documents, explicitly state it's not available
-- **USE THE CORRECT DISTRICT/SOURCE NAMES from the document metadata above, not misspellings from conversation**
 
 Generate a conversational response with proper document references:""")
         ])
         
         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]}...")
             
             # Post-process response to ensure no hallucination
             final_response = self._validate_and_enhance_response(
                 response.content.strip(),
                 documents,
-                query
+                query,
+                filters  # Pass filters for coverage validation
             )
             
             return final_response
         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
+            logger.error(f"❌ RESPONSE GENERATION: Error: {e}")
+            return rag_answer
+    
+    def _generate_conversational_response_without_docs(self, query: str, messages: List[Any]) -> str:
+        """Generate conversational response using only LLM knowledge"""
+        logger.info("💬 RESPONSE GENERATION (NO DOCS): Starting response generation without documents")
+        
+        # Build conversation context
+        conversation_context = ""
+        for msg in messages[-6:]:
+            if isinstance(msg, HumanMessage):
+                conversation_context += f"User: {msg.content}\n"
+            elif isinstance(msg, AIMessage):
+                conversation_context += f"Assistant: {msg.content}\n"
+        
+        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."""),
+            HumanMessage(content=f"""Current Question: {query}
+
+Conversation History:
+{conversation_context}
+
+Generate a conversational response based on your knowledge:""")
+        ])
+        
+        try:
+            response = self.llm.invoke(response_prompt.format_messages())
+            return response.content.strip()
+        except Exception as e:
+            logger.error(f"❌ RESPONSE GENERATION (NO DOCS): Error: {e}")
+            return "I apologize, but I encountered an error. Please try asking your question differently."
+    
+    # ==================== HELPER METHODS ====================
     
     def _build_conversation_context(self, messages: List[Any]) -> str:
         """Build conversation history context for response generation."""
@@ -1138,7 +218,6 @@ Generate a conversational response with proper document references:""")
             return "No previous conversation."
         
         context_lines = []
-        # Show last 6 messages for context (to capture the current exchange)
         for msg in messages[-6:]:
             if isinstance(msg, HumanMessage):
                 context_lines.append(f"User: {msg.content}")
@@ -1153,7 +232,7 @@ Generate a conversational response with proper document references:""")
             return "No documents retrieved."
         
         details = []
-        for i, doc in enumerate(documents[:15], 1):  # Show up to 15 documents
+        for i, doc in enumerate(documents[:15], 1):
             metadata = getattr(doc, 'metadata', {}) if hasattr(doc, 'metadata') else (doc if isinstance(doc, dict) else {})
             content = getattr(doc, 'page_content', '') if hasattr(doc, 'page_content') else (doc.get('content', '') if isinstance(doc, dict) else '')
             
@@ -1171,7 +250,7 @@ Generate a conversational response with proper document references:""")
                 doc_info += f"\n  Source: {source}"
                 if page != 'Unknown':
                     doc_info += f"\n  Page: {page}"
-                doc_info += f"\n  Content: {content[:300]}{'...' if len(content) > 300 else ''}"
+                doc_info += f"\n  Content: {content[:500]}{'...' if len(content) > 500 else ''}"
                 details.append(doc_info)
         
         return "\n\n".join(details) if details else "No document details available."
@@ -1200,16 +279,18 @@ Generate a conversational response with proper document references:""")
         if years:
             result.append(f"Years: {', '.join(sorted(years))}")
         
-        if result:
-            return "\n".join(result) + "\n\nIMPORTANT: Use these EXACT spellings in your response, even if the conversation history has misspellings."
-        return "No metadata available."
+        return "\n".join(result) if result else "No metadata available."
     
-    def _validate_and_enhance_response(self, response: str, documents: List[Any], query: str) -> str:
-        """Validate response and ensure all claims are referenced."""
-        # Extract years and districts from documents
+    def _validate_and_enhance_response(self, response: str, documents: List[Any], query: str, filters: Dict[str, Any] = None) -> str:
+        """Validate response and ensure all claims are referenced.
+        
+        Compares REQUESTED filters against RETRIEVED document metadata to identify gaps.
+        """
+        import re
+        
+        # Extract years and districts from RETRIEVED documents
         doc_years = set()
         doc_districts = set()
-        doc_sources = set()
         
         for doc in documents:
             metadata = getattr(doc, 'metadata', {}) if hasattr(doc, 'metadata') else (doc if isinstance(doc, dict) else {})
@@ -1218,204 +299,67 @@ Generate a conversational response with proper document references:""")
                     doc_years.add(str(metadata['year']))
                 if metadata.get('district'):
                     doc_districts.add(str(metadata['district']))
-                if metadata.get('source'):
-                    doc_sources.add(str(metadata['source']))
-        
-        # Correct misspellings in response using correct names from documents
-        # response = self._correct_misspellings_in_response(response, doc_districts, doc_sources)
         
-        # Check if response mentions years not in documents
-        year_pattern = r'\b(20\d{2})\b'
-        mentioned_years = set(re.findall(year_pattern, response))
-        
-        # Check if user query mentions a year
-        query_years = set(re.findall(year_pattern, query))
-        
-        # If user asks about a year not in documents, add a warning
-        missing_years = query_years - doc_years
-        if missing_years and doc_years:
-            warning = f"\n\n⚠️ Note: The retrieved documents cover years {', '.join(sorted(doc_years))}, but I don't have information for {', '.join(sorted(missing_years))} in the retrieved documents."
-            if warning not in response:
-                response = response + warning
-        
-        # Check if response has document references
-        doc_ref_pattern = r'\[Doc\s+\d+\]'
-        has_refs = bool(re.search(doc_ref_pattern, response))
-        
-        # If response has factual claims but no references, add a note
-        if not has_refs and len(documents) > 0:
-            # Check if response has numbers or specific claims (simple heuristic)
-            has_numbers = bool(re.search(r'\d+', response))
-            if has_numbers and len(response) > 50:
-                logger.warning("⚠️ Response contains factual claims but no document references")
-                # Don't modify response, but log the issue
+        logger.info(f"📊 VALIDATION: Retrieved docs cover years={doc_years}, districts={doc_districts}")
+        
+        warnings = []
+        
+        # Get REQUESTED filters
+        requested_years = set()
+        requested_districts = set()
+        
+        if filters:
+            if filters.get('year'):
+                requested_years = set(str(y) for y in filters['year']) if isinstance(filters['year'], list) else {str(filters['year'])}
+            if filters.get('district'):
+                requested_districts = set(str(d) for d in filters['district']) if isinstance(filters['district'], list) else {str(filters['district'])}
+        
+        logger.info(f"📊 VALIDATION: Requested years={requested_years}, districts={requested_districts}")
+        
+        # Compare requested vs retrieved YEARS
+        if requested_years and doc_years:
+            missing_years = requested_years - doc_years
+            if missing_years:
+                warnings.append(f"You requested data for years {', '.join(sorted(requested_years))}, but the retrieved documents only cover {', '.join(sorted(doc_years))}. Data for {', '.join(sorted(missing_years))} may not be available in the database.")
+        elif requested_years and not doc_years:
+            warnings.append(f"You requested data for years {', '.join(sorted(requested_years))}, but no documents were retrieved with year metadata.")
+        
+        # Compare requested vs retrieved DISTRICTS
+        if requested_districts and doc_districts:
+            # Normalize for comparison (case-insensitive)
+            requested_districts_lower = {d.lower() for d in requested_districts}
+            doc_districts_lower = {d.lower() for d in doc_districts}
+            missing_districts_lower = requested_districts_lower - doc_districts_lower
+            
+            if missing_districts_lower:
+                # Get original case versions for display
+                missing_districts = [d for d in requested_districts if d.lower() in missing_districts_lower]
+                warnings.append(f"You requested data for districts {', '.join(sorted(requested_districts))}, but the retrieved documents only cover {', '.join(sorted(doc_districts))}. Data for {', '.join(sorted(missing_districts))} may not be available in the database.")
+        elif requested_districts and not doc_districts:
+            warnings.append(f"You requested data for districts {', '.join(sorted(requested_districts))}, but no documents were retrieved with district metadata.")
+        
+        # Fallback: Check query text for explicit years not in documents (for cases without filters)
+        if not requested_years:
+            year_pattern = r'\b(20\d{2})\b'
+            query_years = set(re.findall(year_pattern, query))
+            missing_years = query_years - doc_years
+            if missing_years and doc_years:
+                warnings.append(f"The retrieved documents cover years {', '.join(sorted(doc_years))}, but I don't have information for {', '.join(sorted(missing_years))} in the retrieved documents.")
+        
+        # Add warnings to response if any
+        if warnings and "⚠️" not in response:
+            warning_text = "\n\n⚠️ **Note:** " + " ".join(warnings)
+            response = response + warning_text
+            logger.info(f"📊 VALIDATION: Added warning: {warning_text}")
         
         return response
     
-    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:
-            # Ensure the conversations directory exists with proper permissions
-            conversation_file.parent.mkdir(parents=True, mode=0o777, exist_ok=True)
-            
-            # 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}")
-            logger.error(f"Traceback: {traceback.format_exc()}")
-    
 
 def get_multi_agent_chatbot():
     """Get multi-agent chatbot instance"""
     return MultiAgentRAGChatbot()
 
+
 if __name__ == "__main__":
     # Test the multi-agent system
     chatbot = MultiAgentRAGChatbot()