added new code

src/api_client.py

@@ -0,0 +1,254 @@
+"""
+API client for communicating with the backend service.
+"""
+import requests
+import logging
+from typing import Dict, Any, Tuple, Optional, List
+
+
+class APIClient:
+    
+    def __init__(self, base_url: str, timeout: int = 60):
+        self.base_url = base_url.rstrip('/')  # Remove trailing slash if present
+        self.timeout = timeout  # Used only for health checks, not for query processing
+        self.logger = logging.getLogger(__name__)
+        
+        # Define endpoints
+        self.endpoints = {
+            'process_text': '/api/process-text',
+            'health': '/api/health',
+            'models': '/api/models',
+            'change_model': '/api/change-model'
+        }
+    
+    def send_query(self, question: str, model: str = None, agent: str = None) -> Dict[str, Any]:
+        """
+        Send a query to the backend with optional model specification
+        
+        Args:
+            question: The user's question
+            model: Optional AI model name to use for processing
+            agent: Agent type (legacy parameter, not used in current backend)
+        """
+        # Prepare payload with model_name for the new API
+        payload = {"question": question}
+        if model:
+            payload["model_name"] = model
+            
+        self.logger.info(f"Sending API request with: {payload}")
+        
+        # Construct full URL
+        full_url = f"{self.base_url}{self.endpoints['process_text']}"
+        
+        try:
+            # Remove timeout for query processing - wait for actual result
+            # Complex queries may take time and should not be interrupted
+            response = requests.post(full_url, json=payload)
+            return self._process_response(response)
+        except requests.exceptions.ConnectionError as e:
+            self.logger.error(f"Connection error: {e}")
+            return {
+                "message": "❌ Cannot connect to the server. Please check if the backend is running.",
+                "rows": [],
+                "chart": None,
+                "error": True
+            }
+        except requests.exceptions.RequestException as e:
+            self.logger.error(f"Request error: {e}")
+            return {
+                "message": f"❌ Request failed: {str(e)}",
+                "rows": [],
+                "chart": None,
+                "error": True
+            }
+        except Exception as e:
+            self.logger.error(f"Unexpected error: {e}")
+            return {
+                "message": f"⚠️ Exception: {str(e)}",
+                "rows": [],
+                "chart": None,
+                "error": True
+            }
+    
+    def _process_response(self, response: requests.Response) -> Dict[str, Any]:
+        """Process the API response and extract relevant data."""
+        try:
+            result = response.json()
+            self.logger.info(f"API response received: {type(result)} with keys: {list(result.keys()) if isinstance(result, dict) else 'Not a dict'}")
+        except ValueError as e:
+            self.logger.error(f"Failed to parse JSON response: {e}")
+            result = {"detail": response.text}
+        
+        if response.status_code != 200:
+            error_msg = result.get('detail') or result.get('message') or response.text
+            self.logger.error(f"API error response (status {response.status_code}): {error_msg}")
+            return {
+                "message": f"❌ Error: {error_msg}",
+                "rows": [],
+                "chart": None,
+                "error": True,
+                "model_used": result.get('model_used', 'unknown') if isinstance(result, dict) else 'unknown'
+            }
+        
+        # Handle message-only responses (chat responses, errors, etc.)
+        if "message" in result and not ("rows" in result or "chart" in result):
+            self.logger.info("Processing message-only response")
+            return {
+                "message": result["message"],
+                "rows": [],
+                "chart": None,
+                "error": False,
+                "model_used": result.get("model_used", "unknown"),
+                "status": result.get("status", "message")
+            }
+        
+        # Handle data responses (successful SQL queries)
+        rows = result.get("rows", [])
+        chart = result.get("chart", None)
+        heading = result.get("summary", "")
+        sql = result.get("sql", "")
+        
+        self.logger.info(f"Processing data response - Rows: {len(rows) if isinstance(rows, list) else 'Not a list'}, Heading: {heading}, SQL length: {len(sql) if sql else 0}")
+        
+        # Use heading from backend as-is (backend already parsed JSON and model includes record count)
+        if heading and isinstance(heading, str) and heading.strip():
+            message = heading.strip()
+        else:
+            # Fallback message if no heading provided
+            if rows and len(rows) > 0:
+                message = f"Here are the {len(rows)} results I found:"
+            else:
+                message = "I could not find matching records for your query."
+        
+        processed_response = {
+            "message": message,
+            "rows": rows,
+            "chart": chart,
+            "error": False,
+            "model_used": result.get("model_used", "unknown"),
+            "status": result.get("status", "success"),
+            "sql": sql,
+            "heading": heading,  # Clean heading from backend
+            "summary": result.get("summary", "")  # Summary from backend if available
+        }
+        
+        self.logger.info(f"Final processed response keys: {list(processed_response.keys())}")
+        return processed_response
+
+    def check_health(self) -> Tuple[str, str, str]:
+      
+        try:
+            # Construct health URL
+            health_url = f"{self.base_url}{self.endpoints['health']}"
+            response = requests.get(health_url)
+            try:
+                result = response.json()
+                status = result.get("status", "")
+            except Exception:
+                result = {}
+                status = ""
+            
+            if status == "healthy":
+                return "🟢 Active", "Online", "success"
+            elif response.status_code == 503:
+                return "🟡 Degraded", "Some Issues", "warning"
+            else:
+                return "🟡 Limited", f"Status: {response.status_code}", "warning"
+        except requests.exceptions.RequestException:
+            # Fallback to socket check
+            try:
+                import socket
+                # Extract host and port from base_url
+                from urllib.parse import urlparse
+                parsed = urlparse(self.base_url)
+                host = parsed.hostname or "127.0.0.1"
+                port = parsed.port or 8000
+                socket.create_connection((host, port), timeout=1).close()
+                return "🟡 Reachable", "Port Open", "warning"
+            except:
+                return "🔴 Offline", "Connection Failed", "error"
+        except Exception:
+            return "🟡 Unknown", "Check Required", "warning"
+
+    def get_detailed_health(self) -> Dict[str, Any]:
+        """
+        Get detailed health information from the API.
+        
+        Returns:
+            Dictionary containing detailed health status or error information
+        """
+        try:
+            health_url = f"{self.base_url}{self.endpoints['health']}"
+            response = requests.get(health_url, timeout=5)
+            
+            if response.status_code in [200, 503]:
+                return response.json()
+            else:
+                return {
+                    "status": "error",
+                    "message": f"Health endpoint returned status {response.status_code}",
+                    "checks": {}
+                }
+                
+        except requests.exceptions.RequestException as e:
+            return {
+                "status": "error", 
+                "message": f"Connection failed: {str(e)}",
+                "checks": {}
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": f"Health check failed: {str(e)}",
+                "checks": {}
+            }
+    
+    def get(self, endpoint: str, params: Optional[Dict] = None) -> Optional[Dict[str, Any]]:
+        """
+        Generic GET method for API endpoints
+        
+        Args:
+            endpoint: API endpoint path (e.g., "/models")
+            params: Optional query parameters
+            
+        Returns:
+            API response as dictionary or None if failed
+        """
+        try:
+            url = f"{self.base_url}/api{endpoint}" if not endpoint.startswith('/api') else f"{self.base_url}{endpoint}"
+            response = requests.get(url, params=params, timeout=10)
+            
+            if response.status_code == 200:
+                return response.json()
+            else:
+                self.logger.error(f"GET {endpoint} failed with status {response.status_code}")
+                return None
+                
+        except Exception as e:
+            self.logger.error(f"GET {endpoint} error: {str(e)}")
+            return None
+    
+    def post(self, endpoint: str, data: Optional[Dict] = None) -> Optional[Dict[str, Any]]:
+        """
+        Generic POST method for API endpoints
+        
+        Args:
+            endpoint: API endpoint path
+            data: Optional POST data
+            
+        Returns:
+            API response as dictionary or None if failed
+        """
+        try:
+            url = f"{self.base_url}/api{endpoint}" if not endpoint.startswith('/api') else f"{self.base_url}{endpoint}"
+            response = requests.post(url, json=data, timeout=10)
+            
+            if response.status_code in [200, 201]:
+                return response.json()
+            else:
+                self.logger.error(f"POST {endpoint} failed with status {response.status_code}")
+                return None
+                
+        except Exception as e:
+            self.logger.error(f"POST {endpoint} error: {str(e)}")
+            return None

src/chat_components.py

@@ -0,0 +1,151 @@
+"""
+Chat components for rendering messages and handling user interactions.
+"""
+import streamlit as st
+import pandas as pd
+import base64
+from typing import List, Dict, Any
+from config import OUTLINE_INDIGO_USER, DARK_MODE_SLATE_AI, RETRY_BUTTON_TEXT, DOWNLOAD_BUTTON_TEXT
+
+
+class ChatRenderer:
+    """Handles rendering of chat messages and interactions."""
+    
+    def __init__(self):
+        self.user_avatar = OUTLINE_INDIGO_USER
+        self.ai_avatar = DARK_MODE_SLATE_AI
+    
+    def render_chat_history(self, messages: List[Dict[str, Any]]) -> None:
+        """Render the complete chat history."""
+        st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
+        
+        for i, msg in enumerate(messages):
+            if msg["role"] == "user":
+                self._render_user_message(msg, i, messages)
+            elif msg["role"] == "assistant":
+                # Skip error messages that are already shown after user messages
+                if not self._is_error_shown_after_user(i, messages):
+                    self._render_assistant_message(msg)
+        
+        st.markdown("</div>", unsafe_allow_html=True)
+    
+    def _render_user_message(self, msg: Dict[str, Any], index: int, messages: List[Dict[str, Any]]) -> None:
+        """Render a user message with avatar."""
+        st.markdown(
+            f'''<div style="display: flex; align-items: flex-start; justify-content: flex-end; margin-bottom: 0.5em;">
+                    <div style="margin-right: 0.5em;">
+                        <img src="{self.user_avatar}" alt="User" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #e3f2fd; background: #fff; object-fit: cover;" />
+                    </div>
+                    <div class="user-bubble">{msg["content"]}</div>
+               </div>''',
+            unsafe_allow_html=True
+        )
+        
+        # Check if next message is an error and render retry option
+        self._render_error_retry_if_needed(index, messages)
+    
+    def _render_error_retry_if_needed(self, index: int, messages: List[Dict[str, Any]]) -> None:
+        """Render error message and retry button if the next message is an error."""
+        if (
+            index + 1 < len(messages)
+            and messages[index + 1]["role"] == "assistant"
+            and messages[index + 1].get("is_error")
+        ):
+            error_msg = messages[index + 1]["content"]
+            cols = st.columns([0.3, 0.7])
+            
+            with cols[1]:
+                col1, col2 = st.columns([0.8, 0.2])
+                
+                with col1:
+                    st.markdown(
+                        f'<div style="background: #ffebee; color: #b71c1c; font-weight: bold; border-radius: 14px; padding: 8px 14px; max-width: 100%; word-wrap: break-word; box-shadow: 0 2px 8px rgba(0,0,0,0.15); text-align: right;">{error_msg}</div>',
+                        unsafe_allow_html=True
+                    )
+                
+                with col2:
+                    if st.button(RETRY_BUTTON_TEXT, key=f"retry_{index}"):
+                        self._handle_retry(index)
+    
+    def _handle_retry(self, index: int) -> None:
+        """Handle retry button click."""
+        messages = st.session_state["messages"]
+        st.session_state["messages"] = messages[:index + 1]
+        st.session_state["messages"].append({
+            "role": "assistant",
+            "content": "🤔 Thinking...",
+            "is_placeholder": True
+        })
+        st.rerun()
+    
+    def _is_error_shown_after_user(self, index: int, messages: List[Dict[str, Any]]) -> bool:
+        """Check if this error message is already shown after the previous user message."""
+        if not messages[index].get("is_error"):
+            return False
+        
+        # Check if this is an error that follows a user message
+        if index > 0 and messages[index - 1]["role"] == "user":
+            return True
+        
+        return False
+    
+    def _render_assistant_message(self, msg: Dict[str, Any]) -> None:
+        """Render an assistant message with avatar and optional data/charts."""
+        bubble_class = "error-bubble" if msg.get("is_error") else "ai-bubble"
+        
+        # Render message bubble
+        if msg.get("is_placeholder"):
+            self._render_thinking_message(bubble_class)
+        else:
+            self._render_regular_message(msg["content"], bubble_class)
+        
+        # Render data table if present
+        if msg.get("data"):
+            self._render_data_table(msg["data"], msg)
+        
+        # Render chart if present
+        if msg.get("chart"):
+            self._render_chart(msg["chart"])
+    
+    def _render_thinking_message(self, bubble_class: str) -> None:
+        """Render a thinking/loading message with spinner."""
+        st.markdown(
+            f'''<div style="display: flex; align-items: flex-start; margin-bottom: 0.5em;">
+                    <div style="margin-right: 0.5em;">
+                        <img src="{self.ai_avatar}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
+                    </div>
+                    <div class="{bubble_class}"><span class="spinner"></span>Thinking...</div>
+               </div>''',
+            unsafe_allow_html=True
+        )
+    
+    def _render_regular_message(self, content: str, bubble_class: str) -> None:
+        """Render a regular assistant message."""
+        st.markdown(
+            f'''<div style="display: flex; align-items: flex-start; margin-bottom: 0.5em;">
+                    <div style="margin-right: 0.5em;">
+                        <img src="{self.ai_avatar}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
+                    </div>
+                    <div class="{bubble_class}">{content}</div>
+               </div>''',
+            unsafe_allow_html=True
+        )
+    
+    def _render_data_table(self, data: List[Dict], msg: Dict[str, Any]) -> None:
+        """Render data table with download option."""
+        df = pd.DataFrame(data)
+        st.dataframe(df, use_container_width=True)
+        
+        csv = df.to_csv(index=False).encode("utf-8")
+        st.download_button(
+            DOWNLOAD_BUTTON_TEXT,
+            csv,
+            "results.csv",
+            "text/csv",
+            key=f"download_csv_{id(msg)}"
+        )
+    
+    def _render_chart(self, chart_data: str) -> None:
+        """Render chart from base64 data."""
+        img_data = base64.b64decode(chart_data)
+        st.image(img_data, use_column_width=True)

