app.py

import os
import gradio as gr
import pandas as pd
from sqlalchemy import create_engine, text
from langchain.tools import tool
from code_agent import CodeAgent
from hf_api_model import HfApiModel

# Initialize SQLite database engine
engine = create_engine('sqlite:///data.db')

def clear_database():
    """
    Clear all tables from the database.
    """
    with engine.connect() as con:
        # Get all table names
        tables = con.execute(text(
            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
        )).fetchall()
        
        # Drop each table
        for table in tables:
            con.execute(text(f"DROP TABLE IF EXISTS {table[0]}"))

def create_dynamic_table(df):
    """
    Create a table dynamically based on DataFrame structure.
    """
    df.to_sql('data_table', engine, index=False, if_exists='replace')
    return 'data_table'

def insert_rows_into_table(records, table_name):
    """
    Insert records into the specified table.
    """
    with engine.begin() as conn:
        for record in records:
            conn.execute(
                text(f"INSERT INTO {table_name} ({', '.join(record.keys())}) VALUES ({', '.join(['?' for _ in record])})")
                .bindparams(*record.values())
            )

def get_data_table():
    """
    Get the current data table as a DataFrame.
    """
    try:
        # Get list of tables
        with engine.connect() as con:
            tables = con.execute(text(
                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
            )).fetchall()
            
        if not tables:
            return pd.DataFrame()
            
        # Use the first table found
        table_name = tables[0][0]
        
        # Read the table into a DataFrame
        return pd.read_sql_table(table_name, engine)
        
    except Exception as e:
        return pd.DataFrame({"Error": [str(e)]})

def get_table_info():
    """
    Gets the current table name and column information.
    Returns:
        tuple: (table_name, list of column names, column info)
    """
    try:
        # Get list of tables
        with engine.connect() as con:
            tables = con.execute(text(
                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
            )).fetchall()
            
        if not tables:
            return None, [], {}
            
        # Use the first table found
        table_name = tables[0][0]
        
        # Get column information
        with engine.connect() as con:
            columns = con.execute(text(f"PRAGMA table_info({table_name})")).fetchall()
            
        # Extract column names and types
        column_names = [col[1] for col in columns]
        column_info = {
            col[1]: {
                'type': col[2],
                'is_primary': bool(col[5])
            }
            for col in columns
        }
        
        return table_name, column_names, column_info
        
    except Exception as e:
        print(f"Error getting table info: {str(e)}")
        return None, [], {}

def process_sql_file(file_path):
    """
    Process an SQL file and execute its contents.
    """
    try:
        # Read the SQL file
        with open(file_path, 'r') as file:
            sql_content = file.read()
            
        # Replace AUTO_INCREMENT with AUTOINCREMENT for SQLite compatibility
        sql_content = sql_content.replace('AUTO_INCREMENT', 'AUTOINCREMENT')
            
        # Split into individual statements
        statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
        
        # Clear existing database
        clear_database()
        
        # Execute each statement
        with engine.begin() as conn:
            for statement in statements:
                if statement.strip():
                    conn.execute(text(statement))
                    
        return True, "SQL file successfully executed!"
        
    except Exception as e:
        return False, f"Error processing SQL file: {str(e)}"

def process_csv_file(file_path):
    """
    Process a CSV file and load it into the database.
    """
    try:
        # Read the CSV file
        df = pd.read_csv(file_path)
        
        if len(df.columns) == 0:
            return False, "Error: File contains no columns"
            
        # Clear existing database and create new table
        clear_database()
        table = create_dynamic_table(df)
        
        # Convert DataFrame to list of dictionaries and insert
        records = df.to_dict('records')
        insert_rows_into_table(records, table)
        
        return True, "CSV file successfully loaded!"
        
    except Exception as e:
        return False, f"Error processing CSV file: {str(e)}"

def process_uploaded_file(file):
    """
    Process the uploaded file (either SQL or CSV).
    """
    try:
        if file is None:
            return False, "Please upload a file."
            
        # Get file extension
        file_ext = os.path.splitext(file)[1].lower()
        
        if file_ext == '.sql':
            return process_sql_file(file)
        elif file_ext == '.csv':
            return process_csv_file(file)
        else:
            return False, "Error: Unsupported file type. Please upload either a .sql or .csv file."
            
    except Exception as e:
        return False, f"Error processing file: {str(e)}"

@tool
def sql_engine(query: str) -> str:
    """
    Executes an SQL query and returns formatted results.
    
    Args:
        query: The SQL query string to execute on the database. Must be a valid SELECT query.

    Returns:
        str: The formatted query results as a string.
    """
    try:
        with engine.connect() as con:
            rows = con.execute(text(query)).fetchall()

        if not rows:
            return "No results found."

        if len(rows) == 1 and len(rows[0]) == 1:
            return str(rows[0][0])

        return "\n".join([", ".join(map(str, row)) for row in rows])

    except Exception as e:
        return f"Error: {str(e)}"

