Gemini FSA (#6)

- refactor (1154c8ddd84455509a7f915a8588d951f925bbc6)
- refactor (72318ee79a0a0bcafca07cc5be70aace39a25f0c)
- add district Metadata (7c8b7838d143b45686e8360bbf64d2ef3c4a5624)
- refactor and add sample questions (02d7f4f76ed3415d70f279e6ab188435ae60fe92)
- add retrieval visualisations (5262a14ec5cccd4d5e9796cf7c881b211809f9fc)
- add Retrieval Distribution stats (763a8b9f7efaf0b823138df29267434d72edd477)
- Merge branch 'main' of https://huggingface.co/spaces/akryldigital/audit_assistant (264ca849bc4217b186eb0f8de0d23bce39856de7)
- Merge branch 'main' of https://huggingface.co/spaces/akryldigital/audit_assistant (b4984e28f7a41b312d3eba90b68651ab903f7d08)
- refactor + add gemini (72eb0bfa173ea05e8fd8e3b63429ef4678a01663)
- fix use_container_width=False (f8a1d4171fe8d1b3b84938b7b14bbe497e4159fe)
- finalize gemini version (3fc1b5f53b40772ba3c8abf400f1d987c12c4ee1)
- add gemini traceability (6f5999e84c97f19e3d1ff442873b52ef6ca8208e)
- fix gemini chunk extraction (06faccdb62f42a514e9f6b8cc93ca73f5cab5fa1)
- fix gemini chunk extraction (54bf55f7a03022a79da5ead74a289cb893efa88f)
- add upload debug (39edab443db3982dfc9f582868c66b9ab787a208)
- Remove scripts and ignore local_* files (de1d74a230264cf4ae8df516291869434f265d9f)

.gitignore

@@ -109,4 +109,7 @@ pytest_cache/
 tmp/
 temp/
 *.tmp
-*.temp
+*.temp
+
+
+local_*

app.py

@@ -10,6 +10,7 @@ import uuid
 import logging
 import traceback
 from pathlib import Path
+
 from collections import Counter
 from typing import List, Dict, Any, Optional
 
@@ -19,10 +20,11 @@ import streamlit as st
 import plotly.express as px
 from langchain_core.messages import HumanMessage, AIMessage
 
-from multi_agent_chatbot import get_multi_agent_chatbot
-from smart_chatbot import get_chatbot as get_smart_chatbot
-from src.reporting.snowflake_connector import save_to_snowflake
-from src.reporting.feedback_schema import create_feedback_from_dict
+
+from src.agents import get_multi_agent_chatbot, get_smart_chatbot, get_gemini_chatbot
+from src.feedback import FeedbackManager
+from src.ui_components import get_custom_css, display_chunk_statistics_charts, display_chunk_statistics_table, extract_chunk_statistics
+
 from src.config.paths import (
     IS_DEPLOYED,
     PROJECT_DIR,
@@ -31,6 +33,7 @@ from src.config.paths import (
     CONVERSATIONS_DIR,
 )
 
+
 # ===== CRITICAL: Fix OMP_NUM_THREADS FIRST, before ANY other imports =====
 # Some libraries load at import time and will fail if OMP_NUM_THREADS is invalid
 omp_threads = os.environ.get("OMP_NUM_THREADS", "")
@@ -70,6 +73,9 @@ if IS_DEPLOYED and HF_CACHE_DIR:
     except (PermissionError, OSError):
         # If we can't create it, log but continue (might already exist from Dockerfile)
         pass
+else:
+    from dotenv import load_dotenv
+    load_dotenv()
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -90,116 +96,9 @@ st.set_page_config(
     page_title="Intelligent Audit Report Chatbot"
 )
 
-# Custom CSS
-st.markdown("""
-<style>
-    .main-header {
-        font-size: 2.5rem;
-        font-weight: bold;
-        color: #1f77b4;
-        text-align: center;
-        margin-bottom: 1rem;
-        width: 100%;
-        display: block;
-    }
-    
-    .subtitle {
-        font-size: 1.2rem;
-        color: #666;
-        text-align: center;
-        margin-bottom: 2rem;
-        width: 100%;
-        display: block;
-    }
-    
-    .session-info {
-        background-color: #f0f2f6;
-        padding: 10px;
-        border-radius: 5px;
-        margin-bottom: 20px;
-        font-size: 0.9rem;
-    }
-    
-    .user-message {
-        background-color: #007bff;
-        color: white;
-        padding: 12px 16px;
-        border-radius: 18px 18px 4px 18px;
-        margin: 8px 0;
-        margin-left: 20%;
-        word-wrap: break-word;
-    }
-    
-    .bot-message {
-        background-color: #f1f3f4;
-        color: #333;
-        padding: 12px 16px;
-        border-radius: 18px 18px 18px 4px;
-        margin: 8px 0;
-        margin-right: 20%;
-        word-wrap: break-word;
-        border: 1px solid #e0e0e0;
-    }
-    
-    .filter-section {
-        margin-bottom: 20px;
-        padding: 15px;
-        background-color: #f8f9fa;
-        border-radius: 8px;
-        border: 1px solid #e9ecef;
-    }
-    
-    .filter-title {
-        font-weight: bold;
-        margin-bottom: 10px;
-        color: #495057;
-    }
-    
-    .feedback-section {
-        background-color: #f8f9fa;
-        padding: 20px;
-        border-radius: 10px;
-        margin-top: 30px;
-        border: 2px solid #dee2e6;
-    }
-    
-    .retrieval-history {
-        background-color: #ffffff;
-        padding: 15px;
-        border-radius: 5px;
-        margin: 10px 0;
-        border-left: 4px solid #007bff;
-    }
-    
-    .retrieval-distribution-container {
-        background-color: #ffffff;
-        padding: 25px;
-        border-radius: 10px;
-        margin: 20px 0;
-        border: 2px solid #e0e0e0;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1), 0 2px 4px rgba(0, 0, 0, 0.06);
-    }
-    
-    .metric-label {
-        font-size: 0.9rem;
-        color: #555;
-        margin-bottom: 5px;
-        text-align: center;
-    }
-    
-    .metric-value {
-        font-size: 1.8rem;
-        font-weight: bold;
-        color: #000000;
-        text-align: center;
-    }
-    
-    .metric-container {
-        text-align: center;
-        padding: 10px;
-    }
-</style>
-""", unsafe_allow_html=True)
+
+st.markdown(get_custom_css(), unsafe_allow_html=True)
+
 
 def get_system_type():
     """Get the current system type"""
@@ -209,14 +108,17 @@ def get_system_type():
     else:
         return "Multi-Agent System"
 
-def get_chatbot():
-    """Initialize and return the chatbot based on system type"""
-    # Check environment variable for system type
-    system = os.environ.get('CHATBOT_SYSTEM', 'multi-agent')
-    if system == 'smart':
-        return get_smart_chatbot()
+def get_chatbot(version: str = "v1"):
+    """Initialize and return the chatbot based on version"""
+    if version == "beta":
+        return get_gemini_chatbot()
     else:
-        return get_multi_agent_chatbot()
+        # Check environment variable for system type (v1)
+        system = os.environ.get('CHATBOT_SYSTEM', 'multi-agent')
+        if system == 'smart':
+            return get_smart_chatbot()
+        else:
+            return get_multi_agent_chatbot()
 
 def serialize_messages(messages):
     """Serialize LangChain messages to dictionaries"""
@@ -262,368 +164,8 @@ def serialize_documents(sources):
     return serialized
 
 
-def extract_transcript(messages: List[Any]) -> List[Dict[str, str]]:
-    """Extract transcript from messages - only user and bot messages, no extra metadata"""
-    transcript = []
-    for msg in messages:
-        if isinstance(msg, HumanMessage):
-            transcript.append({
-                "role": "user",
-                "content": str(msg.content) if hasattr(msg, 'content') else str(msg)
-            })
-        elif isinstance(msg, AIMessage):
-            transcript.append({
-                "role": "assistant",
-                "content": str(msg.content) if hasattr(msg, 'content') else str(msg)
-            })
-    return transcript
-
-
-def build_retrievals_structure(rag_retrieval_history: List[Dict[str, Any]], messages: List[Any]) -> List[Dict[str, Any]]:
-    """Build retrievals structure from retrieval history"""
-    retrievals = []
-    
-    for entry in rag_retrieval_history:
-        # Get the user message that triggered this retrieval
-        # The entry has conversation_up_to which includes messages up to that point
-        conversation_up_to = entry.get("conversation_up_to", [])
-        
-        # Find the last user message in conversation_up_to (this is the trigger)
-        user_message_trigger = ""
-        for msg_dict in reversed(conversation_up_to):
-            if msg_dict.get("type") == "HumanMessage":
-                user_message_trigger = msg_dict.get("content", "")
-                break
-        
-        # Fallback: if not found in conversation_up_to, get from actual messages
-        # This handles edge cases where conversation_up_to might be incomplete
-        if not user_message_trigger:
-            # Find which retrieval this is (0-indexed)
-            retrieval_idx = rag_retrieval_history.index(entry)
-            # The user message that triggered this retrieval is at position (retrieval_idx * 2)
-            # because each retrieval is preceded by: user message, bot response, user message, ...
-            # But we need to account for the fact that the first retrieval happens after the first user message
-            user_msgs = [msg for msg in messages if isinstance(msg, HumanMessage)]
-            if retrieval_idx < len(user_msgs):
-                user_message_trigger = str(user_msgs[retrieval_idx].content)
-            elif user_msgs:
-                # Fallback to last user message
-                user_message_trigger = str(user_msgs[-1].content)
-        
-        # Get retrieved documents and truncate content to 100 chars
-        docs_retrieved = entry.get("docs_retrieved", [])
-        retrieved_docs = []
-        for doc in docs_retrieved:
-            doc_copy = doc.copy()
-            # Truncate content to 100 characters (keep all other fields)
-            if "content" in doc_copy:
-                doc_copy["content"] = doc_copy["content"][:100]
-            retrieved_docs.append(doc_copy)
-        
-        retrievals.append({
-            "retrieved_docs": retrieved_docs,
-            "user_message_trigger": user_message_trigger
-        })
-    
-    return retrievals
-
-
-def build_feedback_score_related_retrieval_docs(
-    is_feedback_about_last_retrieval: bool,
-    messages: List[Any],
-    rag_retrieval_history: List[Dict[str, Any]]
-) -> Optional[Dict[str, Any]]:
-    """Build feedback_score_related_retrieval_docs structure"""
-    if not rag_retrieval_history:
-        return None
-    
-    # Get the relevant retrieval entry
-    if is_feedback_about_last_retrieval:
-        relevant_entry = rag_retrieval_history[-1]
-    else:
-        # If feedback is about all retrievals, use the last one as default
-        relevant_entry = rag_retrieval_history[-1]
-    
-    # Get conversation up to that point
-    conversation_up_to = relevant_entry.get("conversation_up_to", [])
-    
-    # Convert to transcript format (role/content)
-    conversation_up_to_point = []
-    for msg_dict in conversation_up_to:
-        if msg_dict.get("type") == "HumanMessage":
-            conversation_up_to_point.append({
-                "role": "user",
-                "content": msg_dict.get("content", "")
-            })
-        elif msg_dict.get("type") == "AIMessage":
-            conversation_up_to_point.append({
-                "role": "assistant",
-                "content": msg_dict.get("content", "")
-            })
-    
-    # Get retrieved docs with full content (not truncated)
-    retrieved_docs = relevant_entry.get("docs_retrieved", [])
-    
-    return {
-        "conversation_up_to_point": conversation_up_to_point,
-        "retrieved_docs": retrieved_docs
-    }
+feedback_manager = FeedbackManager()
 
-def extract_chunk_statistics(sources: List[Any]) -> Dict[str, Any]:
-    """Extract statistics from retrieved chunks."""
-    if not sources:
-        return {}
-    
-    sources_list = []
-    years = []
-    filenames = []
-    districts = []
-    
-    for doc in sources:
-        metadata = getattr(doc, 'metadata', {})
-        
-        # Extract source
-        source = metadata.get('source', 'Unknown')
-        sources_list.append(source)
-        
-        # Extract year
-        year = metadata.get('year', 'Unknown')
-        if year and year != 'Unknown':
-            try:
-                # Convert to int first, then back to string to ensure it's a proper year
-                year_int = int(float(year))  # Handle both int and float strings
-                if 1900 <= year_int <= 2030:  # Reasonable year range
-                    years.append(str(year_int))
-                else:
-                    years.append('Unknown')
-            except (ValueError, TypeError):
-                years.append('Unknown')
-        else:
-            years.append('Unknown')
-        
-        # Extract filename
-        filename = metadata.get('filename', 'Unknown')
-        filenames.append(filename)
-        
-        # Extract district
-        district = metadata.get('district', 'Unknown')
-        if district and district != 'Unknown':
-            districts.append(district)
-        else:
-            districts.append('Unknown')
-    
-    # Count occurrences
-    source_counts = Counter(sources_list)
-    year_counts = Counter(years)
-    filename_counts = Counter(filenames)
-    district_counts = Counter(districts)
-    
-    return {
-        'total_chunks': len(sources),
-        'unique_sources': len(source_counts),
-        'unique_years': len([y for y in year_counts.keys() if y != 'Unknown']),
-        'unique_filenames': len(filename_counts),
-        'unique_districts': len([d for d in district_counts.keys() if d != 'Unknown']),
-        'source_distribution': dict(source_counts),
-        'year_distribution': dict(year_counts),
-        'filename_distribution': dict(filename_counts),
-        'district_distribution': dict(district_counts),
-        'sources': sources_list,
-        'years': years,
-        'filenames': filenames,
-        'districts': districts
-    }
-
-def display_chunk_statistics_charts(stats: Dict[str, Any], title: str = "Retrieval Statistics"):
-    """Display statistics as interactive charts for 10+ results."""
-    if not stats or stats.get('total_chunks', 0) == 0:
-        return
-    
-    # Wrap everything in one styled container - open it
-    st.markdown(f"""
-    <div class="retrieval-distribution-container">
-        <h3 style="margin-top: 0;">📊 {title}</h3>
-        <div style="display: flex; justify-content: space-around; align-items: center; padding: 15px 0; border-bottom: 1px solid #e0e0e0; margin-bottom: 20px;">
-            <div class="metric-container">
-                <div class="metric-label">Total Chunks</div>
-                <div class="metric-value">{stats['total_chunks']}</div>
-            </div>
-            <div class="metric-container">
-                <div class="metric-label">Unique Sources</div>
-                <div class="metric-value">{stats['unique_sources']}</div>
-            </div>
-            <div class="metric-container">
-                <div class="metric-label">Unique Years</div>
-                <div class="metric-value">{stats['unique_years']}</div>
-            </div>
-            <div class="metric-container">
-                <div class="metric-label">Unique Files</div>
-                <div class="metric-value">{stats['unique_filenames']}</div>
-            </div>
-        </div>
-    """, unsafe_allow_html=True)
-    
-    # Charts - three columns to include Districts
-    col1, col2, col3 = st.columns(3)
-    
-    with col1:
-        # Source distribution chart
-        if stats['source_distribution']:
-            source_df = pd.DataFrame(
-                list(stats['source_distribution'].items()),
-                columns=['Source', 'Count']
-            )
-            fig_source = px.bar(
-                source_df,
-                x='Count',
-                y='Source',
-                orientation='h',
-                title='Distribution by Source',
-                color='Count',
-                color_continuous_scale='viridis'
-            )
-            fig_source.update_layout(height=400, showlegend=False)
-            st.plotly_chart(fig_source, use_container_width=True)
-    
-    with col2:
-        # Year distribution chart
-        if stats['year_distribution']:
-            # Filter out 'Unknown' years for the chart
-            year_dist_filtered = {k: v for k, v in stats['year_distribution'].items() if k != 'Unknown'}
-            if year_dist_filtered:
-                year_df = pd.DataFrame(
-                    list(year_dist_filtered.items()),
-                    columns=['Year', 'Count']
-                )
-                # Sort by year as integer but keep as string for categorical display
-                year_df['Year_Int'] = year_df['Year'].astype(int)
-                year_df = year_df.sort_values('Year_Int').drop('Year_Int', axis=1)
-                
-                fig_year = px.bar(
-                    year_df,
-                    x='Year',
-                    y='Count',
-                    title='Distribution by Year',
-                    color='Count',
-                    color_continuous_scale='plasma'
-                )
-                # Ensure years are treated as categorical (discrete) not continuous
-                fig_year.update_xaxes(type='category')
-                fig_year.update_layout(height=400, showlegend=False)
-                st.plotly_chart(fig_year, use_container_width=True)
-            else:
-                st.info("No valid years found in the results")
-    
-    with col3:
-        # District distribution chart
-        if stats.get('district_distribution'):
-            district_dist_filtered = {k: v for k, v in stats['district_distribution'].items() if k != 'Unknown'}
-            if district_dist_filtered:
-                district_df = pd.DataFrame(
-                    list(district_dist_filtered.items()),
-                    columns=['District', 'Count']
-                )
-                district_df = district_df.sort_values('Count', ascending=False)
-                
-                fig_district = px.bar(
-                    district_df,
-                    x='Count',
-                    y='District',
-                    orientation='h',
-                    title='Distribution by District',
-                    color='Count',
-                    color_continuous_scale='blues'
-                )
-                fig_district.update_layout(height=400, showlegend=False)
-                st.plotly_chart(fig_district, use_container_width=True)
-            else:
-                st.info("No valid districts found in the results")
-    
-    # Close the container
-    st.markdown('</div>', unsafe_allow_html=True)
-
-def display_chunk_statistics_table(stats: Dict[str, Any], title: str = "Retrieval Distribution"):
-    """Display statistics as tables for smaller results with fixed alignment."""
-    if not stats or stats.get('total_chunks', 0) == 0:
-        return
-    
-    # Wrap in styled container
-    st.markdown('<div class="retrieval-distribution-container">', unsafe_allow_html=True)
-    
-    st.subheader(f"📊 {title}")
-    
-    # Create a container with fixed height for alignment
-    stats_container = st.container()
-    
-    with stats_container:
-        # Create 4 equal columns for consistent alignment
-        col1, col2, col3, col4 = st.columns(4)
-        
-        with col1:
-            st.markdown("**🏘️ Districts**")
-            if stats.get('district_distribution'):
-                district_dist_filtered = {k: v for k, v in stats['district_distribution'].items() if k != 'Unknown'}
-                if district_dist_filtered:
-                    district_data = {
-                        "District": list(district_dist_filtered.keys()),
-                        "Count": list(district_dist_filtered.values())
-                    }
-                    district_df = pd.DataFrame(district_data).sort_values('Count', ascending=False)
-                    st.dataframe(district_df, hide_index=True, use_container_width=True)
-                else:
-                    st.write("No district data")
-            else:
-                st.write("No district data")
-        
-        with col2:
-            st.markdown("**📂 Sources**")
-            if stats['source_distribution']:
-                source_data = {
-                    "Source": list(stats['source_distribution'].keys()),
-                    "Count": list(stats['source_distribution'].values())
-                }
-                source_df = pd.DataFrame(source_data).sort_values('Count', ascending=False)
-                st.dataframe(source_df, hide_index=True, use_container_width=True)
-            else:
-                st.write("No source data")
-        
-        with col3:
-            st.markdown("**📅 Years**")
-            if stats['year_distribution']:
-                year_dist_filtered = {k: v for k, v in stats['year_distribution'].items() if k != 'Unknown'}
-                if year_dist_filtered:
-                    year_data = {
-                        "Year": list(year_dist_filtered.keys()),
-                        "Count": list(year_dist_filtered.values())
-                    }
-                    year_df = pd.DataFrame(year_data)
-                    # Sort by year as integer but display as string
-                    year_df['Year_Int'] = year_df['Year'].astype(int)
-                    year_df = year_df.sort_values('Year_Int')[['Year', 'Count']]
-                    st.dataframe(year_df, hide_index=True, use_container_width=True)
-                else:
-                    st.write("No year data")
-            else:
-                st.write("No year data")
-        
-        with col4:
-            st.markdown("**📄 Files**")
-            if stats['filename_distribution']:
-                filename_items = list(stats['filename_distribution'].items())
-                filename_items.sort(key=lambda x: x[1], reverse=True)
-                
-                # Show top files with truncated names
-                file_data = {
-                    "File": [f[:30] + "..." if len(f) > 30 else f for f, c in filename_items[:5]],
-                    "Count": [c for f, c in filename_items[:5]]
-                }
-                file_df = pd.DataFrame(file_data)
-                st.dataframe(file_df, hide_index=True, use_container_width=True)
-            else:
-                st.write("No file data")
-    
-    # Close container
-    st.markdown('</div>', unsafe_allow_html=True)
 
 @st.cache_data
 def load_filter_options():
@@ -649,11 +191,48 @@ def main():
     # Track RAG retrieval history for feedback
     if 'rag_retrieval_history' not in st.session_state:
         st.session_state.rag_retrieval_history = []
-    # Initialize chatbot only once per app session (cached)
-    if 'chatbot' not in st.session_state:
-        with st.spinner("🔄 Loading AI models and connecting to database..."):
-            st.session_state.chatbot = get_chatbot()
-        st.success("✅ AI system ready!")
+    # Version selection (v1 or beta)
+    if 'chatbot_version' not in st.session_state:
+        st.session_state.chatbot_version = "v1"
+    
+    # Initialize chatbot based on version (only if not already initialized for this version)
+    chatbot_version_key = f"chatbot_{st.session_state.chatbot_version}"
+    
+    # Check if we need to initialize: chatbot doesn't exist OR version changed
+    needs_init = (
+        chatbot_version_key not in st.session_state or 
+        st.session_state.get('_last_version') != st.session_state.chatbot_version
+    )
+    
+    if needs_init:
+        try:
+            # Different spinner messages for different versions
+            if st.session_state.chatbot_version == "beta":
+                spinner_msg = "🔄 Initializing Gemini FSA"
+            else:
+                spinner_msg = "🔄 Loading AI models and connecting to database..."
+            
+            with st.spinner(spinner_msg):
+                st.session_state[chatbot_version_key] = get_chatbot(st.session_state.chatbot_version)
+                st.session_state['_last_version'] = st.session_state.chatbot_version
+                st.session_state.chatbot = st.session_state[chatbot_version_key]
+            print("✅ AI system ready!")
+        except Exception as e:
+            st.error(f"❌ Failed to initialize chatbot: {str(e)}")
+            # Only show Gemini-specific error message for beta version
+            if st.session_state.chatbot_version == "beta":
+                st.error("Please check your environment variables (GEMINI_API_KEY, GEMINI_FILESTORE_NAME for beta)")
+            else:
+                st.error("Please check your configuration and ensure all required models and databases are accessible.")
+            # Reset to v1 to prevent infinite loop
+            st.session_state.chatbot_version = "v1"
+            st.session_state['_last_version'] = "v1"
+            if 'chatbot' in st.session_state:
+                del st.session_state['chatbot']
+            st.stop()  # Stop execution to prevent infinite loop
+    else:
+        # Chatbot already initialized for this version, just use it
+        st.session_state.chatbot = st.session_state[chatbot_version_key]
     
     # Reset conversation history if needed (but keep chatbot cached)
     if 'reset_conversation' in st.session_state and st.session_state.reset_conversation:
@@ -665,9 +244,43 @@ def main():
         st.session_state.reset_conversation = False
         st.rerun()
 
-    # Header - fully center aligned
-    st.markdown('<h1 class="main-header">🤖 Intelligent Audit Report Chatbot</h1>', unsafe_allow_html=True)
-    st.markdown('<p class="subtitle">Ask questions about audit reports. Use the sidebar filters to narrow down your search!</p>', unsafe_allow_html=True)
+
+    # Version selection radio button (top right)
+    col1, col2 = st.columns([3, 1])
+    with col1:
+        st.markdown('<p class="subtitle">Ask questions about audit reports. Use the sidebar filters to narrow down your search!</p>', unsafe_allow_html=True)
+    with col2:
+        st.markdown("<br>", unsafe_allow_html=True)  # Add some spacing
+        selected_version = st.radio(
+            "**Version:**",
+            options=["v1", "beta"],
+            index=0 if st.session_state.chatbot_version == "v1" else 1,
+            horizontal=True,
+            key="version_selector",
+            help="Select v1 (default RAG system) or beta (Gemini FSA)"
+        )
+        
+        # Update version if changed
+        if selected_version != st.session_state.chatbot_version:
+            # Store the old version to check if we need to switch
+            old_version = st.session_state.chatbot_version
+            st.session_state.chatbot_version = selected_version
+            
+            # If chatbot for new version already exists, just switch to it
+            new_chatbot_key = f"chatbot_{selected_version}"
+            if new_chatbot_key in st.session_state:
+                # Chatbot already exists, just switch
+                st.session_state.chatbot = st.session_state[new_chatbot_key]
+                st.session_state['_last_version'] = selected_version
+            else:
+                # Need to initialize new version - will be handled by initialization logic above
+                st.session_state['_last_version'] = old_version  # Set to old to trigger init check
+            
+            st.rerun()
+    
+    # Show version info
+    if st.session_state.chatbot_version == "beta":
+        st.info("🔬 **Beta Mode**: Using Google Gemini FSA")
     
     # Session info
     duration = int(time.time() - st.session_state.session_start_time)
@@ -729,7 +342,7 @@ def main():
         # Determine if filename filter is active
         filename_mode = len(selected_filenames) > 0
         # Sources filter
-        st.markdown('<div class="filter-section">', unsafe_allow_html=True)
+        # st.markdown('<div class="filter-section">', unsafe_allow_html=True)
         st.markdown('<div class="filter-title">📊 Sources</div>', unsafe_allow_html=True)
         selected_sources = st.multiselect(
             "Select sources:",
@@ -826,7 +439,7 @@ def main():
             )
         
         with col2:
-            send_button = st.button("Send", key="send_button", use_container_width=True)
+            send_button = st.button("Send", key="send_button", width='stretch')
         
         # Clear chat button
         if st.button("🗑️ Clear Chat", key="clear_chat_button"):
@@ -878,6 +491,36 @@ def main():
                         if rag_result:
                             sources = rag_result.get('sources', []) if isinstance(rag_result, dict) else (rag_result.sources if hasattr(rag_result, 'sources') else [])
                             
+                            # For Gemini, also check gemini_result for sources
+                            if not sources or len(sources) == 0:
+                                gemini_result = chat_result.get('gemini_result')
+                                print(f"🔍 DEBUG: Checking gemini_result for sources...")
+                                print(f"   gemini_result exists: {gemini_result is not None}")
+                                if gemini_result:
+                                    print(f"   gemini_result type: {type(gemini_result)}")
+                                    print(f"   has sources attr: {hasattr(gemini_result, 'sources')}")
+                                    if hasattr(gemini_result, 'sources'):
+                                        print(f"   sources length: {len(gemini_result.sources) if gemini_result.sources else 0}")
+                                
+                                if gemini_result and hasattr(gemini_result, 'sources'):
+                                    # Format Gemini sources for display
+                                    if hasattr(st.session_state.chatbot, 'gemini_client'):
+                                        sources = st.session_state.chatbot.gemini_client.format_sources_for_display(gemini_result)
+                                        print(f"✅ Formatted {len(sources)} sources from gemini_client")
+                                    elif hasattr(st.session_state.chatbot, '_format_gemini_sources'):
+                                        sources = st.session_state.chatbot._format_gemini_sources(gemini_result)
+                                        print(f"✅ Formatted {len(sources)} sources from _format_gemini_sources")
+                            
+                            # Update rag_result with sources if we found them
+                            if sources and len(sources) > 0:
+                                if isinstance(rag_result, dict):
+                                    rag_result['sources'] = sources
+                                elif hasattr(rag_result, 'sources'):
+                                    rag_result.sources = sources
+                                # Update last_rag_result with sources
+                                st.session_state.last_rag_result = rag_result
+                                print(f"✅ Updated rag_result with {len(sources)} sources")
+                            
                             # Get the actual RAG query
                             actual_rag_query = chat_result.get('actual_rag_query', '')
                             if actual_rag_query:
@@ -887,12 +530,25 @@ def main():
                             else:
                                 formatted_query = "No RAG query available"
                             
+                            # Extract filters from active filters
+                            filters_used = {
+                                "sources": st.session_state.active_filters.get('sources', []),
+                                "years": st.session_state.active_filters.get('years', []),
+                                "districts": st.session_state.active_filters.get('districts', []),
+                                "filenames": st.session_state.active_filters.get('filenames', [])
+                            }
+                            
                             retrieval_entry = {
                                 "conversation_up_to": serialize_messages(st.session_state.messages),
                                 "rag_query_expansion": formatted_query,
-                                "docs_retrieved": serialize_documents(sources)
+                                "docs_retrieved": serialize_documents(sources),
+                                "filters_applied": filters_used,
+                                "timestamp": time.time()
                             }
                             st.session_state.rag_retrieval_history.append(retrieval_entry)
+                            
+                            # Debug logging
+                            print(f"📊 RETRIEVAL TRACKING: {len(sources)} sources stored in retrieval history")
                     else:
                         response = chat_result
                         st.session_state.last_rag_result = None
@@ -922,6 +578,16 @@ def main():
                 # Dictionary format from multi-agent system
                 sources = rag_result['sources']
             
+            # For Gemini, also check if we need to format sources from gemini_result
+            if (not sources or len(sources) == 0) and isinstance(rag_result, dict):
+                gemini_result = rag_result.get('gemini_result')
+                if gemini_result and hasattr(gemini_result, 'sources'):
+                    # Format Gemini sources for display
+                    if hasattr(st.session_state.chatbot, 'gemini_client'):
+                        sources = st.session_state.chatbot.gemini_client.format_sources_for_display(gemini_result)
+                    elif hasattr(st.session_state.chatbot, '_format_gemini_sources'):
+                        sources = st.session_state.chatbot._format_gemini_sources(gemini_result)
+            
             if sources and len(sources) > 0:
                 # Count unique filenames
                 unique_filenames = set()
@@ -951,9 +617,18 @@ def main():
                 for i, doc in enumerate(sources):  # Show all documents
                     # Get relevance score and ID if available
                     metadata = getattr(doc, 'metadata', {})
-                    score = metadata.get('reranked_score', metadata.get('original_score', None))
-                    chunk_id = metadata.get('_id', 'Unknown')
-                    score_text = f" (Score: {score:.3f}, ID: {chunk_id[:8]}...)" if score is not None else f" (ID: {chunk_id[:8]}...)"
+                    # Handle both standard RAG scores and Gemini scores
+                    score = metadata.get('reranked_score') or metadata.get('original_score') or metadata.get('score')
+                    chunk_id = metadata.get('_id') or metadata.get('chunk_id', 'Unknown')
+                    if score is not None:
+                        try:
+                            score_text = f" (Score: {float(score):.3f})"
+                        except (ValueError, TypeError):
+                            score_text = ""
+                    else:
+                        score_text = ""
+                    if chunk_id and chunk_id != 'Unknown':
+                        score_text += f" (ID: {str(chunk_id)[:8]}...)" if score_text else f" (ID: {str(chunk_id)[:8]}...)"
                     
                     with st.expander(f"📄 Document {i+1}: {getattr(doc, 'metadata', {}).get('filename', 'Unknown')[:50]}...{score_text}"):
                         # Display document metadata with emojis
@@ -1031,7 +706,7 @@ def main():
                 
                 submitted = st.form_submit_button(
                     "📤 Submit Feedback", 
-                    use_container_width=True,
+                    width='stretch',
                     disabled=submit_disabled
                 )
                 
@@ -1043,16 +718,18 @@ def main():
                     st.write("🔍 **Debug: Feedback Data Being Submitted:**")
                     
                     # Extract transcript from messages
-                    transcript = extract_transcript(st.session_state.messages)
+                    transcript = feedback_manager.extract_transcript(st.session_state.messages)
                     
                     # Build retrievals structure
-                    retrievals = build_retrievals_structure(
+                    retrievals = feedback_manager.build_retrievals_structure(
+
                         st.session_state.rag_retrieval_history.copy() if st.session_state.rag_retrieval_history else [],
                         st.session_state.messages
                     )
                     
                     # Build feedback_score_related_retrieval_docs
-                    feedback_score_related_retrieval_docs = build_feedback_score_related_retrieval_docs(
+
+                    feedback_score_related_retrieval_docs = feedback_manager.build_feedback_score_related_retrieval_docs(
                         is_feedback_about_last_retrieval,
                         st.session_state.messages,
                         st.session_state.rag_retrieval_history.copy() if st.session_state.rag_retrieval_history else []
@@ -1082,7 +759,7 @@ def main():
                     # Create UserFeedback dataclass instance
                     feedback_obj = None  # Initialize outside try block
                     try:
-                        feedback_obj = create_feedback_from_dict(feedback_dict)
+                        feedback_obj = feedback_manager.create_feedback_from_dict(feedback_dict)
                         print(f"✅ FEEDBACK SUBMISSION: Feedback object created - ID={feedback_obj.feedback_id}")
                         st.write(f"✅ **Feedback Object Created**")
                         st.write(f"- Feedback ID: {feedback_obj.feedback_id}")
@@ -1138,7 +815,11 @@ def main():
                                         logger.info("📤 SNOWFLAKE UI: Attempting to save feedback to Snowflake...")
                                         print("📤 SNOWFLAKE UI: Attempting to save feedback to Snowflake...")
                                         
-                                        snowflake_success = save_to_snowflake(feedback_obj)
+                                        # Show spinner while saving to Snowflake (can take 10-15 seconds)
+                                        # This includes: connection establishment (~5s), data preparation, and SQL execution (~5s)
+                                        with st.spinner("💾 Saving feedback to Snowflake... This may take 10-15 seconds (connecting to database, preparing data, and executing query)"):
+                                            snowflake_success = feedback_manager.save_to_snowflake(feedback_obj)
+                                        
                                         if snowflake_success:
                                             logger.info("✅ SNOWFLAKE UI: Successfully saved to Snowflake")
                                             print("✅ SNOWFLAKE UI: Successfully saved to Snowflake")
@@ -1193,7 +874,7 @@ def main():
             st.success("✅ Feedback already submitted for this conversation!")
             col1, col2 = st.columns([1, 1])
             with col1:
-                if st.button("🔄 Submit New Feedback", key="new_feedback_button", use_container_width=True):
+                if st.button("🔄 Submit New Feedback", key="new_feedback_button", width='stretch'):
                     try:
                         st.session_state.feedback_submitted = False
                         st.rerun()
@@ -1202,7 +883,7 @@ def main():
                         logger.error(f"Error resetting feedback state: {e}")
                         st.error(f"Error resetting feedback. Please refresh the page.")
             with col2:
-                if st.button("📋 View Conversation", key="view_conversation_button", use_container_width=True):
+                if st.button("📋 View Conversation", key="view_conversation_button", width='stretch'):
                     # Scroll to conversation - this is handled by the auto-scroll at bottom
                     pass
     
@@ -1211,20 +892,111 @@ def main():
         st.markdown("---")
         st.markdown("#### 📊 Retrieval History")
         
-        with st.expander(f"View {len(st.session_state.rag_retrieval_history)} retrieval entries", expanded=False):
+        with st.expander(f"View {len(st.session_state.rag_retrieval_history)} retrieval entries", expanded=True):
             for idx, entry in enumerate(st.session_state.rag_retrieval_history, 1):
-                st.markdown(f"**Retrieval #{idx}**")
+                st.markdown(f"### **Retrieval #{idx}**")
+                
+                # Display timestamp if available
+                if entry.get("timestamp"):
+                    timestamp_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(entry["timestamp"]))
+                    st.caption(f"🕐 {timestamp_str}")
                 
                 # Display the actual RAG query
                 rag_query_expansion = entry.get("rag_query_expansion", "No query available")
+                st.markdown("**🔍 RAG Query:**")
                 st.code(rag_query_expansion, language="text")
                 
+                # Display filters used
+                filters_applied = entry.get("filters_applied", {})
+                if filters_applied and any(filters_applied.values()):
+                    st.markdown("**🎯 Filters Applied:**")
+                    filter_display = {}
+                    if filters_applied.get("sources"):
+                        filter_display["Sources"] = filters_applied["sources"]
+                    if filters_applied.get("years"):
+                        filter_display["Years"] = filters_applied["years"]
+                    if filters_applied.get("districts"):
+                        filter_display["Districts"] = filters_applied["districts"]
+                    if filters_applied.get("filenames"):
+                        filter_display["Filenames"] = filters_applied["filenames"]
+                    
+                    if filter_display:
+                        st.json(filter_display)
+                    else:
+                        st.info("No filters applied")
+                else:
+                    st.info("No filters applied")
+                
+                # Display conversation history up to retrieval point
+                conversation_up_to = entry.get("conversation_up_to", [])
+                if conversation_up_to:
+                    st.markdown("**💬 Conversation History (up to retrieval point):**")
+                    with st.expander(f"View {len(conversation_up_to)} messages", expanded=False):
+                        for msg_idx, msg in enumerate(conversation_up_to, 1):
+                            role = msg.get("type", "unknown")
+                            content = msg.get("content", "")
+                            
+                            if role == "HumanMessage" or role == "human":
+                                st.markdown(f"**👤 User {msg_idx}:** {content[:200]}{'...' if len(content) > 200 else ''}")
+                            elif role == "AIMessage" or role == "ai":
+                                st.markdown(f"**🤖 Assistant {msg_idx}:** {content[:200]}{'...' if len(content) > 200 else ''}")
+                else:
+                    st.info("No conversation history available")
+                
+                # Display documents retrieved
+                docs_retrieved = entry.get("docs_retrieved", [])
+                if docs_retrieved:
+                    st.markdown(f"**📄 Documents Retrieved ({len(docs_retrieved)}):**")
+                    with st.expander(f"View {len(docs_retrieved)} documents", expanded=False):
+                        for doc_idx, doc in enumerate(docs_retrieved, 1):
+                            st.markdown(f"**Document {doc_idx}:**")
+                            
+                            # Display metadata
+                            metadata = doc.get("metadata", {})
+                            if metadata:
+                                col1, col2, col3 = st.columns(3)
+                                with col1:
+                                    st.write(f"📄 **File:** {metadata.get('filename', 'Unknown')}")
+                                with col2:
+                                    st.write(f"🏛️ **Source:** {metadata.get('source', 'Unknown')}")
+                                with col3:
+                                    st.write(f"📅 **Year:** {metadata.get('year', 'Unknown')}")
+                                
+                                # Additional metadata
+                                if metadata.get('district'):
+                                    st.write(f"📍 **District:** {metadata.get('district')}")
+                                if metadata.get('page'):
+                                    st.write(f"📖 **Page:** {metadata.get('page')}")
+                                if metadata.get('score') is not None:
+                                    st.write(f"⭐ **Score:** {metadata.get('score'):.3f}" if isinstance(metadata.get('score'), (int, float)) else f"⭐ **Score:** {metadata.get('score')}")
+                            
+                            # Display content preview (first 200 chars)
+                            content = doc.get("content", doc.get("page_content", ""))
+                            if content:
+                                st.markdown("**Content Preview:**")
+                                st.text_area(
+                                    "Content Preview",
+                                    value=content[:200] + ("..." if len(content) > 200 else ""),
+                                    height=100,
+                                    disabled=True,
+                                    label_visibility="collapsed",
+                                    key=f"retrieval_{idx}_doc_{doc_idx}_preview"
+                                )
+                            
+                            if doc_idx < len(docs_retrieved):
+                                st.markdown("---")
+                else:
+                    st.info("No documents retrieved")
+                
                 # Display summary stats
+                st.markdown("**📊 Summary:**")
                 st.json({
-                    "conversation_length": len(entry.get("conversation_up_to", [])),
-                    "documents_retrieved": len(entry.get("docs_retrieved", []))
+                    "conversation_length": len(conversation_up_to),
+                    "documents_retrieved": len(docs_retrieved)
                 })
-                st.markdown("---")
+                
+                if idx < len(st.session_state.rag_retrieval_history):
+                    st.markdown("---")
     
     # Example Questions Section
     st.markdown("---")
@@ -1245,7 +1017,7 @@ def main():
         st.markdown(f"**Example:** `{example_q1}`")
         st.info("💡 **Filter to apply:** Select a Filename from the sidebar panel before asking this question.")
     with col2:
-        if st.button("📋 Use This Question", key="use_example_1", use_container_width=True):
+        if st.button("📋 Use This Question", key="use_example_1", width='stretch'):
             st.session_state.pending_question = example_q1
             st.session_state.input_counter = (st.session_state.get('input_counter', 0) + 1) % 1000
             st.rerun()
@@ -1266,7 +1038,7 @@ def main():
     )
     col1, col2 = st.columns([1, 4])
     with col1:
-        if st.button("📋 Use Question 2", key="use_custom_1", use_container_width=True):
+        if st.button("📋 Use Question 2", key="use_custom_1", width='stretch'):
             if custom_q1.strip():
                 st.session_state.pending_question = custom_q1.strip()
                 st.session_state.custom_question_1 = custom_q1.strip()
@@ -1292,7 +1064,7 @@ def main():
     )
     col1, col2 = st.columns([1, 4])
     with col1:
-        if st.button("📋 Use Question 3", key="use_custom_2", use_container_width=True):
+        if st.button("📋 Use Question 3", key="use_custom_2", width='stretch'):
             if custom_q2.strip():
                 st.session_state.pending_question = custom_q2.strip()
                 st.session_state.custom_question_2 = custom_q2.strip()

src/agents/__init__.py

@@ -0,0 +1,10 @@
+"""
+Agent modules for chatbot implementations
+"""
+
+from .smart_chatbot import get_chatbot as get_smart_chatbot
+from .multi_agent_chatbot import get_multi_agent_chatbot
+from .gemini_chatbot import get_gemini_chatbot
+
+__all__ = ["get_smart_chatbot", "get_multi_agent_chatbot", "get_gemini_chatbot"]
+

src/agents/gemini_chatbot.py

@@ -0,0 +1,392 @@
+"""
+Gemini File Search Chatbot (Beta Version)
+
+This chatbot uses Google Gemini File Search API for RAG.
+It provides a simpler architecture: Main Agent + Gemini Agent
+"""
+
+import os
+import json
+import time
+import logging
+import traceback
+from pathlib import Path
+from typing import Dict, List, Any, Optional, TypedDict
+
+from langgraph.graph import StateGraph, END
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+
+from src.gemini.file_search import GeminiFileSearchClient, GeminiFileSearchResult
+from src.config.paths import CONVERSATIONS_DIR
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+class GeminiState(TypedDict):
+    """State for Gemini chatbot conversation flow"""
+    conversation_id: str
+    messages: List[Any]
+    current_query: str
+    query_context: Optional[Dict[str, Any]]
+    gemini_result: Optional[GeminiFileSearchResult]
+    final_response: Optional[str]
+    agent_logs: List[str]
+    conversation_context: Dict[str, Any]
+    session_start_time: float
+    last_ai_message_time: float
+    filters: Optional[Dict[str, Any]]
+
+
+class GeminiRAGChatbot:
+    """Gemini File Search RAG chatbot (Beta version)"""
+    
+    def __init__(self):
+        """Initialize the Gemini chatbot"""
+        logger.info("🤖 INITIALIZING: Gemini File Search Chatbot (Beta)")
+        
+        # Initialize Gemini File Search client
+        try:
+            self.gemini_client = GeminiFileSearchClient()
+            logger.info("✅ Gemini File Search client initialized")
+        except Exception as e:
+            logger.error(f"❌ Failed to initialize Gemini client: {e}")
+            raise RuntimeError(f"Gemini client initialization failed: {e}")
+        
+        # Build the LangGraph with LangSmith tracing if enabled
+        self.graph = self._build_graph()
+        
+        # Enable LangSmith tracing if configured
+        langsmith_enabled = os.getenv("LANGCHAIN_TRACING_V2", "false").lower() == "true"
+        if langsmith_enabled:
+            logger.info("🔍 LangSmith tracing enabled")
+            langsmith_project = os.getenv("LANGCHAIN_PROJECT", "gemini-chatbot")
+            logger.info(f"📊 LangSmith project: {langsmith_project}")
+        
+        # Conversations directory
+        self.conversations_dir = CONVERSATIONS_DIR
+        try:
+            self.conversations_dir.mkdir(parents=True, mode=0o777, exist_ok=True)
+        except (PermissionError, OSError) as e:
+            logger.warning(f"Could not create conversations directory: {e}")
+            self.conversations_dir = Path("conversations")
+            self.conversations_dir.mkdir(parents=True, mode=0o777, exist_ok=True)
+        
+        logger.info("✅ Gemini File Search Chatbot initialized")
+    
+    def _build_graph(self) -> StateGraph:
+        """Build the LangGraph for Gemini chatbot"""
+        graph = StateGraph(GeminiState)
+        
+        # Add nodes
+        graph.add_node("main_agent", self._main_agent)
+        graph.add_node("gemini_agent", self._gemini_agent)
+        
+        # Define the flow
+        graph.set_entry_point("main_agent")
+        graph.add_edge("main_agent", "gemini_agent")
+        graph.add_edge("gemini_agent", END)
+        
+        return graph.compile()
+    
+    def _main_agent(self, state: GeminiState) -> GeminiState:
+        """Main Agent: Extracts filters and prepares query"""
+        logger.info("🎯 MAIN AGENT: Processing query")
+        
+        query = state["current_query"]
+        messages = state["messages"]
+        
+        # Extract UI filters if present in query
+        ui_filters = self._extract_ui_filters(query)
+        
+        # Extract context from conversation
+        context = self._extract_context_from_conversation(messages, ui_filters)
+        
+        # Store context and filters
+        state["query_context"] = context
+        state["filters"] = context.get("filters", {})
+        
+        logger.info(f"🎯 MAIN AGENT: Filters extracted: {state['filters']}")
+        
+        return state
+    
+    def _gemini_agent(self, state: GeminiState) -> GeminiState:
+        """Gemini Agent: Performs file search and generates response"""
+        logger.info("🔍 GEMINI AGENT: Starting file search")
+        
+        query = state["current_query"]
+        filters = state.get("filters", {})
+        
+        # Perform Gemini file search
+        try:
+            result = self.gemini_client.search(query=query, filters=filters)
+            logger.info(f"✅ GEMINI AGENT: Search completed, {len(result.sources)} sources found")
+            
+            # Enhance response with document references
+            enhanced_response = self._enhance_response_with_references(
+                result.answer, 
+                result.sources,
+                query
+            )
+            
+            state["gemini_result"] = result
+            state["final_response"] = enhanced_response
+            state["last_ai_message_time"] = time.time()
+            
+            state["agent_logs"].append(f"GEMINI AGENT: Found {len(result.sources)} sources")
+            
+        except Exception as e:
+            logger.error(f"❌ GEMINI AGENT ERROR: {e}")
+            traceback.print_exc()
+            state["final_response"] = "I apologize, but I encountered an error while searching. Please try again."
+            state["last_ai_message_time"] = time.time()
+        
+        return state
+    
+    def _enhance_response_with_references(self, answer: str, sources: List[Any], query: str) -> str:
+        """Enhance Gemini response to include document references and format nicely"""
+        if not sources or not answer:
+            return answer
+        
+        # Use LLM to intelligently add document references and format nicely
+        try:
+            from src.llm.adapters import get_llm_client
+            llm = get_llm_client()
+            
+            # Prepare document summaries for the LLM
+            doc_summaries = []
+            for idx, doc in enumerate(sources, 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 '')
+                
+                filename = metadata.get('filename', 'Unknown') if isinstance(metadata, dict) else 'Unknown'
+                year = metadata.get('year', 'Unknown') if isinstance(metadata, dict) else 'Unknown'
+                source = metadata.get('source', 'Unknown') if isinstance(metadata, dict) else 'Unknown'
+                district = metadata.get('district', '') if isinstance(metadata, dict) else ''
+                
+                doc_info = f"{filename}"
+                if year and year != 'Unknown':
+                    doc_info += f" ({year})"
+                if source and source != 'Unknown':
+                    doc_info += f" - {source}"
+                if district:
+                    doc_info += f" - {district}"
+                
+                doc_summaries.append(f"[Doc {idx}] {doc_info}: {content[:300]}...")
+            
+            prompt = f"""You are enhancing a response from a document search system. The original response is:
+
+{answer}
+
+The following documents were retrieved and used to generate this response:
+
+{chr(10).join(doc_summaries)}
+
+CRITICAL RULES:
+1. Format the response nicely with proper paragraphs, bullet points, or structured sections where appropriate
+2. The response should ONLY contain information from the retrieved documents listed above
+3. If the response mentions information NOT found in the retrieved documents, you must REMOVE or CORRECT that information
+4. Add document references [Doc i] at the end of sentences that use information from specific documents
+5. Only reference documents that are actually used in the response
+6. If the response mentions years, sources, or data that don't match the retrieved documents, you must correct it
+7. Keep the response natural, conversational, and well-formatted
+8. Use proper formatting: paragraphs, line breaks, and structure for readability
+9. Don't change the core content that matches the documents, just add references where appropriate and improve formatting
+10. If multiple documents support the same claim, use [Doc i, Doc j] format
+11. If the response contains information that cannot be verified in the retrieved documents, add a note like: "Note: This information may not be in the retrieved documents."
+
+Return ONLY the enhanced, well-formatted response with references added and any corrections made. Do not include any explanation or meta-commentary."""
+
+            enhanced = llm.invoke(prompt).content if hasattr(llm.invoke(prompt), 'content') else str(llm.invoke(prompt))
+            
+            # Fallback: if LLM fails, just return original with basic formatting
+            if not enhanced or len(enhanced) < len(answer) * 0.5:
+                logger.warning("LLM enhancement failed, using original response with basic formatting")
+                # Basic formatting: add line breaks after periods for readability
+                formatted = answer.replace('. ', '.\n\n')
+                if sources:
+                    ref_list = ", ".join([f"[Doc {i+1}]" for i in range(min(len(sources), 5))])
+                    formatted += f"\n\n*Based on documents: {ref_list}*"
+                return formatted
+            
+            return enhanced
+            
+        except Exception as e:
+            logger.warning(f"Failed to enhance response with references: {e}")
+            # Fallback: add basic formatting and references at the end
+            formatted = answer.replace('. ', '.\n\n')  # Basic paragraph formatting
+            if sources:
+                ref_list = ", ".join([f"[Doc {i+1}]" for i in range(min(len(sources), 5))])
+                formatted += f"\n\n*Based on documents: {ref_list}*"
+            return formatted
+    
+    def _extract_ui_filters(self, query: str) -> Dict[str, List[str]]:
+        """Extract UI filters from query if present"""
+        filters = {}
+        
+        if "FILTER CONTEXT:" in 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()
+            
+            if "Sources:" in filter_section:
+                sources_line = [line for line in filter_section.split('\n') if line.strip().startswith('Sources:')]
+                if sources_line:
+                    sources_str = sources_line[0].split("Sources:")[1].strip()
+                    if sources_str and sources_str != "None":
+                        filters["sources"] = [s.strip() for s in sources_str.split(",")]
+            
+            if "Years:" in filter_section:
+                years_line = [line for line in filter_section.split('\n') if line.strip().startswith('Years:')]
+                if years_line:
+                    years_str = years_line[0].split("Years:")[1].strip()
+                    if years_str and years_str != "None":
+                        filters["year"] = [y.strip() for y in years_str.split(",")]
+            
+            if "Districts:" in filter_section:
+                districts_line = [line for line in filter_section.split('\n') if line.strip().startswith('Districts:')]
+                if districts_line:
+                    districts_str = districts_line[0].split("Districts:")[1].strip()
+                    if districts_str and districts_str != "None":
+                        filters["district"] = [d.strip() for d in districts_str.split(",")]
+            
+            if "Filenames:" in filter_section:
+                filenames_line = [line for line in filter_section.split('\n') if line.strip().startswith('Filenames:')]
+                if filenames_line:
+                    filenames_str = filenames_line[0].split("Filenames:")[1].strip()
+                    if filenames_str and filenames_str != "None":
+                        filters["filenames"] = [f.strip() for f in filenames_str.split(",")]
+        
+        return filters
+    
+    def _extract_context_from_conversation(
+        self,
+        messages: List[Any],
+        ui_filters: Dict[str, List[str]]
+    ) -> Dict[str, Any]:
+        """Extract context from conversation history"""
+        # Use UI filters if available
+        filters = ui_filters.copy() if ui_filters else {}
+        
+        # For Gemini, we pass filters directly to the search function
+        # The filters will be used to add context to the query
+        
+        return {
+            "filters": filters,
+            "has_filters": bool(filters)
+        }
+    
+    def chat(self, user_input: str, conversation_id: str = "default") -> Dict[str, Any]:
+        """Main chat interface"""
+        logger.info(f"💬 GEMINI CHAT: Processing '{user_input[:50]}...'")
+        
+        # Load conversation
+        conversation_file = self.conversations_dir / f"{conversation_id}.json"
+        conversation = self._load_conversation(conversation_file)
+        
+        # Add user message
+        conversation["messages"].append(HumanMessage(content=user_input))
+        
+        # Prepare state
+        state = GeminiState(
+            conversation_id=conversation_id,
+            messages=conversation["messages"],
+            current_query=user_input,
+            query_context=None,
+            gemini_result=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"],
+            filters=None
+        )
+        
+        # Run graph
+        final_state = self.graph.invoke(state)
+        
+        # Add AI response to conversation
+        if final_state["final_response"]:
+            conversation["messages"].append(AIMessage(content=final_state["final_response"]))
+        
+        # Update conversation
+        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)
+        
+        # Format sources for display
+        sources = []
+        gemini_result = final_state.get("gemini_result")
+        if gemini_result:
+            sources = self.gemini_client.format_sources_for_display(gemini_result)
+            logger.info(f"📋 GEMINI CHAT: Formatted {len(sources)} sources for display")
+        
+        return {
+            'response': final_state["final_response"] or "I apologize, but I couldn't process your request.",
+            'rag_result': {
+                'sources': sources,
+                'answer': final_state["final_response"]
+            },
+            'agent_logs': final_state["agent_logs"],
+            'actual_rag_query': final_state["current_query"],
+            'gemini_result': gemini_result  # Include raw result for tracking
+        }
+    
+    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)
+                    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 {
+            "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:
+            conversation_file.parent.mkdir(parents=True, mode=0o777, exist_ok=True)
+            
+            messages_data = []
+            for msg in conversation["messages"]:
+                if isinstance(msg, HumanMessage):
+                    messages_data.append({"type": "human", "content": msg.content})
+                elif isinstance(msg, AIMessage):
+                    messages_data.append({"type": "ai", "content": msg.content})
+            
+            conversation_data = {
+                "messages": messages_data,
+                "session_start_time": conversation["session_start_time"],
+                "last_ai_message_time": conversation["last_ai_message_time"],
+                "context": conversation.get("context", {})
+            }
+            
+            with open(conversation_file, 'w') as f:
+                json.dump(conversation_data, f, indent=2)
+                
+        except Exception as e:
+            logger.error(f"Could not save conversation: {e}")
+
+
+def get_gemini_chatbot():
+    """Get Gemini chatbot instance"""
+    return GeminiRAGChatbot()
+

multi_agent_chatbot.py → src/agents/multi_agent_chatbot.py

@@ -208,6 +208,59 @@ class MultiAgentRAGChatbot:
         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
+    
     def _build_graph(self) -> StateGraph:
         """Build the multi-agent LangGraph"""
         graph = StateGraph(MultiAgentState)
@@ -512,6 +565,10 @@ class MultiAgentRAGChatbot:
    - 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)**:
