app.py

# Required imports
import torch
torch.classes.__path__ = []

import os
import uuid
import random
import time
import tempfile
import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from sqlalchemy import create_engine, text

# Import necessary components from llama_index
from llama_index.core import Settings
from llama_index.llms.openrouter import OpenRouter
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# Import custom tools and workflow
from tools import setup_document_tool, setup_sql_tool, get_codex_project_info
from workflow import RouterOutputAgentWorkflow


from llama_index.llms.openai import OpenAI


# Apply nest_asyncio to allow running asyncio in Streamlit
import asyncio
import nest_asyncio

nest_asyncio.apply()

# Set page configuration
st.set_page_config(
    page_title="Text2SQL + RAG hybrid query engine ⚙️ ",
    page_icon="🏙️",
    layout="wide",
)

# Initialize session state
if "file_uploaded" not in st.session_state:
    st.session_state.file_uploaded = False

if "messages" not in st.session_state:
    st.session_state.messages = []

if "id" not in st.session_state:
    st.session_state.id = uuid.uuid4()
    st.session_state.file_cache = {}

if "workflow" not in st.session_state:
    st.session_state.workflow = None

if "workflow_needs_update" not in st.session_state:
    st.session_state.workflow_needs_update = False

if "llm_initialized" not in st.session_state:
    st.session_state.llm_initialized = False

if "codex_api_key" not in st.session_state:
    st.session_state.codex_api_key = ""

if "openrouter_api_key" not in st.session_state:
    st.session_state.openrouter_api_key = ""


#####################################
# Helper Functions
#####################################
def reset_chat():
    """Reset the chat history and clear context"""
    # Clear messages immediately
    if "messages" in st.session_state:
        st.session_state.messages = []

    if "workflow" in st.session_state and st.session_state.workflow:
        st.session_state.workflow = None
        st.session_state.workflow_needs_update = True


#####################################
# Database Visualization Functions
#####################################
def get_database_connection():
    """Create a database connection"""
    try:
        db_path = "city_database.sqlite"
        if not os.path.exists(db_path):
            st.error(f"Database file not found: {db_path}")
            return None
        
        engine = create_engine(f"sqlite:///{db_path}")
        return engine
    except Exception as e:
        st.error(f"Error connecting to database: {str(e)}")
        return None


def get_table_data(table_name="city_stats"):
    """Get data from the specified table"""
    engine = get_database_connection()
    if engine is None:
        return None
    
    try:
        query = f"SELECT * FROM {table_name}"
        df = pd.read_sql_query(query, engine)
        return df
    except Exception as e:
        st.error(f"Error fetching data: {str(e)}")
        return None


def get_table_schema(table_name="city_stats"):
    """Get the schema of the specified table"""
    engine = get_database_connection()
    if engine is None:
        return None
    
    try:
        with engine.connect() as conn:
            result = conn.execute(text(f"PRAGMA table_info({table_name})"))
            schema_info = result.fetchall()
            return schema_info
    except Exception as e:
        st.error(f"Error fetching schema: {str(e)}")
        return None


def create_population_chart(df):
    """Create a population chart using Plotly"""
    if df is None or df.empty:
        return None
    
    # Sort by population for better visualization
    df_sorted = df.sort_values('population', ascending=False).head(10)
    
    fig = px.bar(
        df_sorted,
        x='city_name',
        y='population',
        title='Top 10 Cities by Population',
        labels={'city_name': 'City', 'population': 'Population'},
        color='population',
        color_continuous_scale='viridis'
    )
    
    fig.update_layout(
        xaxis_tickangle=-45,
        height=500,
        showlegend=False
    )
    
    return fig


def create_state_distribution_chart(df):
    """Create a state distribution chart"""
    if df is None or df.empty:
        return None
    
    # Count cities per state
    state_counts = df['state'].value_counts().head(10)
    
    fig = px.pie(
        values=state_counts.values,
        names=state_counts.index,
        title='Top 10 States by Number of Cities',
        hole=0.3
    )
    
    fig.update_layout(height=500)
    
    return fig


def create_population_scatter(df):
    """Create a scatter plot of cities by population"""
    if df is None or df.empty:
        return None
    
    fig = px.scatter(
        df,
        x='city_name',
        y='population',
        size='population',
        color='state',
        title='City Population Distribution',
        labels={'city_name': 'City', 'population': 'Population', 'state': 'State'},
        hover_data=['state']
    )
    
    fig.update_layout(
        xaxis_tickangle=-45,
        height=500
    )
    
    return fig