def process_sql_result(generated_sql, table_name, column_names):
    """
    Process and execute the generated SQL query.
    """
    # Remove any trailing semicolons
    generated_sql = generated_sql.strip().rstrip(';')

    # Fix table names
    for wrong_name in ['table_name', 'customers', 'main']:
        if wrong_name in generated_sql:
            generated_sql = generated_sql.replace(wrong_name, table_name)
    
    # Add quotes around column names that need them
    for col in column_names:
        if ' ' in col:  # If column name contains spaces
            if col in generated_sql and f'"{col}"' not in generated_sql and f'`{col}`' not in generated_sql:
                generated_sql = generated_sql.replace(col, f'"{col}"')

    try:
        # Execute the query
        result = sql_engine(generated_sql)
        
        # Try to format as number if possible
        try:
            float_result = float(result)
            return f"{float_result:,.0f}"  # Format with commas, no decimals
        except ValueError:
            return result
            
    except Exception as e:
        if str(e).startswith("(sqlite3.OperationalError) near"):
            # If it's a SQL syntax error, return the raw result
            return generated_sql
        return f"Error executing query: {str(e)}"

def query_sql(user_query: str, show_full: bool) -> tuple:
    """
    Converts natural language input to an SQL query using CodeAgent.
    Returns both short and full responses based on switch state.
    """
    table_name, column_names, column_info = get_table_info()
    
    if not table_name:
        return "Error: No data table exists. Please upload a file first.", ""

    schema_info = (
        f"The database has a table named '{table_name}' with the following columns:\n"
        + "\n".join([
            f"- {col} ({info['type']}{' primary key' if info['is_primary'] else ''})"
            for col, info in column_info.items()
        ])
        + "\n\nGenerate a valid SQL SELECT query using ONLY these column names.\n"
        "The table name is '" + table_name + "'.\n"
        "If column names contain spaces, they must be quoted.\n"
        "You can use aggregate functions like COUNT, AVG, SUM, etc.\n"
        "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
    )

    # Get full response from the agent
    full_response = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
    
    # Process the short response as before
    if not isinstance(full_response, str):
        return "Error: Invalid query generated", ""

    # Extract and process SQL for short response
    generated_sql = full_response
    if generated_sql.isnumeric():
        short_response = generated_sql
    else:
        sql_lines = [line for line in generated_sql.split('\n') if 'select' in line.lower()]
        if sql_lines:
            generated_sql = sql_lines[0]
        
        # Process the SQL query and get the short result
        short_response = process_sql_result(generated_sql, table_name, column_names)

    return short_response, full_response

def handle_upload(file_obj):
    if file_obj is None:
        return (
            "Please upload a file.",
            None,
            "No schema available",
            gr.update(visible=True),
            gr.update(visible=False)
        )
    
    success, message = process_uploaded_file(file_obj)
    if success:
        df = get_data_table()
        _, _, column_info = get_table_info()
        schema = "\n".join([
            f"- {col} ({info['type']}){'primary key' if info['is_primary'] else ''}"
            for col, info in column_info.items()
        ])
        return (
            message,
            df,
            f"### Current Schema:\n```\n{schema}\n```",
            gr.update(visible=False),
            gr.update(visible=True)
        )
    return (
        message,
        None,
        "No schema available",
        gr.update(visible=True),
        gr.update(visible=False)
    )

def refresh_data():
    df = get_data_table()
    _, _, column_info = get_table_info()
    schema = "\n".join([
        f"- {col} ({info['type']}){'primary key' if info['is_primary'] else ''}"
        for col, info in column_info.items()
    ])
    return df, f"### Current Schema:\n```\n{schema}\n```"

# Initialize the CodeAgent
agent = CodeAgent(
    tools=[sql_engine],
    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
)

# Create the Gradio interface
with gr.Blocks() as demo:
    with gr.Group() as upload_group:
        gr.Markdown("""
        # CSVAgent
        Upload your data file to begin.
        
        ### Supported File Types:
        - CSV (.csv): CSV file with headers that will be automatically converted to a table
        
        ### CSV Requirements:
        - Must include headers
        - First column will be used as the primary key
        - Column types will be automatically detected
        - Sample CSV Files: https://github.com/datablist/sample-csv-files
        ### Based on ZennyKenny's SqlAgent

        ### SQL to CSV File Conversion
        https://tableconvert.com/sql-to-csv
        - Will work on the handling of SQL files soon.
        
        ### Try it out! Upload a CSV file and then ask a question about the data!
        """)
        
        file_input = gr.File(
            label="Upload Data File",
            file_types=[".csv", ".sql"],
            type="filepath"
        )
        status = gr.Textbox(label="Status", interactive=False)

    with gr.Group(visible=False) as query_group:
        with gr.Row():
            with gr.Column(scale=1):
                user_input = gr.Textbox(label="Ask a question about the data")
                query_output = gr.Textbox(label="Result")
                # Add the switch and secondary result box
                full_response_switch = gr.Switch(label="Show Full Response", value=False)
                full_response_output = gr.Textbox(label="Full Response", visible=False)
            
            with gr.Column(scale=2):
                gr.Markdown("### Current Data")
                data_table = gr.Dataframe(
                    value=None,
                    label="Data Table",
                    interactive=False
                )
        
        schema_display = gr.Markdown(value="Loading schema...")
        refresh_btn = gr.Button("Refresh Data")

    # Event handlers
    file_input.upload(
        fn=handle_upload,
        inputs=file_input,
        outputs=[
            status,
            data_table,
            schema_display,
            upload_group,
            query_group
        ]
    )

    user_input.change(
        fn=query_sql,
        inputs=[user_input, full_response_switch],
        outputs=[query_output, full_response_output]
    )
    
    # Add switch change event to control visibility of full response
    full_response_switch.change(
        fn=lambda x: gr.update(visible=x),
        inputs=full_response_switch,
        outputs=full_response_output
    )
    
    refresh_btn.click(
        fn=refresh_data,
        outputs=[data_table, schema_display]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860
    )