@@ -656,13 +713,9 @@ Analyze this query using ONLY the exact values provided above:""")
                     # Validate each district in the array
                     valid_districts = []
                     for district in extracted_district:
-                        if district in self.district_whitelist:
-                            valid_districts.append(district)
-                        else:
-                            # Try removing "District" suffix
-                            district_name = district.replace(" District", "").replace(" district", "")
-                            if district_name in self.district_whitelist:
-                                valid_districts.append(district_name)
+                        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
@@ -671,16 +724,15 @@ Analyze this query using ONLY the exact values provided above:""")
                         logger.warning(f"⚠️ No valid districts found in: '{extracted_district}'")
                         extracted_district = None
                 else:
-                    # Single district validation
-                    if extracted_district not in self.district_whitelist:
-                        # Try removing "District" suffix
-                        district_name = extracted_district.replace(" District", "").replace(" district", "")
-                        if district_name in self.district_whitelist:
-                            logger.info(f"🔍 QUERY ANALYSIS: Normalized district '{extracted_district}' to '{district_name}'")
-                            extracted_district = district_name
-                        else:
-                            logger.warning(f"⚠️ Invalid district extracted: '{extracted_district}' not in whitelist")
-                            extracted_district = None
+                    # 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:
@@ -918,6 +970,23 @@ Rewrite the best retrieval query:""")
                 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):
@@ -926,16 +995,6 @@ Rewrite the best retrieval query:""")
                     filters["year"] = [context.extracted_year]
                 logger.info(f"🔧 FILTER BUILDING: Added extracted year filter (UI missing): {context.extracted_year}")
             