src/chat_ui.py

@@ -1,315 +1,981 @@
-
+"""
+AI Database Assistant - Streamlit Chat Interface
+"""
 import streamlit as st
-import requests
 import pandas as pd
 import base64
 import logging
-import os
-
+from typing import List, Dict, Any
+from api_client import APIClient
+from themes import ThemeManager
+from config import (
+    BASE_URL, PAGE_TITLE, PAGE_LAYOUT, AVAILABLE_MODELS, AVAILABLE_AGENTS,
+    OUTLINE_INDIGO_USER, DARK_MODE_SLATE_AI, AVAILABLE_THEMES,
+    CHAT_INPUT_PLACEHOLDER, THINKING_MESSAGE, WORKING_MESSAGE, RETRY_BUTTON_TEXT, DOWNLOAD_BUTTON_TEXT
+)
 
 # ===============================
-# Config: Dynamic API_URL based on environment
+# Configuration
 # ===============================
+# Configuration is now imported from config.py
+# To change environments, only modify BASE_URL in config.py
 
-# Read API base URL and endpoint from environment variables
-api_base_url = os.getenv("API_BASE_URL", "http://127.0.0.1:8000")
-api_endpoint = os.getenv("API_ENDPOINT", "/api/process-text")
-# Ensure no double slashes
-API_URL = api_base_url.rstrip("/") + "/" + api_endpoint.lstrip("/")
-
-st.set_page_config(page_title="AI Database Assistant", layout="wide")
-
-# Show current API URL in sidebar for debugging
-# st.sidebar.info(f"**API URL:** {API_URL}")
-
-# Setup logger
+# Setup
+st.set_page_config(page_title=PAGE_TITLE, layout=PAGE_LAYOUT)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
-st.title("💬 AI Database Assistant")
-st.write("Ask questions in plain English to generate and run SQL queries.")
-# st.markdown(
-#     "<h1 style='text-align: center; color: white; background: rgba(0,0,0,0.6); "
-#     "padding: 12px; border-radius: 10px;'>🤖 AI Database Assistant</h1>",
-#     unsafe_allow_html=True,
-# )
-# st.write("Ask questions in plain English to generate and run SQL queries.")
+# Initialize API Client with base URL from config
+api_client = APIClient(BASE_URL)
 
 # ===============================
-# Sidebar Settings
+# Theme Management
 # ===============================
-st.sidebar.header("⚙️ Settings")
-
-model = st.sidebar.selectbox("Select Model", ["gpt-3.5-turbo", "gpt-4", "gpt-4o-mini"])
-agent = st.sidebar.selectbox("Select Agent", ["default", "sql-agent", "custom-agent"])
-theme = st.sidebar.selectbox("Theme", ["Light", "Dark", "Custom"])
+# Initialize theme manager
+theme_manager = ThemeManager()
 
 # ===============================
-# Theme CSS
+# Database Status Functions  
 # ===============================
-def inject_css(theme_choice: str):
-    if theme_choice == "Light":
-        css = """
-        body, .stApp { background: #f8fafc !important; }
-        .chat-container { max-width: 900px; margin: auto; }
-        .user-bubble {
-            background: linear-gradient(90deg, #e3f2fd, #ffffff);
-            color: #222;
-            float: right;
-            clear: both;
-        }
-        .ai-bubble {
-            background: linear-gradient(90deg, #f3e5f5, #ffffff);
-            color: #222;
-            float: left;
-            clear: both;
-        }
-        .error-bubble {
-            background: #ffebee; color: #b71c1c; font-weight: bold;
-            float: left; clear: both;
-        }
-        section[data-testid="stSidebar"] { background: #e3f2fd !important; }
-        """
-    elif theme_choice == "Dark":
-        css = """
-        body, .stApp { background: #181825 !important; }
-        .chat-container { max-width: 900px; margin: auto; }
-        .user-bubble {
-            background: linear-gradient(135deg, #3a0ca3, #4361ee);
-            color: #f8fafc;
-            float: right;
-            clear: both;
-        }
-        .ai-bubble {
-            background: linear-gradient(135deg, #14213d, #1d3557);
-            color: #f8fafc;
-            float: left;
-            clear: both;
-        }
-        .error-bubble {
-            background: #ffb4ab; color: #b71c1c; font-weight: bold;
-            float: left; clear: both;
-        }
-        section[data-testid="stSidebar"] { background: #232946 !important; }
-        """
-    else:  # Custom theme
-        css = """
-        body, .stApp { background: #fffbe7 !important; }
-        .chat-container { max-width: 900px; margin: auto; }
-        .user-bubble {
-            background: linear-gradient(90deg, #ffe082, #fffbe7);
-            color: #6d4c00;
-            float: right;
-            clear: both;
-        }
-        .ai-bubble {
-            background: linear-gradient(90deg, #b2dfdb, #fffbe7);
-            color: #004d40;
-            float: left;
-            clear: both;
-        }
-        .error-bubble {
-            background: #ffe0b2; color: #b71c1c; font-weight: bold;
-            float: left; clear: both;
+@st.cache_data(ttl=30)  # Reduced cache time for more responsive health checks
+def check_api_status():
+    """Check if the API is reachable and responsive with a lightweight health check."""
+    try:
+        return api_client.check_health()
+    except Exception as e:
+        logging.error(f"Health check failed: {e}")
+        return "🔴 Error", f"Failed: {str(e)}", "error"
+
+@st.cache_data(ttl=30)  # Reduced cache time for detailed health
+def get_detailed_health_status():
+    """Get detailed health status from the API."""
+    try:
+        return api_client.get_detailed_health()
+    except Exception as e:
+        logging.error(f"Detailed health check failed: {e}")
+        return {
+            "status": "error",
+            "message": f"Health check failed: {str(e)}",
+            "checks": {}
         }
-        section[data-testid="stSidebar"] { background: #ffe082 !important; }
-        """
 
-    # Common bubble styling
-    css += """
-    .user-bubble, .ai-bubble, .error-bubble {
-        border-radius: 14px;
-        padding: 10px 14px;
-        margin: 6px 0;
-        display: inline-block;
-        max-width: 75%;
-        word-wrap: break-word;
-        box-shadow: 0 2px 8px rgba(0,0,0,0.25);
-    }
-    """
-    st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)
+def should_skip_api_call(force_refresh=False):
+    """Enhanced validation to determine if API calls should be skipped - FOR HEALTH CHECKS ONLY."""
+    try:
+        # Always allow API calls when force refresh is requested
+        if force_refresh:
+            logging.info("API call allowed: Force refresh requested")
+            return False
+        
+        # Check for immediate health check flag (when System Status is just enabled)
+        if st.session_state.get("force_immediate_health_check", False):
+            logging.info("API call allowed: System Status just enabled")
+            # Clear the flag after use
+            st.session_state["force_immediate_health_check"] = False
+            return False
+        
+        # Skip if currently processing a query (to prevent concurrent calls)
+        if st.session_state.get("processing_query", False):
+            logging.info("API call skipped: Query processing in progress")
+            return True
+        
+        # Check if System Status section is enabled - only call API if it's visible
+        if not st.session_state.get("sidebar_settings", {}).get("show_system_status", False):
+            logging.info("API call skipped: System Status section is disabled")
+            return True
+        
+        # Block API calls if ANY UI interaction happened in last 5 seconds (FOR HEALTH CHECKS)
+        if "recent_ui_action" in st.session_state:
+            current_time = pd.Timestamp.now().timestamp()
+            time_since_action = current_time - st.session_state.get("recent_ui_action", 0)
+            if time_since_action < 5:  # 5 seconds protection
+                logging.info(f"API call blocked: UI interaction {time_since_action:.1f}s ago")
+                return True
+        
+        # Check rate limiting for regular health checks
+        current_time = pd.Timestamp.now().timestamp()
+        last_update = st.session_state.get("last_status_update", 0)
+        
+        # Allow call if we don't have cached status
+        if "last_api_status" not in st.session_state:
+            logging.info("API call allowed: No cached status available")
+            return False
+            
+        # Otherwise, respect the rate limit
+        time_since_update = current_time - last_update
+        if time_since_update < 30:  # 30 seconds rate limit
+            logging.info(f"API call skipped: Rate limit (updated {time_since_update:.1f}s ago)")
+            return True
+            
+        logging.info("API call allowed: Rate limit passed")
+        return False
+    except Exception as e:
+        logging.error(f"Error in should_skip_api_call: {e}")
+        return False  # Allow API call on error (fail safe)
 
-inject_css(theme)
+def should_skip_query_processing():
+    """Determine if query processing should be skipped - FOR QUERY PROCESSING ONLY."""
+    try:
+        # Never skip if we have a legitimate query
+        if "legitimate_query_time" in st.session_state:
+            current_time = pd.Timestamp.now().timestamp()
+            time_since_query = current_time - st.session_state.get("legitimate_query_time", 0)
+            if time_since_query < 10:  # Allow legitimate queries within 10 seconds
+                logging.info(f"Query processing allowed: Legitimate query {time_since_query:.1f}s ago")
+                return False
+        
+        # Skip if already processing a query (to prevent concurrent calls)
+        if st.session_state.get("processing_query", False):
+            logging.info("Query processing skipped: Already processing a query")
+            return True
+        
+        # Block if recent UI action but no legitimate query flag
+        if "recent_ui_action" in st.session_state and "legitimate_query_time" not in st.session_state:
+            current_time = pd.Timestamp.now().timestamp()
+            time_since_action = current_time - st.session_state.get("recent_ui_action", 0)
+            if time_since_action < 5:  # 5 seconds protection
+                logging.info(f"Query processing blocked: UI interaction {time_since_action:.1f}s ago without legitimate query")
+                return True
+        
+        logging.info("Query processing allowed: No blocking conditions")
+        return False
+    except Exception as e:
+        logging.error(f"Error in should_skip_query_processing: {e}")
+        return False  # Allow processing on error (fail safe)
+
+def silent_status_update():
+    """Update status silently without UI disruption."""
+    try:
+        # Use the validation function
+        if should_skip_api_call():
+            return
+            
+        # Clear cache and update timestamp
+        st.cache_data.clear()
+        st.session_state["last_status_update"] = pd.Timestamp.now().timestamp()
+    except:
+        pass  # Silent failure
+
+def get_query_stats():
+    """Get query statistics from session state."""
+    if "query_stats" not in st.session_state:
+        st.session_state["query_stats"] = {
+            "total_queries": 0,
+            "successful_queries": 0,
+            "failed_queries": 0,
+            "session_start": pd.Timestamp.now()
+        }
+    
+    stats = st.session_state["query_stats"]
+    success_rate = 0
+    if stats["total_queries"] > 0:
+        success_rate = round((stats["successful_queries"] / stats["total_queries"]) * 100, 1)
+    
+    return stats["total_queries"], success_rate, stats["successful_queries"], stats["failed_queries"]
+
+def update_query_stats(success=True):
+    """Update query statistics."""
+    if "query_stats" not in st.session_state:
+        st.session_state["query_stats"] = {
+            "total_queries": 0,
+            "successful_queries": 0,
+            "failed_queries": 0,
+            "session_start": pd.Timestamp.now()
+        }
+    
+    st.session_state["query_stats"]["total_queries"] += 1
+    if success:
+        st.session_state["query_stats"]["successful_queries"] += 1
+    else:
+        st.session_state["query_stats"]["failed_queries"] += 1
 
 # ===============================
-# Avatar Setup
+# Session State Management
 # ===============================
-# Default base64 encoded avatars
-DEFAULT_USER_AVATAR = """
-data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+PHBhdGggZD0iTTEyIDJDNi40OCAyIDIgNi40OCAyIDEyczQuNDggMTAgMTAgMTAgMTAtNC40OCAxMC0xMFMxNy41MiAyIDEyIDJ6bTAgM2MyLjY3IDAgOC0xLjMzIDgtM3Y0LjE1YzAgMS42My0zLjMzIDIuODUtOCAyLjg1cy04LTEuMjItOC0yLjg1VjJjMCAxLjY3IDUuMzMgMyA4IDN6bTAgMTFjLTIuNjcgMC04LTEuMzMtOC0zdjQuMTVjMCAxLjYzIDMuMzMgMi44NSA4IDIuODVzOC0xLjIyIDgtMi44NVYxM2MwIDEuNjctNS4zMyAzLTggM3oiIGZpbGw9IiM0Q0FGNTAiLz48L3N2Zz4=
-"""
+if "messages" not in st.session_state:
+    st.session_state["messages"] = []
 