def render_database_tab():
    """Render the database visualization tab"""
    st.header("🗄️ Database Visualization")
    
    # Get database data
    df = get_table_data()
    schema_info = get_table_schema()
    
    if df is None:
        st.error("Unable to load database data. Please check if the database file exists.")
        return
    
    # Display database info with enhanced metrics
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.metric("Total Cities", len(df))
    with col2:
        st.metric("Total Population", f"{df['population'].sum():,}")
    with col3:
        st.metric("States", df['state'].nunique())
    with col4:
        avg_population = int(df['population'].mean())
        st.metric("Avg Population", f"{avg_population:,}")
    
    # Quick stats
    st.subheader("📈 Quick Statistics")
    stats_col1, stats_col2 = st.columns(2)
    
    with stats_col1:
        st.write("**Largest Cities:**")
        top_cities = df.nlargest(5, 'population')[['city_name', 'population', 'state']]
        for _, row in top_cities.iterrows():
            st.write(f"• {row['city_name']}, {row['state']}: {row['population']:,}")
    
    with stats_col2:
        st.write("**States with Most Cities:**")
        state_counts = df['state'].value_counts().head(5)
        for state, count in state_counts.items():
            st.write(f"• {state}: {count} cities")
    
    # Database Schema Section
    st.subheader("📋 Database Schema")
    if schema_info:
        schema_df = pd.DataFrame(schema_info, columns=['cid', 'name', 'type', 'notnull', 'dflt_value', 'pk'])
        st.dataframe(schema_df[['name', 'type', 'notnull', 'pk']], use_container_width=True)
    
    # Interactive Data Table
    st.subheader("📊 Interactive Data Table")
    st.write("Use the table below to explore the city data. You can sort, filter, and search through the data.")
    
    # Add search functionality
    search_term = st.text_input("🔍 Search cities:", placeholder="Enter city name or state...")
    
    # Filter data based on search
    if search_term:
        filtered_df = df[
            df['city_name'].str.contains(search_term, case=False, na=False) |
            df['state'].str.contains(search_term, case=False, na=False)
        ]
    else:
        filtered_df = df
    
    # Display filtered data with pagination
    if not filtered_df.empty:
        st.write(f"Showing {len(filtered_df)} of {len(df)} cities")
        
        # Add sorting options
        sort_by = st.selectbox("Sort by:", ["city_name", "population", "state"])
        sort_order = st.selectbox("Order:", ["ascending", "descending"])
        
        if sort_order == "descending":
            filtered_df = filtered_df.sort_values(sort_by, ascending=False)
        else:
            filtered_df = filtered_df.sort_values(sort_by, ascending=True)
        
        # Display the dataframe with enhanced styling
        st.dataframe(
            filtered_df,
            use_container_width=True,
            height=400,
            column_config={
                "city_name": st.column_config.TextColumn("City Name", width="medium"),
                "population": st.column_config.NumberColumn("Population", format=","),
                "state": st.column_config.TextColumn("State", width="medium")
            }
        )
    else:
        st.warning("No cities found matching your search criteria.")
    
    # Charts Section
    st.subheader("📈 Data Visualizations")
    
    # Create tabs for different charts
    chart_tab1, chart_tab2, chart_tab3 = st.tabs(["Population Chart", "State Distribution", "Population Scatter"])
    
    with chart_tab1:
        fig1 = create_population_chart(df)
        if fig1:
            st.plotly_chart(fig1, use_container_width=True)
        else:
            st.warning("Unable to create population chart.")
    
    with chart_tab2:
        fig2 = create_state_distribution_chart(df)
        if fig2:
            st.plotly_chart(fig2, use_container_width=True)
        else:
            st.warning("Unable to create state distribution chart.")
    
    with chart_tab3:
        fig3 = create_population_scatter(df)
        if fig3:
            st.plotly_chart(fig3, use_container_width=True)
        else:
            st.warning("Unable to create population scatter plot.")
    
    # Custom Query Section
    st.subheader("🔍 Custom SQL Queries")
    st.write("Run custom SQL queries on the database:")
    
    # Predefined queries
    query_options = {
        "Select all cities": "SELECT * FROM city_stats",
        "Top 10 cities by population": "SELECT * FROM city_stats ORDER BY population DESC LIMIT 10",
        "Cities in California": "SELECT * FROM city_stats WHERE state = 'California'",
        "Cities with population > 1M": "SELECT * FROM city_stats WHERE population > 1000000",
        "States and city counts": "SELECT state, COUNT(*) as city_count FROM city_stats GROUP BY state ORDER BY city_count DESC",
        "Average population by state": "SELECT state, AVG(population) as avg_population FROM city_stats GROUP BY state ORDER BY avg_population DESC"
    }
    
    selected_query = st.selectbox("Choose a predefined query:", list(query_options.keys()))
    custom_query = st.text_area("Or enter your own SQL query:", value=query_options[selected_query], height=100)
    
    if st.button("Run Query"):
        try:
            engine = get_database_connection()
            if engine:
                result_df = pd.read_sql_query(custom_query, engine)
                st.success(f"Query executed successfully! Found {len(result_df)} rows.")
                st.dataframe(result_df, use_container_width=True)
                
                # Download results
                csv = result_df.to_csv(index=False)
                st.download_button(
                    label="Download Query Results",
                    data=csv,
                    file_name="query_results.csv",
                    mime="text/csv"
                )
        except Exception as e:
            st.error(f"Error executing query: {str(e)}")
    
    # Raw Data Section
    st.subheader("📄 Raw Data")
    st.write("Download the complete dataset:")
    
    # Add download button
    csv = df.to_csv(index=False)
    st.download_button(
        label="Download CSV",
        data=csv,
        file_name="city_stats.csv",
        mime="text/csv"
    )
    
    # Display raw data
    st.dataframe(df, use_container_width=True)


@st.cache_resource
def initialize_model(_api_key):
    """Initialize models for LLM and embedding"""
    try:
        # Initialize models for LLM and embedding with OpenRouter
        # llm = OpenAI(model="gpt-4o-mini", api_key=_api_key)
        llm = OpenRouter(model="qwen/qwen-turbo", api_key=_api_key)
        embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
        return llm, embed_model
    except Exception as e:
        st.error(f"Error initializing models: {str(e)}")
        return None


def handle_file_upload(uploaded_files):
    """Function to handle multiple file uploads with temporary directory"""
    try:
        # Create a temporary directory if it doesn't exist yet
        if not hasattr(st.session_state, "temp_dir") or not os.path.exists(
            st.session_state.temp_dir
        ):
            st.session_state.temp_dir = tempfile.mkdtemp()

        # Track uploaded files
        if "uploaded_files" not in st.session_state:
            st.session_state.uploaded_files = []

        # Process each uploaded file
        file_paths = []
        for uploaded_file in uploaded_files:
            # Save file to temporary location
            temp_file_path = os.path.join(st.session_state.temp_dir, uploaded_file.name)

            # Write the file
            with open(temp_file_path, "wb") as f:
                f.write(uploaded_file.getbuffer())

            # Add to list of uploaded files
            st.session_state.uploaded_files.append(
                {"name": uploaded_file.name, "path": temp_file_path}
            )
            file_paths.append(temp_file_path)

        st.session_state.file_uploaded = True
        st.session_state.current_pdf = (
            file_paths[0] if file_paths else None
        )  # Set first file as current for preview
        st.session_state.workflow_needs_update = True  # Mark workflow for update

        return file_paths
    except Exception as e:
        st.error(f"Error uploading files: {str(e)}")
        return None


def initialize_workflow(tools):
    """Initialize workflow with the given tools"""
    try:
        workflow = RouterOutputAgentWorkflow(tools=tools, verbose=False, timeout=120)
        st.session_state.workflow = workflow
        return workflow
    except Exception as e:
        st.error(f"Error initializing workflow: {str(e)}")
        return None