-            if not filters.get("district") and context.extracted_district:
-                # Handle both single values and arrays
-                if isinstance(context.extracted_district, list):
-                    # Normalize district names to title case (match Qdrant metadata format)
-                    normalized = [d.title() for d in context.extracted_district]
-                    filters["district"] = normalized
-                else:
-                    filters["district"] = [context.extracted_district.title()]
-                logger.info(f"🔧 FILTER BUILDING: Added extracted district filter (UI missing): {context.extracted_district}")
-            
             if not filters.get("sources") and context.extracted_source:
                 # Handle both single values and arrays
                 if isinstance(context.extracted_source, list):
@@ -963,12 +1022,21 @@ Rewrite the best retrieval query:""")
                 logger.info(f"🔧 FILTER BUILDING: Added extracted year filter: {context.extracted_year}")
             
             if context.extracted_district:
-                # Handle both single values and arrays
+                # Normalize district names using the normalization function
                 if isinstance(context.extracted_district, list):
-                    filters["district"] = context.extracted_district
+                    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:
-                    filters["district"] = [context.extracted_district]
-                logger.info(f"🔧 FILTER BUILDING: Added extracted district filter: {context.extracted_district}")
+                    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
@@ -978,49 +1046,212 @@ Rewrite the best retrieval query:""")
         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")