-DEFAULT_AI_AVATAR = """
-data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+PHBhdGggZD0iTTEyIDJDNi40OCAyIDIgNi40OCAyIDEyczQuNDggMTAgMTAgMTAgMTAtNC40OCAxMC0xMFMxNy41MiAyIDEyIDJ6bS0xIDE2LjkydjEuMDhjMCAuNTUuNDUgMSAxIDFoMmMuNTUgMCAxLS40NSAxLTF2LTEuMDhjLS42My4xMy0xLjI5LjItMS45Ny4yLS43IDAtMS4zNC0uMDctMS45Ny0uMnptNS43OC0yLjc0QzE1LjI0IDE3LjM0IDEzLjY3IDE4IDEyIDE4cy0zLjI0LS42Ni00Ljc4LTEuODJDNC40NCAxNC4zNCAzIDEyLjA0IDMgOS41IDMgNS45MSA1LjkxIDMgOS41IDNoNUM2LjY1IDMgNCA0LjY1IDQgNi43MWMwIDIuMDYgMS42NyAzLjc0IDMuNzMgMy43NC4zNiAwIC43MS0uMDUgMS4wNS0uMTUgMS4xIDEuMTEgMi41NyAxLjc5IDQuMjIgMS43OSAxLjY0IDAgMy4xMi0uNjkgNC4yMi0xLjc5LjMzLjEuNjkuMTUgMS4wNS4xNSAyLjA2IDAgMy43My0xLjY4IDMuNzMtMy43NCAwLTIuMDYtMS42Ny0zLjcxLTMuNzMtMy43MWgtNUM3LjkxIDMgNSA1LjkxIDUgOS41YzAgMi41NCAxLjQ0IDQuODQgMy43OCA2LjY4eiIgZmlsbD0iIzI5NzlGRiIvPjwvc3ZnPg==
-"""
+if "query_stats" not in st.session_state:
+    st.session_state["query_stats"] = {
+        "total_queries": 0,
+        "successful_queries": 0,
+        "failed_queries": 0,
+        "session_start": pd.Timestamp.now()
+    }
+
+if "processing_query" not in st.session_state:
+    st.session_state["processing_query"] = False
 
 # ===============================
-# Session State
+# Main App Setup
 # ===============================
