Update app.py

app.py

@@ -1,1002 +1,1090 @@
-"""
-Schema-Agnostic Database Chatbot - Streamlit Application
-
-A production-grade chatbot that connects to ANY database
-(MySQL, PostgreSQL, SQLite) and provides intelligent querying 
-through RAG and Text-to-SQL.
-
-Uses Groq for FREE LLM inference!
-"""
-
-import os
-from pathlib import Path
-
-# Load .env FIRST before any other imports
-from dotenv import load_dotenv
-load_dotenv(Path(__file__).parent / ".env")
-
-import streamlit as st
-import uuid
-import time
-import io
-import csv
-import base64
-from datetime import datetime
-
-# Page config must be first
-st.set_page_config(
-    page_title="OnceDataBot",
-    page_icon="🤖",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-# Imports
-from config import config, DatabaseConfig, DatabaseType
-from database import get_db, get_schema, get_introspector
-from database.connection import DatabaseConnection
-from llm import create_llm_client
-from chatbot import create_chatbot, DatabaseChatbot
-from memory import ChatMemory, EnhancedChatMemory
-from viz_utils import render_visualization
-
-
-
-
-
-
-# Groq models (all FREE!)
-GROQ_MODELS = [
-    "llama-3.3-70b-versatile",
-    "llama-3.1-8b-instant",
-    "mixtral-8x7b-32768",
-    "gemma2-9b-it"
-]
-
-# Database types
-DB_TYPES = {
-    "MySQL": "mysql",
-    "PostgreSQL": "postgresql"
-}
-
-# Supported languages for multi-language responses
-SUPPORTED_LANGUAGES = {
-    "English": "en",
-    "हिन्दी (Hindi)": "hi",
-    "Español (Spanish)": "es",
-    "Français (French)": "fr",
-    "Deutsch (German)": "de",
-    "中文 (Chinese)": "zh",
-    "日本語 (Japanese)": "ja",
-    "한국어 (Korean)": "ko",
-    "Português (Portuguese)": "pt",
-    "العربية (Arabic)": "ar",
-    "Русский (Russian)": "ru",
-    "Italiano (Italian)": "it",
-    "Nederlands (Dutch)": "nl",
-    "தமிழ் (Tamil)": "ta",
-    "తెలుగు (Telugu)": "te",
-    "मराठी (Marathi)": "mr",
-    "বাংলা (Bengali)": "bn",
-    "ગુજરાতી (Gujarati)": "gu"
-}
-
-
-def create_custom_db_config(db_type: str, **kwargs) -> DatabaseConfig:
-    """Create a custom database configuration from user input."""
-    return DatabaseConfig(
-        db_type=DatabaseType(db_type),
-        host=kwargs.get("host", ""),
-        port=kwargs.get("port", 3306 if db_type == "mysql" else 5432),
-        database=kwargs.get("database", ""),
-        username=kwargs.get("username", ""),
-        password=kwargs.get("password", ""),
-        ssl_ca=kwargs.get("ssl_ca", None)
-    )
-
-
-def create_custom_memory(session_id: str, user_id: str, db_connection, llm_client=None, 
-                         enable_summarization=True, summary_threshold=10) -> EnhancedChatMemory:
-    """Create enhanced memory with a custom database connection."""
-    return EnhancedChatMemory(
-        session_id=session_id,
-        user_id=user_id,
-        max_messages=20,
-        db_connection=db_connection,
-        llm_client=llm_client,
-        enable_summarization=enable_summarization,
-        summary_threshold=summary_threshold
-    )
-
-
-def init_session_state():
-    """Initialize Streamlit session state."""
-    if "session_id" not in st.session_state:
-        st.session_state.session_id = str(uuid.uuid4())
-    
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-    
-    if "chatbot" not in st.session_state:
-        st.session_state.chatbot = None
-    
-    if "initialized" not in st.session_state:
-        st.session_state.initialized = False
-    
-    if "user_id" not in st.session_state:
-        st.session_state.user_id = "default"
-    
-    if "enable_summarization" not in st.session_state:
-        st.session_state.enable_summarization = True
-    
-    if "summary_threshold" not in st.session_state:
-        st.session_state.summary_threshold = 10
-    
-    if "memory" not in st.session_state:
-        st.session_state.memory = None
-        
-    if "indexed" not in st.session_state:
-        st.session_state.indexed = False
-    
-    if "db_source" not in st.session_state:
-        st.session_state.db_source = "environment"  # "environment" or "custom"
-    
-    if "custom_db_config" not in st.session_state:
-        st.session_state.custom_db_config = None
-    
-    if "custom_db_connection" not in st.session_state:
-        st.session_state.custom_db_connection = None
-        
-    if "ignored_tables" not in st.session_state:
-        st.session_state.ignored_tables = set()
-    
-    if "response_language" not in st.session_state:
-        st.session_state.response_language = "English"
-    
-    if "favorites" not in st.session_state:
-        st.session_state.favorites = []  # List of message indices that are favorited
-
-
-def export_results_to_csv(results: list) -> str:
-    """Convert SQL results to CSV format and return as downloadable string."""
-    if not results:
-        return ""
-    
-    output = io.StringIO()
-    writer = csv.DictWriter(output, fieldnames=results[0].keys())
-    writer.writeheader()
-    writer.writerows(results)
-    return output.getvalue()
-
-
-def export_chat_to_text() -> str:
-    """Export chat messages to text format."""
-    if not st.session_state.messages:
-        return "No messages to export."
-    
-    lines = []
-    lines.append("=" * 50)
-    lines.append(f"OnceDataBot Chat Export")
-    lines.append(f"Exported: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
-    lines.append(f"User: {st.session_state.user_id}")
-    lines.append("=" * 50)
-    lines.append("")
-    
-    for i, msg in enumerate(st.session_state.messages):
-        role = "🧑 User" if msg["role"] == "user" else "🤖 Assistant"
-        is_favorited = "⭐ " if i in st.session_state.favorites else ""
-        lines.append(f"{is_favorited}{role}:")
-        lines.append(msg["content"])
-        
-        if msg["role"] == "assistant" and "metadata" in msg:
-            meta = msg["metadata"]
-            if meta.get("sql_query"):
-                lines.append(f"\n📝 SQL Query: {meta['sql_query']}")
-            if meta.get("query_type"):
-                lines.append(f"📌 Query Type: {meta['query_type']}")
-            if meta.get("execution_time"):
-                lines.append(f"⏱️ Execution Time: {meta['execution_time']:.2f}s")
-        
-        lines.append("-" * 40)
-        lines.append("")
-    
-    return "\n".join(lines)
-
-
-def render_copy_button(text: str, key: str):
-    """Render a copy to clipboard button using Streamlit."""
-    # Using a workaround with st.code which has built-in copy
-    st.code(text, language="sql")
-
-
-def render_database_config():
-    """Render database configuration section in sidebar."""
-    st.subheader("🗄️ Database Configuration")
-    
-    # Database source selection
-    db_source = st.radio(
-        "Database Source",
-        options=["Use Environment Variables", "Custom Database"],
-        index=0 if st.session_state.db_source == "environment" else 1,
-        key="db_source_radio",
-        help="Choose to use .env settings or enter custom credentials"
-    )
-    
-    st.session_state.db_source = "environment" if db_source == "Use Environment Variables" else "custom"
-    
-    if st.session_state.db_source == "environment":
-        # Show current environment config
-        current_db_type = config.database.db_type.value.upper()
-        st.info(f"📌 Using {current_db_type} from environment")
-        st.caption(f"Host: {config.database.host}")
-        return None
-    
-    else:
-        # Custom database configuration
-        st.markdown("##### Enter Database Credentials")
-        
-        # Database type selector
-        db_type_label = st.selectbox(
-            "Database Type",
-            options=list(DB_TYPES.keys()),
-            index=0,
-            key="custom_db_type"
-        )
-        db_type = DB_TYPES[db_type_label]
-        
-        if True:  # MySQL or PostgreSQL (SQLite removed)
-            # MySQL or PostgreSQL
-            col1, col2 = st.columns([3, 1])
-            with col1:
-                host = st.text_input(
-                    "Host",
-                    value="",
-                    key="db_host_input",
-                    placeholder="your-database-host.com"
-                )
-            with col2:
-                default_port = 3306 if db_type == "mysql" else 5432
-                port = st.number_input(
-                    "Port",
-                    value=default_port,
-                    min_value=1,
-                    max_value=65535,
-                    key="db_port_input"
-                )
-            
-            database = st.text_input(
-                "Database Name",
-                value="",
-                key="db_name_input",
-                placeholder="your_database"
-            )
-            
-            username = st.text_input(
-                "Username",
-                value="",
-                key="db_user_input",
-                placeholder="your_username"
-            )
-            
-            password = st.text_input(
-                "Password",
-                value="",
-                type="password",
-                key="db_pass_input"
-            )
-            
-            # Optional SSL
-            with st.expander("🔒 SSL Settings (Optional)"):
-                ssl_ca = st.text_input(
-                    "SSL CA Certificate Path",
-                    value="",
-                    key="ssl_ca_input",
-                    help="Path to SSL CA certificate file (for cloud databases like Aiven)"
-                )
-            
-            return {
-                "db_type": db_type,
-                "host": host,
-                "port": int(port),
-                "database": database,
-                "username": username,
-                "password": password,
-                "ssl_ca": ssl_ca if ssl_ca else None
-            }
-
-
-def render_sidebar():
-    """Render the configuration sidebar."""
-    with st.sidebar:
-        st.title("⚙️ Settings")
-        
-        # Session Dashboard
-        if st.session_state.messages:
-            st.markdown("### 📊 Session Stats")
-            
-            # Calculate stats
-            total_msgs = len(st.session_state.messages)
-            assistant_msgs = [m for m in st.session_state.messages if m.get("role") == "assistant"]
-            sql_queries = sum(1 for m in assistant_msgs if m.get("metadata", {}).get("sql_query"))
-            
-            total_tokens = 0
-            exec_times = []
-            for m in assistant_msgs:
-                meta = m.get("metadata", {})
-                total_tokens += meta.get("token_usage", {}).get("total", 0)
-                if meta.get("execution_time"):
-                    exec_times.append(meta["execution_time"])
-            
-            avg_time = sum(exec_times) / len(exec_times) if exec_times else 0
-            
-            col_s1, col_s2 = st.columns(2)
-            col_s1.metric("Queries", sql_queries)
-            col_s2.metric("Tokens", f"{total_tokens:,}")
-            st.caption(f"⏱️ Avg Time: {avg_time:.2f}s | 💬 Msgs: {total_msgs}")
-            st.divider()
-        
-
-        
-        # User Profile
-        st.subheader("👤 User Profile")
-        user_id = st.text_input(
-            "User ID / Name", 
-            value=st.session_state.get("user_id", "default"),
-            key="user_id_input",
-            help="Your unique ID for private memory storage"
-        )
-        if user_id != st.session_state.get("user_id"):
-            st.session_state.user_id = user_id
-            st.session_state.session_id = str(uuid.uuid4())
-            st.session_state.messages = []
-            
-            # Recreate memory for new user
-            if st.session_state.custom_db_connection:
-                st.session_state.memory = create_custom_memory(
-                    st.session_state.session_id,
-                    user_id,
-                    st.session_state.custom_db_connection,
-                    st.session_state.get("llm"),
-                    st.session_state.enable_summarization,
-                    st.session_state.summary_threshold
-                )
-            elif st.session_state.initialized:
-                from memory import create_enhanced_memory
-                st.session_state.memory = create_enhanced_memory(
-                    st.session_state.session_id,
-                    user_id=user_id,
-                    enable_summarization=st.session_state.enable_summarization,
-                    summary_threshold=st.session_state.summary_threshold
-                )
-            
-            if st.session_state.memory:
-                st.session_state.memory.clear_user_history()
-            st.rerun()
-        
-        st.divider()
-        
-        # Language Selection
-        st.subheader("🌐 Response Language")
-        selected_language = st.selectbox(
-            "Select Language",
-            options=list(SUPPORTED_LANGUAGES.keys()),
-            index=list(SUPPORTED_LANGUAGES.keys()).index(st.session_state.response_language),
-            key="language_selector",
-            help="Choose the language for chatbot responses"
-        )
-        if selected_language != st.session_state.response_language:
-            st.session_state.response_language = selected_language
-            st.toast(f"🌐 Language changed to {selected_language}")
-        
-        st.divider()
-        
-        # Export Chat Button
-        if st.session_state.messages:
-            st.download_button(
-                label="📄 Export Chat",
-                data=export_chat_to_text(),
-                file_name=f"chat_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
-                mime="text/plain",
-                use_container_width=True,
-                help="Download your chat conversation as a text file"
-            )
-        
-        st.divider()
-        
-        # Database Configuration
-        custom_db_params = render_database_config()
-        
-        st.divider()
-        
-        # LLM Configuration
-        st.subheader("🤖 LLM Configuration")
-        
-        # Show status of API key
-        if os.getenv("GROQ_API_KEY"):
-            st.success("✓ API Key configured")
-        else:
-            st.warning("⚠️ GROQ_API_KEY not set in environment")
-        
-        st.divider()
-        
-        # Initialize Button
-        if st.button("🚀 Connect & Initialize", use_container_width=True, type="primary"):
-            with st.spinner("Connecting to database..."):
-                success = initialize_chatbot(custom_db_params, None, None)
-                if success:
-                    st.success("✅ Connected!")
-                    st.rerun()
-        
-        # Index Button (after initialization)
-        if st.session_state.initialized:
-            if st.button("📚 Index Text Data", use_container_width=True):
-                with st.spinner("Indexing text data..."):
-                    index_data()
-                    st.success("✅ Indexed!")
-                    st.rerun()
-        
-        st.divider()
-        
-        # Status
-        st.subheader("📊 Status")
-        if st.session_state.initialized:
-            # Show database type
-            if st.session_state.custom_db_connection:
-                db_type = st.session_state.custom_db_connection.db_type.value.upper()
-            else:
-                db_type = get_db().db_type.value.upper()
-            
-            st.success(f"Database: {db_type} ✓")
-            
-            try:
-                schema = get_schema()
-                st.info(f"Tables: {len(schema.tables)}")
-            except:
-                st.warning("Schema not loaded")
-            
-            if st.session_state.indexed:
-                from rag import get_rag_engine
-                engine = get_rag_engine()
-                st.info(f"Indexed Docs: {engine.document_count}")
-        else:
-            st.warning("Not connected")
-        
-        # New Chat
-        if st.button("➕ New Chat", use_container_width=True, type="secondary"):
-            if st.session_state.memory:
-                st.session_state.memory.clear()
-            
-            st.session_state.messages = []
-            st.session_state.session_id = str(uuid.uuid4())
-            
-            current_user = st.session_state.get("user_id", "default")
-            
-            if st.session_state.custom_db_connection:
-                st.session_state.memory = create_custom_memory(
-                    st.session_state.session_id,
-                    current_user,
-                    st.session_state.custom_db_connection,
-                    st.session_state.get("llm"),
-                    st.session_state.enable_summarization,
-                    st.session_state.summary_threshold
-                )
-            elif st.session_state.initialized:
-                from memory import create_enhanced_memory
-                st.session_state.memory = create_enhanced_memory(
-                    st.session_state.session_id, 
-                    user_id=current_user,
-                    enable_summarization=st.session_state.enable_summarization,
-                    summary_threshold=st.session_state.summary_threshold
-                )
-                if st.session_state.get("llm"):
-                    st.session_state.memory.set_llm_client(st.session_state.llm)
-            
-            st.rerun()
-        
-        # Disconnect button (when using custom DB)
-        if st.session_state.initialized and st.session_state.db_source == "custom":
-            if st.button("🔌 Disconnect", use_container_width=True):
-                if st.session_state.custom_db_connection:
-                    st.session_state.custom_db_connection.close()
-                st.session_state.custom_db_connection = None
-                st.session_state.chatbot = None
-                st.session_state.initialized = False
-                st.session_state.indexed = False
-                st.session_state.memory = None
-                st.success("Disconnected!")
-                st.rerun()
-        
-        st.divider()
-        
-        # Chat History Section
-        if st.session_state.memory:
-            st.subheader("🕰️ Chat History")
-            sessions = st.session_state.memory.get_user_sessions()
-            
-            if not sessions:
-                st.caption("No previous chats found.")
-            else:
-                for session in sessions:
-                    # Highlight current session
-                    is_current = session["id"] == st.session_state.session_id
-                    icon = "🟢" if is_current else "💬"
-                    
-                    if st.button(
-                        f"{icon} {session['title']}", 
-                        key=f"hist_{session['id']}",
-                        use_container_width=True,
-                        type="secondary" if not is_current else "primary"
-                    ):
-                        if not is_current:
-                            # Load selected session
-                            st.session_state.session_id = session["id"]
-                            st.session_state.memory.session_id = session["id"]
-                            st.session_state.memory.messages = [] # Clear current state local cache
-                            
-                            # Load from DB
-                            msgs = st.session_state.memory.load_session(session["id"])
-                            st.session_state.messages = msgs
-                            
-                            # Re-populate memory object messages list for context
-                            # (We need to convert dicts back to ChatMessage objects implicitly or just rely on reload)
-                            # Actually, we should probably re-init the memory to be safe or manually populate
-                            # Let's manually populate to keep the connection valid
-                            from memory import ChatMessage
-                            st.session_state.memory.messages = [
-                                ChatMessage(
-                                    role=m['role'], 
-                                    content=m['content'], 
-                                    metadata=m.get('metadata')
-                                ) for m in msgs
-                            ]
-                            
-                            st.rerun()
-
-
-def initialize_chatbot(custom_db_params=None, api_key=None, model=None) -> bool:
-    """Initialize the chatbot with either environment or custom database."""
-    try:
-        # Get API key
-        groq_api_key = api_key or os.getenv("GROQ_API_KEY", "")
-        groq_model = model or os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
-        
-        if not groq_api_key:
-            st.error("GROQ_API_KEY not configured. Please enter your API key.")
-            return False
-        
-        # Create LLM client
-        llm = create_llm_client("groq", api_key=groq_api_key, model=groq_model)
-        
-        # Create database connection
-        if custom_db_params and st.session_state.db_source == "custom":
-            # Validate custom params
-            db_type = custom_db_params.get("db_type", "mysql")
-            
-            if True:
-                if not all([custom_db_params.get("host"), 
-                           custom_db_params.get("database"),
-                           custom_db_params.get("username")]):
-                    st.error("Please fill in all required database fields.")
-                    return False
-            
-            # Create custom config
-            db_config = create_custom_db_config(**custom_db_params)
-            
-            # Create custom connection
-            custom_connection = DatabaseConnection(db_config)
-            
-            # Test connection
-            success, msg = custom_connection.test_connection()
-            if not success:
-                st.error(f"Connection failed: {msg}")
-                return False
-            
-            st.session_state.custom_db_connection = custom_connection
-            st.session_state.custom_db_config = db_config
-            
-            # Override the global db connection for the chatbot
-            # We need to create a chatbot with this custom connection
-            from chatbot import DatabaseChatbot
-            from database.schema_introspector import SchemaIntrospector
-            from rag import get_rag_engine
-            from sql import get_sql_generator, get_sql_validator
-            from router import get_query_router
-            
-            chatbot = DatabaseChatbot.__new__(DatabaseChatbot)
-            chatbot.db = custom_connection
-            chatbot.introspector = SchemaIntrospector()
-            chatbot.introspector.db = custom_connection
-            chatbot.rag_engine = get_rag_engine()
-            chatbot.sql_generator = get_sql_generator(db_type)
-            chatbot.sql_validator = get_sql_validator()
-            chatbot.router = get_query_router()
-            chatbot.llm_client = llm
-            chatbot._schema_initialized = False
-            chatbot._rag_initialized = False
-            
-            # Set LLM client
-            chatbot.set_llm_client(llm)
-            
-            # Initialize (introspect schema)
-            schema = chatbot.introspector.introspect(force_refresh=True)
-            chatbot.sql_validator.set_allowed_tables(schema.table_names)
-            chatbot._schema_initialized = True
-            
-            st.session_state.chatbot = chatbot
-            
-        else:
-            # Use environment-based connection (existing flow)
-            chatbot = create_chatbot(llm)
-            chatbot.set_llm_client(llm)
-            
-            success, msg = chatbot.initialize()
-            if not success:
-                st.error(f"Initialization failed: {msg}")
-                return False
-            
-            st.session_state.chatbot = chatbot
-            st.session_state.custom_db_connection = None
-        
-        st.session_state.llm = llm
-        st.session_state.initialized = True
-        st.session_state.indexed = False  # Reset index status on new connection
-        
-        # Clear RAG index to ensure no data from previous DB connection persists
-        if hasattr(chatbot, 'rag_engine') and hasattr(chatbot.rag_engine, 'clear_index'):
-            chatbot.rag_engine.clear_index()
-        
-        # Create memory with appropriate connection
-        db_conn = st.session_state.custom_db_connection or get_db()
-        st.session_state.memory = create_custom_memory(
-            st.session_state.session_id,
-            st.session_state.user_id,
-            db_conn,
-            llm,
-            st.session_state.enable_summarization,
-            st.session_state.summary_threshold
-        )
-        
-        return True
-        
-    except Exception as e:
-        st.error(f"Error: {str(e)}")
-        import traceback
-        st.error(traceback.format_exc())
-        return False
-
-
-def index_data():
-    """Index text data from the database."""
-    if st.session_state.chatbot:
-        progress = st.progress(0)
-        status = st.empty()
-        
-        # Get schema from the correct introspector
-        schema = st.session_state.chatbot.introspector.introspect()
-        total_tables = len(schema.tables)
-        indexed = 0
-        
-        def progress_callback(table_name, docs):
-            nonlocal indexed
-            indexed += 1
-            progress.progress(indexed / total_tables)
-            status.text(f"Indexed {table_name}: {docs} documents")
-        
-        total_docs = st.session_state.chatbot.index_text_data(progress_callback)
-        st.session_state.indexed = True
-        status.text(f"Total: {total_docs} documents indexed")
-
-
-def render_schema_explorer():
-    """Render schema explorer in an expander."""
-    if not st.session_state.initialized:
-        return
-    
-    with st.expander("📋 Database Schema", expanded=False):
-        try:
-            schema = st.session_state.chatbot.introspector.introspect()
-            
-            tab_list, tab_erd = st.tabs(["📋 Table List", "🕸️ Schema Diagram"])
-            
-            with tab_list:
-                st.markdown("Uncheck tables to exclude them from the chat context.")
-                
-                for table_name, table_info in schema.tables.items():
-                    col1, col2 = st.columns([0.05, 0.95])
-                    
-                    with col1:
-                        is_active = table_name not in st.session_state.ignored_tables
-                        active = st.checkbox(
-                            "Use", 
-                            value=is_active, 
-                            key=f"use_{table_name}", 
-                            label_visibility="collapsed",
-                            help=f"Include {table_name} in chat analysis"
-                        )
-                        
-                        if not active:
-                            st.session_state.ignored_tables.add(table_name)
-                        else:
-                            st.session_state.ignored_tables.discard(table_name)
-                    
-                    with col2:
-                        with st.container():
-                            st.markdown(f"**{table_name}** ({table_info.row_count or '?'} rows)")
-                            
-                            cols = []
-                            for col in table_info.columns:
-                                pk = "🔑" if col.is_primary_key else ""
-                                txt = "📝" if col.is_text_type else ""
-                                cols.append(f"`{col.name}` {col.data_type} {pk}{txt}")
-                            
-                            st.caption(" | ".join(cols))
-                            st.divider()
-            
-            with tab_erd:
-                if len(schema.tables) > 50:
-                    st.warning("⚠️ Too many tables to visualize effectively (limit: 50).")
-                else:
-                    try:
-                        # Build Graphviz DOT string
-                        dot = ['digraph Database {']
-                        dot.append('  rankdir=LR;')
-                        dot.append('  node [shape=box, style="filled,rounded", fillcolor="#f0f2f6", fontname="Arial", fontsize=10];')
-                        dot.append('  edge [fontname="Arial", fontsize=9, color="#666666"];')
-                        
-                        # Add nodes (tables)
-                        for table_name in schema.tables:
-                            if table_name not in st.session_state.ignored_tables:
-                                dot.append(f'  "{table_name}" [label="{table_name}", fillcolor="#e1effe", color="#1e40af"];')
-                            else:
-                                dot.append(f'  "{table_name}" [label="{table_name} (ignored)", fillcolor="#f3f4f6", color="#9ca3af", fontcolor="#9ca3af"];')
-                        
-                        # Add edges (relationships)
-                        has_edges = False
-                        for table_name, table_info in schema.tables.items():
-                            for col_name, ref_str in table_info.foreign_keys.items():
-                                # ref_str format: "referenced_table.referenced_column"
-                                if "." in ref_str:
-                                    ref_table = ref_str.split(".")[0]
-                                    # specific_col = ref_str.split(".")[1]
-                                    
-                                    # Only draw if both tables exist in our schema list
-                                    if ref_table in schema.tables:
-                                        dot.append(f'  "{table_name}" -> "{ref_table}" [label="{col_name}"];')
-                                        has_edges = True
-                        
-                        dot.append('}')
-                        graph_code = "\n".join(dot)
-                        st.graphviz_chart(graph_code, width="stretch")
-                        
-                        if not has_edges:
-                            st.info("No foreign key relationships detected in the schema metadata.")
-                            
-                    except Exception as e:
-                        st.error(f"Could not render diagram: {e}")
-
-        except Exception as e:
-            st.error(f"Error loading schema: {e}")
-
-
-def render_chat_interface():
-    """Render the main chat interface."""
-    st.title("🤖 OnceDataBot")
-    st.caption("Schema-agnostic chatbot • MySQL | PostgreSQL • Powered by Groq (FREE!)")
-    
-    # Schema explorer
-    render_schema_explorer()
-    
-    # Chat container
-    chat_container = st.container()
-    
-    with chat_container:
-        # Display messages
-        for i, msg in enumerate(st.session_state.messages):
-            with st.chat_message(msg["role"]):
-                # Create columns for message and favorite button
-                msg_col, fav_col = st.columns([0.95, 0.05])
-                
-                with msg_col:
-                    st.markdown(msg["content"])
-                
-                with fav_col:
-                    # Favorite button for assistant messages
-                    if msg["role"] == "assistant":
-                        is_favorited = i in st.session_state.favorites
-                        if st.button(
-                            "⭐" if is_favorited else "☆",
-                            key=f"fav_{i}",
-                            help="Click to favorite/unfavorite this response"
-                        ):
-                            if is_favorited:
-                                st.session_state.favorites.remove(i)
-                            else:
-                                st.session_state.favorites.append(i)
-                            st.rerun()
-                
-                # Show metadata for assistant messages
-                if msg["role"] == "assistant" and "metadata" in msg:
-                    meta = msg["metadata"]
-                    
-                    # Show token usage in a dropdown expander
-                    if "token_usage" in meta:
-                        usage = meta["token_usage"]
-                        total = usage.get('total', 0)
-                        
-                        with st.expander(f"📊 Token Usage ({total:,} total)", expanded=False):
-                            # Create styled token usage boxes using columns
-                            st.markdown("""
-                            <style>
-                            .token-box {
-                                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-                                border-radius: 12px;
-                                padding: 12px 16px;
-                                color: white;
-                                text-align: center;
-                                box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3);
-                                margin: 4px 0;
-                            }
-                            .token-box-input {
-                                background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
-                                box-shadow: 0 4px 15px rgba(17, 153, 142, 0.3);
-                            }
-                            .token-box-output {
-                                background: linear-gradient(135deg, #ee0979 0%, #ff6a00 100%);
-                                box-shadow: 0 4px 15px rgba(238, 9, 121, 0.3);
-                            }
-                            .token-box-total {
-                                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-                                box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3);
-                            }
-                            .token-label {
-                                font-size: 11px;
-                                text-transform: uppercase;
-                                letter-spacing: 1px;
-                                opacity: 0.9;
-                                margin-bottom: 4px;
-                            }
-                            .token-value {
-                                font-size: 20px;
-                                font-weight: 700;
-                            }
-                            </style>
-                            """, unsafe_allow_html=True)
-                            
-                            col1, col2, col3 = st.columns(3)
-                            
-                            with col1:
-                                st.markdown(f"""
-                                <div class="token-box token-box-input">
-                                    <div class="token-label">📥 Input Tokens</div>
-                                    <div class="token-value">{usage.get('input', 0):,}</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                            
-                            with col2:
-                                st.markdown(f"""
-                                <div class="token-box token-box-output">
-                                    <div class="token-label">📤 Output Tokens</div>
-                                    <div class="token-value">{usage.get('output', 0):,}</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                            
-                            with col3:
-                                st.markdown(f"""
-                                <div class="token-box token-box-total">
-                                    <div class="token-label">📊 Total Tokens</div>
-                                    <div class="token-value">{usage.get('total', 0):,}</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                    
-                    if meta.get("query_type"):
-                        # Show query type and execution time on same line
-                        info_text = f"Query type: {meta['query_type']}"
-                        if meta.get("execution_time"):
-                            info_text += f" • ⏱️ {meta['execution_time']:.2f}s"
-                        st.caption(info_text)
-                        
-                    # SQL Query expander
-                    if meta.get("sql_query"):
-                        with st.expander("🛠️ SQL Query & Details"):
-                            st.code(meta["sql_query"], language="sql")
-                            
-                    # Visualizations and CSV export
-                    if meta.get("sql_results"):
-                        # Only render viz if we have results
-                        render_visualization(meta["sql_results"], f"viz_{i}")
-                        
-                        # CSV Export button
-                        csv_data = export_results_to_csv(meta["sql_results"])
-                        if csv_data:
-                            st.download_button(
-                                label="📊 Export to CSV",
-                                data=csv_data,
-                                file_name=f"query_results_{i}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
-                                mime="text/csv",
-                                key=f"csv_export_{i}",
-                                help="Download query results as CSV file"
-                            )
-    
-    # Chat input
-    if prompt := st.chat_input("Ask about your data..."):
-        if not st.session_state.initialized:
-            st.error("Please connect to a database first!")
-            return
-        
-        # Add user message
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        
-        # Calculate memory context for display? No, just render user msg
-        with st.chat_message("user"):
-            st.markdown(prompt)
-        
-        # Get response
-        with st.spinner("Thinking..."):
-            try:
-                # Add memory interaction
-                if st.session_state.memory:
-                    st.session_state.memory.add_message("user", prompt)
-                
-                # Track execution time
-                start_time = time.time()
-                
-                response = st.session_state.chatbot.chat(
-                    prompt, 
-                    st.session_state.memory,
-                    ignored_tables=list(st.session_state.ignored_tables),
-                    language=st.session_state.response_language
-                )
-                
-                execution_time = time.time() - start_time
-                
-
-                
-                # Create metadata dict
-                metadata = {
-                    "query_type": response.query_type,
-                    "sql_query": response.sql_query,
-                    "sql_results": response.sql_results,
-                    "token_usage": response.token_usage,
-                    "execution_time": execution_time
-                }
-                
-                # Save to session state
-                st.session_state.messages.append({
-                    "role": "assistant",
-                    "content": response.answer,
-                    "metadata": metadata
-                })
-                
-                # Set flag to auto-read the latest response
-                st.session_state.auto_read_latest = True
-                
-                # Save to active memory
-                if st.session_state.memory:
-                    st.session_state.memory.add_message("assistant", response.answer)
-                
-                st.rerun()
-                
-            except Exception as e:
-                st.error(f"An error occurred: {e}")
-                import traceback
-                st.error(traceback.format_exc())
-
-
-def main():
-    """Main application entry point."""
-    init_session_state()
-    
-    # Auto-connect to environment database on first load
-    if "auto_connect_attempted" not in st.session_state:
-        st.session_state.auto_connect_attempted = True
-        if st.session_state.db_source == "environment":
-            success = initialize_chatbot()
-            if success:
-                st.toast("✅ Auto-connected to database!")
-
-    render_sidebar()
-    render_chat_interface()
-
-
-if __name__ == "__main__":
-    main()
+"""
+Schema-Agnostic Database Chatbot - Streamlit Application
+
+A production-grade chatbot that connects to ANY database
+(MySQL, PostgreSQL, SQLite) and provides intelligent querying 
+through RAG and Text-to-SQL.
+
+Uses Groq for FREE LLM inference!
+"""
+
+import os
+from pathlib import Path
+
+# Load .env FIRST before any other imports
+from dotenv import load_dotenv
+load_dotenv(Path(__file__).parent / ".env")
+
+import streamlit as st
+import uuid
+import time
+import io
+import csv
+import base64
+import pandas as pd
+from datetime import datetime
+
+# Page config must be first
+st.set_page_config(
+    page_title="OnceDataBot",
+    page_icon="🤖",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# Imports
+from config import config, DatabaseConfig, DatabaseType
+from database import get_db, get_schema, get_introspector
+from database.connection import DatabaseConnection
+from llm import create_llm_client
+from chatbot import create_chatbot, DatabaseChatbot
+from memory import ChatMemory, EnhancedChatMemory
+from viz_utils import render_visualization
+
+
+
+
+
+
+# Groq models (all FREE!)
+GROQ_MODELS = [
+    "llama-3.3-70b-versatile",
+    "llama-3.1-8b-instant",
+    "mixtral-8x7b-32768",
+    "gemma2-9b-it"
+]
+
+# Database types
+DB_TYPES = {
+    "MySQL": "mysql",
+    "PostgreSQL": "postgresql",
+    "SQLite": "sqlite"
+}
+
+# Supported languages for multi-language responses
+SUPPORTED_LANGUAGES = {
+    "English": "en",
+    "हिन्दी (Hindi)": "hi",
+    "Español (Spanish)": "es",
+    "Français (French)": "fr",
+    "Deutsch (German)": "de",
+    "中文 (Chinese)": "zh",
+    "日本語 (Japanese)": "ja",
+    "한국어 (Korean)": "ko",
+    "Português (Portuguese)": "pt",
+    "العربية (Arabic)": "ar",
+    "Русский (Russian)": "ru",
+    "Italiano (Italian)": "it",
+    "Nederlands (Dutch)": "nl",
+    "தமிழ் (Tamil)": "ta",
+    "తెలుగు (Telugu)": "te",
+    "मराठी (Marathi)": "mr",
+    "বাংলা (Bengali)": "bn",
+    "ગુજરાতી (Gujarati)": "gu"
+}
+
+
+def create_custom_db_config(db_type: str, **kwargs) -> DatabaseConfig:
+    """Create a custom database configuration from user input."""
+    return DatabaseConfig(
+        db_type=DatabaseType(db_type),
+        host=kwargs.get("host", ""),
+        port=kwargs.get("port", 3306 if db_type == "mysql" else 5432),
+        database=kwargs.get("database", ""),
+        username=kwargs.get("username", ""),
+        password=kwargs.get("password", ""),
+        ssl_ca=kwargs.get("ssl_ca", None)
+    )
+
+
+def create_custom_memory(session_id: str, user_id: str, db_connection, llm_client=None, 
+                         enable_summarization=True, summary_threshold=10) -> EnhancedChatMemory:
+    """Create enhanced memory with a custom database connection."""
+    return EnhancedChatMemory(
+        session_id=session_id,
+        user_id=user_id,
+        max_messages=20,
+        db_connection=db_connection,
+        llm_client=llm_client,
+        enable_summarization=enable_summarization,
+        summary_threshold=summary_threshold
+    )
+
+
+def init_session_state():
+    """Initialize Streamlit session state."""
+    if "session_id" not in st.session_state:
+        st.session_state.session_id = str(uuid.uuid4())
+    
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    if "chatbot" not in st.session_state:
+        st.session_state.chatbot = None
+    
+    if "initialized" not in st.session_state:
+        st.session_state.initialized = False
+    
+    if "user_id" not in st.session_state:
+        st.session_state.user_id = "default"
+    
+    if "enable_summarization" not in st.session_state:
+        st.session_state.enable_summarization = True
+    
+    if "summary_threshold" not in st.session_state:
+        st.session_state.summary_threshold = 10
+    
+    if "memory" not in st.session_state:
+        st.session_state.memory = None
+        
+    if "indexed" not in st.session_state:
+        st.session_state.indexed = False
+    
+    if "db_source" not in st.session_state:
+        st.session_state.db_source = "environment"  # "environment" or "custom"
+    
+    if "custom_db_config" not in st.session_state:
+        st.session_state.custom_db_config = None
+    
+    if "custom_db_connection" not in st.session_state:
+        st.session_state.custom_db_connection = None
+        
+    if "ignored_tables" not in st.session_state:
+        st.session_state.ignored_tables = set()
+    
+    if "response_language" not in st.session_state:
+        st.session_state.response_language = "English"
+    
+    if "favorites" not in st.session_state:
+        st.session_state.favorites = []  # List of message indices that are favorited
+
+
+def export_results_to_csv(results: list) -> str:
+    """Convert SQL results to CSV format and return as downloadable string."""
+    if not results:
+        return ""
+    
+    output = io.StringIO()
+    writer = csv.DictWriter(output, fieldnames=results[0].keys())
+    writer.writeheader()
+    writer.writerows(results)
+    return output.getvalue()
+
+
+def export_chat_to_text() -> str:
+    """Export chat messages to text format."""
+    if not st.session_state.messages:
+        return "No messages to export."
+    
+    lines = []
+    lines.append("=" * 50)
+    lines.append(f"OnceDataBot Chat Export")
+    lines.append(f"Exported: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
+    lines.append(f"User: {st.session_state.user_id}")
+    lines.append("=" * 50)
+    lines.append("")
+    
+    for i, msg in enumerate(st.session_state.messages):
+        role = "🧑 User" if msg["role"] == "user" else "🤖 Assistant"
+        is_favorited = "⭐ " if i in st.session_state.favorites else ""
+        lines.append(f"{is_favorited}{role}:")
+        lines.append(msg["content"])
+        
+        if msg["role"] == "assistant" and "metadata" in msg:
+            meta = msg["metadata"]
+            if meta.get("sql_query"):
+                lines.append(f"\n📝 SQL Query: {meta['sql_query']}")
+            if meta.get("query_type"):
+                lines.append(f"📌 Query Type: {meta['query_type']}")
+            if meta.get("execution_time"):
+                lines.append(f"⏱️ Execution Time: {meta['execution_time']:.2f}s")
+        
+        lines.append("-" * 40)
+        lines.append("")
+    
+    return "\n".join(lines)
+
+
+def render_copy_button(text: str, key: str):
+    """Render a copy to clipboard button using Streamlit."""
+    # Using a workaround with st.code which has built-in copy
+    st.code(text, language="sql")
+
+
+def render_database_config():
+    """Render database configuration section in sidebar."""
+    st.subheader("🗄️ Database Configuration")
+    
+    # Database source selection
+    db_source = st.radio(
+        "Database Source",
+        options=["Use Environment Variables", "Custom Database"],
+        index=0 if st.session_state.db_source == "environment" else 1,
+        key="db_source_radio",
+        help="Choose to use .env settings or enter custom credentials"
+    )
+    
+    st.session_state.db_source = "environment" if db_source == "Use Environment Variables" else "custom"
+    
+    if st.session_state.db_source == "environment":
+        # Show current environment config
+        current_db_type = config.database.db_type.value.upper()
+        st.info(f"📌 Using {current_db_type} from environment")
+        st.caption(f"Host: {config.database.host}")
+        return None
+    
+    else:
+        # Custom database configuration
+        st.markdown("##### Enter Database Credentials")
+        
+        # Database type selector
+        db_type_label = st.selectbox(
+            "Database Type",
+            options=list(DB_TYPES.keys()),
+            index=0,
+            key="custom_db_type"
+        )
+        db_type = DB_TYPES[db_type_label]
+        
+        if db_type == "sqlite":
+            # SQLite only needs a file path
+            database = st.text_input(
+                "SQLite Database File",
+                value="ingested_data.db",
+                key="db_sqlite_path",
+                help="Path to the .db file (will be created if it doesn't exist)"
+            )
+            return {
+                "db_type": db_type,
+                "database": database
+            }
+        
+        else:  # MySQL or PostgreSQL
+            # MySQL or PostgreSQL
+            col1, col2 = st.columns([3, 1])
+            with col1:
+                host = st.text_input(
+                    "Host",
+                    value="",
+                    key="db_host_input",
+                    placeholder="your-database-host.com"
+                )
+            with col2:
+                default_port = 3306 if db_type == "mysql" else 5432
+                port = st.number_input(
+                    "Port",
+                    value=default_port,
+                    min_value=1,
+                    max_value=65535,
+                    key="db_port_input"
+                )
+            
+            database = st.text_input(
+                "Database Name",
+                value="",
+                key="db_name_input",
+                placeholder="your_database"
+            )
+            
+            username = st.text_input(
+                "Username",
+                value="",
+                key="db_user_input",
+                placeholder="your_username"
+            )
+            
+            password = st.text_input(
+                "Password",
+                value="",
+                type="password",
+                key="db_pass_input"
+            )
+            
+            # Optional SSL
+            with st.expander("🔒 SSL Settings (Optional)"):
+                ssl_ca = st.text_input(
+                    "SSL CA Certificate Path",
+                    value="",
+                    key="ssl_ca_input",
+                    help="Path to SSL CA certificate file (for cloud databases like Aiven)"
+                )
+            
+            return {
+                "db_type": db_type,
+                "host": host,
+                "port": int(port),
+                "database": database,
+                "username": username,
+                "password": password,
+                "ssl_ca": ssl_ca if ssl_ca else None
+            }
+
+
+def render_sidebar():
+    """Render the configuration sidebar."""
+    with st.sidebar:
+        st.title("⚙️ Settings")
+        
+        # Session Dashboard
+        if st.session_state.messages:
+            st.markdown("### 📊 Session Stats")
+            
+            # Calculate stats
+            total_msgs = len(st.session_state.messages)
+            assistant_msgs = [m for m in st.session_state.messages if m.get("role") == "assistant"]
+            sql_queries = sum(1 for m in assistant_msgs if m.get("metadata", {}).get("sql_query"))
+            
+            total_tokens = 0
+            exec_times = []
+            for m in assistant_msgs:
+                meta = m.get("metadata", {})
+                total_tokens += meta.get("token_usage", {}).get("total", 0)
+                if meta.get("execution_time"):
+                    exec_times.append(meta["execution_time"])
+            
+            avg_time = sum(exec_times) / len(exec_times) if exec_times else 0
+            
+            col_s1, col_s2 = st.columns(2)
+            col_s1.metric("Queries", sql_queries)
+            col_s2.metric("Tokens", f"{total_tokens:,}")
+            st.caption(f"⏱️ Avg Time: {avg_time:.2f}s | 💬 Msgs: {total_msgs}")
+            st.divider()
+        
+
+        
+        # User Profile
+        st.subheader("👤 User Profile")
+        user_id = st.text_input(
+            "User ID / Name", 
+            value=st.session_state.get("user_id", "default"),
+            key="user_id_input",
+            help="Your unique ID for private memory storage"
+        )
+        if user_id != st.session_state.get("user_id"):
+            st.session_state.user_id = user_id
+            st.session_state.session_id = str(uuid.uuid4())
+            st.session_state.messages = []
+            
+            # Recreate memory for new user
+            if st.session_state.custom_db_connection:
+                st.session_state.memory = create_custom_memory(
+                    st.session_state.session_id,
+                    user_id,
+                    st.session_state.custom_db_connection,
+                    st.session_state.get("llm"),
+                    st.session_state.enable_summarization,
+                    st.session_state.summary_threshold
+                )
+            elif st.session_state.initialized:
+                from memory import create_enhanced_memory
+                st.session_state.memory = create_enhanced_memory(
+                    st.session_state.session_id,
+                    user_id=user_id,
+                    enable_summarization=st.session_state.enable_summarization,
+                    summary_threshold=st.session_state.summary_threshold
+                )
+            
+            if st.session_state.memory:
+                st.session_state.memory.clear_user_history()
+            st.rerun()
+        
+        st.divider()
+        
+        # Language Selection
+        st.subheader("🌐 Response Language")
+        selected_language = st.selectbox(
+            "Select Language",
+            options=list(SUPPORTED_LANGUAGES.keys()),
+            index=list(SUPPORTED_LANGUAGES.keys()).index(st.session_state.response_language),
+            key="language_selector",
+            help="Choose the language for chatbot responses"
+        )
+        if selected_language != st.session_state.response_language:
+            st.session_state.response_language = selected_language
+            st.toast(f"🌐 Language changed to {selected_language}")
+        
+        st.divider()
+        
+        if st.session_state.messages:
+            st.download_button(
+                label="📄 Export Chat",
+                data=export_chat_to_text(),
+                file_name=f"chat_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
+                mime="text/plain",
+                use_container_width=True,
+                help="Download your chat conversation as a text file"
+            )
+        
+        st.divider()
+        
+        # CSV Ingestion Section
+        st.subheader("📥 Ingest CSV Data")
+        uploaded_files = st.file_uploader(
+            "Upload CSV(s) to create database",
+            type=["csv"],
+            accept_multiple_files=True,
+            help="Your CSVs will be converted to tables in a local SQLite database"
+        )
+        
+        if uploaded_files:
+            if st.button("🚀 Upload & Initialize", use_container_width=True):
+                with st.spinner("Processing CSVs..."):
+                    success_count = 0
+                    table_names = []
+                    for uploaded_file in uploaded_files:
+                        success, name, rows = ingest_csv(uploaded_file)
+                        if success:
+                            success_count += 1
+                            table_names.append(name)
+                        else:
+                            st.error(f"Failed to ingest {uploaded_file.name}: {name}")
+                    
+                    if success_count > 0:
+                        st.success(f"Successfully ingested {success_count} file(s) as tables: {', '.join(table_names)}")
+                        
+                        # Now initialize chatbot with this SQLite DB
+                        sqlite_params = {
+                            "db_type": "sqlite",
+                            "database": "ingested_data.db"
+                        }
+                        # Temporarily set db_source to custom for initialization
+                        old_source = st.session_state.db_source
+                        st.session_state.db_source = "custom"
+                        init_success = initialize_chatbot(sqlite_params, None, None)
+                        if not init_success:
+                            st.session_state.db_source = old_source
+                        else:
+                            st.rerun()
+        
+        st.divider()
+        
+        # Database Configuration
+        custom_db_params = render_database_config()
+        
+        st.divider()
+        
+        # LLM Configuration
+        st.subheader("🤖 LLM Configuration")
+        
+        # Show status of API key
+        if os.getenv("GROQ_API_KEY"):
+            st.success("✓ API Key configured")
+        else:
+            st.warning("⚠️ GROQ_API_KEY not set in environment")
+        
+        st.divider()
+        
+        # Initialize Button
+        if st.button("🚀 Connect & Initialize", use_container_width=True, type="primary"):
+            with st.spinner("Connecting to database..."):
+                success = initialize_chatbot(custom_db_params, None, None)
+                if success:
+                    st.success("✅ Connected!")
+                    st.rerun()
+        
+        # Index Button (after initialization)
+        if st.session_state.initialized:
+            if st.button("📚 Index Text Data", use_container_width=True):
+                with st.spinner("Indexing text data..."):
+                    index_data()
+                    st.success("✅ Indexed!")
+                    st.rerun()
+        
+        st.divider()
+        
+        # Status
+        st.subheader("📊 Status")
+        if st.session_state.initialized:
+            # Show database type
+            if st.session_state.custom_db_connection:
+                db_type = st.session_state.custom_db_connection.db_type.value.upper()
+            else:
+                db_type = get_db().db_type.value.upper()
+            
+            st.success(f"Database: {db_type} ✓")
+            
+            try:
+                schema = get_schema()
+                st.info(f"Tables: {len(schema.tables)}")
+            except:
+                st.warning("Schema not loaded")
+            
+            if st.session_state.indexed:
+                from rag import get_rag_engine
+                engine = get_rag_engine()
+                st.info(f"Indexed Docs: {engine.document_count}")
+        else:
+            st.warning("Not connected")
+        
+        # New Chat
+        if st.button("➕ New Chat", use_container_width=True, type="secondary"):
+            if st.session_state.memory:
+                st.session_state.memory.clear()
+            
+            st.session_state.messages = []
+            st.session_state.session_id = str(uuid.uuid4())
+            
+            current_user = st.session_state.get("user_id", "default")
+            
+            if st.session_state.custom_db_connection:
+                st.session_state.memory = create_custom_memory(
+                    st.session_state.session_id,
+                    current_user,
+                    st.session_state.custom_db_connection,
+                    st.session_state.get("llm"),
+                    st.session_state.enable_summarization,
+                    st.session_state.summary_threshold
+                )
+            elif st.session_state.initialized:
+                from memory import create_enhanced_memory
+                st.session_state.memory = create_enhanced_memory(
+                    st.session_state.session_id, 
+                    user_id=current_user,
+                    enable_summarization=st.session_state.enable_summarization,
+                    summary_threshold=st.session_state.summary_threshold
+                )
+                if st.session_state.get("llm"):
+                    st.session_state.memory.set_llm_client(st.session_state.llm)
+            
+            st.rerun()
+        
+        # Disconnect button (when using custom DB)
+        if st.session_state.initialized and st.session_state.db_source == "custom":
+            if st.button("🔌 Disconnect", use_container_width=True):
+                if st.session_state.custom_db_connection:
+                    st.session_state.custom_db_connection.close()
+                st.session_state.custom_db_connection = None
+                st.session_state.chatbot = None
+                st.session_state.initialized = False
+                st.session_state.indexed = False
+                st.session_state.memory = None
+                st.success("Disconnected!")
+                st.rerun()
+        
+        st.divider()
+        
+        # Chat History Section
+        if st.session_state.memory:
+            st.subheader("🕰️ Chat History")
+            sessions = st.session_state.memory.get_user_sessions()
+            
+            if not sessions:
+                st.caption("No previous chats found.")
+            else:
+                for session in sessions:
+                    # Highlight current session
+                    is_current = session["id"] == st.session_state.session_id
+                    icon = "🟢" if is_current else "💬"
+                    
+                    if st.button(
+                        f"{icon} {session['title']}", 
+                        key=f"hist_{session['id']}",
+                        use_container_width=True,
+                        type="secondary" if not is_current else "primary"
+                    ):
+                        if not is_current:
+                            # Load selected session
+                            st.session_state.session_id = session["id"]
+                            st.session_state.memory.session_id = session["id"]
+                            st.session_state.memory.messages = [] # Clear current state local cache
+                            
+                            # Load from DB
+                            msgs = st.session_state.memory.load_session(session["id"])
+                            st.session_state.messages = msgs
+                            
+                            # Re-populate memory object messages list for context
+                            # (We need to convert dicts back to ChatMessage objects implicitly or just rely on reload)
+                            # Actually, we should probably re-init the memory to be safe or manually populate
+                            # Let's manually populate to keep the connection valid
+                            from memory import ChatMessage
+                            st.session_state.memory.messages = [
+                                ChatMessage(
+                                    role=m['role'], 
+                                    content=m['content'], 
+                                    metadata=m.get('metadata')
+                                ) for m in msgs
+                            ]
+                            
+                            st.rerun()
+
+
+def initialize_chatbot(custom_db_params=None, api_key=None, model=None) -> bool:
+    """Initialize the chatbot with either environment or custom database."""
+    try:
+        # Get API key
+        groq_api_key = api_key or os.getenv("GROQ_API_KEY", "")
+        groq_model = model or os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
+        
+        if not groq_api_key:
+            st.error("GROQ_API_KEY not configured. Please enter your API key.")
+            return False
+        
+        # Create LLM client
+        llm = create_llm_client("groq", api_key=groq_api_key, model=groq_model)
+        
+        # Create database connection
+        if custom_db_params and st.session_state.db_source == "custom":
+            # Validate custom params
+            db_type = custom_db_params.get("db_type", "mysql")
+            
+            if db_type != "sqlite":
+                if not all([custom_db_params.get("host"), 
+                           custom_db_params.get("database"),
+                           custom_db_params.get("username")]):
+                    st.error("Please fill in all required database fields.")
+                    return False
+            else:
+                if not custom_db_params.get("database"):
+                    st.error("Please specify a SQLite database file path.")
+                    return False
+            
+            # Create custom config
+            db_config = create_custom_db_config(**custom_db_params)
+            
+            # Create custom connection
+            custom_connection = DatabaseConnection(db_config)
+            
+            # Test connection
+            success, msg = custom_connection.test_connection()
+            if not success:
+                st.error(f"Connection failed: {msg}")
+                return False
+            
+            st.session_state.custom_db_connection = custom_connection
+            st.session_state.custom_db_config = db_config
+            
+            # Override the global db connection for the chatbot
+            # We need to create a chatbot with this custom connection
+            from chatbot import DatabaseChatbot
+            from database.schema_introspector import SchemaIntrospector
+            from rag import get_rag_engine
+            from sql import get_sql_generator, get_sql_validator
+            from router import get_query_router
+            
+            chatbot = DatabaseChatbot.__new__(DatabaseChatbot)
+            chatbot.db = custom_connection
+            chatbot.introspector = SchemaIntrospector()
+            chatbot.introspector.db = custom_connection
+            chatbot.rag_engine = get_rag_engine()
+            chatbot.sql_generator = get_sql_generator(db_type)
+            chatbot.sql_validator = get_sql_validator()
+            chatbot.router = get_query_router()
+            chatbot.llm_client = llm
+            chatbot._schema_initialized = False
+            chatbot._rag_initialized = False
+            
+            # Set LLM client
+            chatbot.set_llm_client(llm)
+            
+            # Initialize (introspect schema)
+            schema = chatbot.introspector.introspect(force_refresh=True)
+            chatbot.sql_validator.set_allowed_tables(schema.table_names)
+            chatbot._schema_initialized = True
+            
+            st.session_state.chatbot = chatbot
+            
+        else:
+            # Use environment-based connection (existing flow)
+            chatbot = create_chatbot(llm)
+            chatbot.set_llm_client(llm)
+            
+            success, msg = chatbot.initialize()
+            if not success:
+                st.error(f"Initialization failed: {msg}")
+                return False
+            
+            st.session_state.chatbot = chatbot
+            st.session_state.custom_db_connection = None
+        
+        st.session_state.llm = llm
+        st.session_state.initialized = True
+        st.session_state.indexed = False  # Reset index status on new connection
+        
+        # Clear RAG index to ensure no data from previous DB connection persists
+        if hasattr(chatbot, 'rag_engine') and hasattr(chatbot.rag_engine, 'clear_index'):
+            chatbot.rag_engine.clear_index()
+        
+        # Create memory with appropriate connection
+        db_conn = st.session_state.custom_db_connection or get_db()
+        st.session_state.memory = create_custom_memory(
+            st.session_state.session_id,
+            st.session_state.user_id,
+            db_conn,
+            llm,
+            st.session_state.enable_summarization,
+            st.session_state.summary_threshold
+        )
+        
+        return True
+        
+    except Exception as e:
+        st.error(f"Error: {str(e)}")
+        import traceback
+        st.error(traceback.format_exc())
+        return False
+
+
+def ingest_csv(uploaded_file):
+    """Ingest a CSV file into a SQLite database."""
+    from sqlalchemy import create_engine
+    
+    try:
+        # 1. Read CSV
+        # Reset file pointer to beginning in case it was read before
+        uploaded_file.seek(0)
+        df = pd.read_csv(uploaded_file)
+        
+        # 2. Clean table name from filename
+        table_name = Path(uploaded_file.name).stem.replace(" ", "_").replace("-", "_").lower()
+        # Ensure it starts with a letter and only contains alphanumeric/underscore
+        table_name = "".join([c for c in table_name if c.isalnum() or c == "_"])
+        if not table_name[0].isalpha():
+            table_name = "t_" + table_name
+            
+        # 3. Create/Connect to SQLite DB
+        db_path = "ingested_data.db"
+        engine = create_engine(f"sqlite:///{db_path}")
+        
+        # 4. Write to DB
+        df.to_sql(table_name, engine, if_exists='replace', index=False)
+        
+        return True, table_name, len(df)
+    except Exception as e:
+        return False, str(e), 0
+
+
+def index_data():
+    """Index text data from the database."""
+    if st.session_state.chatbot:
+        progress = st.progress(0)
+        status = st.empty()
+        
+        # Get schema from the correct introspector
+        schema = st.session_state.chatbot.introspector.introspect()
+        total_tables = len(schema.tables)
+        indexed = 0
+        
+        def progress_callback(table_name, docs):
+            nonlocal indexed
+            indexed += 1
+            progress.progress(indexed / total_tables)
+            status.text(f"Indexed {table_name}: {docs} documents")
+        
+        total_docs = st.session_state.chatbot.index_text_data(progress_callback)
+        st.session_state.indexed = True
+        status.text(f"Total: {total_docs} documents indexed")
+
+
+def render_schema_explorer():
+    """Render schema explorer in an expander."""
+    if not st.session_state.initialized:
+        return
+    
+    with st.expander("📋 Database Schema", expanded=False):
+        try:
+            schema = st.session_state.chatbot.introspector.introspect()
+            
+            tab_list, tab_erd = st.tabs(["📋 Table List", "🕸️ Schema Diagram"])
+            
+            with tab_list:
+                st.markdown("Uncheck tables to exclude them from the chat context.")
+                
+                for table_name, table_info in schema.tables.items():
+                    col1, col2 = st.columns([0.05, 0.95])
+                    
+                    with col1:
+                        is_active = table_name not in st.session_state.ignored_tables
+                        active = st.checkbox(
+                            "Use", 
+                            value=is_active, 
+                            key=f"use_{table_name}", 
+                            label_visibility="collapsed",
+                            help=f"Include {table_name} in chat analysis"
+                        )
+                        
+                        if not active:
+                            st.session_state.ignored_tables.add(table_name)
+                        else:
+                            st.session_state.ignored_tables.discard(table_name)
+                    
+                    with col2:
+                        with st.container():
+                            st.markdown(f"**{table_name}** ({table_info.row_count or '?'} rows)")
+                            
+                            cols = []
+                            for col in table_info.columns:
+                                pk = "🔑" if col.is_primary_key else ""
+                                txt = "📝" if col.is_text_type else ""
+                                cols.append(f"`{col.name}` {col.data_type} {pk}{txt}")
+                            
+                            st.caption(" | ".join(cols))
+                            st.divider()
+            
+            with tab_erd:
+                if len(schema.tables) > 50:
+                    st.warning("⚠️ Too many tables to visualize effectively (limit: 50).")
+                else:
+                    try:
+                        # Build Graphviz DOT string
+                        dot = ['digraph Database {']
+                        dot.append('  rankdir=LR;')
+                        dot.append('  node [shape=box, style="filled,rounded", fillcolor="#f0f2f6", fontname="Arial", fontsize=10];')
+                        dot.append('  edge [fontname="Arial", fontsize=9, color="#666666"];')
+                        
+                        # Add nodes (tables)
+                        for table_name in schema.tables:
+                            if table_name not in st.session_state.ignored_tables:
+                                dot.append(f'  "{table_name}" [label="{table_name}", fillcolor="#e1effe", color="#1e40af"];')
+                            else:
+                                dot.append(f'  "{table_name}" [label="{table_name} (ignored)", fillcolor="#f3f4f6", color="#9ca3af", fontcolor="#9ca3af"];')
+                        
+                        # Add edges (relationships)
+                        has_edges = False
+                        for table_name, table_info in schema.tables.items():
+                            for col_name, ref_str in table_info.foreign_keys.items():
+                                # ref_str format: "referenced_table.referenced_column"
+                                if "." in ref_str:
+                                    ref_table = ref_str.split(".")[0]
+                                    # specific_col = ref_str.split(".")[1]
+                                    
+                                    # Only draw if both tables exist in our schema list
+                                    if ref_table in schema.tables:
+                                        dot.append(f'  "{table_name}" -> "{ref_table}" [label="{col_name}"];')
+                                        has_edges = True
+                        
+                        dot.append('}')
+                        graph_code = "\n".join(dot)
+                        st.graphviz_chart(graph_code, width="stretch")
+                        
+                        if not has_edges:
+                            st.info("No foreign key relationships detected in the schema metadata.")
+                            
+                    except Exception as e:
+                        st.error(f"Could not render diagram: {e}")
+
+        except Exception as e:
+            st.error(f"Error loading schema: {e}")
+
+
+def render_chat_interface():
+    """Render the main chat interface."""
+    st.title("🤖 OnceDataBot")
+    st.caption("Schema-agnostic chatbot • MySQL | PostgreSQL • Powered by Groq (FREE!)")
+    
+    # Schema explorer
+    render_schema_explorer()
+    
+    # Chat container
+    chat_container = st.container()
+    
+    with chat_container:
+        # Display messages
+        for i, msg in enumerate(st.session_state.messages):
+            with st.chat_message(msg["role"]):
+                # Create columns for message and favorite button
+                msg_col, fav_col = st.columns([0.95, 0.05])
+                
+                with msg_col:
+                    st.markdown(msg["content"])
+                
+                with fav_col:
+                    # Favorite button for assistant messages
+                    if msg["role"] == "assistant":
+                        is_favorited = i in st.session_state.favorites
+                        if st.button(
+                            "⭐" if is_favorited else "☆",
+                            key=f"fav_{i}",
+                            help="Click to favorite/unfavorite this response"
+                        ):
+                            if is_favorited:
+                                st.session_state.favorites.remove(i)
+                            else:
+                                st.session_state.favorites.append(i)
+                            st.rerun()
+                
+                # Show metadata for assistant messages
+                if msg["role"] == "assistant" and "metadata" in msg:
+                    meta = msg["metadata"]
+                    
+                    # Show token usage in a dropdown expander
+                    if "token_usage" in meta:
+                        usage = meta["token_usage"]
+                        total = usage.get('total', 0)
+                        
+                        with st.expander(f"📊 Token Usage ({total:,} total)", expanded=False):
+                            # Create styled token usage boxes using columns
+                            st.markdown("""
+                            <style>
+                            .token-box {
+                                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                                border-radius: 12px;
+                                padding: 12px 16px;
+                                color: white;
+                                text-align: center;
+                                box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3);
+                                margin: 4px 0;
+                            }
+                            .token-box-input {
+                                background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
+                                box-shadow: 0 4px 15px rgba(17, 153, 142, 0.3);
+                            }
+                            .token-box-output {
+                                background: linear-gradient(135deg, #ee0979 0%, #ff6a00 100%);
+                                box-shadow: 0 4px 15px rgba(238, 9, 121, 0.3);
+                            }
+                            .token-box-total {
+                                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                                box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3);
+                            }
+                            .token-label {
+                                font-size: 11px;
+                                text-transform: uppercase;
+                                letter-spacing: 1px;
+                                opacity: 0.9;
+                                margin-bottom: 4px;
+                            }
+                            .token-value {
+                                font-size: 20px;
+                                font-weight: 700;
+                            }
+                            </style>
+                            """, unsafe_allow_html=True)
+                            
+                            col1, col2, col3 = st.columns(3)
+                            
+                            with col1:
+                                st.markdown(f"""
+                                <div class="token-box token-box-input">
+                                    <div class="token-label">📥 Input Tokens</div>
+                                    <div class="token-value">{usage.get('input', 0):,}</div>
+                                </div>
+                                """, unsafe_allow_html=True)
+                            
+                            with col2:
+                                st.markdown(f"""
+                                <div class="token-box token-box-output">
+                                    <div class="token-label">📤 Output Tokens</div>
+                                    <div class="token-value">{usage.get('output', 0):,}</div>
+                                </div>
+                                """, unsafe_allow_html=True)
+                            
+                            with col3:
+                                st.markdown(f"""
+                                <div class="token-box token-box-total">
+                                    <div class="token-label">📊 Total Tokens</div>
+                                    <div class="token-value">{usage.get('total', 0):,}</div>
+                                </div>
+                                """, unsafe_allow_html=True)
+                    
+                    if meta.get("query_type"):
+                        # Show query type and execution time on same line
+                        info_text = f"Query type: {meta['query_type']}"
+                        if meta.get("execution_time"):
+                            info_text += f" • ⏱️ {meta['execution_time']:.2f}s"
+                        st.caption(info_text)
+                        
+                    # SQL Query expander
+                    if meta.get("sql_query"):
+                        with st.expander("🛠️ SQL Query & Details"):
+                            st.code(meta["sql_query"], language="sql")
+                            
+                    # Visualizations and CSV export
+                    if meta.get("sql_results"):
+                        # Only render viz if we have results
+                        render_visualization(meta["sql_results"], f"viz_{i}")
+                        
+                        # CSV Export button
+                        csv_data = export_results_to_csv(meta["sql_results"])
+                        if csv_data:
+                            st.download_button(
+                                label="📊 Export to CSV",
+                                data=csv_data,
+                                file_name=f"query_results_{i}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
+                                mime="text/csv",
+                                key=f"csv_export_{i}",
+                                help="Download query results as CSV file"
+                            )
+    
+    # Chat input
+    if prompt := st.chat_input("Ask about your data..."):
+        if not st.session_state.initialized:
+            st.error("Please connect to a database first!")
+            return
+        
+        # Add user message
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        
+        # Calculate memory context for display? No, just render user msg
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        
+        # Get response
+        with st.spinner("Thinking..."):
+            try:
+                # Add memory interaction
+                if st.session_state.memory:
+                    st.session_state.memory.add_message("user", prompt)
+                
+                # Track execution time
+                start_time = time.time()
+                
+                response = st.session_state.chatbot.chat(
+                    prompt, 
+                    st.session_state.memory,
+                    ignored_tables=list(st.session_state.ignored_tables),
+                    language=st.session_state.response_language
+                )
+                
+                execution_time = time.time() - start_time
+                
+
+                
+                # Create metadata dict
+                metadata = {
+                    "query_type": response.query_type,
+                    "sql_query": response.sql_query,
+                    "sql_results": response.sql_results,
+                    "token_usage": response.token_usage,
+                    "execution_time": execution_time
+                }
+                
+                # Save to session state
+                st.session_state.messages.append({
+                    "role": "assistant",
+                    "content": response.answer,
+                    "metadata": metadata
+                })
+                
+                # Set flag to auto-read the latest response
+                st.session_state.auto_read_latest = True
+                
+                # Save to active memory
+                if st.session_state.memory:
+                    st.session_state.memory.add_message("assistant", response.answer)
+                
+                st.rerun()
+                
+            except Exception as e:
+                st.error(f"An error occurred: {e}")
+                import traceback
+                st.error(traceback.format_exc())
+
+
+def main():
+    """Main application entry point."""
+    init_session_state()
+    
+    # Auto-connect to environment database on first load
+    if "auto_connect_attempted" not in st.session_state:
+        st.session_state.auto_connect_attempted = True
+        if st.session_state.db_source == "environment":
+            success = initialize_chatbot()
+            if success:
+                st.toast("✅ Auto-connected to database!")
+
+    render_sidebar()
+    render_chat_interface()
+
+
+if __name__ == "__main__":
+    main()