+        
+        # Build conversation history context
+        conversation_context = self._build_conversation_context(messages)
+        
+        # Build detailed document information
+        document_details = self._build_document_details(documents)
+        
+        # Extract correct district/source/year names from documents (to correct misspellings)
+        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.
 
+CRITICAL RULES - NO HALLUCINATION:
+1. **ONLY use information from the retrieved documents provided below**
+2. **EVERY sentence with facts, numbers, or specific claims MUST have a [Doc i] reference**
+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.
+
 RULES:
 1. Answer the user's question directly and clearly
-2. Use the retrieved documents as evidence
+2. Use ONLY the retrieved documents as evidence - DO NOT use your training data
 3. Be conversational, not technical
 4. Don't mention scores, retrieval details, or technical implementation
 5. If relevant documents were found, reference them naturally
-6. If no relevant documents, explain based on your knowledge (if you have it) or just say you do not have enough information.
-7. If the passages have useful facts or numbers, use them in your answer.
-8. When you use information from a passage, mention where it came from by using [Doc i] at the end of the sentence. i stands for the number of the document.
+6. If no relevant documents, say you do not have enough information - DO NOT hallucinate
+7. If the passages have useful facts or numbers, use them in your answer WITH references
+8. **MANDATORY**: When you use information from a passage, mention where it came from by using [Doc i] at the end of the sentence. i stands for the number of the document.
 9. Do not use the sentence 'Doc i says ...' to say where information came from.
 10. If the same thing is said in more than one document, you can mention all of them like this: [Doc i, Doc j, Doc k]
 11. Do not just summarize each passage one by one. Group your summaries to highlight the key parts in the explanation.
 12. If it makes sense, use bullet points and lists to make your answers easier to understand.
 13. You do not need to use every passage. Only use the ones that help answer the question.
-14. If the documents do not have the information needed to answer the question, just say you do not have enough information.
-
+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."""),
-            HumanMessage(content=f"""User Question: {query}
+            HumanMessage(content=f"""Conversation History:
+{conversation_context}
+
+Current User Question: {query}
 
 Retrieved Documents: {len(documents)} documents found
 
+CORRECT NAMES TO USE (from document metadata - use these exact spellings):
+{correct_names}
+
+Full Document Details:
+{document_details}
+
 RAG Answer: {rag_answer}
 
-Generate a conversational response:""")
+CRITICAL: 
+- Responses should be grounded to what is available in the retrieved documents
+- 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]}...")
-            return response.content.strip()
+            
+            # Post-process response to ensure no hallucination
+            final_response = self._validate_and_enhance_response(
+                response.content.strip(),
+                documents,
+                query
+            )
+            
+            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
     
+    def _build_conversation_context(self, messages: List[Any]) -> str:
+        """Build conversation history context for response generation."""
+        if not messages:
+            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}")
+            elif isinstance(msg, AIMessage):
+                context_lines.append(f"Assistant: {msg.content}")
+        
+        return "\n".join(context_lines) if context_lines else "No previous conversation."
+    
+    def _build_document_details(self, documents: List[Any]) -> str:
+        """Build detailed document information for response generation."""
+        if not documents:
+            return "No documents retrieved."
+        
+        details = []
+        for i, doc in enumerate(documents[:15], 1):  # Show up to 15 documents
+            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 '')
+            
+            if isinstance(metadata, dict):
+                filename = metadata.get('filename', 'Unknown')
+                year = metadata.get('year', 'Unknown')
+                district = metadata.get('district', 'Unknown')
+                source = metadata.get('source', 'Unknown')
+                page = metadata.get('page', metadata.get('page_label', 'Unknown'))
+                
+                doc_info = f"[Doc {i}]"
+                doc_info += f"\n  Filename: {filename}"
+                doc_info += f"\n  Year: {year}"
+                doc_info += f"\n  District: {district}"
+                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 ''}"
+                details.append(doc_info)
+        
+        return "\n\n".join(details) if details else "No document details available."
+    
+    def _extract_correct_names_from_documents(self, documents: List[Any]) -> str:
+        """Extract correct district/source names from documents to correct misspellings."""
+        districts = set()
+        sources = set()
+        years = set()
+        
+        for doc in documents:
+            metadata = getattr(doc, 'metadata', {}) if hasattr(doc, 'metadata') else (doc if isinstance(doc, dict) else {})
+            if isinstance(metadata, dict):
+                if metadata.get('district'):
+                    districts.add(str(metadata['district']))
+                if metadata.get('source'):
+                    sources.add(str(metadata['source']))
+                if metadata.get('year'):
+                    years.add(str(metadata['year']))
+        
+        result = []
+        if districts:
+            result.append(f"Districts: {', '.join(sorted(districts))}")
+        if sources:
+            result.append(f"Sources: {', '.join(sorted(sources))}")
+        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."
+    
+    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
+        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 {})
+            if isinstance(metadata, dict):
+                if metadata.get('year'):
+                    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
+        
+        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")