-if "messages" not in st.session_state:
-    st.session_state["messages"] = []
+# Professional header positioned at top-left of chat area
+st.markdown(f"""
+<div style=" display: flex; align-items: center; padding-left: 0.5rem;">
+    <h2 style="margin: 0; font-size: 1.5rem; font-weight: 600; color: inherit;">
+        💬 {PAGE_TITLE}
+    </h2>
+</div>
+""", unsafe_allow_html=True)
+
+st.markdown("""
+<div style=" padding-left: 0.5rem;">
+    <p style="margin: 0; font-size: 0.9rem; opacity: 0.7; color: inherit;">
+        Ask questions in plain English to generate and run SQL queries.
+    </p>
+</div>
+""", unsafe_allow_html=True)
+
+
+# Sidebar
+with st.sidebar:
+    # Quick status indicator at the top
+    if "sidebar_settings" in st.session_state:
+        visible_count = sum([
+            st.session_state["sidebar_settings"]["show_model_selection"],
+            st.session_state["sidebar_settings"]["show_agent_selection"],
+            st.session_state["sidebar_settings"]["show_theme_selection"],
+            st.session_state["sidebar_settings"]["show_system_status"],
+            st.session_state["sidebar_settings"]["show_tips"]
+        ])
+        
+        if visible_count == 5:
+            st.caption("🟢 All sections visible")
+        elif visible_count > 0:
+            st.caption(f"🟡 {visible_count}/5 sections visible")
+        else:
+            st.caption("🔴 No sections visible")
+    
+    # Sidebar Display Settings - User Configurable
+    # Initialize settings state tracking
+    if "settings_interaction_count" not in st.session_state:
+        st.session_state["settings_interaction_count"] = 0
+    
+    # Track if user has recently interacted with settings
+    keep_expanded = st.session_state.get("settings_interaction_count", 0) > 0
+    
+    # Create expander that stays open for a few interactions
+    settings_expander = st.expander(
+        "⚙️ Sidebar Settings", 
+        expanded=keep_expanded
+    )
+    
+    with settings_expander:
+        st.markdown("**Choose what to display:**")
+        
+        # Initialize sidebar visibility settings in session state with new defaults
+        if "sidebar_settings" not in st.session_state:
+            st.session_state["sidebar_settings"] = {
+                "show_model_selection": True,  # Default unchecked
+                "show_agent_selection": False,  # Default unchecked
+                "show_theme_selection": True,
+                "show_system_status": False,    # Default unchecked - only load health checks when enabled
+                "show_tips": True
+            }
+        
+        # Store previous values to detect changes
+        prev_settings = st.session_state["sidebar_settings"].copy()
+        
+        # Configurable checkboxes
+        col1, col2 = st.columns(2)
+        with col1:
+            show_model = st.checkbox(
+                "🤖 AI Model", 
+                value=st.session_state["sidebar_settings"]["show_model_selection"],
+                help="Show/hide AI model selection",
+                key="settings_model"
+            )
+            show_agent = st.checkbox(
+                "🎯 Agent Type", 
+                value=st.session_state["sidebar_settings"]["show_agent_selection"],
+                help="Show/hide agent selection",
+                key="settings_agent"
+            )
+            show_theme = st.checkbox(
+                "🎨 Theme", 
+                value=st.session_state["sidebar_settings"]["show_theme_selection"],
+                help="Show/hide theme selection",
+                key="settings_theme"
+            )
+        
+        with col2:
+            show_status = st.checkbox(
+                "📊 System Status", 
+                value=st.session_state["sidebar_settings"]["show_system_status"],
+                help="Show/hide system status",
+                key="settings_status"
+            )
+            show_tips = st.checkbox(
+                "💡 Tips & Help", 
+                value=st.session_state["sidebar_settings"]["show_tips"],
+                help="Show/hide tips section",
+                key="settings_tips"
+            )
+        
+        # Detect if any setting changed
+        new_settings = {
+            "show_model_selection": show_model,
+            "show_agent_selection": show_agent,
+            "show_theme_selection": show_theme,
+            "show_system_status": show_status,
+            "show_tips": show_tips
+        }
+        
+        # Check if settings changed
+        settings_changed = any(
+            prev_settings.get(key) != new_settings[key] 
+            for key in new_settings.keys()
+        )
+        
+        # Special handling for System Status being enabled
+        status_just_enabled = (
+            not prev_settings.get("show_system_status", False) and 
+            new_settings["show_system_status"]
+        )
+        
+        # Update session state
+        st.session_state["sidebar_settings"].update(new_settings)
+        
+        # Increment interaction count when settings change
+        if settings_changed:
+            st.session_state["settings_interaction_count"] += 1
+            # Mark UI action to prevent unnecessary API calls for most settings
+            if not status_just_enabled:  # Don't block API calls when System Status is enabled
+                st.session_state["recent_ui_action"] = pd.Timestamp.now().timestamp()
+            # Reset counter after 10 interactions to prevent it growing indefinitely
+            if st.session_state["settings_interaction_count"] > 10:
+                st.session_state["settings_interaction_count"] = 5
+        
+        # If System Status was just enabled, clear cache and allow immediate health check
+        if status_just_enabled:
+            st.cache_data.clear()
+            # Clear any cached status to force fresh call
+            for key in ["last_api_status", "last_api_delta", "last_api_type", "last_detailed_health"]:
+                if key in st.session_state:
+                    del st.session_state[key]
+            st.session_state["last_status_update"] = 0
+            # Set flag to trigger immediate health check
+            st.session_state["force_immediate_health_check"] = True
+            # Remove any recent UI action timestamp to allow API call
+            if "recent_ui_action" in st.session_state:
+                del st.session_state["recent_ui_action"]
+        
+        # Show current settings status
+        visible_count = sum([show_model, show_agent, show_theme, show_status, show_tips])
+        
+        if visible_count == 5:
+            st.success(f"✅ All {visible_count} sections visible")
+        elif visible_count > 0:
+            st.info(f"ℹ️ {visible_count}/5 sections visible")
+        else:
+            st.warning("⚠️ No sections visible")
+            
+    # Use the current session state values for conditional rendering
+    show_model = st.session_state["sidebar_settings"]["show_model_selection"]
+    show_agent = st.session_state["sidebar_settings"]["show_agent_selection"] 
+    show_theme = st.session_state["sidebar_settings"]["show_theme_selection"]
+    show_status = st.session_state["sidebar_settings"]["show_system_status"]
+    show_tips = st.session_state["sidebar_settings"]["show_tips"]
+    
+    st.markdown("---")
+    
+    # Model Selection (conditional display)
+    if show_model:
+        st.markdown("### 🤖 AI Model")
+        
+        # Static list of available models - no API calls needed
+        available_models = [
+            "gpt-4o-mini",
+            "gpt-3.5-turbo", 
+            "gemini-1.5-pro",
+            "gemini-1.5-flash",
+            "claude-3-haiku",
+            "claude-3-sonnet", 
+            "mistral-small",
+            "mistral-medium",
+            "mistral-large",
+        ]
+        
+        model = st.selectbox(
+            "Choose your AI model:",
+            available_models,
+            index=0,  # Default to first model
+            help="Select the AI model for processing your queries",
+            key="model_selector"
+        )
+        
+        # Show model descriptions - static information
+        model_descriptions = {
+            "gpt-4o-mini": "⚡ Fast & cost-effective OpenAI model",
+            "gpt-3.5-turbo": "🔥 Reliable & quick OpenAI model", 
+            "gemini-pro": "💎 Google's powerful Gemini model",
+            "gemini-1.5-pro": "🔬 Google's latest Gemini model",
+            "gemini-1.5-flash": "⚡ Google's fast Gemini model",
+            "claude-3-haiku": "🌸 Anthropic's efficient Claude model",
+            "claude-3-sonnet": "🎵 Anthropic's balanced Claude model", 
+            "mistral-small": "🎯 Mistral's efficient model",
+            "mistral-medium": "⚖️ Mistral's balanced model",
+            "mistral-large": "🦾 Mistral's most capable model",
+           
+        }
+        
+        description = model_descriptions.get(model, "🤖 Advanced AI model")
+        st.info(description)
+        
+        # Show provider information - static
+        provider_info = {
+            "gpt-4o-mini": "🏢 OpenAI",
+            "gpt-3.5-turbo": "🏢 OpenAI",
+            "gemini-pro": "🔍 Google",
+            "gemini-1.5-pro": "🔍 Google",
+            "gemini-1.5-flash": "🔍 Google",
+            "claude-3-haiku": "🤖 Anthropic",
+            "claude-3-sonnet": "🤖 Anthropic",
+            "mistral-small": "⚡ Mistral AI",
+            "mistral-medium": "⚡ Mistral AI",
+            "mistral-large": "⚡ Mistral AI",
+        }
+        
+        provider = provider_info.get(model, "🤖 AI Provider")
+        st.caption(f"Provider: {provider}")
+        
+        # Show setup hints for different providers
+        if model.startswith("gemini"):
+            st.caption("💡 Requires GOOGLE_API_KEY in .env")
+        elif model.startswith("claude"):
+            st.caption("💡 Requires ANTHROPIC_API_KEY in .env")
+        elif model.startswith("mistral"):
+            st.caption("💡 Requires MISTRAL_API_KEY in .env")
+        elif model in ["llama3.2", "llama3.1", "codellama", "phi3"]:
+            st.caption("💡 Requires Ollama installed locally")
+        else:
+            st.caption("💡 Requires OPENAI_API_KEY in .env")
+        
+        # Mark UI action to prevent unnecessary API calls for other interactions
+        if model:
+            st.session_state["recent_ui_action"] = pd.Timestamp.now().timestamp()
+        
+        st.markdown("---")
+    else:
+        # Use default model when hidden
+        model = AVAILABLE_MODELS[0]
+    
+    # Agent Selection (conditional display)
+    if show_agent:
+        st.markdown("### 🎯 Agent Type")
+        agent = st.selectbox(
+            "Choose your agent:",
+            AVAILABLE_AGENTS,
+            help="Select the specialized agent for your database tasks",
+            key="agent_selector"
+        )
+        
+        # Mark UI action to prevent unnecessary API calls
+        if agent:
+            st.session_state["recent_ui_action"] = pd.Timestamp.now().timestamp()
+        
+        # Show agent info
+        agent_info = {
+            "default": "🔧 General-purpose database assistant",
+            "sql-agent": "💾 Specialized in SQL optimization",
+            "custom-agent": "🎨 Customized for specific workflows"
+        }
+        st.info(agent_info.get(agent, "Specialized database agent"))
+        st.markdown("---")
+    else:
+        # Use default agent when hidden
+        agent = AVAILABLE_AGENTS[0]
+    
+    # Theme Selection (conditional display)
+    if show_theme:
+        st.markdown("### 🎨 Appearance")
+        theme = st.selectbox(
+            "Choose your theme:",
+            AVAILABLE_THEMES,
+            help="Customize the visual appearance",
+            key="theme_selector"
+        )
+        
+        # Mark UI action to prevent unnecessary API calls
+        if theme:
+            st.session_state["recent_ui_action"] = pd.Timestamp.now().timestamp()
+        
+        st.markdown("---")
+    else:
+        # Use default theme when hidden
+        theme = AVAILABLE_THEMES[0]
+    
+    # System Status (conditional display)
+    if show_status:
+        st.markdown("### 📊 System Status")
+        
+        # Skip status check if currently processing a query to avoid slowdown
+        if st.session_state.get("processing_query", False):
+            st.info("⏳ Status check paused during query processing")
+            # Show cached stats only
+            total_queries, success_rate, successful, failed = get_query_stats()
+            col1, col2 = st.columns(2)
+            with col1:
+                st.metric("API Status", "⏳ Processing", delta="Query in progress")
+            with col2:
+                st.metric("Total Queries", total_queries, delta=f"+{total_queries}")
+        else:
+            # Check if this is a force refresh scenario
+            is_force_refresh = st.session_state.get("force_refresh_requested", False)
+            
+            # Use centralized validation to avoid unnecessary API calls
+            if not should_skip_api_call(force_refresh=is_force_refresh):
+                logging.info("=== HEALTH CHECK API CALL STARTING ===")
+                try:
+                    # Get real-time API status only when validation passes
+                    status_text, status_delta, status_type = check_api_status()
+                    
+                    # Get detailed health information
+                    detailed_health = get_detailed_health_status()
+                    
+                    # Store the status for future use
+                    st.session_state["last_api_status"] = status_text
+                    st.session_state["last_api_delta"] = status_delta  
+                    st.session_state["last_api_type"] = status_type
+                    st.session_state["last_detailed_health"] = detailed_health
+                    st.session_state["last_status_update"] = pd.Timestamp.now().timestamp()
+                    
+                    logging.info("=== HEALTH CHECK API CALL COMPLETED ===")
+                    
+                    # Clear force refresh flag after successful update
+                    if is_force_refresh:
+                        st.session_state["force_refresh_requested"] = False
+                    
+                except Exception as e:
+                    logging.error(f"Status check failed: {e}")
+                    # Use error values if API fails
+                    status_text = "🔴 Failed"
+                    status_delta = f"Error: {str(e)[:50]}..."
+                    status_type = "error"
+                    detailed_health = {
+                        "status": "error",
+                        "message": f"Status check failed: {str(e)}",
+                        "checks": {}
+                    }
+                    
+                    # Clear force refresh flag even on error
+                    if is_force_refresh:
+                        st.session_state["force_refresh_requested"] = False
+            else:
+                # Use cached values to avoid API calls, with better defaults
+                status_text = st.session_state.get("last_api_status", "🟡 Loading...")
+                status_delta = st.session_state.get("last_api_delta", "Initializing...")
+                status_type = st.session_state.get("last_api_type", "normal")
+                detailed_health = st.session_state.get("last_detailed_health", {
+                    "status": "unknown",
+                    "message": "Loading system status...",
+                    "checks": {}
+                })
+            
+            # Get query statistics
+            total_queries, success_rate, successful, failed = get_query_stats()
+            
+            # Display main metrics
+            col1, col2 = st.columns(2)
+            with col1:
+                st.metric("API Status", status_text, delta=status_delta)
+            with col2:
+                st.metric("Total Queries", total_queries, delta=f"+{total_queries}")
+            
+            # Status indicator with more details
+            if status_type == "success":
+                st.success("✅ All systems operational")
+            elif status_type == "warning":
+                st.warning("⚠️ Limited functionality - some features may be slow")
+            else:
+                st.error("❌ API connection failed - please check your backend service")
+            
+            # Show detailed health checks in an expander
+            with st.expander("🔍 Detailed System Health", expanded=False):
+                if detailed_health.get("status") in ["healthy", "unhealthy"]:
+                    # Display timestamp
+                    if "timestamp" in detailed_health:
+                        st.caption(f"Last checked: {detailed_health['timestamp']}")
+                    
+                    # Display individual checks
+                    checks = detailed_health.get("checks", {})
+                    
+                    if "database" in checks:
+                        db_check = checks["database"]
+                        db_status = db_check.get("status", "unknown")
+                        db_message = db_check.get("message", "No information")
+                        
+                        if db_status == "healthy":
+                            st.success(f"🗃️ Database: {db_message}")
+                        elif db_status == "unhealthy":
+                            st.error(f"🗃️ Database: {db_message}")
+                        else:
+                            st.warning(f"🗃️ Database: {db_message}")
+                    
+                    if "openai_api" in checks:
+                        api_check = checks["openai_api"]
+                        api_status = api_check.get("status", "unknown")
+                        api_message = api_check.get("message", "No information")
+                        
+                        if api_status == "configured":
+                            st.success(f"🤖 OpenAI API: {api_message}")
+                        elif api_status == "error":
+                            st.error(f"🤖 OpenAI API: {api_message}")
+                        else:
+                            st.warning(f"🤖 OpenAI API: {api_message}")
+                    
+                    # Display version if available
+                    if "version" in detailed_health:
+                        st.info(f"📦 Version: {detailed_health['version']}")
+                        
+                else:
+                    # Show error information
+                    st.error(f"❌ Health check failed: {detailed_health.get('message', 'Unknown error')}")
+                    st.caption("Unable to retrieve detailed system status")
+        
+        # Additional detailed stats (always show regardless of processing state)
+        if total_queries > 0:
+            col3, col4 = st.columns(2)
+            with col3:
+                st.metric("Success Rate", f"{success_rate}%", delta=f"{successful} successful")
+            with col4:
+                session_duration = pd.Timestamp.now() - st.session_state["query_stats"]["session_start"]
+                hours = int(session_duration.total_seconds() // 3600)
+                minutes = int((session_duration.total_seconds() % 3600) // 60)
+                st.metric("Session Time", f"{hours}h {minutes}m", delta="Active")
+        
+        # Show last update time and refresh button
+        col_refresh1, col_refresh2 = st.columns([2, 1])
+        with col_refresh1:
+            st.caption(f"🔄 Last updated: {pd.Timestamp.now().strftime('%H:%M:%S')}")
+        with col_refresh2:
+            if st.button("🔄", help="Force Refresh Status", key="refresh_status"):
+                # Set force refresh flag to bypass validation
+                st.session_state["force_refresh_requested"] = True
+                # Clear all cached data and force fresh API calls
+                st.cache_data.clear()
+                # Reset validation timestamps to allow immediate API calls
+                st.session_state["last_status_update"] = 0
+                if "recent_ui_action" in st.session_state:
+                    del st.session_state["recent_ui_action"]
+                # Clear cached status values
+                for key in ["last_api_status", "last_api_delta", "last_api_type", "last_detailed_health"]:
+                    if key in st.session_state:
+                        del st.session_state[key]
+                st.rerun()
+        
+        st.markdown("---")
+    
+    # Tips and Help (conditional display)
+    if show_tips:
+        st.markdown("### 💡 Quick Tips")
+        with st.expander("📝 How to ask questions"):
+            st.markdown("""
+            - **"Show me all customers from Chicago"**
+            - **"What are the top 5 branches by transactions?"**
+            - **"Calculate total transactions by month"**
+            - **"Find customers who haven't done any transactions recently"**
+            """)
+        
+        with st.expander("⚡ Pro Tips"):
+            st.markdown("""
+            - Be specific about what data you want
+            - Mention date ranges when relevant
+            - Ask for summaries or aggregations
+            - Use natural language - no SQL needed!
+            """)
+        
+        with st.expander("🔧 Troubleshooting"):
+            st.markdown("""
+            - **No results?** Try rephrasing your question
+            - **Error message?** Use the retry button
+            - **Slow response?** Check your connection
+            - **Wrong data?** Be more specific in your query
+            """)
+        
+        st.markdown("---")
+    
+    # Quick actions
+    st.markdown("### 🚀 Quick Actions")
+    
+    # Check if query is currently being processed
+    is_processing = st.session_state.get("processing_query", False)
+    
+    # Show processing indicator if query is running
+    if is_processing:
+        st.info("⚡ Query in progress... All action buttons are temporarily disabled.")
+    
+    col_action1, col_action2 = st.columns(2)
+    with col_action1:
+        # Disable Clear Chat button when query is processing
+        clear_chat_disabled = is_processing
+        clear_chat_help = "Cannot clear chat while query is processing" if is_processing else "Clear all chat messages"
+        
+        if st.button("🗑️ Clear Chat", 
+                    use_container_width=True, 
+                    disabled=clear_chat_disabled,
+                    help=clear_chat_help):
+            # Simple UI action - block API calls for 5 seconds
+            st.session_state["recent_ui_action"] = pd.Timestamp.now().timestamp()
+            st.session_state["messages"] = []
+            logging.info("Clear Chat button clicked - blocking API calls for 5 seconds")
+            st.rerun()
+    
+    with col_action2:
+        # Disable Reset Stats button when query is processing to prevent any API calls
+        reset_stats_disabled = is_processing
+        reset_stats_help = "Cannot reset stats while query is processing" if is_processing else "Reset query statistics"
+        
+        if st.button("📊 Reset Stats", 
+                    use_container_width=True,
+                    disabled=reset_stats_disabled,
+                    help=reset_stats_help):
+            # Simple UI action - block API calls for 5 seconds
+            st.session_state["recent_ui_action"] = pd.Timestamp.now().timestamp()
+            st.session_state["query_stats"] = {
+                "total_queries": 0,
+                "successful_queries": 0,
+                "failed_queries": 0,
+                "session_start": pd.Timestamp.now()
+            }
+            logging.info("Reset Stats button clicked - blocking API calls for 5 seconds")
+            st.rerun()
+    
+    # Sample Queries button - also disabled during processing to avoid confusion
+    sample_disabled = is_processing
+    sample_help = "Cannot show samples while query is processing" if is_processing else "Show sample queries you can copy"
+    
+    if st.button("📋 Sample Queries", 
+                use_container_width=True,
+                disabled=sample_disabled,
+                help=sample_help):
+        # Simple UI action - block API calls for 5 seconds
+        st.session_state["recent_ui_action"] = pd.Timestamp.now().timestamp()
+        logging.info("Sample Queries button clicked - blocking API calls for 5 seconds")
+        sample_queries = [
+            "Show me the top 10 customers by transactions",
+            "Which branches collected more transactions last month?",
+            "Calculate average transactions value",
+            "List all active customers"
+        ]
+        # Display sample queries in the sidebar instead of adding to chat
+        st.markdown("**Sample queries you can copy and paste:**")
+        for query in sample_queries:
+            st.code(query)
+
+   # Footer
+    st.markdown("### ℹ️ About")
+    st.markdown("""
+    **AI Database Assistant** v2.0  
+    🚀 Powered by advanced AI  
+    💬 Natural language to SQL  
+    📈 Real-time analytics  
+    """)
+    
+
+# Apply theme
+theme_manager.inject_theme(theme)
 
 # ===============================
-# Helper: Render Chat History
+# Chat Rendering Function
 # ===============================
 def render_chat():
+    """Render chat history with proper error handling."""
     st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
-    for msg in st.session_state["messages"]:
+    messages = st.session_state["messages"]
+    
+    for i, msg in enumerate(messages):
         if msg["role"] == "user":
-            st.markdown(f'<div class="user-bubble">{msg["content"]}</div>', unsafe_allow_html=True)
+            # User message with avatar
+            st.markdown(
+                f'''<div style="display: flex; align-items: flex-start; justify-content: flex-end; margin-bottom: 0.5em;">
+                        <div style="margin-right: 0.5em;">
+                            <img src="{OUTLINE_INDIGO_USER}" alt="User" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #e3f2fd; background: #fff; object-fit: cover;" />
+                        </div>
+                        <div class="user-bubble">{msg["content"]}</div>
+                   </div>''',
+                unsafe_allow_html=True)
+        
         elif msg["role"] == "assistant":
+            # Show all assistant messages normally (including errors)
             bubble_class = "error-bubble" if msg.get("is_error") else "ai-bubble"
-            st.markdown(f'<div class="{bubble_class}">{msg["content"]}</div>', unsafe_allow_html=True)
-
-            if msg.get("data"):
-                df = pd.DataFrame(msg["data"])
-                st.dataframe(df, use_container_width=True)
-
-                # Add download button for results
-                csv = df.to_csv(index=False).encode("utf-8")
-                st.download_button("📥 Download CSV", csv, "results.csv", "text/csv", key=f"download_csv_{id(msg)}")
-
-            if msg.get("chart"):
-                img_data = base64.b64decode(msg["chart"])
-                st.image(img_data, use_column_width=True)
-    st.markdown("</div>", unsafe_allow_html=True)
-
-
-    def render_chat():
-        # Spinner CSS (only inject once)
-        st.markdown('''
-            <style>
-            .spinner {
-                display: inline-block;
-                width: 1.3em;
-                height: 1.3em;
-                border: 3px solid #e0e0e0;
-                border-top: 3px solid #2193b0;
-                border-radius: 50%;
-                animation: spin 0.8s linear infinite;
-                margin-right: 0.7em;
-            }
-            @keyframes spin {
-                0% { transform: rotate(0deg); }
-                100% { transform: rotate(360deg); }
-            }
-            </style>
-        ''', unsafe_allow_html=True)
-        st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
-        user_avatar = DEFAULT_USER_AVATAR
-        ai_avatar = DEFAULT_AI_AVATAR
-        for msg in st.session_state["messages"]:
-            if msg["role"] == "user":
-                st.markdown(
-                    f'''<div style="display: flex; align-items: flex-start; justify-content: flex-end; margin-bottom: 0.5em;">
-                            <div style="margin-right: 0.5em;">
-                                <img src="{user_avatar}" alt="User" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #e3f2fd; background: #fff; object-fit: cover;" />
-                            </div>
-                            <div class="user-bubble">{msg["content"]}</div>
-                       </div>''',
-                    unsafe_allow_html=True)
-            elif msg["role"] == "assistant":
-                bubble_class = "error-bubble" if msg.get("is_error") else "ai-bubble"
-                # Show spinner if this is a placeholder 'Thinking...' message
-                if msg.get("is_placeholder"):
+            
+            # Skip rendering placeholder messages (we use Streamlit spinner instead)
+            if msg.get("is_placeholder"):
+                continue
+            else:
+                # For error messages, show retry button below the message
+                if msg.get("is_error"):
                     st.markdown(
                         f'''<div style="display: flex; align-items: flex-start; margin-bottom: 0.5em;">
                                 <div style="margin-right: 0.5em;">
-                                    <img src="{ai_avatar_url}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
+                                    <img src="{DARK_MODE_SLATE_AI}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
                                 </div>
-                                <div class="{bubble_class}"><span class="spinner"></span>Thinking...</div>
+                                <div class="{bubble_class}">{msg["content"]}</div>
                            </div>''',
                         unsafe_allow_html=True)
+                    
+                    # Show retry button below the error message, aligned with the error message
+                    # Use same layout structure as the error message for alignment
+                    cols = st.columns([0.03, 0.85])
+                    with cols[0]:
+                        st.empty()  # Empty space where avatar would be
+                    with cols[1]:
+                        if st.button(RETRY_BUTTON_TEXT, key=f"retry_error_{i}"):
+                            # Use the stored user query and index for reliable retry
+                            stored_user_query = msg.get("user_query")
+                            stored_user_index = msg.get("user_query_index")
+                            
+                            if stored_user_query and stored_user_index is not None:
+                                # Remove all messages after the user query and retry
+                                st.session_state["messages"] = st.session_state["messages"][:stored_user_index+1]
+                                st.session_state["messages"].append({"role": "assistant", "content": "� Processing your query...", "is_placeholder": True})
+                                # Mark this as a legitimate retry, not a UI interaction
+                                st.session_state["legitimate_query_time"] = pd.Timestamp.now().timestamp()
+                                # Remove any recent UI action flag to allow this query to process
+                                if "recent_ui_action" in st.session_state:
+                                    del st.session_state["recent_ui_action"]
+                                st.rerun()
                 else:
+                    # Regular AI message
                     st.markdown(
                         f'''<div style="display: flex; align-items: flex-start; margin-bottom: 0.5em;">
                                 <div style="margin-right: 0.5em;">
-                                    <img src="{ai_avatar_url}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
+                                    <img src="{DARK_MODE_SLATE_AI}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
                                 </div>
                                 <div class="{bubble_class}">{msg["content"]}</div>
                            </div>''',
                         unsafe_allow_html=True)
+            
+            # Show data table and download if present
+            if msg.get("data"):
+                df = pd.DataFrame(msg["data"])
+                st.dataframe(df, use_container_width=True)
+                csv = df.to_csv(index=False).encode("utf-8")
+                st.download_button(DOWNLOAD_BUTTON_TEXT, csv, "results.csv", "text/csv", key=f"download_csv_{id(msg)}")
+            
+            # Show chart if present
+            if msg.get("chart"):
+                img_data = base64.b64decode(msg["chart"])
+                st.image(img_data, use_column_width=True)
+    
+    st.markdown("</div>", unsafe_allow_html=True)
 
-                if msg.get("data"):
-                    df = pd.DataFrame(msg["data"])
-                    st.dataframe(df, use_container_width=True)
-
-                    # Add download button for results
-                    csv = df.to_csv(index=False).encode("utf-8")
-                    st.download_button("\U0001F4C5 Download CSV", csv, "results.csv", "text/csv", key=f"download_csv_{id(msg)}")
-
-                if msg.get("chart"):
-                    img_data = base64.b64decode(msg["chart"])
-                    st.image(img_data, use_column_width=True)
-        st.markdown("</div>", unsafe_allow_html=True)
-
+# Render chat
 render_chat()
 
 # ===============================
-# User Input
+# User Input Handling
 # ===============================
-
-# Disable chat input while AI is working
+# Check if AI is thinking
 pending = False
-if st.session_state["messages"]:  # Check if there are any messages
+if st.session_state["messages"]:
     if st.session_state["messages"][-1]["role"] == "assistant":
         pending = st.session_state["messages"][-1].get("is_placeholder", False)
 
-user_query = st.chat_input(
-    "Ask me something about your database...",
-    disabled=pending or False  # Ensure disabled is always boolean
-)
+# Get user input
+user_query = st.chat_input(CHAT_INPUT_PLACEHOLDER, disabled=pending)
+
 if user_query and not pending:
     st.session_state["messages"].append({"role": "user", "content": user_query})
-    st.session_state["messages"].append({"role": "assistant", "content": "🤔 Thinking...", "is_placeholder": True})
+    st.session_state["messages"].append({"role": "assistant", "content": "� Processing ...", "is_placeholder": True})
+    # Mark this as a legitimate user query, not a UI interaction
+    st.session_state["legitimate_query_time"] = pd.Timestamp.now().timestamp()
+    # Remove any recent UI action flag to allow this query to process
+    if "recent_ui_action" in st.session_state:
+        del st.session_state["recent_ui_action"]
     st.rerun()
 
 # ===============================
-# Handle Pending Assistant Response
+# API Response Handling
 # ===============================
-if (
+# CRITICAL: Only process API calls for legitimate user queries, not UI interactions
+# Check if we have a placeholder message from a user query
+has_placeholder = (
     st.session_state["messages"]
     and st.session_state["messages"][-1].get("is_placeholder")
     and len(st.session_state["messages"]) >= 2
     and st.session_state["messages"][-2]["role"] == "user"
-):
-    user_query = st.session_state["messages"][-2]["content"]
-
-    try:
-        with st.spinner("Working..."):
-            payload = {"question": user_query, "model": model, "agent": agent}
-            logging.info(f"Sending API request with: {payload}")
-            response = requests.post(API_URL, json=payload, timeout=60)
+)
 
-            try:
-                result = response.json()
-            except ValueError:
-                result = {"detail": response.text}
+# Check if this is a legitimate query that should be processed
+# Block if recent UI action (sidebar interactions) triggered this rerun
+should_process_query = (
+    has_placeholder
+    and not should_skip_query_processing()  # Use query-specific validation
+)
 
-        if response.status_code != 200:
-            answer_text = f"❌ Error: {result.get('detail') or result.get('message') or response.text}"
-            rows, chart, is_error = [], None, True
-        elif "message" in result and not ("rows" in result or "chart" in result):
-            answer_text, rows, chart, is_error = result["message"], [], None, False
-        else:
-            rows = result.get("rows", [])
-            chart = result.get("chart", None)
-            if rows:
-                answer_text = f"Here are the top {len(rows)} results I found:"
-            else:
-                answer_text = "I could not find matching records for your query."
-            is_error = False
+if has_placeholder:
+    logging.info(f"=== QUERY PROCESSING CHECK ===")
+    logging.info(f"Has placeholder: {has_placeholder}")
+    logging.info(f"Should process query: {should_process_query}")
+    if not should_process_query:
+        logging.info("BLOCKED: Query processing blocked by should_skip_api_call")
+    else:
+        logging.info("ALLOWED: Query processing allowed")
 
+if should_process_query:
+    user_query = st.session_state["messages"][-2]["content"]
+    user_query_index = len(st.session_state["messages"]) - 2  # Store the user query index
+    
+    # Set processing flag to pause status checks
+    st.session_state["processing_query"] = True
+    
+    try:
+        with st.spinner(THINKING_MESSAGE):
+            logging.info(f"Sending query to API: {user_query}")
+            
+            try:
+                # Use the API client instead of direct requests
+                result = api_client.send_query(user_query, model, agent)
+
+                if result:
+                    answer_text = result.get("message", "No response received")
+                    rows = result.get("rows", [])
+                    chart = result.get("chart", None)
+                    is_error = result.get("error", False)
+                    model_used = result.get("model_used", "unknown")
+                    status = result.get("status", "unknown")
+                    
+                    # Add model information to the response for successful queries
+                    if not is_error and model_used != "unknown":
+                        answer_text += f"\n\n*Powered by: {model_used}*"
+                else:
+                    answer_text = "❌ Error: Unable to process your request. Please try again."
+                    rows, chart, is_error = [], None, True
+                    model_used, status = "unknown", "error"
+                    
+            except Exception as e:
+                logging.error(f"API connectivity error: {e}")
+                answer_text, rows, chart, is_error = (
+                    "❌ API is not available. Please check your connection or try again later.", [], None, True
+                )
+                model_used, status = "unknown", "error"
+                
     except Exception as e:
         logging.error(f"Exception: {e}")
         answer_text, rows, chart, is_error = f"⚠️ Exception: {str(e)}", [], None, True
-
-    # Replace placeholder with final assistant message
+        model_used, status = "unknown", "error"
+
+    # Update query statistics
+    update_query_stats(success=not is_error)
+    
+    # Clear processing flag
+    st.session_state["processing_query"] = False
+    
+    # Clear legitimate query flag after processing
+    if "legitimate_query_time" in st.session_state:
+        del st.session_state["legitimate_query_time"]
+    
+    # Replace placeholder with final response
     st.session_state["messages"][-1] = {
         "role": "assistant",
         "content": answer_text,
         "data": rows,
         "chart": chart,
         "is_error": is_error,
+        "model_used": model_used,
+        "status": status,
+        "user_query": user_query if is_error else None,  # Store user query for retry
+        "user_query_index": user_query_index if is_error else None,  # Store user query index for retry
     }
-
     st.rerun()

src/config.py

@@ -0,0 +1,54 @@
+"""
+Configuration settings for the AI Database Assistant.
+"""
+import logging
+from typing import List
+import os
+
+# API Configuration
+# Change only the BASE_URL for different deployment environments
+BASE_URL = os.environ.get("BASE_URL", "http://127.0.0.1:8000")  # Development Environment
+
+# UI Configuration
+PAGE_TITLE = "AI Assistant"
+PAGE_LAYOUT = "wide"
+
+# Model Options - Simplified for better UX
+AVAILABLE_MODELS: List[str] = [
+    "gpt-4o-mini",     # Fast & cost-effective
+    "gpt-3.5-turbo",   # Reliable & quick
+    "gpt-4"            # Most capable
+]
+AVAILABLE_AGENTS: List[str] = [
+    "default",         # General database assistant
+    "sql-agent"        # SQL optimization expert
+]
+AVAILABLE_THEMES: List[str] = [
+    "Light",           # Clean & bright
+    "Dark"             # Easy on eyes
+]
+
+# Logging Configuration
+LOG_LEVEL = logging.INFO
+LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"
+
+# UI Messages
+CHAT_INPUT_PLACEHOLDER = "Ask me something about your database..."
+THINKING_MESSAGE = "🧠 Processing..."
+WORKING_MESSAGE = "Working..."
+RETRY_BUTTON_TEXT = "🔄 Try Again"
+DOWNLOAD_BUTTON_TEXT = "📥 Download CSV"
+
+# Default Avatars (Base64 encoded SVGs)
+DEFAULT_USER_AVATAR = """
+data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+CiAgPCEtLSBkYXRhYmFzZSAtLT4KICA8ZWxsaXBzZSBjeD0iMTciIGN5PSI2IiByeD0iNS41IiByeT0iMi41IiBmaWxsPSIjV0ZGM0UwIiBzdHJva2U9IiNFRjZD MDAiIHN0cm9rZS13aWR0aD0iMSIvPgogIDxwYXRoIGQ9Ik0xMS41IDZ2NWMwIDEuMyAyLjQ2IDIuNCA1LjUgMi40czUuNS0xLjEgNS41LTIuNFY2IiBmaWxsPSIjV0ZGM0UwIiBzdHJva2U9IiNFRjZD MDAiIHN0cm9rZS13aWR0aD0iMSIvPgogIDxwYXRoIGQ9Ik0xMS41IDguNWMwIDEuMyAyLjQ2IDIuNCA1LjUgMi40czUuNS0xLjEgNS41LTIuNCIgZmlsbD0ibm9uZSIgc3Ryb2tlPSIjRUY2QzAwIiBzdHJva2Utd2lkdGg9IjEiLz4KICA8IS0tIHVzZXIgYnVzdCAtLT4KICA8Y2lyY2xlIGN4PSI3IiBjeT0iOSIgcj0iMyIgZmlsbD0iI0ZGQTcyNiIvPgogIDxwYXRoIGQ9Ik0yLjUgMTdjMC0yLjUgMi42LTQuNSA0LjUtNC41UzExLjUgMTQuNSAxMS41IDE3djEuNUgyLjVWMTd6IiBmaWxsPSIjRkZDQzgwIi8+Cjwvc3ZnPg==
+"""
+OUTLINE_INDIGO_USER ="""
+data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+CiAgPGcgZmlsbD0ibm9uZSIgc3Ryb2tlPSIjNjM2NkYxIiBzdHJva2Utd2lkdGg9IjEuMyIgc3Ryb2tlLWxpbmVjYXA9InJvdW5kIiBzdHJva2UtbGluZWpvaW49InJvdW5kIj4KICAgIDxlbGxpcHNlIGN4PSIxMiIgY3k9IjUuNSIgcng9IjciIHJ5PSIyLjgiLz4KICAgIDxwYXRoIGQ9Ik01IDUuNXY3YzAgMS43IDMuMSAzIDcgM3M3LTEuMyA3LTMgMHYtNyIvPgogICAgPHBhdGggZD0iTSA1IDguOGMwIDEuNyAzLjEgMyA3IDNzNy0xLjMgNy0zIi8+CiAgICA8Y2lyY2xlIGN4PSI4LjIiIGN5PSIxMy44IiByPSIyLjIiLz4KICAgIDxwYXRoIGQ9Ik0zLjUgMTkuNWMuNy0yLjQgMy4xLTQgNC43LTRzNCAxLjYgNC43IDQiLz4KICA8L2c+Cjwvc3ZnPg==
+"""
+DARK_MODE_SLATE_AI="""
+data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+CiAgPHJlY3Qgd2lkdGg9IjI0IiBoZWlnaHQ9IjI0IiByeD0iNiIgZmlsbD0iIzBGMTE3QSIvPgogIDwhLS0gc3R5bGl6ZWQgY3lsaW5kZXIgLS0+CiAgPGVsbGlwc2UgY3g9IjEyIiBjeT0iNyIgcng9IjciIHJ5PSIzIiBmaWxsPSIjMUYyOTM3IiBzdHJva2U9IiM5NEEzQjgiIHN0cm9rZS13aWR0aD0iMSIvPgogIDxwYXRoIGQ9Ik01IDd2N2MwIDEuNyAzLjEgMyA3IDNzNy0xLjMgNy0zVjciIGZpbGw9IiMxMTE4MjciIHN0cm9rZT0iIzk0QTNCOCIgc3Ryb2tlLXdpZHRoPSIxIi8+CiAgPHBhdGggZD0iTSA1IDEwYzAgMS43IDMuMSAzIDcgM3M3LTEuMyA3LTMiIGZpbGw9Im5vbmUiIHN0cm9rZT0iIzZCNzI4MCIgc3Ryb2tlLXdpZHRoPSIxIi8+CiAgPCEtLSBib3QgZmFjZSAtLT4KICA8Y2lyY2xlIGN4PSIxMiIgY3k9IjE1LjUiIHI9IjMiIGZpbGw9IiMxMTE4MjciIHN0cm9rZT0iIzk0QTNCOCIvPgogIDxjaXJjbGUgY3g9IjExIiBjeT0iMTUuNSIgcj0iMC43IiBmaWxsPSIjOTRBM0I4Ii8+CiAgPGNpcmNsZSBjeD0iMTMiIGN5PSIxNS41IiByPSIwLjciIGZpbGw9IiM5NEEzQjgiLz4KPC9zdmc+
+"""
+DEFAULT_AI_AVATAR = """
+data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+CiAgPGRlZnM+CiAgICA8bGluZWFyR3JhZGllbnQgaWQ9ImciIHgxPSIwIiB5MT0iMCIgeDI9IjI0IiB5Mj0iMjQiIGdyYWRpZW50VW5pdHM9InVzZXJTcGFjZSI+CiAgICAgIDxzdG9wIG9mZnNldD0iMCIgc3RvcC1jb2xvcj0iIzRDQUY1MCIvPgogICAgICA8c3RvcCBvZmZzZXQ9IjEiIHN0b3AtY29sb3I9IiMyRTdEMzIiLz4KICAgIDwvbGluZWFyR3JhZGllbnQ+CiAgPC9kZWZzPgogIDwhLS0gZGF0YWJhc2UgLS0+CiAgPGVsbGlwc2UgY3g9IjEyIiBjeT0iNiIgcng9IjciIHJ5PSIzIiBmaWxsPSIjRThGNUU5IiBzdHJva2U9IiMyRTdEMzIiIHN0cm9rZS13aWR0aD0iMSIvPgogIDxwYXRoIGQ9Ik01IDZ2NmMwIDEuNjYgMy4xMyAzIDcgM3M3LTEuMzQgNy0zVjYiIGZpbGw9IiNFOEY1RTkiIHN0cm9rZT0iIzJFN0QzMiIgc3Ryb2tlLXdpZHRoPSIxIi8+CiAgPHBhdGggZD0iTSA1IDljMCAxLjY2IDMuMTMgMyA3IDNzNy0xLjM0IDctMyIgZmlsbD0ibm9uZSIgc3Ryb2tlPSIjMkU3RDMyIiBzdHJva2Utd2lkdGg9IjEiLz4KICA8IS0tIGJvdCBoZWFkIC0tPgogIDxjaXJjbGUgY3g9IjEyIiBjeT0iMTUuNSIgcj0iMyIgZmlsbD0idXJsKCNnKSIgLz4KICA8cmVjdCB4PSI5IiB5PSIxNC4yIiB3aWR0aD0iNiIgaGVpZ2h0PSIyLjYiIHJ4PSIxLjMiIGZpbGw9IiNGRkZGRkYiIG9wYWNpdHk9IjAuOSIvPgogIDxjaXJjbGUgY3g9IjEwLjciIGN5PSIxNS41IiByPSIwLjciIGZpbGw9IiMyRTdEMzIiLz4KICA8Y2lyY2xlIGN4PSIxMy4zIiBjeT0iMTUuNSIgcj0iMC43IiBmaWxsPSIjMkU3RDMyIi8+CiAgPCEtLSBhbnRlbm5hIC0tPgogIDxsaW5lIHgxPSIxMiIgeTE9IjEyLjMiIHgyPSIxMiIgeTI9IjEwLjMiIHN0cm9rZT0iIzJFN0QzMiIgc3Ryb2tlLXdpZHRoPSIxIi8+CiAgPGNpcmNsZSBjeD0iMTIiIGN5PSI5LjciIHI9IjAuNiIgZmlsbD0iIzJFN0QzMiIvPgo8L3N2Zz4=
+"""

src/run_api_client_tests.py

@@ -0,0 +1,23 @@
+"""
+Test runner for API Client unit tests
+"""
+import sys
+import os
+
+# Add the current directory to path so we can import api_client
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import pytest
+
+if __name__ == "__main__":
+    # Run the tests with verbose output and coverage
+    exit_code = pytest.main([
+        "test_api_client.py",
+        "-v",  # Verbose output
+        "--tb=short",  # Short traceback format
+        "--durations=10",  # Show 10 slowest tests
+        "-x",  # Stop on first failure
+    ])
+    
+    print(f"\nTest execution completed with exit code: {exit_code}")
+    sys.exit(exit_code)

src/session_manager.py

@@ -0,0 +1,81 @@
+"""
+Session management for the AI Database Assistant.
+"""
+import streamlit as st
+from typing import List, Dict, Any
+
+
+class SessionManager:
+    """Manages session state and message handling."""
+    
+    @staticmethod
+    def initialize_session() -> None:
+        """Initialize session state variables."""
+        if "messages" not in st.session_state:
+            st.session_state["messages"] = []
+    
+    @staticmethod
+    def get_messages() -> List[Dict[str, Any]]:
+        """Get all messages from session state."""
+        return st.session_state.get("messages", [])
+    
+    @staticmethod
+    def add_user_message(content: str) -> None:
+        """Add a user message to the session."""
+        st.session_state["messages"].append({
+            "role": "user",
+            "content": content
+        })
+    
+    @staticmethod
+    def add_thinking_message() -> None:
+        """Add a thinking placeholder message."""
+        st.session_state["messages"].append({
+            "role": "assistant",
+            "content": "🤔 Thinking...",
+            "is_placeholder": True
+        })
+    
+    @staticmethod
+    def replace_last_message(content: str, data: List[Dict] = None, 
+                           chart: str = None, is_error: bool = False) -> None:
+        """Replace the last message with final response."""
+        st.session_state["messages"][-1] = {
+            "role": "assistant",
+            "content": content,
+            "data": data or [],
+            "chart": chart,
+            "is_error": is_error,
+        }
+    
+    @staticmethod
+    def is_ai_thinking() -> bool:
+        """Check if AI is currently thinking (has placeholder message)."""
+        messages = SessionManager.get_messages()
+        if not messages:
+            return False
+        
+        last_message = messages[-1]
+        return (
+            last_message["role"] == "assistant" and 
+            last_message.get("is_placeholder", False)
+        )
+    
+    @staticmethod
+    def get_last_user_message() -> str:
+        """Get the content of the last user message."""
+        messages = SessionManager.get_messages()
+        if len(messages) >= 2:
+            return messages[-2]["content"]
+        return ""
+    
+    @staticmethod
+    def has_pending_response() -> bool:
+        """Check if there's a pending response to process."""
+        messages = SessionManager.get_messages()
+        return (
+            messages and
+            messages[-1].get("is_placeholder") and
+            len(messages) >= 2 and
+            messages[-2]["role"] == "user"
+        )

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/test_api_client.py

@@ -0,0 +1,347 @@
+"""
+Unit tests for APIClient class
+"""
+import pytest
+import requests
+import json
+from unittest.mock import Mock, patch, MagicMock
+from api_client import APIClient
+
+
+class TestAPIClient:
+    """Test cases for APIClient class"""
+
+    def setup_method(self):
+        """Set up test fixtures before each test method."""
+        self.base_url = "http://localhost:8000"
+        self.client = APIClient(self.base_url, timeout=30)
+
+    def test_init(self):
+        """Test APIClient initialization"""
+        client = APIClient("http://localhost:8000/", timeout=45)
+        assert client.base_url == "http://localhost:8000"  # Trailing slash removed
+        assert client.timeout == 45
+        assert client.endpoints['process_text'] == '/api/process-text'
+        assert client.endpoints['health'] == '/api/health'
+
+    @patch('requests.post')
+    def test_send_query_success_with_data(self, mock_post):
+        """Test successful query with data response"""
+        # Mock successful response
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "sql": "SELECT * FROM employees",
+            "rows": [{"id": 1, "name": "John Doe"}],
+            "heading": "Employee Details",
+            "summary": "List of employees",
+            "chart": None
+        }
+        mock_post.return_value = mock_response
+
+        result = self.client.send_query("Show me all employees")
+
+        assert result["error"] == False
+        assert result["message"] == "Employee Details"
+        assert len(result["rows"]) == 1
+        assert result["rows"][0]["name"] == "John Doe"
+        assert result["sql"] == "SELECT * FROM employees"
+
+    @patch('requests.post')
+    def test_send_query_success_with_message_only(self, mock_post):
+        """Test successful query with message-only response"""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "message": "I understand your question about banking services.",
+            "model_used": "gpt-4"
+        }
+        mock_post.return_value = mock_response
+
+        result = self.client.send_query("What services do you offer?")
+
+        assert result["error"] == False
+        assert result["message"] == "I understand your question about banking services."
+        assert result["rows"] == []
+        assert result["model_used"] == "gpt-4"
+        assert result["status"] == "message"
+
+    @patch('requests.post')
+    def test_send_query_with_model_parameter(self, mock_post):
+        """Test query with specific model parameter"""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"message": "Success"}
+        mock_post.return_value = mock_response
+
+        self.client.send_query("Test question", model="gpt-4-turbo")
+
+        # Verify the payload includes model_name
+        args, kwargs = mock_post.call_args
+        payload = kwargs['json']
+        assert payload["question"] == "Test question"
+        assert payload["model_name"] == "gpt-4-turbo"
+
+    @patch('requests.post')
+    def test_send_query_connection_error(self, mock_post):
+        """Test handling of connection errors"""
+        mock_post.side_effect = requests.exceptions.ConnectionError("Connection failed")
+
+        result = self.client.send_query("Test question")
+
+        assert result["error"] == True
+        assert "Cannot connect to the server" in result["message"]
+        assert result["rows"] == []
+
+    @patch('requests.post')
+    def test_send_query_request_exception(self, mock_post):
+        """Test handling of general request exceptions"""
+        mock_post.side_effect = requests.exceptions.RequestException("Request failed")
+
+        result = self.client.send_query("Test question")
+
+        assert result["error"] == True
+        assert "Request failed" in result["message"]
+        assert result["rows"] == []
+
+    @patch('requests.post')
+    def test_send_query_unexpected_exception(self, mock_post):
+        """Test handling of unexpected exceptions"""
+        mock_post.side_effect = ValueError("Unexpected error")
+
+        result = self.client.send_query("Test question")
+
+        assert result["error"] == True
+        assert "Exception" in result["message"]
+        assert "Unexpected error" in result["message"]
+
+    def test_process_response_success_with_data(self):
+        """Test _process_response with successful data response"""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "sql": "SELECT * FROM branches",
+            "rows": [{"id": 1, "name": "Main Branch"}, {"id": 2, "name": "Downtown"}],
+            "heading": "Branch Information",
+            "summary": "List of all branches"
+        }
+
+        result = self.client._process_response(mock_response)
+
+        assert result["error"] == False
+        assert result["message"] == "Branch Information"
+        assert len(result["rows"]) == 2
+        assert result["heading"] == "Branch Information"
+
+    def test_process_response_with_empty_heading(self):
+        """Test _process_response with empty heading"""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "rows": [{"id": 1, "name": "Test"}],
+            "heading": "",
+            "summary": ""
+        }
+
+        result = self.client._process_response(mock_response)
+
+        assert result["message"] == "Here are the 1 results I found:"
+
+    def test_process_response_no_data_no_heading(self):
+        """Test _process_response with no data and no heading"""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "rows": [],
+            "heading": "",
+            "summary": ""
+        }
+
+        result = self.client._process_response(mock_response)
+
+        assert result["message"] == "I could not find matching records for your query."
+
+    def test_process_response_error_status(self):
+        """Test _process_response with error status code"""
+        mock_response = Mock()
+        mock_response.status_code = 400
+        mock_response.json.return_value = {
+            "detail": "Bad request error"
+        }
+
+        result = self.client._process_response(mock_response)
+
+        assert result["error"] == True
+        assert "Error: Bad request error" in result["message"]
+
+    def test_process_response_invalid_json(self):
+        """Test _process_response with invalid JSON"""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.side_effect = ValueError("Invalid JSON")
+        mock_response.text = "Invalid response text"
+
+        result = self.client._process_response(mock_response)
+
+        # Should handle the JSON parsing error gracefully
+        assert "detail" in result or "message" in result
+
+    @patch('requests.get')
+    def test_check_health_healthy(self, mock_get):
+        """Test health check with healthy status"""
+        mock_response = Mock()
+        mock_response.json.return_value = {"status": "healthy"}
+        mock_get.return_value = mock_response
+
+        status, message, level = self.client.check_health()
+
+        assert status == "🟢 Active"
+        assert message == "Online"
+        assert level == "success"
+
+    @patch('requests.get')
+    def test_check_health_degraded(self, mock_get):
+        """Test health check with degraded status"""
+        mock_response = Mock()
+        mock_response.status_code = 503
+        mock_response.json.return_value = {"status": "degraded"}
+        mock_get.return_value = mock_response
+
+        status, message, level = self.client.check_health()
+
+        assert status == "🟡 Degraded"
+        assert message == "Some Issues"
+        assert level == "warning"
+
+    @patch('requests.get')
+    @patch('socket.create_connection')
+    def test_check_health_connection_error_with_socket_fallback(self, mock_socket, mock_get):
+        """Test health check with connection error but socket reachable"""
+        mock_get.side_effect = requests.exceptions.RequestException("Connection failed")
+        mock_socket.return_value.close.return_value = None
+
+        status, message, level = self.client.check_health()
+
+        assert status == "🟡 Reachable"
+        assert message == "Port Open"
+        assert level == "warning"
+
+    @patch('requests.get')
+    @patch('socket.create_connection')
+    def test_check_health_completely_offline(self, mock_socket, mock_get):
+        """Test health check when completely offline"""
+        mock_get.side_effect = requests.exceptions.RequestException("Connection failed")
+        mock_socket.side_effect = Exception("Socket connection failed")
+
+        status, message, level = self.client.check_health()
+
+        assert status == "🔴 Offline"
+        assert message == "Connection Failed"
+        assert level == "error"
+
+    @patch('requests.get')
+    def test_get_detailed_health_success(self, mock_get):
+        """Test get_detailed_health with successful response"""
+        expected_health = {
+            "status": "healthy",
+            "checks": {"database": "ok", "api": "ok"}
+        }
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = expected_health
+        mock_get.return_value = mock_response
+
+        result = self.client.get_detailed_health()
+
+        assert result == expected_health
+
+    @patch('requests.get')
+    def test_get_detailed_health_error(self, mock_get):
+        """Test get_detailed_health with error response"""
+        mock_response = Mock()
+        mock_response.status_code = 500
+        mock_get.return_value = mock_response
+
+        result = self.client.get_detailed_health()
+
+        assert result["status"] == "error"
+        assert "500" in result["message"]
+
+    @patch('requests.get')
+    def test_get_method_success(self, mock_get):
+        """Test generic GET method with success"""
+        expected_data = {"models": ["gpt-4", "gpt-3.5"]}
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = expected_data
+        mock_get.return_value = mock_response
+
+        result = self.client.get("/models")
+
+        assert result == expected_data
+
+    @patch('requests.get')
+    def test_get_method_error(self, mock_get):
+        """Test generic GET method with error"""
+        mock_response = Mock()
+        mock_response.status_code = 404
+        mock_get.return_value = mock_response
+
+        result = self.client.get("/models")
+
+        assert result is None
+
+    @patch('requests.post')
+    def test_post_method_success(self, mock_post):
+        """Test generic POST method with success"""
+        expected_data = {"result": "success"}
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = expected_data
+        mock_post.return_value = mock_response
+
+        result = self.client.post("/change-model", {"model": "gpt-4"})
+
+        assert result == expected_data
+
+    @patch('requests.post')
+    def test_post_method_error(self, mock_post):
+        """Test generic POST method with error"""
+        mock_response = Mock()
+        mock_response.status_code = 400
+        mock_post.return_value = mock_response
+
+        result = self.client.post("/change-model", {"model": "invalid"})
+
+        assert result is None
+
+    def test_url_construction(self):
+        """Test URL construction for different endpoints"""
+        client = APIClient("http://localhost:8000/")
+        
+        # Test that trailing slash is removed
+        assert client.base_url == "http://localhost:8000"
+        
+        # Test endpoint URLs
+        process_url = f"{client.base_url}{client.endpoints['process_text']}"
+        assert process_url == "http://localhost:8000/api/process-text"
+
+    @patch('requests.post')
+    def test_send_query_with_legacy_agent_parameter(self, mock_post):
+        """Test that legacy agent parameter is handled correctly"""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"message": "Success"}
+        mock_post.return_value = mock_response
+
+        # Agent parameter should be ignored in payload
+        self.client.send_query("Test", model="gpt-4", agent="legacy_agent")
+
+        args, kwargs = mock_post.call_args
+        payload = kwargs['json']
+        assert "agent" not in payload  # Agent should not be in payload
+        assert payload["model_name"] == "gpt-4"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

