app.py

"""
Self-Correcting SQL Agent with Streamlit + LangGraph
Fixed for SQLite thread safety
"""

import streamlit as st
import pandas as pd
import sqlite3
import os
import re
from typing import TypedDict, Literal
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv
import threading

# Load environment variables
load_dotenv()

# ========== THREAD-SAFE DATABASE MANAGER ==========
class DatabaseManager:
    """Manages SQLite connections per thread"""
    
    def __init__(self, db_path: str):
        self.db_path = db_path
        self._local = threading.local()
    
    def get_connection(self):
        """Get or create connection for current thread"""
        if not hasattr(self._local, 'connection') or self._local.connection is None:
            self._local.connection = sqlite3.connect(self.db_path, check_same_thread=False)
            self._local.connection.row_factory = sqlite3.Row  # Enable named columns
        return self._local.connection
    
    def close_connection(self):
        """Close connection for current thread if exists"""
        if hasattr(self._local, 'connection') and self._local.connection:
            self._local.connection.close()
            self._local.connection = None

# Global database manager instance
db_manager = DatabaseManager("sample_company.db")

# ========== PAGE CONFIG ==========
st.set_page_config(
    page_title="Self-Correcting SQL Agent",
    page_icon="🐘",
    layout="wide"
)

# ========== INITIALIZE SESSION STATE ==========
if "messages" not in st.session_state:
    st.session_state.messages = []
if "db_initialized" not in st.session_state:
    st.session_state.db_initialized = False
if "current_schema" not in st.session_state:
    st.session_state.current_schema = None

# ========== DATABASE SETUP ==========
def init_sample_database():
    """Create a sample SQLite database for demonstration"""
    try:
        # Use a temporary connection for initialization
        conn = sqlite3.connect("sample_company.db")
        cursor = conn.cursor()
        
        # Create employees table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS employees (
                id INTEGER PRIMARY KEY,
                name TEXT NOT NULL,
                department TEXT,
                salary REAL,
                hire_date TEXT,
                manager_id INTEGER
            )
        """)
        
        # Create departments table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS departments (
                id INTEGER PRIMARY KEY,
                name TEXT NOT NULL,
                budget REAL,
                location TEXT
            )
        """)
        
        # Check if data exists
        cursor.execute("SELECT COUNT(*) FROM employees")
        count = cursor.fetchone()[0]
        
        if count == 0:
            # Insert sample employees
            sample_employees = [
                (1, "Alice Johnson", "Engineering", 95000, "2020-01-15", None),
                (2, "Bob Smith", "Sales", 75000, "2019-03-20", None),
                (3, "Carol Davis", "Engineering", 110000, "2018-11-01", 1),
                (4, "David Brown", "Marketing", 68000, "2021-06-10", None),
                (5, "Eve Wilson", "Sales", 82000, "2020-09-05", 2),
                (6, "Frank Miller", "Engineering", 125000, "2017-04-12", 1),
                (7, "Grace Lee", "HR", 72000, "2022-01-30", None),
            ]
            cursor.executemany("INSERT INTO employees VALUES (?,?,?,?,?,?)", sample_employees)
            
            # Insert sample departments
            sample_depts = [
                (1, "Engineering", 500000, "Building A"),
                (2, "Sales", 300000, "Building B"),
                (3, "Marketing", 200000, "Building C"),
                (4, "HR", 150000, "Building A"),
            ]
            cursor.executemany("INSERT INTO departments VALUES (?,?,?,?,?)", sample_depts)
        
        conn.commit()
        conn.close()
        
        # Reset any existing connections
        db_manager.close_connection()
        
        return True
    except Exception as e:
        st.error(f"Database initialization error: {str(e)}")
        return False