src/feedback/__init__.py

@@ -0,0 +1,152 @@
+"""
+Feedback Management Module
+
+This module provides a unified interface for handling user feedback,
+including data preparation, validation, and Snowflake storage.
+"""
+
+from typing import Dict, Any, List, Optional
+from langchain_core.messages import HumanMessage, AIMessage
+
+from .feedback_schema import UserFeedback, create_feedback_from_dict, generate_snowflake_schema_sql
+from .snowflake_connector import SnowflakeFeedbackConnector, save_to_snowflake, get_snowflake_connector_from_env
+
+
+class FeedbackManager:
+    """
+    Unified manager for feedback operations.
+    
+    This class provides a single interface for all feedback-related functionality,
+    including data preparation, validation, and storage.
+    """
+    
+    def __init__(self):
+        """Initialize the FeedbackManager"""
+        pass
+    
+    @staticmethod
+    def extract_transcript(messages: List[Any]) -> List[Dict[str, str]]:
+        """Extract transcript from messages - only user and bot messages, no extra metadata"""
+        transcript = []
+        for msg in messages:
+            if isinstance(msg, HumanMessage):
+                transcript.append({
+                    "role": "user",
+                    "content": str(msg.content) if hasattr(msg, 'content') else str(msg)
+                })
+            elif isinstance(msg, AIMessage):
+                transcript.append({
+                    "role": "assistant",
+                    "content": str(msg.content) if hasattr(msg, 'content') else str(msg)
+                })
+        return transcript
+    
+    @staticmethod
+    def build_retrievals_structure(rag_retrieval_history: List[Dict[str, Any]], messages: List[Any]) -> List[Dict[str, Any]]:
+        """Build retrievals structure from retrieval history"""
+        retrievals = []
+        
+        for entry in rag_retrieval_history:
+            # Get the user message that triggered this retrieval
+            # The entry has conversation_up_to which includes messages up to that point
+            conversation_up_to = entry.get("conversation_up_to", [])
+            
+            # Find the last user message in conversation_up_to (this is the trigger)
+            user_message_trigger = ""
+            for msg_dict in reversed(conversation_up_to):
+                if msg_dict.get("type") == "HumanMessage":
+                    user_message_trigger = msg_dict.get("content", "")
+                    break
+            
+            # Fallback: if not found in conversation_up_to, get from actual messages
+            # This handles edge cases where conversation_up_to might be incomplete
+            if not user_message_trigger:
+                # Find which retrieval this is (0-indexed)
+                retrieval_idx = rag_retrieval_history.index(entry)
+                # The user message that triggered this retrieval is at position (retrieval_idx * 2)
+                # because each retrieval is preceded by: user message, bot response, user message, ...
+                # But we need to account for the fact that the first retrieval happens after the first user message
+                user_msgs = [msg for msg in messages if isinstance(msg, HumanMessage)]
+                if retrieval_idx < len(user_msgs):
+                    user_message_trigger = str(user_msgs[retrieval_idx].content)
+                elif user_msgs:
+                    # Fallback to last user message
+                    user_message_trigger = str(user_msgs[-1].content)
+            
+            # Get retrieved documents and truncate content to 100 chars
+            docs_retrieved = entry.get("docs_retrieved", [])
+            retrieved_docs = []
+            for doc in docs_retrieved:
+                doc_copy = doc.copy()
+                # Truncate content to 100 characters (keep all other fields)
+                if "content" in doc_copy:
+                    doc_copy["content"] = doc_copy["content"][:100]
+                retrieved_docs.append(doc_copy)
+            
+            retrievals.append({
+                "retrieved_docs": retrieved_docs,
+                "user_message_trigger": user_message_trigger
+            })
+        
+        return retrievals
+    
+    @staticmethod
+    def build_feedback_score_related_retrieval_docs(
+        is_feedback_about_last_retrieval: bool,
+        messages: List[Any],
+        rag_retrieval_history: List[Dict[str, Any]]
+    ) -> Optional[Dict[str, Any]]:
+        """Build feedback_score_related_retrieval_docs structure"""
+        if not rag_retrieval_history:
+            return None
+        
+        # Get the relevant retrieval entry
+        if is_feedback_about_last_retrieval:
+            relevant_entry = rag_retrieval_history[-1]
+        else:
+            # If feedback is about all retrievals, use the last one as default
+            relevant_entry = rag_retrieval_history[-1]
+        
+        # Get conversation up to that point
+        conversation_up_to = relevant_entry.get("conversation_up_to", [])
+        
+        # Convert to transcript format (role/content)
+        conversation_up_to_point = []
+        for msg_dict in conversation_up_to:
+            if msg_dict.get("type") == "HumanMessage":
+                conversation_up_to_point.append({
+                    "role": "user",
+                    "content": msg_dict.get("content", "")
+                })
+            elif msg_dict.get("type") == "AIMessage":
+                conversation_up_to_point.append({
+                    "role": "assistant",
+                    "content": msg_dict.get("content", "")
+                })
+        
+        # Get retrieved docs with full content (not truncated)
+        retrieved_docs = relevant_entry.get("docs_retrieved", [])
+        
+        return {
+            "conversation_up_to_point": conversation_up_to_point,
+            "retrieved_docs": retrieved_docs
+        }
+    
+    @staticmethod
+    def create_feedback_from_dict(data: Dict[str, Any]) -> UserFeedback:
+        """Create UserFeedback instance from dictionary"""
+        return create_feedback_from_dict(data)
+    
+    @staticmethod
+    def save_to_snowflake(feedback: UserFeedback, table_name: Optional[str] = None) -> bool:
+        """Save feedback to Snowflake"""
+        return save_to_snowflake(feedback, table_name)
+    
+    @staticmethod
+    def generate_snowflake_schema_sql(table_name: Optional[str] = None) -> str:
+        """Generate Snowflake schema SQL"""
+        return generate_snowflake_schema_sql(table_name)
+
+
+__all__ = ["FeedbackManager", "UserFeedback", "save_to_snowflake", "SnowflakeFeedbackConnector"]
+

src/feedback/feedback_schema.py