src/test_heading_parsing.py

@@ -0,0 +1,22 @@
+"""
+Quick test to check heading parsing in API client
+"""
+import logging
+from api_client import APIClient
+
+# Set up logging to see debug info
+logging.basicConfig(level=logging.INFO)
+
+# Create API client
+client = APIClient("http://localhost:8000")
+
+# Test with a simple query
+print("Testing heading parsing...")
+response = client.send_query("Show me all employees")
+
+print(f"\nResponse keys: {list(response.keys())}")
+print(f"Message: {response.get('message')}")
+print(f"Heading: {response.get('heading')}")
+print(f"Summary: {response.get('summary')}")
+print(f"Original heading: {response.get('original_heading')}")
+print(f"Rows count: {len(response.get('rows', []))}")

src/themes.py

@@ -0,0 +1,740 @@
+"""
+Theme management for the AI Database Assistant.
+Handles CSS injection and theme-specific styling.
+"""
+import streamlit as st
+from typing import Dict
+
+
+class ThemeManager:
+    """Manages themes and CSS styling for the application."""
+    
+    def __init__(self):
+        self.themes = {
+            "Light": self._get_light_theme(),
+            "Dark": self._get_dark_theme(),
+            "Custom": self._get_custom_theme()
+        }
+    
+    def _get_light_theme(self) -> str:
+        """Returns CSS for light theme."""
+        return """
+        /* Light Theme - Clean and Professional */
+        body, .stApp { 
+            background: #f8fafc !important; 
+            color: #1f2937 !important;
+        }
+        
+        /* Ensure all text is dark in light mode */
+        div, p, span, label, h1, h2, h3, h4, h5, h6 {
+            color: #1f2937 !important;
+        }
+        
+        /* Professional header colors for light theme */
+        .main-header h1 {
+            color: #1f2937 !important;
+        }
+        
+        .main-header p {
+            color: #6b7280 !important;
+        }
+        
+        .chat-container { max-width: 900px; margin: auto; }
+        
+        .user-bubble {
+            background: linear-gradient(90deg, #e3f2fd, #ffffff);
+            color: #222 !important; float: right; clear: both;
+        }
+        .ai-bubble {
+            background: linear-gradient(90deg, #f3e5f5, #ffffff);
+            color: #222 !important; float: left; clear: both;
+        }
+        .error-bubble {
+            background: #ffebee; color: #b71c1c !important; font-weight: bold;
+            float: left; clear: both;
+        }
+        
+        /* Sidebar styling */
+        section[data-testid="stSidebar"] { 
+            background: #e3f2fd !important; 
+            border-right: 2px solid #1976d2;
+        }
+        
+        /* Sidebar text should be dark */
+        section[data-testid="stSidebar"] * {
+            color: #1f2937 !important;
+        }
+        
+        .sidebar-content { padding: 1rem; }
+        
+        /* Input fields in light mode */
+        .stTextInput > div > div > input {
+            background-color: #ffffff !important;
+            color: #1f2937 !important;
+            border: 1px solid #d1d5db !important;
+        }
+        
+        /* Select boxes in light mode */
+        .stSelectbox > div > div > select {
+            background-color: #ffffff !important;
+            color: #1f2937 !important;
+            border: 1px solid #d1d5db !important;
+        }
+        
+        /* Dropdown options in light mode */
+        .stSelectbox > div > div > select > option {
+            background-color: #ffffff !important;
+            color: #1f2937 !important;
+        }
+        
+        /* Buttons in light mode - Strong selectors for all buttons */
+        .stButton > button {
+            background: linear-gradient(45deg, #667eea, #764ba2) !important;
+            color: #ffffff !important;
+            border: none !important;
+            font-weight: 600 !important;
+            border-radius: 8px !important;
+            padding: 0.6rem 1.2rem !important;
+            min-height: 2.8rem !important;
+            font-size: 14px !important;
+            transition: all 0.2s ease !important;
+        }
+        
+        .stButton > button:hover {
+            background: linear-gradient(45deg, #5a67d8, #6b46c1) !important;
+            color: #ffffff !important;
+            transform: translateY(-1px) !important;
+        }
+        
+        /* Ensure button text stays white - stronger selectors */
+        .stButton > button span,
+        .stButton > button div,
+        .stButton > button *,
+        .stButton > button p {
+            color: #ffffff !important;
+            font-weight: 600 !important;
+        }
+        
+        /* Specific targeting for all button states and text elements */
+        button[kind="primary"],
+        button[kind="secondary"],
+        div[data-testid="stButton"] > button,
+        .stButton button {
+            background: linear-gradient(45deg, #667eea, #764ba2) !important;
+            color: #ffffff !important;
+            border: none !important;
+            font-weight: 600 !important;
+        }
+        
+        /* All text inside buttons must be white */
+        button[kind="primary"] *,
+        button[kind="primary"] span,
+        button[kind="primary"] div,
+        button[kind="primary"] p,
+        button[kind="secondary"] *,
+        button[kind="secondary"] span,
+        button[kind="secondary"] div,
+        button[kind="secondary"] p,
+        div[data-testid="stButton"] > button *,
+        div[data-testid="stButton"] > button span,
+        div[data-testid="stButton"] > button div,
+        div[data-testid="stButton"] > button p,
+        .stButton button *,
+        .stButton button span,
+        .stButton button div,
+        .stButton button p {
+            color: #ffffff !important;
+            font-weight: 600 !important;
+        }
+        
+        button[kind="primary"]:hover,
+        button[kind="secondary"]:hover,
+        div[data-testid="stButton"] > button:hover,
+        .stButton button:hover {
+            background: linear-gradient(45deg, #5a67d8, #6b46c1) !important;
+            color: #ffffff !important;
+        }
+        
+        /* Hover state text colors */
+        button[kind="primary"]:hover *,
+        button[kind="secondary"]:hover *,
+        div[data-testid="stButton"] > button:hover *,
+        .stButton button:hover * {
+            color: #ffffff !important;
+        }
+        
+        /* Force all button text elements to be white */
+        .stButton > button * {
+            color: #ffffff !important;
+        }
+        
+        /* Disabled button styling */
+        .stButton > button:disabled {
+            background: #9ca3af !important;
+            color: #ffffff !important;
+            opacity: 0.6 !important;
+        }
+        
+        .stButton > button:disabled * {
+            color: #ffffff !important;
+        }
+        
+        /* Copy button styling in light mode */
+        .copy-button {
+            background: #6b7280 !important;
+            color: #ffffff !important;
+            border: none !important;
+        }
+        
+        .copy-button:hover {
+            background: #4b5563 !important;
+            color: #ffffff !important;
+        }
+        
+        /* Labels and form elements */
+        .stTextInput label, .stSelectbox label, .stRadio label, .stCheckbox label {
+            color: #1f2937 !important;
+            font-weight: 500 !important;
+        }
+        
+        /* Metrics in light mode */
+        [data-testid="metric-container"] {
+            background: #ffffff !important;
+            border: 1px solid #e5e7eb !important;
+            color: #1f2937 !important;
+        }
+        
+        [data-testid="metric-container"] * {
+            color: #1f2937 !important;
+        }
+        
+        /* Additional button overrides for light theme */
+        div[data-testid="column"] .stButton > button {
+            background: linear-gradient(45deg, #667eea, #764ba2) !important;
+            color: #ffffff !important;
+        }
+        
+        /* Force override of any inherited text colors */
+        .stButton > button p, .stButton > button div, .stButton > button span {
+            color: #ffffff !important;
+        }
+        
+        /* Hide Streamlit toolbar and menu */
+        .stToolbar {
+            display: none !important;
+        }
+        
+        /* Hide Streamlit header menu */
+        header[data-testid="stHeader"] {
+            display: none !important;
+        }
+        
+        /* Hide the running/rerun indicator */
+        .stAppView > .main .block-container {
+            padding-top: 1rem !important;
+        }
+        """
+    
+    def _get_dark_theme(self) -> str:
+        """Returns CSS for dark theme."""
+        return """
+        /* Modern Dark Theme - High Contrast & Clean */
+        body, .stApp { 
+            background: #1a1a1a !important; 
+            color: #ffffff !important;
+        }
+        
+        /* Override all text colors for visibility */
+        div, p, span, label, h1, h2, h3, h4, h5, h6 {
+            color: #ffffff !important;
+        }
+        
+        /* Professional header colors for dark theme */
+        .main-header h1 {
+            color: #ffffff !important;
+        }
+        
+        .main-header p {
+            color: #9ca3af !important;
+        }
+        
+        /* Main container */
+        .chat-container { 
+            max-width: 900px; 
+            margin: auto; 
+        }
+        
+        /* User message bubble - Clean Blue */
+        .user-bubble {
+            background: #2563eb !important;
+            color: #ffffff !important; 
+            float: right; 
+            clear: both;
+            border: none !important;
+            box-shadow: 0 2px 10px rgba(37, 99, 235, 0.4) !important;
+        }
+        
+        /* AI message bubble - Clean Dark Gray */
+        .ai-bubble {
+            background: #374151 !important;
+            color: #ffffff !important; 
+            float: left; 
+            clear: both;
+            border: 1px solid #4b5563 !important;
+            box-shadow: 0 2px 10px rgba(0, 0, 0, 0.3) !important;
+        }
+        
+        /* Error bubble - Clean Red */
+        .error-bubble {
+            background: #dc2626 !important;
+            color: #ffffff !important; 
+            font-weight: bold;
+            float: left; 
+            clear: both;
+            border: none !important;
+            box-shadow: 0 2px 10px rgba(220, 38, 38, 0.4) !important;
+        }
+        
+        /* Sidebar - Clean Dark */
+        section[data-testid="stSidebar"] { 
+            background: #2d2d2d !important; 
+            border-right: 2px solid #4b5563 !important;
+        }
+        
+        /* Sidebar content visibility */
+        section[data-testid="stSidebar"] * {
+            color: #ffffff !important;
+        }
+        
+        /* Input fields - High contrast */
+        .stTextInput > div > div > input {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+            border: 2px solid #6b7280 !important;
+        }
+        
+        /* Select boxes - High contrast */
+        .stSelectbox > div > div > select {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+            border: 2px solid #6b7280 !important;
+        }
+        
+        /* Dropdown options styling */
+        .stSelectbox > div > div > select > option {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+        }
+        
+        /* Alternative dropdown styling for better browser support */
+        div[data-baseweb="select"] {
+            background-color: #374151 !important;
+        }
+        
+        div[data-baseweb="select"] > div {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+            border: 2px solid #6b7280 !important;
+        }
+        
+        /* Dropdown menu styling */
+        ul[role="listbox"] {
+            background-color: #374151 !important;
+            border: 1px solid #6b7280 !important;
+        }
+        
+        li[role="option"] {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+        }
+        
+        li[role="option"]:hover {
+            background-color: #4b5563 !important;
+            color: #ffffff !important;
+        }
+        
+        /* Buttons - Clean and visible - Strong selectors for all buttons */
+        .stButton > button {
+            background: #2563eb !important;
+            color: #ffffff !important;
+            border: none !important;
+            font-weight: 600 !important;
+            border-radius: 8px !important;
+            padding: 0.6rem 1.2rem !important;
+            min-height: 2.8rem !important;
+            font-size: 14px !important;
+            transition: all 0.2s ease !important;
+        }
+        
+        .stButton > button:hover {
+            background: #1d4ed8 !important;
+            color: #ffffff !important;
+            transform: translateY(-1px) !important;
+        }
+        
+        /* Ensure button text stays white - stronger selectors */
+        .stButton > button span,
+        .stButton > button div,
+        .stButton > button *,
+        .stButton > button p {
+            color: #ffffff !important;
+            font-weight: 600 !important;
+        }
+        
+        /* Specific targeting for all button states and text elements */
+        button[kind="primary"],
+        button[kind="secondary"],
+        div[data-testid="stButton"] > button,
+        .stButton button {
+            background: #2563eb !important;
+            color: #ffffff !important;
+            border: none !important;
+            font-weight: 600 !important;
+        }
+        
+        /* All text inside buttons must be white */
+        button[kind="primary"] *,
+        button[kind="primary"] span,
+        button[kind="primary"] div,
+        button[kind="primary"] p,
+        button[kind="secondary"] *,
+        button[kind="secondary"] span,
+        button[kind="secondary"] div,
+        button[kind="secondary"] p,
+        div[data-testid="stButton"] > button *,
+        div[data-testid="stButton"] > button span,
+        div[data-testid="stButton"] > button div,
+        div[data-testid="stButton"] > button p,
+        .stButton button *,
+        .stButton button span,
+        .stButton button div,
+        .stButton button p {
+            color: #ffffff !important;
+            font-weight: 600 !important;
+        }
+        
+        button[kind="primary"]:hover,
+        button[kind="secondary"]:hover,
+        div[data-testid="stButton"] > button:hover,
+        .stButton button:hover {
+            background: #1d4ed8 !important;
+            color: #ffffff !important;
+        }
+        
+        /* Hover state text colors */
+        button[kind="primary"]:hover *,
+        button[kind="secondary"]:hover *,
+        div[data-testid="stButton"] > button:hover *,
+        .stButton button:hover * {
+            color: #ffffff !important;
+        }
+        
+        /* Expander - Clean styling */
+        .streamlit-expanderHeader {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+            border: 1px solid #6b7280 !important;
+        }
+        
+        /* Metrics - High visibility */
+        [data-testid="metric-container"] {
+            background: #374151 !important;
+            border: 1px solid #6b7280 !important;
+            color: #ffffff !important;
+        }
+        
+        [data-testid="metric-container"] * {
+            color: #ffffff !important;
+        }
+        
+        /* Code blocks - GitHub dark style */
+        pre, code {
+            background-color: #0d1117 !important;
+            color: #f0f6fc !important;
+            border: 1px solid #30363d !important;
+        }
+        
+        /* Tables - Clean dark */
+        .dataframe, .dataframe * {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+        }
+        
+        /* Table headers */
+        .dataframe th {
+            background-color: #2563eb !important;
+            color: #ffffff !important;
+        }
+        
+        /* Streamlit widgets override */
+        .stRadio > div, .stCheckbox > div {
+            color: #ffffff !important;
+        }
+        
+        .stRadio > div > label, .stCheckbox > div > label {
+            color: #ffffff !important;
+        }
+        
+        /* Spinner for dark theme */
+        .stSpinner > div {
+            border-top-color: #2563eb !important;
+        }
+        
+        /* Success/Info messages */
+        .stSuccess, .stInfo {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+            border: 1px solid #6b7280 !important;
+        }
+        
+        /* Warning messages */
+        .stWarning {
+            background-color: #f59e0b !important;
+            color: #000000 !important;
+        }
+        
+        /* Error messages */
+        .stError {
+            background-color: #dc2626 !important;
+            color: #ffffff !important;
+        }
+        
+        /* Additional dropdown fixes */
+        .stSelectbox [data-testid="stMarkdownContainer"] {
+            color: #ffffff !important;
+        }
+        
+        /* Streamlit's custom selectbox */
+        div[data-testid="stSelectbox"] div[data-testid="stMarkdownContainer"] p {
+            color: #ffffff !important;
+        }
+        
+        /* Multi-select components */
+        .stMultiSelect > div > div > div {
+            background-color: #374151 !important;
+            color: #ffffff !important;
+            border: 2px solid #6b7280 !important;
+        }
+        
+        /* Radio button labels */
+        .stRadio div[role="radiogroup"] label {
+            color: #ffffff !important;
+        }
+        
+        /* Hide Streamlit toolbar and menu */
+        .stToolbar {
+            display: none !important;
+        }
+        
+        /* Hide Streamlit header menu */
+        header[data-testid="stHeader"] {
+            display: none !important;
+        }
+        
+        /* Hide the running/rerun indicator */
+        .stAppView > .main .block-container {
+            padding-top: 1rem !important;
+        }
+        """
+    
+    def _get_custom_theme(self) -> str:
+        """Returns CSS for custom theme."""
+        return """
+        body, .stApp { background: #fffbe7 !important; }
+        .chat-container { max-width: 900px; margin: auto; }
+        .user-bubble {
+            background: linear-gradient(90deg, #ffe082, #fffbe7);
+            color: #6d4c00;
+            float: right;
+            clear: both;
+        }
+        .ai-bubble {
+            background: linear-gradient(90deg, #b2dfdb, #fffbe7);
+            color: #004d40;
+            float: left;
+            clear: both;
+        }
+        .error-bubble {
+            background: #ffe0b2; color: #b71c1c; font-weight: bold;
+            float: left; clear: both;
+        }
+        section[data-testid="stSidebar"] { background: #ffe082 !important; }
+        
+        /* Hide Streamlit toolbar and menu */
+        .stToolbar {
+            display: none !important;
+        }
+        
+        /* Hide Streamlit header menu */
+        header[data-testid="stHeader"] {
+            display: none !important;
+        }
+        
+        /* Hide the running/rerun indicator */
+        .stAppView > .main .block-container {
+            padding-top: 1rem !important;
+        }
+        """
+    
+    def _get_common_styles(self) -> str:
+        """Returns common CSS styles for all themes."""
+        return """
+        /* Professional header styling */
+        .main-header h1 {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', sans-serif !important;
+            font-weight: 600 !important;
+            line-height: 1.2 !important;
+            margin: 0 !important;
+        }
+        
+        .main-header p {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', sans-serif !important;
+            font-weight: 400 !important;
+            line-height: 1.4 !important;
+            margin: 0 !important;
+        }
+        
+        .user-bubble, .ai-bubble, .error-bubble {
+            border-radius: 14px;
+            padding: 10px 14px;
+            margin: 6px 0;
+            display: inline-block;
+            max-width: 75%;
+            word-wrap: break-word;
+            box-shadow: 0 2px 8px rgba(0,0,0,0.25);
+        }
+        .spinner {
+            display: inline-block;
+            width: 1.3em;
+            height: 1.3em;
+            border: 3px solid #e0e0e0;
+            border-top: 3px solid #2193b0;
+            border-radius: 50%;
+            animation: spin 0.8s linear infinite;
+            margin-right: 0.7em;
+        }
+        @keyframes spin {
+            0% { transform: rotate(0deg); }
+            100% { transform: rotate(360deg); }
+        }
+        
+        /* Enhanced button styling for all themes */
+        .stButton > button {
+            border-radius: 8px !important;
+            font-weight: 600 !important;
+            transition: all 0.2s ease !important;
+            padding: 0.6rem 1.2rem !important;
+            min-height: 2.8rem !important;
+            font-size: 14px !important;
+            border: none !important;
+        }
+        
+        .stButton > button:hover {
+            transform: translateY(-1px) !important;
+        }
+        
+        /* Force all button text elements to maintain proper colors - Super strong selectors */
+        .stButton > button *, 
+        .stButton > button p, 
+        .stButton > button div, 
+        .stButton > button span,
+        .stButton > button strong,
+        .stButton > button em,
+        .stButton > button i,
+        .stButton > button b,
+        div[data-testid="stButton"] > button *,
+        div[data-testid="stButton"] > button p,
+        div[data-testid="stButton"] > button div,
+        div[data-testid="stButton"] > button span,
+        div[data-testid="stButton"] > button strong,
+        div[data-testid="stButton"] > button em,
+        div[data-testid="stButton"] > button i,
+        div[data-testid="stButton"] > button b,
+        button[kind="primary"] *,
+        button[kind="primary"] p,
+        button[kind="primary"] div,
+        button[kind="primary"] span,
+        button[kind="primary"] strong,
+        button[kind="primary"] em,
+        button[kind="primary"] i,
+        button[kind="primary"] b,
+        button[kind="secondary"] *,
+        button[kind="secondary"] p,
+        button[kind="secondary"] div,
+        button[kind="secondary"] span,
+        button[kind="secondary"] strong,
+        button[kind="secondary"] em,
+        button[kind="secondary"] i,
+        button[kind="secondary"] b,
+        .stButton button *,
+        .stButton button p,
+        .stButton button div,
+        .stButton button span,
+        .stButton button strong,
+        .stButton button em,
+        .stButton button i,
+        .stButton button b {
+            font-weight: 600 !important;
+            /* Note: color is set by individual themes */
+        }
+        
+        /* Additional strong selectors for buttons */
+        div[data-testid="stButton"] button,
+        .stButton button,
+        button[data-testid],
+        [data-testid="stButton"] button {
+            border-radius: 8px !important;
+            font-weight: 600 !important;
+            transition: all 0.2s ease !important;
+            padding: 0.6rem 1.2rem !important;
+            min-height: 2.8rem !important;
+            font-size: 14px !important;
+            border: none !important;
+        }
+        
+        /* Ensure ALL button content inherits button color */
+        .stButton > button,
+        div[data-testid="stButton"] > button,
+        button[kind="primary"],
+        button[kind="secondary"],
+        .stButton button {
+            /* Background and text colors are set by individual themes */
+        }
+        
+        /* Critical: Force text color inheritance from button */
+        .stButton > button *,
+        div[data-testid="stButton"] > button *,
+        button[kind="primary"] *,
+        button[kind="secondary"] *,
+        .stButton button * {
+            color: inherit !important;
+        }
+        
+        /* Scrollbar styling */
+        ::-webkit-scrollbar {
+            width: 8px;
+            height: 8px;
+        }
+        ::-webkit-scrollbar-track {
+            background: rgba(0,0,0,0.1);
+            border-radius: 4px;
+        }
+        ::-webkit-scrollbar-thumb {
+            background: rgba(128,128,128,0.4);
+            border-radius: 4px;
+        }
+        ::-webkit-scrollbar-thumb:hover {
+            background: rgba(128,128,128,0.6);
+        }
+        """
+    
+    def inject_theme(self, theme_name: str) -> None:
+        """Inject the selected theme CSS into the Streamlit app."""
+        if theme_name not in self.themes:
+            theme_name = "Light"  # Default fallback
+        
+        theme_css = self.themes[theme_name]
+        common_css = self._get_common_styles()
+        combined_css = theme_css + common_css
+        
+        st.markdown(f"<style>{combined_css}</style>", unsafe_allow_html=True)