def get_table_schema() -> str:
    """Extract schema information from database (thread-safe)"""
    try:
        conn = db_manager.get_connection()
        cursor = conn.cursor()
        schema_parts = []
        
        # Get all tables
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
        tables = cursor.fetchall()
        
        for table in tables:
            table_name = table[0]
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = cursor.fetchall()
            
            # Get sample data
            cursor.execute(f"SELECT * FROM {table_name} LIMIT 2")
            sample_rows = cursor.fetchall()
            
            schema_parts.append(f"\nTable: {table_name}")
            schema_parts.append(f"Columns: {', '.join([col[1] + ' (' + col[2] + ')' for col in columns])}")
            if sample_rows:
                schema_parts.append(f"Sample rows: {[dict(row) for row in sample_rows[:2]]}")
        
        return "\n".join(schema_parts) if schema_parts else "No tables found. Please initialize the database."
    except Exception as e:
        return f"Error reading schema: {str(e)}"

def execute_sql(query: str) -> tuple[bool, str, pd.DataFrame]:
    """Execute SQL query safely with error handling (thread-safe)"""
    try:
        conn = db_manager.get_connection()
        cursor = conn.cursor()
        cursor.execute(query)
        
        # Check if SELECT query
        if query.strip().upper().startswith("SELECT"):
            rows = cursor.fetchall()
            if rows:
                # Get column names
                columns = [description[0] for description in cursor.description]
                # Convert rows to list of lists
                data = [list(row) for row in rows]
                df = pd.DataFrame(data, columns=columns)
                return True, f"✓ Query executed successfully. Returned {len(rows)} rows.", df
            else:
                return True, "✓ Query executed successfully. Returned 0 rows.", pd.DataFrame()
        else:
            conn.commit()
            return True, f"✓ Query executed successfully. Affected rows: {cursor.rowcount}", None
    except Exception as e:
        return False, f"✗ Error: {str(e)}", None

# ========== AGENT STATE DEFINITION ==========
class AgentState(TypedDict):
    question: str
    sql_query: str
    execution_result: str
    error_message: str
    attempts: int
    max_attempts: int
    schema_info: str
    reflection: str
    final_answer: str

# ========== AGENT NODES ==========
class SQLAgent:
    def __init__(self, llm):
        self.llm = llm
        self.schema = get_table_schema()
    
    def generate_sql(self, state: AgentState) -> AgentState:
        """Generate SQL from natural language"""
        error_context = ""
        if state.get('error_message'):
            error_context = f"""
Previous attempt failed with error: {state['error_message']}
Reflection: {state.get('reflection', 'No reflection provided')}
Please generate a corrected SQL query.
"""
        
        prompt = f"""You are a SQL expert. Generate a SQLite query for the following question.

Database Schema:
{self.schema}

User Question: {state['question']}

{error_context}

Important rules:
- Use valid SQLite syntax
- Only use tables and columns from the schema
- For string comparisons, use single quotes
- Return ONLY the SQL query, no explanation

SQL Query:"""
        
        response = self.llm.invoke([HumanMessage(content=prompt)])
        sql = response.content.strip()
        
        # Clean up markdown if present
        sql = re.sub(r'```sql\n?|```', '', sql)
        sql = sql.replace('`', '')  # Remove backticks
        
        return {**state, "sql_query": sql, "attempts": state.get("attempts", 0) + 1}
    
    def execute_sql_node(self, state: AgentState) -> AgentState:
        """Execute SQL and capture results"""
        success, result, df = execute_sql(state["sql_query"])
        
        if success:
            return {
                **state, 
                "execution_result": result,
                "error_message": "",
                "final_answer": df.to_string() if df is not None and not df.empty else result
            }
        else:
            return {**state, "error_message": result, "execution_result": ""}
    
    def reflect_on_error(self, state: AgentState) -> AgentState:
        """Analyze error and provide reflection for correction"""
        if not state.get("error_message"):
            return {**state, "reflection": "Query executed successfully"}
        
        prompt = f"""The SQL query failed. Analyze the error and suggest specific corrections.

Question: {state['question']}
Failed Query: {state['sql_query']}
Error: {state['error_message']}
Schema: {self.schema}

Provide a brief, specific guidance on what went wrong and how to fix it.
Keep it to 2-3 sentences.

Reflection:"""
        
        response = self.llm.invoke([HumanMessage(content=prompt)])
        reflection = response.content.strip()
        
        return {**state, "reflection": reflection}
    
    def should_continue(self, state: AgentState) -> Literal["generate_sql", "finalize"]:
        """Determine if we should retry or finish"""
        if state.get("error_message") and state["attempts"] < state["max_attempts"]:
            return "generate_sql"
        return "finalize"
    
    def finalize(self, state: AgentState) -> AgentState:
        """Generate final answer with explanation"""
        if state.get("error_message"):
            prompt = f"""The system failed to generate a working SQL query after {state['attempts']} attempts.

Question: {state['question']}
Last Attempt SQL: {state['sql_query']}
Last Error: {state['error_message']}

Explain what went wrong and what information would be needed to answer this question."""
            
            response = self.llm.invoke([HumanMessage(content=prompt)])
            return {**state, "final_answer": f"❌ Failed after {state['attempts']} attempts.\n\n{response.content}"}
        else:
            prompt = f"""Based on the SQL query result, provide a natural language answer.

Question: {state['question']}
SQL Used: {state['sql_query']}
Result: {state['final_answer']}

Provide a clear, concise answer in natural language (1-2 sentences)."""
            
            response = self.llm.invoke([HumanMessage(content=prompt)])
            return {**state, "final_answer": response.content}