@@ -0,0 +1,161 @@
+"""
+Feedback Schema for RAG Chatbot
+
+This module defines dataclasses for feedback data structures
+and provides Snowflake schema generation.
+"""
+import os
+from datetime import datetime
+from dataclasses import dataclass, asdict, field
+from typing import List, Optional, Dict, Any, Union
+
+
+
+@dataclass
+class RetrievedDocument:
+    """Single retrieved document metadata"""
+    doc_id: str
+    filename: str
+    page: int
+    score: float
+    content: str
+    metadata: Dict[str, Any]
+    
+
+@dataclass
+class RetrievalEntry:
+    """Single retrieval operation metadata"""
+    rag_query: str
+    documents_retrieved: List[RetrievedDocument]
+    conversation_length: int
+    filters_applied: Optional[Dict[str, Any]] = None
+    timestamp: Optional[float] = None
+    _raw_data: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class UserFeedback:
+    """User feedback submission data"""
+    feedback_id: str
+    open_ended_feedback: Optional[str]
+    score: int
+    is_feedback_about_last_retrieval: bool
+    conversation_id: str
+    timestamp: float
+    message_count: int
+    has_retrievals: bool
+    retrieval_count: int
+    transcript: List[Dict[str, str]]  # List of {"role": "user"/"assistant", "content": "..."}
+    retrievals: List[Dict[str, Any]]  # List of retrieval objects with retrieved_docs and user_message_trigger
+    feedback_score_related_retrieval_docs: Optional[Dict[str, Any]] = None  # Conversation subset + retrieved docs
+    retrieved_data: Optional[List[Dict[str, Any]]] = None  # Preserved old column for backward compatibility
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary with nested data structures"""
+        result = asdict(self)
+        return result
+    
+    def to_snowflake_schema(self) -> Dict[str, Any]:
+        """Generate Snowflake schema for this dataclass"""
+        schema = {
+            "feedback_id": "VARCHAR(255)",
+            "open_ended_feedback": "VARCHAR(16777216)",  # Large text
+            "score": "INTEGER",
+            "is_feedback_about_last_retrieval": "BOOLEAN",
+            "conversation_id": "VARCHAR(255)",
+            "timestamp": "NUMBER(20, 0)",
+            "message_count": "INTEGER",
+            "has_retrievals": "BOOLEAN",
+            "retrieval_count": "INTEGER",
+            "transcript": "VARCHAR(16777216)",  # JSON string of ARRAY of {"role": "user"/"assistant", "content": "..."}
+            "retrievals": "VARCHAR(16777216)",  # JSON string of ARRAY of retrieval objects
+            "feedback_score_related_retrieval_docs": "VARCHAR(16777216)",  # JSON string of OBJECT with conversation subset + retrieved docs
+            "retrieved_data": "VARCHAR(16777216)",  # JSON string - preserved old column for backward compatibility
+            "created_at": "TIMESTAMP_NTZ",
+            # transcript structure: [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}, ...]
+            # retrievals structure: [
+            #   {
+            #     "retrieved_docs": [{"content": "...", "metadata": {...}, ...}],  # content truncated to 100 chars
+            #     "user_message_trigger": "final user message that triggered this retrieval"
+            #   },
+            #   ...
+            # ]
+            # feedback_score_related_retrieval_docs structure: {
+            #   "conversation_up_to_point": [{"role": "user", "content": "..."}, ...],  # subset of transcript
+            #   "retrieved_docs": [{"content": "...", "metadata": {...}, ...}]  # full chunks with all info
+            # }
+        }
+        return schema
+    
+    @classmethod
+    def get_snowflake_create_table_sql(cls, table_name: str = "USER_FEEDBACK_V3") -> str:
+        """Generate CREATE TABLE SQL for Snowflake"""
+        schema = cls.to_snowflake_schema(None)
+        
+        columns = []
+        for col_name, col_type in schema.items():
+            nullable = "NULL" if col_name not in ["feedback_id", "score", "timestamp"] else "NOT NULL"
+            columns.append(f"  {col_name} {col_type} {nullable}")
+        
+        # Build SQL string properly
+        columns_str = ",\n".join(columns)
+        
+        sql = f"""CREATE TABLE IF NOT EXISTS {table_name} (
+{columns_str},
+  PRIMARY KEY (feedback_id)
+)
+CLUSTER BY (timestamp, conversation_id, score);
+-- Note: Snowflake doesn't support traditional indexes on regular tables.
+-- Instead, we use CLUSTER BY to optimize queries on these columns.
+-- Snowflake automatically maintains clustering for efficient querying.
+-- Note: transcript, retrievals, and feedback_score_related_retrieval_docs are stored as VARCHAR (JSON strings),
+-- same approach as the old retrieved_data column. This allows easy storage and retrieval without VARIANT type complexity.
+"""
+        return sql
+
+
+# Snowflake variant schema for retrieved_data array
+RETRIEVAL_ENTRY_SCHEMA = {
+    "rag_query": "VARCHAR",
+    "documents_retrieved": "ARRAY",  # Array of document objects
+    "conversation_length": "INTEGER",
+    "filters_applied": "OBJECT",
+    "timestamp": "NUMBER"
+}
+
+DOCUMENT_SCHEMA = {
+    "doc_id": "VARCHAR",
+    "filename": "VARCHAR",
+    "page": "INTEGER",
+    "score": "DOUBLE",
+    "content": "VARCHAR(16777216)",
+    "metadata": "OBJECT"
+}
+
+
+def generate_snowflake_schema_sql(table_name: Optional[str] = None) -> str:
+    """Generate complete Snowflake schema SQL for feedback system"""
+    if table_name is None:
+        table_name = os.getenv("SNOWFLAKE_FEEDBACK_TABLE", "USER_FEEDBACK_V3")
+    return UserFeedback.get_snowflake_create_table_sql(table_name)
+
+
+def create_feedback_from_dict(data: Dict[str, Any]) -> UserFeedback:
+    """Create UserFeedback instance from dictionary"""
+    return UserFeedback(
+        feedback_id=data.get("feedback_id", f"feedback_{data.get('timestamp', 'unknown')}"),
+        open_ended_feedback=data.get("open_ended_feedback"),
+        score=data["score"],
+        is_feedback_about_last_retrieval=data["is_feedback_about_last_retrieval"],
+        conversation_id=data["conversation_id"],
+        timestamp=data["timestamp"],
+        message_count=data["message_count"],
+        has_retrievals=data["has_retrievals"],
+        retrieval_count=data["retrieval_count"],
+        transcript=data.get("transcript", []),
+        retrievals=data.get("retrievals", []),
+        feedback_score_related_retrieval_docs=data.get("feedback_score_related_retrieval_docs"),
+        retrieved_data=data.get("retrieved_data")
+    )
+

src/feedback/snowflake_connector.py

@@ -0,0 +1,331 @@
+"""
+Snowflake Connector for Feedback System
+
+This module handles inserting user feedback into Snowflake.
+"""
+
+import os
+import json
+import logging
+from typing import Dict, Any, Optional
+from .feedback_schema import UserFeedback
+
+# Try to import snowflake connector
+try:
+    import snowflake.connector
+    SNOWFLAKE_AVAILABLE = True
+except ImportError:
+    SNOWFLAKE_AVAILABLE = False
+    logging.warning("⚠️ snowflake-connector-python not installed. Install with: pip install snowflake-connector-python")
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+class SnowflakeFeedbackConnector:
+    """Connector for inserting feedback into Snowflake"""
+    
+    def __init__(
+        self,
+        user: str,
+        password: str,
+        account: str,
+        warehouse: str,
+        database: str = "SNOWFLAKE_LEARNING",
+        schema: str = "PUBLIC"
+    ):
+        self.user = user
+        self.password = password
+        self.account = account
+        self.warehouse = warehouse
+        self.database = database
+        self.schema = schema
+        self._connection = None
+    
+    def connect(self):
+        """Establish Snowflake connection"""
+        if not SNOWFLAKE_AVAILABLE:
+            raise ImportError("snowflake-connector-python is not installed. Install with: pip install snowflake-connector-python")
+        
+        logger.info("=" * 80)
+        logger.info("🔌 SNOWFLAKE CONNECTION: Attempting to connect...")
+        logger.info(f"   - Account: {self.account}")
+        logger.info(f"   - Warehouse: {self.warehouse}")
+        logger.info(f"   - Database: {self.database}")
+        logger.info(f"   - Schema: {self.schema}")
+        logger.info(f"   - User: {self.user}")
+        
+        try:
+            self._connection = snowflake.connector.connect(
+                user=self.user,
+                password=self.password,
+                account=self.account,
+                warehouse=self.warehouse
+                # Don't set database/schema in connection - we'll do it per query
+            )
+            logger.info("✅ SNOWFLAKE CONNECTION: Successfully connected")
+            logger.info("=" * 80)
+            print(f"✅ Connected to Snowflake: {self.database}.{self.schema}")
+        except Exception as e:
+            logger.error(f"❌ SNOWFLAKE CONNECTION FAILED: {e}")
+            logger.error("=" * 80)
+            print(f"❌ Failed to connect to Snowflake: {e}")
+            raise
+    
+    def disconnect(self):
+        """Close Snowflake connection"""
+        if self._connection:
+            self._connection.close()
+            print("✅ Disconnected from Snowflake")
+    
+    def insert_feedback(self, feedback: UserFeedback, table_name: Optional[str] = None) -> bool:
+        """Insert a single feedback record into Snowflake"""
+        logger.info("=" * 80)
+        logger.info("🔄 SNOWFLAKE INSERT: Starting feedback insertion process")
+        logger.info(f"📝 Feedback ID: {feedback.feedback_id}")
+        
+        # Get table name from parameter, env var, or default
+        if table_name is None:
+            table_name = os.getenv("SNOWFLAKE_FEEDBACK_TABLE", "USER_FEEDBACK_V3")
+        
+        if not self._connection:
+            logger.error("❌ Not connected to Snowflake. Call connect() first.")
+            raise RuntimeError("Not connected to Snowflake. Call connect() first.")
+        
+        try:
+            logger.info("📊 VALIDATION: Validating feedback data structure...")
+            
+            # Validate feedback object
+            validation_errors = []
+            if not feedback.feedback_id:
+                validation_errors.append("Missing feedback_id")
+            if feedback.score is None:
+                validation_errors.append("Missing score")
+            if feedback.timestamp is None:
+                validation_errors.append("Missing timestamp")
+                
+            if validation_errors:
+                logger.error(f"❌ VALIDATION FAILED: {validation_errors}")
+                return False
+            else:
+                logger.info("✅ VALIDATION PASSED: All required fields present")
+            
+            logger.info("📋 Data Summary:")
+            logger.info(f"   - Feedback ID: {feedback.feedback_id}")
+            logger.info(f"   - Score: {feedback.score}")
+            logger.info(f"   - Conversation ID: {feedback.conversation_id}")
+            logger.info(f"   - Has Retrievals: {feedback.has_retrievals}")
+            logger.info(f"   - Retrieval Count: {feedback.retrieval_count}")
+            logger.info(f"   - Message Count: {feedback.message_count}")
+            logger.info(f"   - Timestamp: {feedback.timestamp}")
+            
+            cursor = self._connection.cursor()
+            logger.info("✅ SNOWFLAKE CONNECTION: Cursor created")
+            
+            # Set database and schema context
+            logger.info(f"🔧 SETTING CONTEXT: Database={self.database}, Schema={self.schema}")
+            try:
+                cursor.execute(f'USE DATABASE "{self.database}"')
+                cursor.execute(f'USE SCHEMA "{self.schema}"')
+                cursor.execute("SELECT CURRENT_DATABASE(), CURRENT_SCHEMA()")
+                current_db, current_schema = cursor.fetchone()
+                logger.info(f"✅ Current context verified: Database={current_db}, Schema={current_schema}")
+            except Exception as e:
+                logger.error(f"❌ Could not set context: {e}")
+                raise
+            
+            # Prepare data - convert to JSON strings for VARIANT columns (same approach as old retrieved_data)
+            logger.info("🔧 DATA PREPARATION: Preparing VARIANT columns...")
+            feedback_dict = feedback.to_dict()
+            
+            # Prepare transcript (ARRAY) - convert to JSON string
+            transcript_raw = feedback_dict.get('transcript', [])
+            if transcript_raw:
+                # Convert to JSON string (same approach as old retrieved_data)
+                transcript_for_db = json.dumps(transcript_raw)
+                logger.info(f"   - Transcript: {len(transcript_raw)} messages, JSON length: {len(transcript_for_db)}")
+            else:
+                transcript_for_db = None
+                logger.info("   - Transcript: None")
+            
+            # Prepare retrievals (ARRAY) - convert to JSON string
+            retrievals_raw = feedback_dict.get('retrievals', [])
+            if retrievals_raw:
+                # Convert to JSON string (same approach as old retrieved_data)
+                retrievals_for_db = json.dumps(retrievals_raw)
+                logger.info(f"   - Retrievals: {len(retrievals_raw)} entries, JSON length: {len(retrievals_for_db)}")
+            else:
+                retrievals_for_db = None
+                logger.info("   - Retrievals: None")
+            
+            # Prepare feedback_score_related_retrieval_docs (OBJECT) - convert to JSON string
+            feedback_score_related_raw = feedback_dict.get('feedback_score_related_retrieval_docs')
+            if feedback_score_related_raw:
+                # Convert to JSON string (same approach as old retrieved_data)
+                feedback_score_related_for_db = json.dumps(feedback_score_related_raw)
+                logger.info(f"   - Feedback score related docs: present, JSON length: {len(feedback_score_related_for_db)}")
+            else:
+                feedback_score_related_for_db = None
+                logger.info("   - Feedback score related docs: None")
+            
+            # Prepare retrieved_data (preserved old column) - convert to JSON string
+            retrieved_data_raw = feedback_dict.get('retrieved_data')
+            if retrieved_data_raw:
+                # Convert to JSON string (same approach as old retrieved_data)
+                retrieved_data_for_db = json.dumps(retrieved_data_raw)
+                logger.info(f"   - Retrieved data (preserved): present, JSON length: {len(retrieved_data_for_db)}")
+            else:
+                retrieved_data_for_db = None
+                logger.info("   - Retrieved data (preserved): None")
+            
+            # Build SQL with new column structure
+            # Columns are VARCHAR (storing JSON strings), same approach as old retrieved_data
+            sql = f"""INSERT INTO {table_name} (
+                feedback_id,
+                open_ended_feedback,
+                score,
+                is_feedback_about_last_retrieval,
+                conversation_id,
+                timestamp,
+                message_count,
+                has_retrievals,
+                retrieval_count,
+                transcript,
+                retrievals,
+                feedback_score_related_retrieval_docs,
+                retrieved_data,
+                created_at
+            ) VALUES (
+                %(feedback_id)s, %(open_ended_feedback)s, %(score)s, %(is_feedback_about_last_retrieval)s,
+                %(conversation_id)s, %(timestamp)s, %(message_count)s, %(has_retrievals)s,
+                %(retrieval_count)s, %(transcript)s, %(retrievals)s, %(feedback_score_related_retrieval_docs)s,
+                %(retrieved_data)s, %(created_at)s
+            )"""
+            
+            logger.info("📝 SQL PREPARATION: Building INSERT statement...")
+            logger.info(f"   - Target table: {table_name}")
+            logger.info(f"   - Database: {self.database}")
+            logger.info(f"   - Schema: {self.schema}")
+            
+            # Prepare parameters
+            # Pass JSON strings for VARIANT columns (same approach as old retrieved_data)
+            params = {
+                'feedback_id': feedback.feedback_id,
+                'open_ended_feedback': feedback.open_ended_feedback,
+                'score': feedback.score,
+                'is_feedback_about_last_retrieval': feedback.is_feedback_about_last_retrieval,
+                'conversation_id': feedback.conversation_id,
+                'timestamp': int(feedback.timestamp),
+                'message_count': feedback.message_count,
+                'has_retrievals': feedback.has_retrievals,
+                'retrieval_count': feedback.retrieval_count,
+                'transcript': transcript_for_db,  # JSON string
+                'retrievals': retrievals_for_db,  # JSON string
+                'feedback_score_related_retrieval_docs': feedback_score_related_for_db,  # JSON string
+                'retrieved_data': retrieved_data_for_db,  # JSON string - preserved old column
+                'created_at': feedback.created_at
+            }
+            
+            # Execute insert
+            logger.info("🚀 SQL EXECUTION: Executing INSERT query...")
+            cursor.execute(sql, params)
+            
+            logger.info("✅ SQL EXECUTION: Query executed successfully")
+            logger.info(f"   - Rows affected: 1")
+            logger.info(f"   - Status: SUCCESS")
+            
+            cursor.close()
+            logger.info("✅ SNOWFLAKE INSERT: Feedback inserted successfully")
+            logger.info(f"📝 Inserted feedback: {feedback.feedback_id}")
+            logger.info("=" * 80)
+            return True
+            
+        except Exception as e:
+            # Check if it's a Snowflake error
+            if SNOWFLAKE_AVAILABLE and "ProgrammingError" in str(type(e)):
+                logger.error(f"❌ SQL EXECUTION ERROR: {e}")
+                logger.error(f"   - Error code: {getattr(e, 'errno', 'Unknown')}")
+                logger.error(f"   - SQL state: {getattr(e, 'sqlstate', 'Unknown')}")
+            else:
+                logger.error(f"❌ SNOWFLAKE INSERT FAILED: {type(e).__name__}")
+                logger.error(f"   - Error: {e}")
+            logger.error("=" * 80)
+            return False
+    
+    def __enter__(self):
+        """Context manager entry"""
+        self.connect()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit"""
+        self.disconnect()
+
+
+def get_snowflake_connector_from_env() -> Optional[SnowflakeFeedbackConnector]:
+    """Create Snowflake connector from environment variables"""
+    user = os.getenv("SNOWFLAKE_USER")
+    password = os.getenv("SNOWFLAKE_PASSWORD")
+    account = os.getenv("SNOWFLAKE_ACCOUNT")
+    warehouse = os.getenv("SNOWFLAKE_WAREHOUSE")
+    database = os.getenv("SNOWFLAKE_DATABASE", "SNOWFLAKE_LEARN")
+    schema = os.getenv("SNOWFLAKE_SCHEMA", "PUBLIC")
+    
+    if not all([user, password, account, warehouse]):
+        print("⚠️ Snowflake credentials not found in environment variables")
+        print("Required variables: SNOWFLAKE_USER, SNOWFLAKE_PASSWORD, SNOWFLAKE_ACCOUNT, SNOWFLAKE_WAREHOUSE")
+        return None
+    
+    return SnowflakeFeedbackConnector(
+        user=user,
+        password=password,
+        account=account,
+        warehouse=warehouse,
+        database=database,
+        schema=schema
+    )
+
+
+def save_to_snowflake(feedback: UserFeedback, table_name: Optional[str] = None) -> bool:
+    """Helper function to save feedback to Snowflake"""
+    logger.info("=" * 80)
+    logger.info("🔵 SNOWFLAKE SAVE: Starting save process")
+    logger.info(f"📝 Feedback ID: {feedback.feedback_id}")
+    
+    # Get table name from parameter or env var
+    if table_name is None:
+        table_name = os.getenv("SNOWFLAKE_FEEDBACK_TABLE", "USER_FEEDBACK_V3")
+    
+    connector = get_snowflake_connector_from_env()
+    
+    if not connector:
+        logger.warning("⚠️ SNOWFLAKE SAVE: Skipping insertion (credentials not configured)")
+        logger.warning("   Required variables: SNOWFLAKE_USER, SNOWFLAKE_PASSWORD, SNOWFLAKE_ACCOUNT, SNOWFLAKE_WAREHOUSE")
+        logger.info("=" * 80)
+        return False
+    
+    try:
+        logger.info("📡 SNOWFLAKE SAVE: Establishing connection...")
+        connector.connect()
+        logger.info("✅ SNOWFLAKE SAVE: Connection established")
+        
+        logger.info("📥 SNOWFLAKE SAVE: Attempting to insert feedback...")
+        success = connector.insert_feedback(feedback, table_name=table_name)
+        
+        logger.info("🔌 SNOWFLAKE SAVE: Disconnecting...")
+        connector.disconnect()
+        
+        if success:
+            logger.info("✅ SNOWFLAKE SAVE: Successfully saved feedback")
+        else:
+            logger.error("❌ SNOWFLAKE SAVE: Failed to save feedback")
+        
+        logger.info("=" * 80)
+        return success
+    except Exception as e:
+        logger.error(f"❌ SNOWFLAKE SAVE ERROR: {type(e).__name__}")
+        logger.error(f"   - Error: {e}")
+        logger.info("=" * 80)
+        return False
+

src/gemini/__init__.py

@@ -0,0 +1,11 @@
+"""
+Gemini File Search Integration Module
+
+This module provides integration with Google Gemini File Search API
+for RAG functionality using Gemini's built-in file search capabilities.
+"""
+
+from .file_search import GeminiFileSearchClient, GeminiFileSearchResult
+
+__all__ = ["GeminiFileSearchClient", "GeminiFileSearchResult"]
+

src/gemini/file_search.py