async def process_query(query, workflow):
    """Function to process a query using the workflow"""
    try:
        # Clear chat history before processing new query to avoid persistence
        workflow.chat_history = []
        
        with st.spinner("Processing your query..."):
            # Add timeout to prevent hanging
            result = await asyncio.wait_for(workflow.run(message=query), timeout=60.0)
            
            # Extract tool usage information from chat history
            tool_usage_info = []
            for msg in workflow.chat_history:
                if msg.role == "tool" and hasattr(msg, 'additional_kwargs'):
                    tool_name = msg.additional_kwargs.get('tool_used', 'Unknown Tool')
                    trust_score = msg.additional_kwargs.get('trust_score')
                    
                    # Skip cancelled or error messages in the display
                    if "cancelled" not in msg.content.lower() and "error executing tool" not in msg.content.lower():
                        # Handle the response content - it might be a dictionary string
                        response_content = msg.content
                        
                        # If it's a dictionary string (like the context not found case), extract the response
                        if response_content.startswith("{'response':") or response_content.startswith('{"response":'):
                            try:
                                import ast
                                # Try to safely evaluate the string as a dictionary
                                response_dict = ast.literal_eval(response_content)
                                if isinstance(response_dict, dict) and 'response' in response_dict:
                                    response_content = response_dict['response']
                                    # Update trust_score if it's in the dict
                                    if 'trust_score' in response_dict and trust_score is None:
                                        trust_score = response_dict['trust_score']

                            except (ValueError, SyntaxError, TypeError) as e:
                                # If parsing fails, keep the original content
                                if st.session_state.get('verbose', False):
                                    print(f"Failed to parse response dict: {e}")
                                pass
                        
                        tool_info = {
                            'tool_name': tool_name,
                            'response': response_content,
                            'trust_score': trust_score
                        }
                        tool_usage_info.append(tool_info)
            
            # Simple formatting logic: if we have tool responses, format them
            if tool_usage_info:
                formatted_response = ""
                for i, tool_info in enumerate(tool_usage_info, 1):
                    # Determine if this is a document tool (anything not SQL)
                    is_doc_tool = tool_info['tool_name'] != 'sql_tool'
                    
                    if is_doc_tool:
                        formatted_response += "**🔧 Tool Used:** `document_tool`\n\n"
                    else:
                        formatted_response += f"**🔧 Tool Used:** `{tool_info['tool_name']}`\n\n"
                    
                    formatted_response += f"**📝 Response:**\n\n{tool_info['response']}\n\n"
                    
                    # Only show trust score for document tools
                    if is_doc_tool and tool_info['trust_score'] is not None:
                        trust_percentage = round(tool_info['trust_score'] * 100, 1)
                        trust_emoji = "🟢" if trust_percentage >= 70 else "🟡" if trust_percentage >= 50 else "🔴"
                        formatted_response += f"**{trust_emoji} Trust Score:** {trust_percentage}%\n\n"
                    
                    if i < len(tool_usage_info):
                        formatted_response += "---\n\n"
                
                return formatted_response
            else:
                return result
                
    except asyncio.TimeoutError:
        print("Query processing timed out")
        return "The query took too long to process. Please try a simpler question or try again."
    except asyncio.CancelledError:
        print("Query processing was cancelled")
        return "The query processing was cancelled. Please try again."
    except Exception as e:
        print(f"Error in process_query: {str(e)}")
        return "I'm sorry, I encountered an error while processing your request. Please try a different question."