# ========== BUILD GRAPH ==========
def build_agent_graph(llm):
    """Build the LangGraph state machine"""
    agent = SQLAgent(llm)
    
    workflow = StateGraph(AgentState)
    
    # Add nodes
    workflow.add_node("generate_sql", agent.generate_sql)
    workflow.add_node("execute_sql", agent.execute_sql_node)
    workflow.add_node("reflect", agent.reflect_on_error)
    workflow.add_node("finalize", agent.finalize)
    
    # Set entry point
    workflow.set_entry_point("generate_sql")
    
    # Add edges
    workflow.add_edge("generate_sql", "execute_sql")
    workflow.add_edge("execute_sql", "reflect")
    workflow.add_conditional_edges(
        "reflect",
        agent.should_continue,
        {
            "generate_sql": "generate_sql",
            "finalize": "finalize"
        }
    )
    workflow.add_edge("finalize", END)
    
    return workflow.compile()

# ========== MAIN STREAMLIT UI ==========
def main():
    st.title("🐘 Self-Correcting SQL Agent")
    st.markdown("""
    An agent that writes SQL from natural language, executes it, parses errors, and **self-corrects** 
    using reflection + schema-aware prompting.
    """)
    
    # Sidebar for configuration
    with st.sidebar:
        st.header("🔧 Configuration")
        
        # API Key input
        default_key = os.getenv("OPENAI_API_KEY", "")
        api_key = st.text_input("OpenAI API Key", type="password", value=default_key)
        if api_key:
            os.environ["OPENAI_API_KEY"] = api_key
        
        # Check if API key is provided
        if not os.getenv("OPENAI_API_KEY"):
            st.error("⚠️ Please provide your OpenAI API key")
            st.stop()
        
        # Model selection
        model_name = st.selectbox("Model", ["gpt-3.5-turbo", "gpt-4"], index=0)
        
        # Max attempts
        max_attempts = st.slider("Max Self-Correction Attempts", 1, 5, 3)
        
        st.divider()
        
        # Database initialization
        st.header("🗄️ Database")
        if st.button("🔄 Initialize Sample Database", use_container_width=True):
            with st.spinner("Creating sample database..."):
                if init_sample_database():
                    st.session_state.db_initialized = True
                    st.session_state.current_schema = get_table_schema()
                    st.success("✅ Database initialized successfully!")
                    st.rerun()
                else:
                    st.error("Failed to initialize database")
        
        # Show schema if available
        if st.session_state.db_initialized:
            if st.button("🔄 Refresh Schema", use_container_width=True):
                st.session_state.current_schema = get_table_schema()
                st.rerun()
            
            with st.expander("📊 Database Schema", expanded=False):
                if st.session_state.current_schema:
                    st.code(st.session_state.current_schema, language="sql")
                else:
                    st.info("No schema available. Click 'Refresh Schema'")
    
    # Main content area
    if not st.session_state.db_initialized:
        st.info("👈 Click 'Initialize Sample Database' in the sidebar to start")
        return
    
    # Chat interface
    st.subheader("💬 Ask Questions About Your Data")
    
    # Display chat history
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])
            if "details" in msg and msg["details"]:
                with st.expander("🔍 View SQL & Details"):
                    st.code(msg["details"].get("sql", "No SQL available"), language="sql")
                    st.json({k: v for k, v in msg["details"].items() if k != "sql"})
    
    # Chat input
    if prompt := st.chat_input("Example: 'Show me all employees in Engineering department'"):
        # Add user message
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Initialize LLM
        llm = ChatOpenAI(model=model_name, temperature=0)
        
        # Run agent
        with st.chat_message("assistant"):
            with st.spinner("🤔 Agent analyzing and self-correcting..."):
                # Create initial state
                initial_state = AgentState(
                    question=prompt,
                    sql_query="",
                    execution_result="",
                    error_message="",
                    attempts=0,
                    max_attempts=max_attempts,
                    schema_info=st.session_state.current_schema or "",
                    reflection="",
                    final_answer=""
                )
                
                # Build and run graph
                graph = build_agent_graph(llm)
                
                # Stream execution
                status_placeholder = st.empty()
                final_state = None
                
                for output in graph.stream(initial_state):
                    final_state = output
                    for node_name, node_state in output.items():
                        if node_name == "generate_sql" and node_state.get("sql_query"):
                            status_placeholder.info(f"📝 **Step {node_state['attempts']}**: Generating SQL...")
                        elif node_name == "execute_sql":
                            if node_state.get("error_message"):
                                status_placeholder.warning(f"❌ **Error**: {node_state['error_message'][:200]}...")
                            else:
                                status_placeholder.success(f"✅ **Success**: {node_state.get('execution_result', '')}")
                        elif node_name == "reflect" and node_state.get("reflection") and "successfully" not in node_state["reflection"]:
                            status_placeholder.info(f"💡 **Reflection**: {node_state['reflection'][:200]}...")
                
                # Get final answer
                if final_state:
                    last_state = list(final_state.values())[-1]
                    answer = last_state.get("final_answer", "No answer generated")
                    st.markdown(answer)
                    
                    # Show execution details
                    with st.expander("🔍 Execution Details"):
                        col1, col2 = st.columns(2)
                        with col1:
                            st.metric("Attempts", last_state.get("attempts", 0))
                            st.metric("Max Attempts", max_attempts)
                        with col2:
                            st.metric("Success", "✅" if not last_state.get("error_message") else "❌")
                        
                        if last_state.get("sql_query"):
                            st.subheader("Final SQL Query")
                            st.code(last_state["sql_query"], language="sql")
                        
                        if last_state.get("reflection") and "successfully" not in last_state["reflection"]:
                            st.subheader("Reflections")
                            st.info(last_state["reflection"])
                    
                    # Save to session
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": answer,
                        "details": {
                            "attempts": last_state.get("attempts", 0),
                            "sql": last_state.get("sql_query", ""),
                            "success": not bool(last_state.get("error_message"))
                        }
                    })
                else:
                    st.error("Agent failed to produce a response")
        
        # Clear status and rerun to update UI
        st.rerun()

if __name__ == "__main__":
    main()