@@ -0,0 +1,427 @@
+"""
+Gemini File Search Client
+
+Handles interaction with Google Gemini File Search API for RAG.
+"""
+
+import os
+import logging
+from typing import List, Dict, Any, Optional
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+try:
+    from google import genai
+    from google.genai import types
+    GEMINI_AVAILABLE = True
+except ImportError:
+    GEMINI_AVAILABLE = False
+
+
+@dataclass
+class GeminiFileSearchResult:
+    """Result from Gemini File Search query"""
+    answer: str
+    sources: List[Dict[str, Any]]  # List of document references
+    grounding_metadata: Optional[Dict[str, Any]] = None
+    query: str = ""
+
+
+class GeminiFileSearchClient:
+    """Client for interacting with Gemini File Search API"""
+    
+    def __init__(self, api_key: Optional[str] = None, store_name: Optional[str] = None):
+        """
+        Initialize Gemini File Search client.
+        
+        Args:
+            api_key: Gemini API key (defaults to GEMINI_API_KEY env var)
+            store_name: File search store name (defaults to GEMINI_FILESTORE_NAME env var)
+        """
+        if not GEMINI_AVAILABLE:
+            raise ImportError("google-genai package not installed. Install with: pip install google-genai")
+        
+        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
+        if not self.api_key:
+            raise ValueError("GEMINI_API_KEY not found. Set it in .env file or pass as argument.")
+        
+        store_name_raw = store_name or os.getenv("GEMINI_FILESTORE_NAME")
+        if not store_name_raw:
+            raise ValueError("GEMINI_FILESTORE_NAME not found. Set it in .env file or pass as argument.")
+        
+        # Normalize store name: API expects the FULL path format (fileSearchStores/xxx)
+        # If just the ID is provided, construct the full path
+        if store_name_raw.startswith("fileSearchStores/"):
+            self.store_name = store_name_raw  # Already full path
+        else:
+            # Just the ID provided, construct full path
+            self.store_name = f"fileSearchStores/{store_name_raw}"
+        
+        logger.info(f"📦 Using file search store: {self.store_name}")
+        
+        self.client = genai.Client(api_key=self.api_key)
+        self.model = "gemini-2.5-flash"  # or "gemini-2.5-pro"
+    
+    def search(
+        self,
+        query: str,
+        filters: Optional[Dict[str, Any]] = None,
+        model: Optional[str] = None
+    ) -> GeminiFileSearchResult:
+        """
+        Search using Gemini File Search.
+        
+        Args:
+            query: User query
+            filters: Optional filters (year, source, district, etc.)
+            model: Model to use (defaults to gemini-2.5-flash)
+        
+        Returns:
+            GeminiFileSearchResult with answer and sources
+        """
+        model = model or self.model
+        
+        # Build filter context for the query if filters are provided
+        # Gemini File Search doesn't support explicit filters in the API,
+        # so we add them as context in the query
+        filter_context = ""
+        if filters:
+            filter_parts = []
+            if filters.get("year"):
+                years = filters["year"] if isinstance(filters["year"], list) else [filters["year"]]
+                filter_parts.append(f"Year: {', '.join(years)}")
+            if filters.get("sources"):
+                sources = filters["sources"] if isinstance(filters["sources"], list) else [filters["sources"]]
+                filter_parts.append(f"Source: {', '.join(sources)}")
+            if filters.get("district"):
+                districts = filters["district"] if isinstance(filters["district"], list) else [filters["district"]]
+                filter_parts.append(f"District: {', '.join(districts)}")
+            if filters.get("filenames"):
+                filenames = filters["filenames"] if isinstance(filters["filenames"], list) else [filters["filenames"]]
+                filter_parts.append(f"Filename: {', '.join(filenames)}")
+            
+            if filter_parts:
+                filter_context = f"\n\nPlease focus on documents matching these criteria: {', '.join(filter_parts)}"
+        
+        # Combine query with filter context
+        # Add comprehensive system instructions similar to multi-agent system
+        system_instructions = """You are a helpful audit report assistant specialized in analyzing government audit reports from Uganda's Office of the Auditor General.
+
+CRITICAL RULES:
+1. **NO HALLUCINATION**: Only use information that is explicitly stated in the retrieved documents. Do not make up facts, numbers, or details.
+2. **Document References**: Always cite which documents you're using with [Doc i] references at the end of sentences that use specific information.
+3. **Formatting**: Structure your response with clear paragraphs, bullet points, or sections for readability.
+4. **Accuracy**: If the retrieved documents don't contain the requested information, explicitly state "The retrieved documents do not contain information about [topic]."
+5. **Years and Data**: Pay careful attention to years mentioned in documents. If a user asks about a specific year but documents show different years, explicitly state this.
+6. **District/Source Names**: Use the exact district and source names as they appear in the document metadata (e.g., "Kalangala" not "Kalagala").
+7. **Financial Data**: When providing financial figures, include the currency (UGX) and be precise about amounts.
+8. **Conversational Tone**: Be helpful, clear, and conversational while maintaining accuracy.
+
+IMPORTANT: Only use information from the retrieved documents. Do not use information from your training data unless it's explicitly mentioned in the retrieved documents."""
+        
+        # Combine system instructions with query
+        full_query = f"{system_instructions}\n\nUser Question: {query}{filter_context}\n\nPlease provide a detailed, well-formatted response with proper document references."
+        
+        try:
+            # Generate content with file search
+            # Based on Gemini API docs: https://ai.google.dev/gemini-api/docs/file-search
+            # Try with full path format first, then fallback to just ID if needed
+            store_name_to_try = self.store_name
+            
+            try:
+                # Try the documented format first with full path
+                response = self.client.models.generate_content(
+                    model=model,
+                    contents=full_query,
+                    config=types.GenerateContentConfig(
+                        tools=[
+                            types.Tool(
+                                file_search=types.FileSearch(
+                                    file_search_store_names=[store_name_to_try]
+                                )
+                            )
+                        ]
+                    )
+                )
+            except Exception as api_error:
+                error_str = str(api_error).lower()
+                # If format error, try with just the ID (without fileSearchStores/ prefix)
+                if 'format' in error_str or 'invalid' in error_str or 'too long' in error_str:
+                    logger.warning(f"Full path format failed, trying with just store ID: {api_error}")
+                    # Extract just the ID part
+                    if store_name_to_try.startswith("fileSearchStores/"):
+                        store_id = store_name_to_try.split("/", 1)[1]
+                        store_name_to_try = store_id
+                    
+                    try:
+                        response = self.client.models.generate_content(
+                            model=model,
+                            contents=full_query,
+                            config=types.GenerateContentConfig(
+                                tools=[
+                                    types.Tool(
+                                        file_search=types.FileSearch(
+                                            file_search_store_names=[store_name_to_try]
+                                        )
+                                    )
+                                ]
+                            )
+                        )
+                    except Exception as e2:
+                        raise Exception(f"Failed to call Gemini API with both formats. Full path error: {api_error}, ID-only error: {e2}")
+                else:
+                    # Try alternative dict format
+                    logger.warning(f"Primary API format failed, trying alternative: {api_error}")
+                    try:
+                        response = self.client.models.generate_content(
+                            model=model,
+                            contents=full_query,
+                            tools=[{
+                                "file_search": {
+                                    "file_search_store_names": [store_name_to_try]
+                                }
+                            }]
+                        )
+                    except Exception as e2:
+                        raise Exception(f"Failed to call Gemini API: {e2}")
+            
+            # Extract answer
+            answer = ""
+            if hasattr(response, 'text'):
+                answer = response.text
+            elif hasattr(response, 'candidates') and response.candidates:
+                # Try to get text from first candidate
+                candidate = response.candidates[0]
+                if hasattr(candidate, 'content') and candidate.content:
+                    if hasattr(candidate.content, 'parts'):
+                        text_parts = []
+                        for part in candidate.content.parts:
+                            if hasattr(part, 'text'):
+                                text_parts.append(part.text)
+                        answer = " ".join(text_parts)
+                    elif isinstance(candidate.content, str):
+                        answer = candidate.content
+            else:
+                answer = str(response)
+            
+            # Extract grounding metadata (document references)
+            sources = []
+            grounding_metadata = None
+            
+            logger.info(f"🔍 Extracting sources from Gemini response...")
+            
+            if hasattr(response, 'candidates') and response.candidates:
+                candidate = response.candidates[0]
+                logger.info(f"   Found candidate, checking for grounding_metadata...")
+                
+                # Get grounding metadata
+                if hasattr(candidate, 'grounding_metadata'):
+                    grounding_metadata = candidate.grounding_metadata
+                    logger.info(f"   Found grounding_metadata: {type(grounding_metadata)}")
+                    
+                    # Extract source documents from grounding metadata
+                    # Handle different response formats
+                    grounding_chunks = None
+                    if hasattr(grounding_metadata, 'grounding_chunks'):
+                        grounding_chunks = grounding_metadata.grounding_chunks
+                        logger.info(f"   Found grounding_chunks (attr): {len(grounding_chunks) if grounding_chunks else 0}")
+                    elif isinstance(grounding_metadata, dict) and 'grounding_chunks' in grounding_metadata:
+                        grounding_chunks = grounding_metadata['grounding_chunks']
+                        logger.info(f"   Found grounding_chunks (dict): {len(grounding_chunks) if grounding_chunks else 0}")
+                    elif hasattr(grounding_metadata, '__dict__'):
+                        # Try to access as object attributes
+                        metadata_dict = grounding_metadata.__dict__
+                        if 'grounding_chunks' in metadata_dict:
+                            grounding_chunks = metadata_dict['grounding_chunks']
+                            logger.info(f"   Found grounding_chunks (__dict__): {len(grounding_chunks) if grounding_chunks else 0}")
+                    
+                    if grounding_chunks:
+                        logger.info(f"   Processing {len(grounding_chunks)} grounding chunks...")
+                        for idx, chunk in enumerate(grounding_chunks):
+                            # Handle both object and dict formats
+                            try:
+                                if isinstance(chunk, dict):
+                                    chunk_data = chunk
+                                else:
+                                    # Object format - convert to dict-like access
+                                    chunk_data = {}
+                                    if hasattr(chunk, 'chunk'):
+                                        chunk_obj = chunk.chunk
+                                        chunk_data['chunk'] = {
+                                            'text': getattr(chunk_obj, 'text', ''),
+                                            'file_name': getattr(chunk_obj, 'file_name', '')
+                                        }
+                                    if hasattr(chunk, 'relevance_score'):
+                                        score_obj = chunk.relevance_score
+                                        chunk_data['relevance_score'] = {
+                                            'score': getattr(score_obj, 'score', 0.0)
+                                        }
+                                
+                                chunk_info = chunk_data.get('chunk', {})
+                                text = chunk_info.get('text', '') if isinstance(chunk_info, dict) else ''
+                                file_name = chunk_info.get('file_name', '') if isinstance(chunk_info, dict) else ''
+                                
+                                # Try to extract file URI and parse metadata from it
+                                file_uri = chunk_info.get('file_uri', '') if isinstance(chunk_info, dict) else ''
+                                
+                                # Also check for 'web' attribute (GroundingChunkData structure)
+                                if hasattr(chunk, 'web') and chunk.web:
+                                    web_data = chunk.web
+                                    file_uri = getattr(web_data, 'file_uri', '') or file_uri
+                                    file_name = getattr(web_data, 'title', '') or getattr(web_data, 'filename', '') or file_name
+                                    text = getattr(web_data, 'text', '') or getattr(web_data, 'content', '') or text
+                                
+                                # Check retrieved_context - this is where the actual data seems to be!
+                                if hasattr(chunk, 'retrieved_context') and chunk.retrieved_context:
+                                    rc = chunk.retrieved_context
+                                    # Get text content
+                                    if hasattr(rc, 'text'):
+                                        text = getattr(rc, 'text', '') or text
+                                    # Get document name
+                                    if hasattr(rc, 'document_name'):
+                                        doc_name = getattr(rc, 'document_name', '')
+                                        if doc_name:
+                                            file_name = doc_name or file_name
+                                
+                                # Fallback: Parse from string representation if we still don't have filename
+                                if not file_name:
+                                    chunk_str = str(chunk)
+                                    import re
+                                    # Look for PDF filenames
+                                    pdf_match = re.search(r"([A-Za-z0-9\s_-]+\.pdf)", chunk_str)
+                                    if pdf_match:
+                                        file_name = pdf_match.group(1)
+                                    # Or look for title= pattern
+                                    if not file_name and 'title=' in chunk_str:
+                                        title_match = re.search(r"title=['\"]([^'\"]+)['\"]", chunk_str)
+                                        if title_match:
+                                            file_name = title_match.group(1)
+                                
+                                if not file_name and file_uri:
+                                    # Extract filename from URI if available
+                                    file_name = file_uri.split('/')[-1] if '/' in file_uri else file_uri
+                                
+                                score_data = chunk_data.get('relevance_score', {})
+                                score = score_data.get('score', 0.0) if isinstance(score_data, dict) else 0.0
+                                
+                                if text or file_name:  # Only add if we have content
+                                    source_info = {
+                                        "content": text,
+                                        "filename": file_name,
+                                        "score": score,
+                                        "file_uri": file_uri,
+                                    }
+                                    sources.append(source_info)
+                                    logger.info(f"📄 Extracted source {idx+1}: {file_name} (score: {score:.3f}, content length: {len(text)})")
+                            except Exception as e:
+                                logger.warning(f"Error extracting chunk {idx+1} info: {e}")
+                                import traceback
+                                logger.debug(traceback.format_exc())
+                                continue
+                    else:
+                        logger.warning(f"   No grounding_chunks found in grounding_metadata")
+                else:
+                    logger.warning(f"   Candidate does not have grounding_metadata attribute")
+                    
+                # Also try to get file references from other parts of the response
+                # Sometimes Gemini includes file references in the response itself
+                if not sources or len(sources) == 0:
+                    logger.info(f"   No sources from grounding_metadata, trying alternative extraction...")
+                    # Check if response has file references in other attributes
+                    if hasattr(candidate, 'content') and candidate.content:
+                        if hasattr(candidate.content, 'parts'):
+                            for part in candidate.content.parts:
+                                if hasattr(part, 'file_data'):
+                                    file_data = part.file_data
+                                    if hasattr(file_data, 'file_uri') or (isinstance(file_data, dict) and 'file_uri' in file_data):
+                                        file_uri = getattr(file_data, 'file_uri', None) or (file_data.get('file_uri') if isinstance(file_data, dict) else None)
+                                        if file_uri:
+                                            file_name = file_uri.split('/')[-1] if '/' in file_uri else file_uri
+                                            sources.append({
+                                                "content": "",
+                                                "filename": file_name,
+                                                "score": 0.0,
+                                                "file_uri": file_uri,
+                                            })
+                                            logger.info(f"📄 Extracted source from file_data: {file_name}")
+            
+            logger.info(f"✅ Total sources extracted: {len(sources)}")
+            
+            return GeminiFileSearchResult(
+                answer=answer,
+                sources=sources,
+                grounding_metadata=grounding_metadata,
+                query=query
+            )
+            
+        except Exception as e:
+            # Return error result
+            return GeminiFileSearchResult(
+                answer=f"I apologize, but I encountered an error: {str(e)}",
+                sources=[],
+                query=query
+            )
+    
+    def format_sources_for_display(self, result: GeminiFileSearchResult) -> List[Any]:
+        """
+        Format Gemini sources to match the format expected by the UI.
+        
+        Returns list of document-like objects compatible with existing display code.
+        """
+        from langchain.docstore.document import Document
+        
+        formatted_sources = []
+        
+        for i, source in enumerate(result.sources):
+            filename = source.get("filename", "Unknown")
+            
+            # Try to extract metadata from filename (e.g., "Kalangala DLG Report of Auditor General 2021.pdf")
+            year = None
+            district = None
+            source_name = "Gemini File Search"
+            
+            # Parse filename for year
+            import re
+            year_match = re.search(r'\b(20\d{2})\b', filename)
+            if year_match:
+                year = int(year_match.group(1))
+            
+            # Parse filename for district/source
+            if "Kalangala" in filename:
+                district = "Kalangala"
+                source_name = "Kalangala DLG"
+            elif "Gulu" in filename:
+                district = "Gulu"
+                source_name = "Gulu DLG"
+            elif "KCCA" in filename:
+                district = "Kampala"
+                source_name = "KCCA"
+            elif "MAAIF" in filename:
+                source_name = "MAAIF"
+            elif "MWTS" in filename:
+                source_name = "MWTS"
+            elif "Consolidated" in filename:
+                source_name = "Consolidated"
+            
+            # Create a Document object compatible with existing code
+            doc = Document(
+                page_content=source.get("content", ""),
+                metadata={
+                    "filename": filename,
+                    "source": source_name,
+                    "score": source.get("score"),
+                    "chunk_index": i,
+                    "page": None,  # Gemini doesn't provide page numbers
+                    "year": year,
+                    "district": district,
+                    "chunk_id": f"gemini_{i}",
+                    "_id": f"gemini_{i}",
+                }
+            )
+            formatted_sources.append(doc)
+            logger.info(f"📋 Formatted source {i+1}: {filename} ({year}, {source_name})")
+        
+        logger.info(f"✅ Formatted {len(formatted_sources)} sources for display")
+        return formatted_sources
+

src/pipeline.py

@@ -14,7 +14,7 @@ except ModuleNotFoundError as me:
 
 from .logging import log_error
 
-from .loader import chunks_to_documents
+from .llm.loader import chunks_to_documents
 from .vectorstore import VectorStoreManager
 from .reporting.service import ReportService
 from .retrieval.context import ContextRetriever

src/reporting/__init__.py

@@ -1,4 +1,8 @@
-"""Report metadata and utilities."""
+"""Report metadata and utilities.
+
+This module is kept for backward compatibility with pipeline.py.
+For feedback-related functionality, use src.feedback instead.
+"""
 
 from .metadata import get_report_metadata, get_available_sources
 from .service import ReportService

src/streamlit_app.py

@@ -1,40 +0,0 @@
-import altair as alt
-import numpy as np
-import pandas as pd
-import streamlit as st
-
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))

src/ui_components/__init__.py

@@ -0,0 +1,21 @@
+"""
+UI Components Module
+
+This module contains UI-related components including styles, visualizations,
+and utility functions for the Streamlit application.
+"""
+
+from .styles import get_custom_css
+from .components import (
+    display_chunk_statistics_charts,
+    display_chunk_statistics_table
+)
+from .utils import extract_chunk_statistics
+
+__all__ = [
+    "get_custom_css",
+    "display_chunk_statistics_charts",
+    "display_chunk_statistics_table",
+    "extract_chunk_statistics"
+]
+

src/ui_components/components.py

