add single smart chatbot

smart_chatbot.py

@@ -0,0 +1,1098 @@
+"""
+Intelligent RAG Chatbot with Smart Query Analysis and Conversation Management
+
+This chatbot provides intelligent conversation flow with:
+- Smart query analysis and expansion
+- Single LangSmith conversation traces
+- Local conversation logging
+- Context-aware RAG retrieval
+- Natural conversation without technical jargon
+"""
+
+import os
+import json
+import time
+import logging
+from pathlib import Path
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Dict, List, Any, Optional, TypedDict
+
+
+import re
+from langgraph.graph import StateGraph, END
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+
+from src.pipeline import PipelineManager
+from src.config.loader import load_config
+
+
+@dataclass
+class QueryAnalysis:
+    """Analysis result of a user query"""
+    has_district: bool
+    has_source: bool
+    has_year: bool
+    extracted_district: Optional[str]
+    extracted_source: Optional[str]
+    extracted_year: Optional[str]
+    confidence_score: float
+    can_answer_directly: bool
+    missing_filters: List[str]
+    suggested_follow_up: Optional[str]
+    expanded_query: Optional[str] = None  # Query expansion for better RAG
+
+
+class ConversationState(TypedDict):
+    """State for the conversation flow"""
+    conversation_id: str
+    messages: List[Any]
+    current_query: str
+    query_analysis: Optional[QueryAnalysis]
+    rag_query: Optional[str]
+    rag_result: Optional[Any]
+    final_response: Optional[str]
+    conversation_context: Dict[str, Any]  # Store conversation context
+    session_start_time: float
+    last_ai_message_time: float
+
+
+class IntelligentRAGChatbot:
+    """Intelligent chatbot with smart query analysis and conversation management"""
+    
+    def __init__(self, suppress_logs=False):
+        """Initialize the intelligent chatbot"""
+        # Setup logger to avoid cluttering UI
+        self.logger = logging.getLogger(__name__)
+        if suppress_logs:
+            self.logger.setLevel(logging.CRITICAL)  # Suppress all logs
+        else:
+            self.logger.setLevel(logging.INFO)
+            if not self.logger.handlers:
+                handler = logging.StreamHandler()
+                formatter = logging.Formatter('%(message)s')
+                handler.setFormatter(formatter)
+                self.logger.addHandler(handler)
+        
+        self.logger.info("🤖 INITIALIZING: Intelligent RAG Chatbot")
+        
+        # Load configuration first
+        self.config = load_config()
+        
+        # Use the same LLM configuration as the existing system
+        from auditqa.llm.adapters import get_llm_client
+        
+        # Get LLM client using the same configuration
+        reader_config = self.config.get("reader", {})
+        default_type = reader_config.get("default_type", "INF_PROVIDERS")
+        
+        # Convert to lowercase as that's how it's registered
+        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):
+                    # Convert LangChain messages to dict format
+                    message_dicts = []
+                    for msg in messages:
+                        if hasattr(msg, 'content'):
+                            role = "user" if isinstance(msg, HumanMessage) else "assistant"
+                            message_dicts.append({"role": role, "content": msg.content})
+                        else:
+                            message_dicts.append({"role": "user", "content": str(msg)})
+                else:
+                    # Single message
+                    message_dicts = [{"role": "user", "content": str(messages)}]
+                
+                # Use the adapter to generate response
+                llm_response = self.adapter.generate(message_dicts)
+                
+                # Return in LangChain format
+                class MockResponse:
+                    def __init__(self, content):
+                        self.content = content
+                
+                return MockResponse(llm_response.content)
+        
+        self.llm = LLMWrapper(self.llm_adapter)
+        
+        # Initialize pipeline manager for RAG
+        self.logger.info("🔧 PIPELINE: Initializing PipelineManager...")
+        self.pipeline_manager = PipelineManager(self.config)
+        
+        # Ensure vectorstore is connected
+        self.logger.info("🔗 VECTORSTORE: Connecting to Qdrant...")
+        try:
+            vectorstore = self.pipeline_manager.vectorstore_manager.connect_to_existing()
+            self.logger.info("✅ VECTORSTORE: Connected successfully")
+        except Exception as e:
+            self.logger.error(f"❌ VECTORSTORE: Connection failed: {e}")
+        
+        # Fix LLM client to use the same provider as chatbot
+        self.logger.info("🔧 LLM: Fixing PipelineManager LLM client...")
+        self.pipeline_manager.llm_client = self.llm_adapter
+        self.logger.info("✅ LLM: PipelineManager now uses same LLM as chatbot")
+        
+        self.logger.info("✅ PIPELINE: PipelineManager initialized")
+        
+        # Available metadata for filtering
+        self.available_metadata = {
+            'sources': [
+                'KCCA', 'MAAIF', 'MWTS', 'Gulu DLG', 'Kalangala DLG', 'Namutumba DLG',
+                'Lwengo DLG', 'Kiboga DLG', 'Annual Consolidated OAG', 'Consolidated',
+                'Hospital', 'Local Government', 'Ministry, Department and Agency', 
+                'Project', 'Thematic', 'Value for Money'
+            ],
+            'years': ['2018', '2019', '2020', '2021', '2022', '2023', '2024', '2025'],
+            'districts': [
+                'Gulu', 'Kalangala', 'Kampala', 'Namutumba', 'Lwengo', 'Kiboga',
+                'Fort Portal', 'Arua', 'Kasese', 'Kabale', 'Masindi', 'Mbale', 'Jinja', 'Masaka', 'Mbarara',
+                'KCCA'
+            ]
+        }
+
+        # Try to load district whitelist from filter_options.json
+        try:
+            fo = Path("filter_options.json")
+            if fo.exists():
+                with open(fo) as f:
+                    data = json.load(f)
+                    if isinstance(data, dict) and data.get("districts"):
+                        self.district_whitelist = [d.strip() for d in data["districts"] if d]
+                    else:
+                        self.district_whitelist = self.available_metadata['districts']
+            else:
+                self.district_whitelist = self.available_metadata['districts']
+        except Exception:
+            self.district_whitelist = self.available_metadata['districts']
+
+        # Enrich whitelist 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:
+                # Merge while preserving order: existing first, then new ones not present
+                merged = list(self.district_whitelist)
+                for n in sorted(names):
+                    if n not in merged:
+                        merged.append(n)
+                self.district_whitelist = merged
+                self.logger.info(f"🧭 District whitelist enriched: {len(self.district_whitelist)} entries")
+        except Exception as e:
+            self.logger.info(f"ℹ️ Could not enrich districts from add_district_metadata: {e}")
+        
+        # Get dynamic year list from filter_options.json
+        try:
+            fo = Path("filter_options.json")
+            if fo.exists():
+                with open(fo) as f:
+                    data = json.load(f)
+                    if isinstance(data, dict) and data.get("years"):
+                        self.year_whitelist = [str(y).strip() for y in data["years"] if y]
+                    else:
+                        self.year_whitelist = self.available_metadata['years']
+            else:
+                self.year_whitelist = self.available_metadata['years']
+        except Exception:
+            self.year_whitelist = self.available_metadata['years']
+        
+        # Calculate current year dynamically
+        from datetime import datetime
+        self.current_year = str(datetime.now().year)
+        self.previous_year = str(datetime.now().year - 1)
+        
+        # Data context for system prompt
+        self.data_context = self._load_data_context()
+        
+        # Build the LangGraph
+        self.graph = self._build_graph()
+        
+        # Conversation logging
+        self.conversations_dir = Path("conversations")
+        self.conversations_dir.mkdir(exist_ok=True)
+    
+    def _load_data_context(self) -> str:
+        """Load and analyze data context for system prompt"""
+        try:
+            # Try to load from generated context file
+            context_file = Path("data_context.md")
+            if context_file.exists():
+                with open(context_file) as f:
+                    return f.read()
+            
+            # Fallback to basic analysis
+            reports_dir = Path("reports")
+            testset_dir = Path("outputs/datasets/testset")
+            
+            context_parts = []
+            
+            # Report analysis
+            if reports_dir.exists():
+                report_folders = [d for d in reports_dir.iterdir() if d.is_dir()]
+                context_parts.append(f"📊 Available Reports: {len(report_folders)} audit report folders")
+                
+                # Get year range
+                years = []
+                for folder in report_folders:
+                    if "2018" in folder.name:
+                        years.append("2018")
+                    elif "2019" in folder.name:
+                        years.append("2019")
+                    elif "2020" in folder.name:
+                        years.append("2020")
+                    elif "2021" in folder.name:
+                        years.append("2021")
+                    elif "2022" in folder.name:
+                        years.append("2022")
+                    elif "2023" in folder.name:
+                        years.append("2023")
+                
+                if years:
+                    context_parts.append(f"📅 Years covered: {', '.join(sorted(set(years)))}")
+            
+            # Test dataset analysis
+            if testset_dir.exists():
+                test_files = list(testset_dir.glob("*.json"))
+                context_parts.append(f"🧪 Test dataset: {len(test_files)} files with sample questions")
+            
+            return "\n".join(context_parts) if context_parts else "📊 Audit report database with comprehensive coverage"
+            
+        except Exception as e:
+            self.logger.warning(f"⚠️ Could not load data context: {e}")
+            return "📊 Comprehensive audit report database"
+    
+    def _build_graph(self) -> StateGraph:
+        """Build the LangGraph for intelligent conversation flow"""
+        
+        # Define the graph
+        workflow = StateGraph(ConversationState)
+        
+        # Add nodes
+        workflow.add_node("analyze_query", self._analyze_query)
+        workflow.add_node("decide_action", self._decide_action)
+        workflow.add_node("perform_rag", self._perform_rag)
+        workflow.add_node("ask_follow_up", self._ask_follow_up)
+        workflow.add_node("generate_response", self._generate_response)
+        
+        # Add edges
+        workflow.add_edge("analyze_query", "decide_action")
+        
+        # Conditional edges from decide_action
+        workflow.add_conditional_edges(
+            "decide_action",
+            self._should_perform_rag,
+            {
+                "rag": "perform_rag",
+                "follow_up": "ask_follow_up"
+            }
+        )
+        
+        # From perform_rag, go to generate_response
+        workflow.add_edge("perform_rag", "generate_response")
+        
+        # From ask_follow_up, end
+        workflow.add_edge("ask_follow_up", END)
+        
+        # From generate_response, end
+        workflow.add_edge("generate_response", END)
+        
+        # Set entry point
+        workflow.set_entry_point("analyze_query")
+        
+        return workflow.compile()
+
+    def _extract_districts_list(self, text: str) -> List[str]:
+        """Extract one or more districts from free text using whitelist matching.
+        - Case-insensitive substring match for each known district name
+        - Handles multi-district inputs like "Lwengo Kiboga District & Namutumba"
+        """
+        if not text:
+            return []
+        q = text.lower()
+        found: List[str] = []
+        for name in self.district_whitelist:
+            n = name.lower()
+            if n in q:
+                # Map Kampala -> KCCA canonical
+                canonical = 'KCCA' if name.lower() == 'kampala' else name
+                if canonical not in found:
+                    found.append(canonical)
+        return found
+
+    def _extract_years_list(self, text: str) -> List[str]:
+        """Extract year list from text, supporting forms like '2022 / 23', '2022-2023', '2022–23'."""
+        if not text:
+            return []
+        years: List[str] = []
+        q = text
+        # Full 4-digit years
+        for y in re.findall(r"\b(20\d{2})\b", q):
+            if y not in years:
+                years.append(y)
+        # Shorthand like 2022/23 or 2022-23
+        for m in re.finditer(r"\b(20\d{2})\s*[\-/–]\s*(\d{2})\b", q):
+            y1 = m.group(1)
+            y2_short = int(m.group(2))
+            y2 = f"20{y2_short:02d}"
+            for y in [y1, y2]:
+                if y not in years:
+                    years.append(y)
+        return years
+    
+    def _analyze_query(self, state: ConversationState) -> ConversationState:
+        """Analyze the user query with conversation context"""
+        
+        query = state["current_query"]
+        conversation_context = state.get("conversation_context", {})
+        
+        self.logger.info(f"🧠 QUERY ANALYSIS: Starting analysis for: '{query[:50]}...'")
+        
+        # Build conversation context for analysis
+        context_info = ""
+        if conversation_context:
+            context_info = f"\n\nConversation context:\n"
+            for key, value in conversation_context.items():
+                if value:
+                    context_info += f"- {key}: {value}\n"
+        
+        # Also include recent conversation messages for better context
+        recent_messages = state.get("messages", [])
+        if recent_messages and len(recent_messages) > 1:
+            context_info += f"\nRecent conversation:\n"
+            # Get last 3 messages for context
+            for msg in recent_messages[-3:]:
+                if hasattr(msg, 'content'):
+                    role = "User" if isinstance(msg, HumanMessage) else "Assistant"
+                    context_info += f"- {role}: {msg.content[:100]}...\n"
+        
+        # Create analysis prompt with data context
+        analysis_prompt = ChatPromptTemplate.from_messages([
+            SystemMessage(content=f"""You are an expert at analyzing audit report queries. Your job is to extract specific information and determine if a query can be answered directly.
+
+{self.data_context}
+
+DISTRICT RECOGNITION RULES:
+- Kampala = KCCA (Kampala Capital City Authority)
+- Available districts: {', '.join(self.district_whitelist[:15])}... (and {len(self.district_whitelist)-15} more)
+- DLG = District Local Government
+- Uganda has {len(self.district_whitelist)} districts - recognize common ones
+
+SOURCE RECOGNITION RULES:
+- KCCA = Kampala Capital City Authority
+- MAAIF = Ministry of Agriculture, Animal Industry and Fisheries
+- MWTS = Ministry of Works and Transport
+- OAG = Office of the Auditor General
+- Consolidated = Annual Consolidated reports
+
+YEAR RECOGNITION RULES:
+- Available years: {', '.join(self.year_whitelist)}
+- Current year is {self.current_year} - use this to reason about relative years
+- If user mentions "last year", "previous year" - infer {self.previous_year}
+- If user mentions "this year", "current year" - infer {self.current_year}
+
+Analysis rules:
+1. Be SMART - if you have enough context to search, do it
+2. Use conversation context to fill in missing information
+3. For budget/expenditure queries, try to infer missing details from context
+4. Current year is {self.current_year} - use this to reason about relative years
+5. If user mentions "last year", "previous year" - infer {self.previous_year}
+6. If user mentions "this year", "current year" - infer {self.current_year}
+7. If user mentions a department/ministry, infer the source
+8. If user is getting frustrated or asking for results, proceed with RAG even if not perfect
+9. Recognize Kampala as a district (KCCA)
+
+IMPORTANT: You must respond with ONLY valid JSON. No additional text.
+
+Return your analysis as JSON with these exact fields:
+{{
+  "has_district": boolean,
+  "has_source": boolean,
+  "has_year": boolean,
+  "extracted_district": "string or null",
+  "extracted_source": "string or null", 
+  "extracted_year": "string or null",
+  "confidence_score": 0.0-1.0,
+  "can_answer_directly": boolean,
+  "missing_filters": ["list", "of", "missing", "filters"],
+  "suggested_follow_up": "string or null",
+  "expanded_query": "string or null"
+}}
+
+The expanded_query should be a natural language query that combines the original question with any inferred context for better RAG retrieval."""),
+            HumanMessage(content=f"Analyze this query: '{query}'{context_info}")
+        ])
+        
+        # Get analysis from LLM
+        response = self.llm.invoke(analysis_prompt.format_messages())
+        
+        try:
+            # Clean the response content to extract JSON
+            content = response.content.strip()
+            
+            # Try to find JSON in the response
+            if content.startswith('{') and content.endswith('}'):
+                json_content = content
+            else:
+                # Try to extract JSON from the response
+                import re
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    json_content = json_match.group()
+                else:
+                    raise json.JSONDecodeError("No JSON found in response", content, 0)
+            
+            # Parse JSON response
+            analysis_data = json.loads(json_content)
+            
+            query_analysis = QueryAnalysis(
+                has_district=analysis_data.get("has_district", False),
+                has_source=analysis_data.get("has_source", False),
+                has_year=analysis_data.get("has_year", False),
+                extracted_district=analysis_data.get("extracted_district"),
+                extracted_source=analysis_data.get("extracted_source"),
+                extracted_year=analysis_data.get("extracted_year"),
+                confidence_score=analysis_data.get("confidence_score", 0.0),
+                can_answer_directly=analysis_data.get("can_answer_directly", False),
+                missing_filters=analysis_data.get("missing_filters", []),
+                suggested_follow_up=analysis_data.get("suggested_follow_up"),
+                expanded_query=analysis_data.get("expanded_query")
+            )
+            
+        except (json.JSONDecodeError, KeyError, AttributeError) as e:
+            self.logger.info(f"⚠️ JSON parsing failed: {e}")
+            # Fallback analysis - be more permissive
+            query_lower = query.lower()
+            
+            # Simple keyword matching - improved district recognition
+            has_district = any(district.lower() in query_lower for district in [
+                'gulu', 'kalangala', 'kampala', 'namutumba', 'lwengo', 'kiboga', 'kcca', 'maaif', 'mwts'
+            ])
+            
+            # Special case: Kampala = KCCA
+            if 'kampala' in query_lower and not has_district:
+                has_district = True
+            
+            has_source = any(source.lower() in query_lower for source in [
+                'kcca', 'maaif', 'mwts', 'gulu', 'kalangala', 'consolidated', 'oag', 'government'
+            ])
+            
+            # Check for year mentions using dynamic year list
+            has_year = any(year in query_lower for year in self.year_whitelist)
+            
+            # Also check for explicit relative year terms
+            has_year = has_year or any(term in query_lower for term in [
+                'this year', 'last year', 'previous year', 'current year'
+            ])
+            
+            # Extract specific values
+            extracted_district = None
+            extracted_source = None
+            extracted_year = None
+            
+            # Extract districts using comprehensive whitelist
+            for district_name in self.district_whitelist:
+                if district_name.lower() in query_lower:
+                    extracted_district = district_name
+                    break
+            
+            # Also check common aliases
+            district_aliases = {
+                'kampala': 'Kampala',
+                'kcca': 'Kampala', 
+                'gulu': 'Gulu',
+                'kalangala': 'Kalangala'
+            }
+            for alias, full_name in district_aliases.items():
+                if alias in query_lower and not extracted_district:
+                    extracted_district = full_name
+                    break
+            
+            for source in ['kcca', 'maaif', 'mwts', 'consolidated', 'oag']:
+                if source in query_lower:
+                    extracted_source = source.upper()
+                    break
+            
+            # Extract year using dynamic year list
+            for year in self.year_whitelist:
+                if year in query_lower:
+                    extracted_year = year
+                    has_year = True
+                    break
+            
+            # Only handle relative year terms if explicitly mentioned
+            if not extracted_year:
+                if 'last year' in query_lower or 'previous year' in query_lower:
+                    extracted_year = self.previous_year
+                    has_year = True
+                elif 'this year' in query_lower or 'current year' in query_lower:
+                    extracted_year = self.current_year
+                    has_year = True
+                elif 'recent' in query_lower and 'year' in query_lower:
+                    # Use the most recent year from available data
+                    extracted_year = max(self.year_whitelist) if self.year_whitelist else self.previous_year
+                    has_year = True
+            
+            # Be more permissive - if we have some context, try to answer
+            missing_filters = []
+            if not has_district:
+                missing_filters.append("district")
+            if not has_source:
+                missing_filters.append("source")
+            if not has_year:
+                missing_filters.append("year")
+            
+            # If user seems frustrated or asking for results, be more permissive
+            frustration_indicators = ['already', 'just said', 'specified', 'provided', 'crazy', 'answer']
+            is_frustrated = any(indicator in query_lower for indicator in frustration_indicators)
+            
+            can_answer_directly = len(missing_filters) <= 1 or is_frustrated  # More permissive
+            confidence_score = 0.8 if can_answer_directly else 0.3
+            
+            # Generate follow-up suggestion
+            if missing_filters and not is_frustrated:
+                if "district" in missing_filters and "source" in missing_filters:
+                    suggested_follow_up = "I'd be happy to help you with that information! Could you please specify which district and department/ministry you're asking about?"
+                elif "district" in missing_filters:
+                    suggested_follow_up = "Thanks for your question! Could you please specify which district you're asking about?"
+                elif "source" in missing_filters:
+                    suggested_follow_up = "I can help you with that! Could you please specify which department or ministry you're asking about?"
+                elif "year" in missing_filters:
+                    suggested_follow_up = "Great question! Could you please specify which year you're interested in?"
+                else:
+                    suggested_follow_up = "Could you please provide more specific details to help me give you a precise answer?"
+            else:
+                suggested_follow_up = None
+            
+            # Create expanded query
+            expanded_query = query
+            if extracted_district:
+                expanded_query += f" for {extracted_district} district"
+            if extracted_source:
+                expanded_query += f" from {extracted_source}"
+            if extracted_year:
+                expanded_query += f" in {extracted_year}"
+            
+            query_analysis = QueryAnalysis(
+                has_district=has_district,
+                has_source=has_source,
+                has_year=has_year,
+                extracted_district=extracted_district,
+                extracted_source=extracted_source,
+                extracted_year=extracted_year,
+                confidence_score=confidence_score,
+                can_answer_directly=can_answer_directly,
+                missing_filters=missing_filters,
+                suggested_follow_up=suggested_follow_up,
+                expanded_query=expanded_query
+            )
+        
+        # Update conversation context
+        if query_analysis.extracted_district:
+            conversation_context["district"] = query_analysis.extracted_district
+        if query_analysis.extracted_source:
+            conversation_context["source"] = query_analysis.extracted_source
+        if query_analysis.extracted_year:
+            conversation_context["year"] = query_analysis.extracted_year
+        
+        state["query_analysis"] = query_analysis
+        state["conversation_context"] = conversation_context
+        
+        self.logger.info(f"✅ ANALYSIS COMPLETE: district={query_analysis.has_district}, source={query_analysis.has_source}, year={query_analysis.has_year}")
+        self.logger.info(f"📈 Confidence: {query_analysis.confidence_score:.2f}, Can answer directly: {query_analysis.can_answer_directly}")
+        if query_analysis.expanded_query:
+            self.logger.info(f"🔄 Expanded query: {query_analysis.expanded_query}")
+        
+        return state
+    
+    def _decide_action(self, state: ConversationState) -> ConversationState:
+        """Decide what action to take based on query analysis"""
+        
+        analysis = state["query_analysis"]
+        
+        # Add decision reasoning
+        if analysis.can_answer_directly and analysis.confidence_score > 0.7:
+            self.logger.info(f"🚀 DECISION: Query is complete, proceeding with RAG")
+            self.logger.info(f"📊 REASONING: Confidence={analysis.confidence_score:.2f}, Missing filters={len(analysis.missing_filters or [])}")
+            if analysis.missing_filters:
+                self.logger.info(f"📋 Missing: {', '.join(analysis.missing_filters)}")
+            else:
+                self.logger.info(f"✅ All required information available")
+        else:
+            self.logger.info(f"❓ DECISION: Query incomplete, asking follow-up")
+            self.logger.info(f"📊 REASONING: Confidence={analysis.confidence_score:.2f}, Missing filters={len(analysis.missing_filters or [])}")
+            if analysis.missing_filters:
+                self.logger.info(f"📋 Missing: {', '.join(analysis.missing_filters)}")
+            self.logger.info(f"💡 Follow-up needed: {analysis.suggested_follow_up}")
+        
+        return state
+    
+    def _should_perform_rag(self, state: ConversationState) -> str:
+        """Determine whether to perform RAG or ask follow-up"""
+        
+        analysis = state["query_analysis"]
+        conversation_context = state.get("conversation_context", {})
+        recent_messages = state.get("messages", [])
+        
+        # Check if we have enough context from conversation history
+        has_district_context = analysis.has_district or conversation_context.get("district")
+        has_source_context = analysis.has_source or conversation_context.get("source") 
+        has_year_context = analysis.has_year or conversation_context.get("year")
+        
+        # Count how many context pieces we have
+        context_count = sum([bool(has_district_context), bool(has_source_context), bool(has_year_context)])
+        
+        # For PDM queries, we need more specific information
+        current_query = state["current_query"].lower()
+        recent_messages = state.get("messages", [])
+        
+        # Check if this is a PDM query by looking at current query OR recent conversation
+        is_pdm_query = "pdm" in current_query or "parish development" in current_query
+        
+        # Also check recent messages for PDM context
+        if not is_pdm_query and recent_messages:
+            for msg in recent_messages[-3:]:  # Check last 3 messages
+                if isinstance(msg, HumanMessage) and ("pdm" in msg.content.lower() or "parish development" in msg.content.lower()):
+                    is_pdm_query = True
+                    break
+        
+        if is_pdm_query:
+            # For PDM queries, we need district AND year to be specific enough
+            # But we need them to be explicitly provided in the current conversation, not just inferred
+            if has_district_context and has_year_context:
+                # Check if both district and year are explicitly mentioned in recent messages
+                explicit_district = False
+                explicit_year = False
+                
+                for msg in recent_messages[-3:]:  # Check last 3 messages
+                    if isinstance(msg, HumanMessage):
+                        content = msg.content.lower()
+                        if any(district in content for district in ["gulu", "kalangala", "kampala", "namutumba"]):
+                            explicit_district = True
+                        if any(year in content for year in ["2022", "2023", "2022/23", "2023/24"]):
+                            explicit_year = True
+                
+                if explicit_district and explicit_year:
+                    self.logger.info(f"🚀 DECISION: PDM query with explicit district and year, proceeding with RAG")
+                    self.logger.info(f"📊 REASONING: PDM query - explicit_district={explicit_district}, explicit_year={explicit_year}")
+                    return "rag"
+                else:
+                    self.logger.info(f"❓ DECISION: PDM query needs explicit district and year, asking follow-up")
+                    self.logger.info(f"📊 REASONING: PDM query - explicit_district={explicit_district}, explicit_year={explicit_year}")
+                    return "follow_up"
+            else:
+                self.logger.info(f"❓ DECISION: PDM query needs more specific info, asking follow-up")
+                self.logger.info(f"📊 REASONING: PDM query - district={has_district_context}, year={has_year_context}")
+                return "follow_up"
+        
+        # For general queries, be more conservative - need at least 2 pieces AND high confidence
+        if context_count >= 2 and analysis.confidence_score > 0.8:
+            self.logger.info(f"🚀 DECISION: Sufficient context with high confidence, proceeding with RAG")
+            self.logger.info(f"📊 REASONING: Context pieces: district={has_district_context}, source={has_source_context}, year={has_year_context}, confidence={analysis.confidence_score}")
+            return "rag"
+        
+        # If user seems frustrated (short responses like "no"), proceed with RAG
+        if recent_messages and len(recent_messages) >= 3:  # Need more messages to detect frustration
+            last_user_message = None
+            for msg in reversed(recent_messages):
+                if isinstance(msg, HumanMessage):
+                    last_user_message = msg.content.lower().strip()
+                    break
+            
+            if last_user_message and len(last_user_message) < 10 and any(word in last_user_message for word in ["no", "yes", "ok", "sure"]):
+                self.logger.info(f"🚀 DECISION: User seems frustrated with short response, proceeding with RAG")
+                return "rag"
+        
+        # Original logic for direct answers
+        if analysis.can_answer_directly and analysis.confidence_score > 0.7:
+            return "rag"
+        else:
+            return "follow_up"
+    
+    def _ask_follow_up(self, state: ConversationState) -> ConversationState:
+        """Generate a follow-up question to clarify missing information"""
+        
+        analysis = state["query_analysis"]
+        current_query = state["current_query"].lower()
+        conversation_context = state.get("conversation_context", {})
+        
+        # Check if this is a PDM query
+        is_pdm_query = "pdm" in current_query or "parish development" in current_query
+        
+        if is_pdm_query:
+            # Generate PDM-specific follow-up questions
+            missing_info = []
+            
+            if not analysis.has_district and not conversation_context.get("district"):
+                missing_info.append("district (e.g., Gulu, Kalangala)")
+            
+            if not analysis.has_year and not conversation_context.get("year"):
+                missing_info.append("year (e.g., 2022, 2023)")
+            
+            if missing_info:
+                follow_up_message = f"For PDM administrative costs information, I need to know the {', '.join(missing_info)}. Could you please specify these details?"
+            else:
+                follow_up_message = "Could you please provide more specific details about the PDM administrative costs you're looking for?"
+        else:
+            # Use the original follow-up logic
+            if analysis.suggested_follow_up:
+                follow_up_message = analysis.suggested_follow_up
+            else:
+                follow_up_message = "Could you please provide more specific details to help me give you a precise answer?"
+        
+        state["final_response"] = follow_up_message
+        state["last_ai_message_time"] = time.time()
+        
+        return state
+    
+    def _build_comprehensive_query(self, current_query: str, analysis, conversation_context: dict, recent_messages: list) -> str:
+        """Build a better RAG query from conversation.
+        - If latest message is a short modifier (e.g., "financial"), merge it into the last substantive question.
+        - If latest message looks like filters (district/year), keep the last question unchanged.
+        - Otherwise, use the current message.
+        """
+
+        def is_interrogative(text: str) -> bool:
+            t = text.lower().strip()
+            return any(t.startswith(w) for w in ["what", "how", "why", "when", "where", "which", "who"]) or t.endswith("?")
+
+        def is_filter_like(text: str) -> bool:
+            t = text.lower()
+            if "district" in t:
+                return True
+            if re.search(r"\b20\d{2}\b", t) or re.search(r"20\d{2}\s*[\-/–]\s*\d{2}\b", t):
+                return True
+            if self._extract_districts_list(text):
+                return True
+            return False
+
+        # Find last substantive user question
+        last_question = None
+        for msg in reversed(recent_messages[:-1] if recent_messages else []):
+            if isinstance(msg, HumanMessage):
+                if is_interrogative(msg.content) and len(msg.content.strip()) > 15:
+                    last_question = msg.content.strip()
+                    break
+
+        cq = current_query.strip()
+        words = cq.split()
+        is_short_modifier = (not is_interrogative(cq)) and (len(words) <= 3)
+
+        if is_filter_like(cq) and last_question:
+            comprehensive_query = last_question
+        elif is_short_modifier and last_question:
+            modifier = cq
+            if modifier.lower() in last_question.lower():
+                comprehensive_query = last_question
+            else:
+                if last_question.endswith('?'):
+                    comprehensive_query = last_question[:-1] + f" for {modifier}?"
+                else:
+                    comprehensive_query = last_question + f" for {modifier}"
+        else:
+            comprehensive_query = current_query
+        
+        self.logger.info(f"🔄 COMPREHENSIVE QUERY: '{comprehensive_query}'")
+        return comprehensive_query
+
+    def _rewrite_query_with_llm(self, recent_messages: list, draft_query: str) -> str:
+        """Use the LLM to rewrite a clean, focused RAG query from the conversation.
+        Rules enforced in prompt:
+        - Keep the user's main information need from the last substantive question
+        - Integrate short modifiers (e.g., "financial") into that question when appropriate
+        - Do NOT include filter text (years/districts/sources) in the query; those are handled separately
+        - Return a single plain sentence only (no quotes, no markdown)
+        """
+        try:
+            # Build a compact conversation transcript (last 6 messages max)
+            convo_lines = []
+            for msg in recent_messages[-6:]:
+                if isinstance(msg, HumanMessage):
+                    convo_lines.append(f"User: {msg.content}")
+                elif isinstance(msg, AIMessage):
+                    convo_lines.append(f"Assistant: {msg.content}")
+
+            convo_text = "\n".join(convo_lines)
+
+            """
+                                "DECISION GUIDANCE:\n"
+                    "- If the latest user message looks like a modifier (e.g., 'financial'), merge it into the best prior question.\n"
+                    "- If the latest message provides filters (e.g., districts, years), DO NOT embed them; keep the base question.\n"
+                    "- If the latest message itself is a full, clear question, use it.\n"
+                    "- If the draft query is already good, you may refine its clarity but keep the same intent.\n\n"
+            """
+        
+        
+            prompt = ChatPromptTemplate.from_messages([
+                SystemMessage(content=(
+                    "ROLE: Query Rewriter for a RAG system.\n\n"
+                    "PRIMARY OBJECTIVE:\n- Produce ONE retrieval-focused sentence that best represents the user's information need.\n"
+                    "- Maximize recall of relevant evidence; be specific but not overconstrained.\n\n"
+                    "INPUTS:\n- Conversation with User and Assistant turns (latest last).\n- A draft query (heuristic).\n\n"
+                    "OPERATING PRINCIPLES:\n"
+                    "1) Use the last substantive USER question as the backbone of intent.\n"
+                    "2) Merge helpful domain modifiers from any USER turns (financial, procurement, risk) when they sharpen focus; ignore if not helpful.\n"
+                    "3) Treat Assistant messages as guidance only; if the user later provided filters (years, districts, sources), DO NOT embed them in the query (filters are applied separately).\n"
+                    "4) Remove meta-verbs like 'summarize', 'list', 'explain', 'compare' from the query.\n"
+                    "5) Prefer content-bearing terms (topics, programs, outcomes) over task phrasing.\n"
+                    "6) If the latest user message is filters-only, keep the prior substantive question unchanged.\n"
+                    "7) If the draft query is already strong, refine wording for clarity but keep the same intent.\n\n"
+                    "EXAMPLES (multi-turn):\n"
+                    "A)\nUser: What are the top 5 priorities for improving audit procedures?\nAssistant: Could you specify the scope (e.g., financial, procurement)?\nUser: Financial\n→ Output: Top priorities for improving financial audit procedures.\n\n"
+                    "B)\nUser: How were PDM administrative costs utilized and what was the impact of shortfalls?\nAssistant: Please specify district/year for precision.\nUser: Namutumba and Lwengo Districts (2022/23)\n→ Output: How were PDM administrative costs utilized and what was the impact of shortfalls.\n(Exclude districts/years; they are filters.)\n\n"
+                    "C)\nUser: Summarize risk management issues in audit reports.\n→ Output: Key risk management issues in audit reports.\n\n"
+                    "CONSTRAINTS:\n- Do NOT include filters (years, districts, sources, filenames).\n- Do NOT include quotes/markdown/bullets or multiple sentences.\n- Return exactly one plain sentence."
+                )),
+                HumanMessage(content=(
+                    f"Conversation (most recent last):\n{convo_text}\n\n"
+                    f"Draft query: {draft_query}\n\n"
+                    "Rewrite the single best retrieval query sentence now:"
+                )),
+            ])
+
+            # Add timeout for LLM call
+            import signal
+            
+            def timeout_handler(signum, frame):
+                raise TimeoutError("LLM rewrite timeout")
+            
+            # Set 10 second timeout
+            signal.signal(signal.SIGALRM, timeout_handler)
+            signal.alarm(10)
+            
+            try:
+                resp = self.llm.invoke(prompt.format_messages())
+                signal.alarm(0)  # Cancel timeout
+                
+                rewritten = getattr(resp, 'content', '').strip()
+                # Basic sanitization: keep it one line
+                rewritten = rewritten.replace('\n', ' ').strip()
+                if rewritten and len(rewritten) > 5:  # Basic quality check
+                    self.logger.info(f"🛠️ LLM REWRITER: '{rewritten}'")
+                    return rewritten
+                else:
+                    self.logger.info(f"⚠️ LLM rewrite too short/empty, using draft query")
+                    return draft_query
+            except TimeoutError:
+                signal.alarm(0)
+                self.logger.info(f"⚠️ LLM rewrite timeout after 10s, using draft query")
+                return draft_query
+            except Exception as e:
+                signal.alarm(0)
+                self.logger.info(f"⚠️ LLM rewrite failed, using draft query. Error: {e}")
+                return draft_query
+        except Exception as e:
+            self.logger.info(f"⚠️ LLM rewrite setup failed, using draft query. Error: {e}")
+        return draft_query
+    
+    def _perform_rag(self, state: ConversationState) -> ConversationState:
+        """Perform RAG retrieval with smart query expansion"""
+        
+        query = state["current_query"]
+        analysis = state["query_analysis"]
+        conversation_context = state.get("conversation_context", {})
+        recent_messages = state.get("messages", [])
+        
+        # Build comprehensive query from conversation history
+        draft_query = self._build_comprehensive_query(query, analysis, conversation_context, recent_messages)
+        # Let LLM rewrite a clean, focused search query
+        search_query = self._rewrite_query_with_llm(recent_messages, draft_query)
+        
+        self.logger.info(f"🔍 RAG RETRIEVAL: Starting for query: '{search_query[:50]}...'")
+        self.logger.info(f"📊 Analysis: district={analysis.has_district}, source={analysis.has_source}, year={analysis.has_year}")
+        
+        try:
+            # Build filters from analysis and conversation context
+            filters = {}
+            
+            # Use conversation context to fill in missing filters
+            source = analysis.extracted_source or conversation_context.get("source")
+            district = analysis.extracted_district or conversation_context.get("district")
+            year = analysis.extracted_year or conversation_context.get("year")
+            
+            if source:
+                filters["source"] = [source]  # Qdrant expects lists
+                self.logger.info(f"🎯 Filter: source={source}")
+            
+            if year:
+                filters["year"] = [year]
+                self.logger.info(f"🎯 Filter: year={year}")
+            
+            if district:
+                # Map district to source if needed
+                if district.upper() == "KAMPALA":
+                    filters["source"] = ["KCCA"]
+                    self.logger.info(f"🎯 Filter: district={district} -> source=KCCA")
+                elif district.upper() in ["GULU", "KALANGALA"]:
+                    filters["source"] = [f"{district.upper()} DLG"]
+                    self.logger.info(f"🎯 Filter: district={district} -> source={district.upper()} DLG")
+            
+            # Run RAG pipeline with correct parameters
+            result = self.pipeline_manager.run(
+                query=search_query,  # Use expanded query
+                sources=filters.get("source") if filters.get("source") else None,
+                auto_infer_filters=False,  # Our agent already handled filter inference
+                filters=filters if filters else None
+            )
+            
+            self.logger.info(f"✅ RAG completed: Found {len(result.sources)} sources")
+            self.logger.info(f"⏱️ Execution time: {result.execution_time:.2f}s")
+            
+            # Store RAG result in state
+            state["rag_result"] = result
+            state["rag_query"] = search_query
+            
+        except Exception as e:
+            self.logger.info(f"❌ RAG retrieval failed: {e}")
+            state["rag_result"] = None
+        
+        return state
+    
+    def _generate_response(self, state: ConversationState) -> ConversationState:
+        """Generate final response using RAG results"""
+        
+        rag_result = state["rag_result"]
+        
+        self.logger.info(f"📝 RESPONSE: Using RAG result ({len(rag_result.answer)} chars)")
+        
+        # Store the final response directly from RAG
+        state["final_response"] = rag_result.answer
+        state["last_ai_message_time"] = time.time()
+        
+        return state
+    
+    def chat(self, user_input: str, conversation_id: str = "default") -> str:
+        """Main chat interface with conversation management"""
+        
+        self.logger.info(f"💬 CHAT: Processing user input: '{user_input[:50]}...'")
+        self.logger.info(f"📊 Session: {conversation_id}")
+        
+        # Load conversation history
+        conversation_file = self.conversations_dir / f"{conversation_id}.json"
+        conversation = self._load_conversation(conversation_file)
+        
+        # Add user message to conversation
+        conversation["messages"].append(HumanMessage(content=user_input))
+        
+        self.logger.info(f"🔄 LANGGRAPH: Starting graph execution")
+        
+        # Prepare state for LangGraph with conversation context
+        state = ConversationState(
+            conversation_id=conversation_id,
+            messages=conversation["messages"],
+            current_query=user_input,
+            query_analysis=None,
+            conversation_context=conversation.get("context", {}),
+            rag_result=None,
+            final_response=None,
+            session_start_time=conversation["session_start_time"],
+            last_ai_message_time=conversation["last_ai_message_time"]
+        )
+        
+        # Run the graph
+        final_state = self.graph.invoke(state)
+        
+        # Add the AI response to conversation
+        if final_state["final_response"]:
+            conversation["messages"].append(AIMessage(content=final_state["final_response"]))
+        
+        # Update conversation state
+        conversation["last_ai_message_time"] = final_state["last_ai_message_time"]
+        conversation["context"] = final_state["conversation_context"]
+        
+        # Save conversation
+        self._save_conversation(conversation_file, conversation)
+        
+        self.logger.info(f"✅ LANGGRAPH: Graph execution completed")
+        self.logger.info(f"🎯 CHAT COMPLETE: Response ready")
+        
+        # Return both response and RAG result for UI
+        return {
+            'response': final_state["final_response"] or "I apologize, but I couldn't process your request.",
+            'rag_result': final_state["rag_result"],
+            '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:
+                self.logger.info(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 LangChain 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})
+            
+            data = {
+                "messages": messages_data,
+                "session_start_time": conversation["session_start_time"],
+                "last_ai_message_time": conversation["last_ai_message_time"],
+                "context": conversation.get("context", {}),
+                "last_updated": datetime.now().isoformat()
+            }
+            
+            with open(conversation_file, "w") as f:
+                json.dump(data, f, indent=2)
+                
+        except Exception as e:
+            self.logger.info(f"⚠️ Could not save conversation: {e}")
+
+
+def get_chatbot():
+    """Get chatbot instance"""
+    return IntelligentRAGChatbot()
+
+if __name__ == "__main__":
+    # Test the chatbot
+    chatbot = IntelligentRAGChatbot()
+    
+    # Test conversation
+    test_queries = [
+        "How much was the budget allocation for government salary payroll management?",
+        "Namutumba district in 2023",
+        "KCCA"
+    ]
+    
+    for query in test_queries:
+        self.logger.info(f"\n{'='*50}")
+        self.logger.info(f"User: {query}")
+        response = chatbot.chat(query)
+        self.logger.info(f"Bot: {response}")