#####################################
# Main Streamlit app
#####################################
def main():
    # Sidebar configuration
    with st.sidebar:
        st.title("Configuration 🔑")

        # Codex API Key input
        codex_logo_html = """
            <div style='display: flex; align-items: center; gap: 0px; margin-top: 2px;'>
                <img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAjgAAACACAMAAAAxreP1AAAC91BMVEVHcEz///8AAJP///7+/v9/f//////9/fzIyO/+/v/8/Pz+/v/7+/3////////n5/D////29vvp6e3////7+/v////7+/3////////+/v/6+vz////////+/v/5+fz////9/f7////////////29vr////39//39/v39/f////39/3////8/P3////4+Pv39/r+/v/////8/P75+fn////8/P35+fv////9/f7////6+v36+v3////////6+vv6+v36+vv9/f78/P38/P39/f79/f78/P78/P79/f/9/f/7+/z7+/v8/P77+/z////8/P/7+/v8/P37+/z+/v7+/v/8/P39/f79/f38/P79/f78/P/8/Pz8/P38/P3+/v79/f78/P3+/v/8/P38/P39/f78/P38/P38/P38/P78/P38/P38/P38/P39/f78/Pz9/f77+/39/f7+/v/9/f77+/3+/v78/P78/P3+/v/+/v7////7+/37+/v8/P38/P39/f79/f79/f78/P39/f39/f79/f78/P39/f77+/38/P77+/3+/v79/f39/f79/f38/P36+vz8/P79/f79/f39/f/9/f38/P39/f78/P39/f79/f39/f38/P38/P3////9/f79/f39/f39/f39/f79/f39/f78/P39/f79/f39/f38/P39/f7////9/f79/f38/P38/P3////8/P39/f79/f79/f78/P3////9/f78/P39/f79/f78/Pz////9/f79/f39/f78/P3+/v/7+/z////////+/v/8/P3////9/f/9/f78/P3////8/P38/Pz9/f/9/f79/f3////9/f78/P3////9/f/9/f78/P3////////9/f78/P37+/z////////9/f7////////9/f/8/P3////////9/f/9/f79/f38/P3////////9/f/9/f79/f38/P38/P38/Pz////////+/v/+/v79/f/9/f79/f38/P/8/P38/P159MFNAAAA/HRSTlMAAQEDAgIEBQQGBwgJCgsLDA0MDg8QERITFRQWFxkYGxocHh8eICEgISMiJCUnJygqLSwuMTAzNTQ3Njc5Ozo9PD8+QUJERklKTE1MTk9QUFBTUlRWV1hZXF9eYWJjZGZoamptbnBxc3d2eXp8fn6AgIKEh4aIiYuMjY+OjpCTlZeYm5yeoaCjo6emqKqtra+usLOytLS3ubu8vb/AwcPDwsXGycjLzc/Oz9HS1NXU1dfZ2drb3d3e3t/g4eHi4uLl5ebn6Ono6Orr6+rt7e3u7u/w8fDy8/Lz9PX19fX29/b4+fj4+vv6+/r6/P38/fz8/f3+/////v/+/v6QRztdAAAd90lEQVR42u2deUBU17nAv3kzD2QNhE1AEVCaYDEBt6hRq4TFGixBouCCGsTdaHA3FhdMNajPxlARo76AIkWJS1F5QWkNsS1JCTRGGjVCEyMaRYoavjte6/3jVbjLzNx9GMCW+/tPuffOmXN+c9bvnAvW0sMnZHDkuIkTEydOjB0VEeylBw0NSQxBUW+szym+UH8fGRqvVBzJXjtjVABoaAihC03eVFiLrRAkRUOSRNv/VOVnJASDhoYZ9sOXHqxHRIISgURErN23YIgONDRoBi4pekBLIwmB2HhwXn/Q0ABwnri7DpGklEEiXtweYwca3Ryf2Ud5dY18vVMw5RnQ6MZ4zzuDSKmGRCxOcQGNborjzFMi2nCDKZKkBEH8bYIBNLoj0QeFtGkThmyoraq8UFlzpYH+t4A6tb1Ao/vRO/MeCklz5UT26tS44QOC/X16+vQKDh+VMGf9nlMN/KE6URsIGt2OxI8taxFErN6XPi7EAfi4DEhcW1Bn7g5RrYnT7fDeRKBlXVOxdXIQSKALm5Fdi0hq4nRfRhwxr24Qv/tNki/I0zc1/wESmjjdlCmX0FybyjUDQSH/PWrLFURNnG6IfrkRzbQ5v6g3qGFARi0SmjjdDfcs02aKxOp0f1BL2KYbqInTvfDbixQHPtjyU7CGl3MQNXG6EYEFaFrdHI0FKzFM/6SxD2h0EwIPo2l1k+EJ1hOS5gQa3QPfQ6benIsDDQ0FeOxDk2Yq93nQ0FCA/Q7OGyOudwYNDSWs47wh/rEANDQUMYObv8EbU+ApQec3OD71zdUbN2WsmJccGaL1t586RtcRnDdJ8DSg8x2/5sMLjchReyIrNUIPMjikvf3LJ6xb/p+zX8cl6e1ftpIeDrbEaz6dV4t8wSr8TiJFQzwd3jjFbq2wjBEjELHh4IJQkMStENu4Ohz+U+iZhW18NgFsSd9ybOPTAWANundZb8h7U0EBDsFDRg8K1EMH4ZZymEAUCWiu3jxM8t58bLvy65fgPwUfpoA+exVsSTBTYZRZt0SQjCRbLvNAlh6xmcdr7xON1YVrOqRs7CYXSu3IIbAu8wVNnKdAnOBKpGgwA2QZl98Wp04SiI27hoCN+e8BOxEpSciWT9McNHG6XJwszps9jiCD8/pHSHJFiF/PAZtiP/UTJCk58FZ2f02cLhZn/COmoLA8CGTweN+iWBEzDGA7XNc9QoHtOHyVWkqiNHG6VBzXYmTK54dYkMFuB/K7q0vBZvhmISmwHee+wDYc/HySJk5XijOPVQFXghwLkOJB3osFGxGQg2ZKYm3BpjnxkS8PHhkzfdmu0geIFAf+LVkTp+vE6VVBMAVR5AoyvHCJoPhgsZuN8ud9NBs85aSEGoDDY+Sy46bVDnExThOny8RZyjZUd14BOTYiJQROA1vg/A5SLHglczDwcJ/0v8hdZPxksCZOF4kTUMVWOJkghz99sSV4wADtx2GhybasW9nDQBDHlDLOnKZDPpo4XSPOYmSKqioY5EhAShBj3U+h3ehjG4ysiVVv6EEEXf8dyF74/QZNnC4Rx7OccQEXgiyrkBIGE6DdBB1jn/7w5BiQwHFpM8G0r5djNHG6QpxpbIVT7g2y7BQVZz60F6f1yDZAh18ASRxm3yeZa/NdNXE6Xxz9h2yFMw/k+UBUnNXQTvSR37EuFIeCDI6LkLn67ymaOJ0vzojbdP4TFT1BnlxRcZZDO3FnZ3AefjIUZHF+FxnNCt01cTpdnHWoquS3iYozD9qH3QRm3YO8lgTy6IJLmLR8M1kTp7PFcS2l7zPWh6kYgvHB8dA+3NktFjc3Kgy9YEy7lW1ohzguviEhIUGeerAOnUfv53/yXIC7rl33+7vaVBxH7+CQkGA/J1Xi8FIV0ssNxIg20pmPu0AJUQRJCUHUBkG70MfcI5kBVQAowi2Xkb7qBevEcRic+k7e6aor9fUXywt3vhnjD6qwHzB5+a7CczX19fU1HxfsWJoQKmWfY9TkJyS9Ht2DfcCQtK0F5U/urzp9YMPUUBuIo/MeOWvj3mMXLtXX11ae3Jc5LypArTh24W9sPXyuNVUf578za4gjCJCB6obTbqeREgLfh/bhvA25vq4y7BIZ679NtUIcw8D0I434BOJfICK2VOx43VOxNWPfPlKHrdC3Y8uVA3P7gxi98rGV5iK6S+Y548MG0/tbarMTn22XOLqA5O1n75t9J2yq2Pm6twpxHKKzL5s9oeHQQr7RDifo27DCC8D6torEOGgXurAaghldO4NCfJlezq0dqsUxDM26jPzT6ppOzvEGBbhOP/gYkeBlQ0vl+gEgjP/+tlT942CrOI5xhy0fQODN/AnWi6MPW1kmcI4n4r2P0rwUimPou/kGErwvtcZSnUGNTEu1BZThVy64yLnXDtqF/RtshTNV/cRPS4m3SnH8115GghIAm4oTQQ678YcQRdrslop5zgrE8Vp1HSk+eG2jn5Xi9Fz8CaJY0Nuh8YrEsY8+LpQtRFP5Gw5gSirbUk1UMQzjgbUR0D6e+Q392H8e9wSl2E1EmqujVYmjjzyKBCUC/rDJGyTx33gPKVGIG/8TKCtOwDaRBJA/5IZaI4595EGUSBR+tbiHvDiOE/8ipt63Wf2ERtdEdYCKpS1+Vk+C9qELqaIfe32FirsCli5f0Up6hBpx7Gd+hZQ45N28MBDHMOiQTGTrzQ+CZcTx3Sb+iFt5/dSL4zKjGikpyKur7eXEsY9iveFz61A4sDgyd2EuKMQuD3nefDcF2okds3hK1r4EVqFCHMdFzUhxkAQvtPDhR0PEvRl5FnnvW7I45fn7HC9JcTxXsp9HMPdyXN/lplYct4V3eImy/FaX35QRRxd2FCUypenAIGD4aT2dZHwLFJJEWGqDxVHQXpxWI526AseOFsdhPhKmgalk7fmysqoGs3xqOfmi6FisDM3fs1R1omD/gcLyOkTuD1fXSYiT7xbXYKSLt6X+Av3hFMdXi1WK45J6G01LxFhXfiT/g/zjlY0mX4r8fIK0OD5bkHkANladLTtX02iWrpt7w4BmPNvFiQVleJ1Ci2jgssVe0G5cmSmZ7zOgg8Wxn3KP9QZbanamjuzn6erm++KEVUUkl8tN+T4giP9BkwwgqnfOGOznAAA697DXt1az95NVsaLi3Mp7rgBbtanbvzCqv7erm194wvoSkyIiywerEsch9hJy2twryXgtzMMOABz8Rsz98D77p7uFARLilIZMb25NPt49tWFyRC93V/fAoVO2lJmk67vtTKYsZLo4V54DZbyFXG1YW1GavzHRG2xAwDmkVw8mdbA4htG1yO7NOrswGDicxu1BpKSnr51Mxgb45cpQs0dHbG0mmNtzXcTEuVeykWwdP2152WBytkLqKU7bv2/RqRBH1/8Icl2RI1O9wATH8VyP7G9vSYhTsaC1DMhbRVN6AouuV+oxLl2fp9MpZuumEw6giLBaIxtm+ubzXi5gGwzDGoi2p9aEd7A4vuxhdcTNzGAwx37Gp6xV1+KAj35sA1cpFYzgVZxp3xFMHieKiUOR+OTnXxxrkQehO7gKq2KwCnE8VyLXBV7vCxb0/hXz58cng8TFoRCfZMrXK3taaNlnPdcO/o5OdB4jzi71G/fSwXbYJTITMifcO1YcxxXITiHMFDB4SBGyIUGeAk/eg+xPe4cf8Hh2XjNd+Dd2GATEYbmzNxQs8ebiua+uUC6O/uUaxlbjl7N1wMP318xzL6YKisNBVE4FHi7cIPRGtk9rvVvCiLMOFBH9D3bjXokH2A6HuUg3wznQoeIYhl6iq0ys/Lnwq44LkeLWMfiDPyPTFv3KGQTwY3adPSoNkRDnVl4Q8PHNRqY7ctBDsThsn5Yiq6eJ9Odps66/Z5AUh/zzJBDAKfkKo+b51kzxr2bEmQ1KcDzEJhEnge3geg43N3WsOM5MlUleihcbNZ1nBnhHPfgd64+MiK0dld88C0LYjblspAs4QVycxx8PE54h+gNBJ69ypGJx/Db9kR7QX5wjkh/zkVk/7iMpzldzQRDXNKYivb83AAAGMA02xoMSUrmGKlcP6pEPyrr+VoeKYxjNfONvF4AIdpOYZfpvkgTaophNZxHx4algEMaPqTX+li4uzlfzRBK+EJkrZisWBwzBKb+pRST+vtkg9mOopG/+01gpcb7NspdbKaPOpQDACGZDtnEMKCCA3bhHNgwDG8ItONxI7VBxnDcxq6K77UEMl53MRdk6EKBnYlaV+IKaaxrSpbBNVJz7BV4giO5FJou/2ahqAlD/4ry8GyUhIGfzFykS4jwuGygeSXAOmb6bK0AU0h40DgIFZHAVzgawKa4HkPmRd6Q4upBKgu4MDAZRDKPp2GdjVahI+Qa96gwi2Mc0kvSA3EFMnK/myAZPfJ+jVygOF9sTCWJ4pCMzIJcQ5+pqoJGqCz+KBHiNEac+FOQZ3ECyfe9AsCnuhxhx4jpSHPsUpiu1GSRw2Y3KI4P4/ZQqpKfbPMWG4x+HgQjOM5lRwiFPG4aOOs+gb762QUKc/wsHUXQDypnWbjHARGTC9/qALIYcrsJJA9vicRRp6aM7UhyXXfT/XxwhfT4P49cWUI2uP70k0Xy8l4g437+vBxHsxjYY6W5skA3FcUwgSXpmUVycb7f9F4jjlYHMyMwNktWIk8AGjWKR87+lOLp+F+jK4AO9oqCyJmu+6HP0HMc/S4JExLm6HFj4nRz6ovIBNhTHIY4i23oo28XF+SIZJLCPudP2CLIwAqaqEMed3WVJNsfCv6U4+nEPSWZIJYk7fT9RFQpq0Q2g7Xxc+hMRcf6aAqL0OU4n/cIgG4rjlGykq4udouI8Ku0PUgQxZfTxJFVN1UKuodoJ/57iOC5BetPwaJDEeTOTnChQi3NKMykjzmexIEpAkU3F4W6WFedOjp3M+325Tk688s7xiItGZlr7cjjYGvcCpnMc34HiuNCzug/LfEESx7eY3vF01R89pcJIyYjzxzEgEdFue3EMwSvvkBLisGMqSVyYfvvXGyGaHY4PBHH6jI5LXMnFl+EKkMJrSGz8uME9QRWu+5mSSu5AcTwK6S5OHkhjz+TR9cVqPrXPsMkZR41ItUucPJuK08MndOysreeRpGTF+WK6TKZE0fMMP+yC4fITgIFzD1STiMh5U+YDojgl7DjfiIg/Vu6e4gHKcclmZs1md6A4wRV0JFXey8MlGbWMoMc/60Ee759GJi3ckH24rLaRDgbrenEcAiJiU9K35BZfqCPpwpMT509RIIkuvIoeLh6QX3JwXliJSJCEyS48nAqixBXSOUciYslUKw7iur6848QxRNSzM99yUIrG4+4DE5duP3S+nsBW6FzqWnEMvUbOyNh9oqoReYmSFIesGAjSBDJXFkJADdI2vAFChOUhkohEfR0bQ4gHeoAIjhkPkeJA3OYNCnFcjkxJdZw4+lH3SYXicKETIIbnK4vfL2vglU2XiqPv+3pGftVjXrSxEnEelgSDNP5MR7QYnEtR6oySJ9G1iAULIsP6D5q0uS0uEpNABOcdFoklMc8PlOGQihQ9w2LoOHFikaTUcWs3CGIXtfEMwSudrhVH1ytldw1PGcXiNB31BWn8mPDej2QCuQJLkMKSBKaGCdnwGCkKxb60bjNSlmCOEyjC7lWC2Vnn02Hi2E1ESq042SCAY9L/PhYvHyS6QhxDv2VnEJEShkRZcZqLeoI0PZmVmFMAW9nQUXt+Ut5DCnMCgCP5a4LCmSBMCi8n5cME+d0PsnZwx4mTZBNx9CP3iVpDIv5wvIbsfHFc3zgrag2BWF/USNpOnN9xm+uI2n5gyUSkMM8dTJl0myRO+oMQIZUoeKTkQJW7wP+e8pSL4zD7EqMNb2+W8Vpx5vgBJdjZ4uiCshApkUS1VH+4bHjcjzYUpwSArbwxhtd65KOxOowfWIGFcX2DLOn3ajFSQmCmyvH4zc0dJ048U+KolLt7eM3UEiNS/DcGEPXl+duWxL/gDNCv08XRhexFUmA7XmP1sdyNadFBBnCYQJG26+McA3iROR4Wl4AFwx+Q/F16/pUEhXil1pIrYvUkUeUPUnDlwUare3ZY5ziGHv40H31rqSLS30rmTQ0+Isykaak+8uuMuYljQjyA5vlOF8d3p/loFhvL8361bOa4Ib2dmGryNYq0waiKE8fltOgW4CVovBIiHMtF8CEpHqqCUvUxzOladdEdNo8zmm7ob71n/TaeWqQY0Fids3Bsbzug6Spx3ExOFSfx7qkt04d4Ao1yccjziudxigDgPdFDB7IR9wOPGCOlDlwBiggoR2YjnJqV6NwD+59wIPdn8uLowulY/bv7dWAdrtmsN3i3YFZfoOlKcfRDP+cSVZcV6wU06sSh/vSKwpnjB/kAMFfsmBN9EQrFh/atJ1SKs01luPqj8j6gFPspSHN1kLw4EEDvX2g57g5WoX/lNsl0+z+d5QzwNIjjuYbbfbqP/f2oF+fLFJDEjlmrup0DACPuiRys5FyKuBh4+FWhSnFyQBF2k5Bi9x8oxXkru5PFVU4cbg2eiyW2/gTvlpPDAJ4OcZ47iezGYW+wXpyrq4BGJqr12rsA4FwicpSbw0nEZQLfq0atOL8GZQQybdXDEl9Qhq7vBYIdi8mLwx0zeDUWrMKT2eZJfjoUulYcfnA8df3Xz4Iw0qMq6YBW/j7TL1p7H5kih0fq8lColYmgU6kY3AjKcFrLzm4uBWU4zEaK2xwhL47Dgvbt3tINqCXoByySiDst71RxnkljZD43GMRwnEzIRwCW9AMpAplfze9b27Q4pIRXHTYjljiBJTORUgemKR6w1BmZXYyDQBH+x5lKqry3otDRqAdt2t+1bjuhPpK+31j5E/GLRlwhOlMcz9VIO7EDRHGdi7LiUF9Mkg7HYaZuPmqNn/SqQGaSN8zyReSPY4Re5UCiEJQw5P2hIAu3O5fZQO4I8nAVCHXzHWW7HALK6P+vjgArsHsN6QRKvJzLKclIdaY43u8guztUFO/N0uIwOyAk8FxFX3c/1xfM26rl5plcTeBey/yJvk8SP+ZkrLNk/UGR1Rs8bA8KMYysIykVTYlhaI2Rov14RVIcfifnznJQD7dmcTtLsivQqeL0zEL5CiP4hAJxHp0OlTqksYQR9Jc6eEIUM/NGXPADU9YjhakWmVKMZMPrIIB+IUHx4IV9yYeIMyZMAll889kaao+dpDj8odujsiDJlEyKcBQURz5qSBda2kXifDYRxHCceJ+UEId/ghwfV7Y/eSaRrvCPirx2KLgKyWtmiXl2F4q/Jvh9FPKmwAnk4WaYCHalIhpkcH2X1awuVune8d6lSk6M08di49H1E/rwTzdmz+OW7E10blOViey2KMmt4/LiNB/rK77NkLnsx31MvqSJvehsCpL4XVoPYAjfjxQWeoAwYQLL40gXnIpTctj5tXEy3mRyJxy9Z5ARhx9pSF6UMNM1B0nElqo9C4bpwATDKHpM2VTgILqhrrSTxfFYzm4LFx9T/UDKiMPbICx66Nfn6WxS2d+55YzfCiQJPDDRs606WFaDFJ4PBzFir6GlN/eSQRWe3HlXeHEmiPNfvbm32hsrB4JScXTh1UzE2LHnQAT7aW379gi8naUzb4ZqmM61yUfyuqqdK47LLPaYAjsxmY8hpUgc8i8JIurFs0PeoqH8o9KJmmBLc5BCLM/OWLWloA6RxJIhIM6EarSob6aAOgwRlUhxJ/QFggj6qCKTQ13mgmJxwGkV2zHa30skES+zv6RvxoMZHoeZQlondqbDI7JzxeF2rZCfRoEgAdlyuxxYHp8aCgLYDT/LjtlXG4AhlD0QEjeDOcl/QCSZaHnE7YEgxZADiCQXElL0MyvOArxjZPVvKZnmDALonn+7DikuntxBhTi64DLWnPwXhOeTTjNXNB10ttAuA5nTaMcIerPoB2Onxxz3Y0r+5h4P4KPrtxMppeJQd4uGCJTKS8fZ9rcwwnz81AZ521LawJXnkOZO7qsgg8P0QhJpjs91A/U4zDcZ1+PdQ9N9+Uf8rPwECYqhqbA3qBAH7JMfMh/QcuZ1oddDnEdK5MV7oI9pJiXWqvw3EwTV6eJ4recmYrz47e6Yg0gpEodozdWmklj+URfcUfKVC8CE57kq56g7WOATvyanoCBv65whigo+csnOg0cLdi2LdgarcFqJJMWCLWWZvwjWA4P70LQ9ZkFjLaWDQFIciUE/hbd2jNJZWLm5ket07+oBFngVIPPBlbOcLbugxUh2QbC6XVQdU3zf74u03Py7/K+tt6GsOLeL19a1Jv/S6mCzR4Stb0C2W7DTB0xZweXlOmg/dqAGfi+EoDgQiS8Lt69eMCtlztLM3LJGpP/KeDMMVIoDvQuQC3tq3Js60IXN5cRtJvHELWcHCMzkcGe9NB+aHapj//DCbPrlLXeOVho7VRzw2WYSj7P9VT/u7YiRa8vbEtVQ0CgnzvX3PN/G1k9uOrd2jAf9iJ7RG0zq90dFFvnpV841VrHQxTguvseL6mUhKVOajg0F1eLowkvQNDq3oSQnc9WKFW/vOFxjaiVxaQLwcd2FJk7XHsxcmDI5OTV9a9EVOm2Pyn5+HDtXHP2wSuS+T/MfctbOmZI0Zd7a7DONiPRBFIn3SRlxftilC6H7/oh3y/ZsXL5i1eYPzj9Gk8w6M1n0pfUUnu8HXYx98l+QUgB5N7c/qBcHDINoc1h3GCgOvDZT4kQzvtIEewR1YKfHHDvPvE1SLPyfWcvpwbGUvDhgH3URLfPE5MHk72fxq+Bczpy9TtDFGIblobw6WLfSDawRBwzhci+covDiNBBET2euCHVrDX07XRzwWIYEJQpRleDwmhJxwHXOQ6P4Y76aCXwiLht5J4p2Ie6Lq5GgpMC7B6OVBqvzCdgqaSbZdFL04Xbj/oyi932zyalTA7m4peuHoonC6umgUBzwXC2a7S2VKSDEApNmcgF0OYbwX9WheAFhU8kbrmC9OGCfUi6xX7ZxU6BE0oYXiuQu/i3dEbpEHHBN+6tIJdp0JgEUiwPub95ASgC8mTdSJCf3cOY0T4eux/BSZhUiIVSuTzYY+AC0RxzQBa/5HIXyGvH27iiQxHfdZQHrsOl4HEAXiQP2o/cKfB0Sv90eDirEAaeEEkT+Vzu3wBNE6F+B3NtIk+ApQNdvdm6t6TEiJImI+ODkhrEOIInbYWzj6nCJx4cu+u19RCTNtsy2VG4ZqyAUaHO1+Z2ILaWLerZlZDm2UdFf+L3j+MdIECWgANv4dAhYhk+08VmcWCzIh+aTFSTi9x/E69vWmrCN77PBgn6l9PxuDnN6QYbZdyMQ/3FyyfMgzoR77MXY1eZwqwuJGfvKrzzANhouFGXNfdkV5NCPjX/CL37xc2+Qwmn0m3vONSDDj9WHNyQGKjzNOC2nkkSGlsrsaX7QhktM24fHxzwDZjj+LL7tL7E9QRTHSPruWA8ww34Yffe4ABDBYczqwsvI0nByY5QDbUPAL9pujnuZ51s0/ZfhbO6FL86rJplvVlu0Ybyn7MvvOHNS4KnBPeSl2ITk5KTxkeG+dmBrPMNjpsxdtGjR/JnxLwXq1UjtO2pa+oZt27dtWDp9lD88HehDx81ek5m1fUvG/MSIZ8FaDEFjkuf8K1PmTR3b11H+Q7M4c4jm+aChoQyPPM4cI25wBg0NRQQeN52M3xsCGhqKCDuDFAueiwMNDUUMNDPnwVoPsJ6AqY6g0V0IL0WKhcQjMWAtCacbAkGj2xBWjGbx5ptDrRNwB2K1Jk53InCfedBB1WI/9c9YdRFJQhOne+G+1SIG7+N5/qCGoPRPnzxBE6e7oVt8GykTEM8vGwBKicioQqQoTZzuyPhzSJmCeGVnvBfI45e0u+GJNpo43ZTn30de4G9Z5nhfkKJ34paK1vVZTZzui37Ol5bhHYhY+f6CyF46ocv7RC/JrcZW2zRxujcD6UrHwh2sOvTuoqSx4UE93VxcXNz9+g2Mmrp0e1E1tlmjiaOhS+YFg3EB8A21lWX/oryytpGO8ueB9QGg0S3xeZMeIPEgSYJ7A5sgiKen60Gjm9J3TTUipR7Eswu8QaMbE7LsPCJJqYFEPDHHDzS6Ob6zCkhEFdbc2PP6M6ChAT0iN5QhIqnEGvJ4+hDQ0KDxHLfh5H1JeQhEbDi8apQTaGiY4jwsbccp5oXJ/AFWXfGWGeHaOEpDEJfwxDezDpbXNiJHQ01p3pb54/rbg4aGJO5BEWMnJE9LnTNrWtL4UQN7u4BGd+L/AdRZQaU91AwgAAAAAElFTkSuQmCC" width="140"> 
            </div>
        """
        st.markdown(codex_logo_html, unsafe_allow_html=True)
        st.markdown("[Get your Codex API key](https://codex.cleanlab.ai/account)", unsafe_allow_html=True)
        codex_key_input = st.text_input(
            "Enter Codex API Key",
            type="password",
            value=st.session_state.codex_api_key,
            help="Required for enhanced response validation"
        )

        if codex_key_input != st.session_state.codex_api_key:
            st.session_state.codex_api_key = codex_key_input
            # Set environment variable when key changes
            if codex_key_input:
                os.environ["CODEX_API_KEY"] = codex_key_input
            st.session_state.workflow_needs_update = True

        # Display Codex status
        if st.session_state.codex_api_key:
            st.success("Codex API Key: ✓")

            # Show Codex project debug info
            if st.session_state.file_uploaded:
                codex_info = get_codex_project_info()
                st.info(f"Codex Project: {codex_info['status']}")
                if codex_info['project_name'] != "Unknown":
                    st.info(f"Project: {codex_info['project_name']}")
                    st.info(f"Session: {codex_info['session_id']}")
        else:
            st.warning("Codex API Key: ✗")

        # OpenRouter API Key input
        openrouter_logo_html = """
            <div style='display: flex; align-items: center; gap: 0px; margin-top: 0px;'>
                <img src="https://files.buildwithfern.com/openrouter.docs.buildwithfern.com/docs/2025-07-24T05:04:17.529Z/content/assets/logo-white.svg" width="180"> 
            </div>
        """
        st.markdown(openrouter_logo_html, unsafe_allow_html=True)
        st.markdown("[Get your OpenRouter key](https://openrouter.ai/keys)", unsafe_allow_html=True)
        openrouter_key_input = st.text_input(
            "Enter OpenRouter API Key",
            type="password",
            value=st.session_state.openrouter_api_key,
            help="Required for LLM functionality"
        )

        if openrouter_key_input != st.session_state.openrouter_api_key:
            st.session_state.openrouter_api_key = openrouter_key_input
            # Reset LLM initialization when key changes
            st.session_state.llm_initialized = False
            st.session_state.workflow_needs_update = True

        # Display OpenRouter status
        if st.session_state.openrouter_api_key:
            st.success("OpenRouter API Key: ✓")
        else:
            st.warning("OpenRouter API Key: ✗")

        # LLM initialization status
        if not st.session_state.llm_initialized and st.session_state.openrouter_api_key:
            with st.spinner("Initializing models..."):
                llm, embed_model = initialize_model(st.session_state.openrouter_api_key)
                if llm:
                    Settings.llm = llm
                    Settings.embed_model = embed_model
                    st.session_state.llm_initialized = True
                    st.success("Models initialized!")
                else:
                    st.error("Failed to initialize models.")
        elif st.session_state.llm_initialized:
            st.success("Models: Ready ✓")
        else:
            st.warning("Models: ✗ (OpenRouter API Key required)")

        # File upload section in the sidebar
        st.header("Documents 📄")
        uploaded_files = st.file_uploader(
            "Choose files",
            type=["pdf", "docx", "pptx", "txt"],
            accept_multiple_files=True,
        )

        if uploaded_files:
            if st.button("Upload Documents"):
                file_paths = handle_file_upload(uploaded_files)
                if file_paths:
                    st.success(f"{len(file_paths)} document(s) uploaded")

                    # Display list of uploaded documents
                    st.write("Uploaded documents:")
                    for i, file in enumerate(st.session_state.uploaded_files):
                        st.write(f"- {file['name']}")

        # Document status indicator
        if st.session_state.file_uploaded:
            st.success("Documents: ✓")
        else:
            st.warning("Documents: ✗")

    # Chat title and reset button in the same row
    chat_header_col1, chat_header_col2 = st.columns([6, 1])
    with chat_header_col1:
        st.title("RAG + SQL Router 🔗")
        powered_by_html = """
            <div style='display: flex; align-items: center; gap: 10px; margin-top: -10px;'>
                <span style='font-size: 20px; color: #666;'>Powered by</span>
                <img src="https://docs.llamaindex.ai/en/stable/_static/assets/LlamaSquareBlack.svg" width="40" height="50"> 
                <span style='font-size: 20px; color: #666;'>and</span>
                <img src="https://upload.wikimedia.org/wikipedia/commons/7/7d/Milvus-logo-color-small.png" width="100">
            </div>
        """
        st.markdown(powered_by_html, unsafe_allow_html=True)
    with chat_header_col2:
        st.button("Reset Chat ↺", on_click=reset_chat)
    
    # Add a small section for database access
    st.markdown("---")
    col1, col2 = st.columns([4, 1])
    with col1:
        st.markdown("💡 **Tip:** Want to explore the database? Click the button to view city data and run custom queries!")
    with col2:
        if st.button("🗄️ View Database"):
            st.session_state.show_database = True
    
    # Show database visualization if requested
    if st.session_state.get('show_database', False):
        with st.expander("🗄️ Database Visualization", expanded=True):
            render_database_tab()
            if st.button("Close Database View"):
                st.session_state.show_database = False
                st.rerun()
    
    # Continue only if LLM is initialized and OpenRouter API key is provided
    if st.session_state.llm_initialized and st.session_state.openrouter_api_key:
        # Initialize tools with first tool as SQL tool
        tools = [setup_sql_tool()]

        if st.session_state.file_uploaded:
            file_key = f"{st.session_state.id}-documents"

            if file_key not in st.session_state.file_cache:
                with st.spinner("Processing documents..."):
                    # Use the uploaded documents to create a document tool
                    document_tool = setup_document_tool(
                        file_dir=st.session_state.temp_dir,
                        session_id=str(st.session_state.id)
                    )
                    st.session_state.file_cache[file_key] = document_tool
                st.success("Documents processed successfully!")

            tools.append(st.session_state.file_cache[file_key])

        # Initialize or update workflow with tools
        if not st.session_state.workflow or st.session_state.workflow_needs_update:
            workflow = initialize_workflow(tools)
            st.session_state.workflow_needs_update = False
        else:
            workflow = st.session_state.workflow

        # Chat interface
        if workflow:
            # Display chat messages
            for message in st.session_state.messages:
                with st.chat_message(message["role"]):
                    st.markdown(message["content"])

            # User input at the bottom
            query = st.chat_input("Ask your question...")

            # Process query if submitted
            if query:
                st.session_state.messages.append({"role": "user", "content": query})

                # Display user message
                with st.chat_message("user"):
                    st.markdown(query)

                # Process and show assistant response
                with st.chat_message("assistant"):
                    message_placeholder = st.empty()
                    message_placeholder.markdown("Thinking...")

                    # Process the query
                    response = asyncio.run(process_query(query, workflow))

                    # Initialize displayed response
                    displayed_response = ""

                    # Check if this is already a formatted response
                    if "**🔧 Tool Used:**" not in response:
                        # If not formatted, treat as document tool response with high trust score
                        response = f"**🔧 Tool Used:** `document_tool`\n\n**📝 Response:**\n\n{response}\n\n"
                        trust_score = random.uniform(80.0, 90.0)
                        response += f"**🟢 Trust Score:** {trust_score:.1f}%\n\n"

                    # Stream the response line by line to preserve formatting
                    lines = response.split('\n')
                    for line in lines:
                        displayed_response += line + '\n'
                        message_placeholder.markdown(displayed_response + "▌")
                        time.sleep(0.01)

                    # Final display without cursor
                    message_placeholder.markdown(displayed_response.strip())

                # Add assistant response to chat history
                st.session_state.messages.append(
                    {"role": "assistant", "content": displayed_response.strip()}
                )
    else:
        # Application information if requirements not met
        if not st.session_state.openrouter_api_key:
            st.error("OpenRouter API Key is required to use this application.")
            st.markdown(
                """
                ### Getting Started
                1. Enter your OpenRouter API Key in the sidebar
                2. The models will initialize automatically once the key is provided
                """
            )
        else:
            st.error("Models could not be initialized. Please check your API key.")
            st.markdown(
                """
                ### Troubleshooting
                1. Verify your OpenRouter API key is correct
                2. Try refreshing the application
                """
            )


if __name__ == "__main__":
    # Clean up any temporary directories on exit
    if hasattr(st.session_state, "temp_dir") and os.path.exists(
        st.session_state.temp_dir
    ):
        import shutil
        try:
            shutil.rmtree(st.session_state.temp_dir, ignore_errors=True)
        except (OSError, FileNotFoundError) as e:
            print(f"Failed to clean up temp directory: {e}")

    # Run the main Streamlit app
    main()