@@ -0,0 +1,202 @@
+"""
+UI components for displaying statistics and visualizations
+"""
+
+import streamlit as st
+import pandas as pd
+import plotly.express as px
+from typing import Dict, Any
+
+
+def display_chunk_statistics_charts(stats: Dict[str, Any], title: str = "Retrieval Statistics"):
+    """Display statistics as interactive charts for 10+ results."""
+    if not stats or stats.get('total_chunks', 0) == 0:
+        return
+    
+    # Wrap everything in one styled container - open it
+    st.markdown(f"""
+    <div class="retrieval-distribution-container">
+        <h3 style="margin-top: 0;">📊 {title}</h3>
+        <div style="display: flex; justify-content: space-around; align-items: center; padding: 15px 0; border-bottom: 1px solid #e0e0e0; margin-bottom: 20px;">
+            <div class="metric-container">
+                <div class="metric-label">Total Chunks</div>
+                <div class="metric-value">{stats['total_chunks']}</div>
+            </div>
+            <div class="metric-container">
+                <div class="metric-label">Unique Sources</div>
+                <div class="metric-value">{stats['unique_sources']}</div>
+            </div>
+            <div class="metric-container">
+                <div class="metric-label">Unique Years</div>
+                <div class="metric-value">{stats['unique_years']}</div>
+            </div>
+            <div class="metric-container">
+                <div class="metric-label">Unique Files</div>
+                <div class="metric-value">{stats['unique_filenames']}</div>
+            </div>
+        </div>
+    """, unsafe_allow_html=True)
+    
+    # Charts - three columns to include Districts
+    col1, col2, col3 = st.columns(3)
+    
+    with col1:
+        # Source distribution chart
+        if stats['source_distribution']:
+            source_df = pd.DataFrame(
+                list(stats['source_distribution'].items()),
+                columns=['Source', 'Count']
+            )
+            fig_source = px.bar(
+                source_df,
+                x='Count',
+                y='Source',
+                orientation='h',
+                title='Distribution by Source',
+                color='Count',
+                color_continuous_scale='viridis'
+            )
+            fig_source.update_layout(height=400, showlegend=False)
+            st.plotly_chart(fig_source, use_container_width=True)  # Note: plotly_chart still uses use_container_width
+    
+    with col2:
+        # Year distribution chart
+        if stats['year_distribution']:
+            # Filter out 'Unknown' years for the chart
+            year_dist_filtered = {k: v for k, v in stats['year_distribution'].items() if k != 'Unknown'}
+            if year_dist_filtered:
+                year_df = pd.DataFrame(
+                    list(year_dist_filtered.items()),
+                    columns=['Year', 'Count']
+                )
+                # Sort by year as integer but keep as string for categorical display
+                year_df['Year_Int'] = year_df['Year'].astype(int)
+                year_df = year_df.sort_values('Year_Int').drop('Year_Int', axis=1)
+                
+                fig_year = px.bar(
+                    year_df,
+                    x='Year',
+                    y='Count',
+                    title='Distribution by Year',
+                    color='Count',
+                    color_continuous_scale='plasma'
+                )
+                # Ensure years are treated as categorical (discrete) not continuous
+                fig_year.update_xaxes(type='category')
+                fig_year.update_layout(height=400, showlegend=False)
+                st.plotly_chart(fig_year, use_container_width=True)  # Note: plotly_chart still uses use_container_width
+            else:
+                st.info("No valid years found in the results")
+    
+    with col3:
+        # District distribution chart
+        if stats.get('district_distribution'):
+            district_dist_filtered = {k: v for k, v in stats['district_distribution'].items() if k != 'Unknown'}
+            if district_dist_filtered:
+                district_df = pd.DataFrame(
+                    list(district_dist_filtered.items()),
+                    columns=['District', 'Count']
+                )
+                district_df = district_df.sort_values('Count', ascending=False)
+                
+                fig_district = px.bar(
+                    district_df,
+                    x='Count',
+                    y='District',
+                    orientation='h',
+                    title='Distribution by District',
+                    color='Count',
+                    color_continuous_scale='blues'
+                )
+                fig_district.update_layout(height=400, showlegend=False)
+                st.plotly_chart(fig_district, use_container_width=True)  # Note: plotly_chart still uses use_container_width
+            else:
+                st.info("No valid districts found in the results")
+    
+    # Close the container
+    st.markdown('</div>', unsafe_allow_html=True)
+
+
+def display_chunk_statistics_table(stats: Dict[str, Any], title: str = "Retrieval Distribution"):
+    """Display statistics as tables for smaller results with fixed alignment."""
+    if not stats or stats.get('total_chunks', 0) == 0:
+        return
+    
+    # Wrap in styled container
+    st.markdown('<div class="retrieval-distribution-container">', unsafe_allow_html=True)
+    
+    st.subheader(f"📊 {title}")
+    
+    # Create a container with fixed height for alignment
+    stats_container = st.container()
+    
+    with stats_container:
+        # Create 4 equal columns for consistent alignment
+        col1, col2, col3, col4 = st.columns(4)
+        
+        with col1:
+            st.markdown("**🏘️ Districts**")
+            if stats.get('district_distribution'):
+                district_dist_filtered = {k: v for k, v in stats['district_distribution'].items() if k != 'Unknown'}
+                if district_dist_filtered:
+                    district_data = {
+                        "District": list(district_dist_filtered.keys()),
+                        "Count": list(district_dist_filtered.values())
+                    }
+                    district_df = pd.DataFrame(district_data).sort_values('Count', ascending=False)
+                    st.dataframe(district_df, hide_index=True, width='stretch')
+                else:
+                    st.write("No district data")
+            else:
+                st.write("No district data")
+        
+        with col2:
+            st.markdown("**📂 Sources**")
+            if stats['source_distribution']:
+                source_data = {
+                    "Source": list(stats['source_distribution'].keys()),
+                    "Count": list(stats['source_distribution'].values())
+                }
+                source_df = pd.DataFrame(source_data).sort_values('Count', ascending=False)
+                st.dataframe(source_df, hide_index=True, width='stretch')
+            else:
+                st.write("No source data")
+        
+        with col3:
+            st.markdown("**📅 Years**")
+            if stats['year_distribution']:
+                year_dist_filtered = {k: v for k, v in stats['year_distribution'].items() if k != 'Unknown'}
+                if year_dist_filtered:
+                    year_data = {
+                        "Year": list(year_dist_filtered.keys()),
+                        "Count": list(year_dist_filtered.values())
+                    }
+                    year_df = pd.DataFrame(year_data)
+                    # Sort by year as integer but display as string
+                    year_df['Year_Int'] = year_df['Year'].astype(int)
+                    year_df = year_df.sort_values('Year_Int')[['Year', 'Count']]
+                    st.dataframe(year_df, hide_index=True, width='stretch')
+                else:
+                    st.write("No year data")
+            else:
+                st.write("No year data")
+        
+        with col4:
+            st.markdown("**📄 Files**")
+            if stats['filename_distribution']:
+                filename_items = list(stats['filename_distribution'].items())
+                filename_items.sort(key=lambda x: x[1], reverse=True)
+                
+                # Show top files with truncated names
+                file_data = {
+                    "File": [f[:30] + "..." if len(f) > 30 else f for f, c in filename_items[:5]],
+                    "Count": [c for f, c in filename_items[:5]]
+                }
+                file_df = pd.DataFrame(file_data)
+                st.dataframe(file_df, hide_index=True, width='stretch')
+            else:
+                st.write("No file data")
+    
+    # Close container
+    st.markdown('</div>', unsafe_allow_html=True)
+

src/ui_components/styles.py

@@ -0,0 +1,117 @@
+"""
+Custom CSS styles for Streamlit application
+"""
+
+
+def get_custom_css() -> str:
+    """Get custom CSS styles as a string"""
+    return """
+<style>
+    .main-header {
+        font-size: 2.5rem;
+        font-weight: bold;
+        color: #1f77b4;
+        text-align: center;
+        margin-bottom: 1rem;
+        width: 100%;
+        display: block;
+    }
+    
+    .subtitle {
+        font-size: 1.2rem;
+        color: #666;
+        text-align: center;
+        margin-bottom: 2rem;
+        width: 100%;
+        display: block;
+    }
+    
+    .session-info {
+        background-color: #f0f2f6;
+        padding: 10px;
+        border-radius: 5px;
+        margin-bottom: 20px;
+        font-size: 0.9rem;
+    }
+    
+    .user-message {
+        background-color: #007bff;
+        color: white;
+        padding: 12px 16px;
+        border-radius: 18px 18px 4px 18px;
+        margin: 8px 0;
+        margin-left: 20%;
+        word-wrap: break-word;
+    }
+    
+    .bot-message {
+        background-color: #f1f3f4;
+        color: #333;
+        padding: 12px 16px;
+        border-radius: 18px 18px 18px 4px;
+        margin: 8px 0;
+        margin-right: 20%;
+        word-wrap: break-word;
+        border: 1px solid #e0e0e0;
+    }
+    
+    .filter-section {
+        margin-bottom: 20px;
+        padding: 15px;
+        background-color: #f8f9fa;
+        border-radius: 8px;
+        border: 1px solid #e9ecef;
+    }
+    
+    .filter-title {
+        font-weight: bold;
+        margin-bottom: 10px;
+        color: #495057;
+    }
+    
+    .feedback-section {
+        background-color: #f8f9fa;
+        padding: 20px;
+        border-radius: 10px;
+        margin-top: 30px;
+        border: 2px solid #dee2e6;
+    }
+    
+    .retrieval-history {
+        background-color: #ffffff;
+        padding: 15px;
+        border-radius: 5px;
+        margin: 10px 0;
+        border-left: 4px solid #007bff;
+    }
+    
+    .retrieval-distribution-container {
+        background-color: #ffffff;
+        padding: 25px;
+        border-radius: 10px;
+        margin: 20px 0;
+        border: 2px solid #e0e0e0;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1), 0 2px 4px rgba(0, 0, 0, 0.06);
+    }
+    
+    .metric-label {
+        font-size: 0.9rem;
+        color: #555;
+        margin-bottom: 5px;
+        text-align: center;
+    }
+    
+    .metric-value {
+        font-size: 1.8rem;
+        font-weight: bold;
+        color: #000000;
+        text-align: center;
+    }
+    
+    .metric-container {
+        text-align: center;
+        padding: 10px;
+    }
+</style>
+"""
+

src/ui_components/utils.py

@@ -0,0 +1,73 @@
+"""
+UI utility functions for data processing and statistics
+"""
+
+from typing import Dict, Any, List
+from collections import Counter
+
+
+def extract_chunk_statistics(sources: List[Any]) -> Dict[str, Any]:
+    """Extract statistics from retrieved chunks."""
+    if not sources:
+        return {}
+    
+    sources_list = []
+    years = []
+    filenames = []
+    districts = []
+    
+    for doc in sources:
+        metadata = getattr(doc, 'metadata', {})
+        
+        # Extract source
+        source = metadata.get('source', 'Unknown')
+        sources_list.append(source)
+        
+        # Extract year
+        year = metadata.get('year', 'Unknown')
+        if year and year != 'Unknown':
+            try:
+                # Convert to int first, then back to string to ensure it's a proper year
+                year_int = int(float(year))  # Handle both int and float strings
+                if 1900 <= year_int <= 2030:  # Reasonable year range
+                    years.append(str(year_int))
+                else:
+                    years.append('Unknown')
+            except (ValueError, TypeError):
+                years.append('Unknown')
+        else:
+            years.append('Unknown')
+        
+        # Extract filename
+        filename = metadata.get('filename', 'Unknown')
+        filenames.append(filename)
+        
+        # Extract district
+        district = metadata.get('district', 'Unknown')
+        if district and district != 'Unknown':
+            districts.append(district)
+        else:
+            districts.append('Unknown')
+    
+    # Count occurrences
+    source_counts = Counter(sources_list)
+    year_counts = Counter(years)
+    filename_counts = Counter(filenames)
+    district_counts = Counter(districts)
+    
+    return {
+        'total_chunks': len(sources),
+        'unique_sources': len(source_counts),
+        'unique_years': len([y for y in year_counts.keys() if y != 'Unknown']),
+        'unique_filenames': len(filename_counts),
+        'unique_districts': len([d for d in district_counts.keys() if d != 'Unknown']),
+        'source_distribution': dict(source_counts),
+        'year_distribution': dict(year_counts),
+        'filename_distribution': dict(filename_counts),
+        'district_distribution': dict(district_counts),
+        'sources': sources_list,
+        'years': years,
+        'filenames': filenames,
+        'districts': districts
+    }
+

src/vectorstore.py

@@ -1,9 +1,20 @@
 """Vector store management and operations."""
+import os
+# Disable MPS before importing torch to prevent meta tensor issues on Mac
+os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
+
 from pathlib import Path
 from typing import Dict, Any, List, Optional
 
 
 import torch
+# Disable MPS backend explicitly to prevent meta tensor issues
+if hasattr(torch.backends, 'mps'):
+    # Monkey patch to disable MPS
+    original_mps_available = torch.backends.mps.is_available
+    torch.backends.mps.is_available = lambda: False
+
 from langchain_qdrant import QdrantVectorStore
 from langchain.docstore.document import Document
 from langchain_core.embeddings import Embeddings
@@ -28,11 +39,23 @@ class MatryoshkaEmbeddings(Embeddings):
 
         if truncate_dim and "matryoshka" in model_name.lower():
             # Use SentenceTransformer directly for Matryoshka models
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            self.model = SentenceTransformer(model_name, truncate_dim=truncate_dim, device=device)
+            # Fix for meta tensor issue: Explicitly force CPU
+            # MPS is already disabled at module level
+            # Explicitly pass device="cpu" to prevent MPS/CUDA detection
+            self.model = SentenceTransformer(
+                model_name, 
+                truncate_dim=truncate_dim,
+                device="cpu"  # Force CPU to prevent meta tensor issues
+            )
             print(f"🔧 Matryoshka model configured for {truncate_dim} dimensions")
         else:
             # Use standard HuggingFaceEmbeddings
+            # Don't pass device parameter - let it load naturally on CPU
+            # This prevents the meta tensor error
+            if "model_kwargs" not in kwargs:
+                kwargs["model_kwargs"] = {}
+            # Remove device from model_kwargs if present to prevent meta tensor issues
+            kwargs["model_kwargs"].pop("device", None)
             self.model = HuggingFaceEmbeddings(model_name=model_name, **kwargs)
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
@@ -76,12 +99,17 @@ class VectorStoreManager:
 
     def _create_embeddings(self) -> HuggingFaceEmbeddings:
         """Create embeddings model from configuration."""
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        
         model_name = self.config["retriever"]["model"]
         normalize = self.config["retriever"]["normalize"]
         
-        model_kwargs = {"device": device}
+        # Fix for meta tensor issue: Force CPU usage to prevent MPS/CUDA detection
+        # The error occurs when SentenceTransformer detects MPS/CUDA and tries to move meta tensors
+        # MPS is already disabled at module level, now we explicitly force CPU in model_kwargs
+        model_kwargs = {
+            "device": "cpu",  # Explicitly force CPU to prevent MPS/CUDA detection
+            "trust_remote_code": True,  # Some models need this
+        }
+        
         encode_kwargs = {
             "normalize_embeddings": normalize,
             "batch_size": 100,
@@ -108,6 +136,8 @@ class VectorStoreManager:
                 return embeddings
         
         # Use standard HuggingFaceEmbeddings for non-Matryoshka models
+        # Don't pass device in model_kwargs - let HuggingFaceEmbeddings handle it
+        # but ensure we're not using meta device
         embeddings = HuggingFaceEmbeddings(
             model_name=model_name,
             model_kwargs=model_kwargs,