Upload 6 files

Dockerfile

@@ -0,0 +1,20 @@
+# Base image
+FROM python:3.12-slim
+
+# Set working directory
+WORKDIR /app
+
+# Copy project files
+COPY . /app
+
+# Upgrade pip
+RUN pip install --upgrade pip
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Expose FastAPI port
+EXPOSE 7860
+
+# Set command for HF Spaces to run your app
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,122 @@
+import os
+import tempfile
+import sqlite3
+from fastapi import FastAPI, UploadFile, File, Form
+from pydantic import BaseModel
+from typing import Optional
+
+# Import schema extraction logic
+from schema import (
+    get_sqlite_schema,
+    get_uploaded_sqlite_schema,
+    get_mysql_schema,
+    get_postgresql_schema,
+)
+
+# Import your AI SQL logic
+from sql_analyst import generate_and_execute_sql
+
+app = FastAPI(title="AI SQL Analyst API")
+
+# ----------------------------
+# Health / Root Endpoint
+# ----------------------------
+@app.get("/")
+def root():
+    return {"status": "AI SQL Analyst API is running"}
+
+# ----------------------------
+# POST /sqlite
+# Accept a SQLite database file + user query
+# ----------------------------
+@app.post("/sqlite")
+async def analyze_sqlite(
+    file: UploadFile = File(...),
+    user_query: str = Form(...)
+):
+    """
+    Upload a SQLite file and run a natural‑language query against it.
+    """
+    # Save uploaded file to a temp path
+    temp_dir = tempfile.gettempdir()
+    sqlite_path = os.path.join(temp_dir, file.filename)
+
+    with open(sqlite_path, "wb") as f:
+        f.write(await file.read())
+
+    # Extract schema
+    schema = get_uploaded_sqlite_schema(sqlite_path)
+
+    # Build SQLite config
+    db_info = {"type": "sqlite", "path": sqlite_path}
+
+    # Generate + execute SQL
+    sql, result = generate_and_execute_sql(user_query, schema, db_info)
+
+    # Optionally remove the file (cleanup)
+    try:
+        os.remove(sqlite_path)
+    except:
+        pass
+
+    return {"sql": sql, "result": result}
+
+# ----------------------------
+# POST /builtinsqlite
+# Run query on built‑in business SQLite
+# ----------------------------
+@app.post("/builtinsqlite")
+async def analyze_business_sqlite(user_query: str = Form(...)):
+    """
+    Run a query against the built‑in business SQLite database.
+    """
+    # Get business SQLite schema
+    schema = get_sqlite_schema()
+    db_info = {"type": "sqlite", "path": "business_database.db"}
+
+    sql, result = generate_and_execute_sql(user_query, schema, db_info)
+    return {"sql": sql, "result": result}
+
+# ----------------------------
+# POST /remote
+# Accept remote DB connection details + user query
+# ----------------------------
+class RemoteRequest(BaseModel):
+    host: str
+    port: str
+    user: str
+    password: str
+    database: str
+    db_type: str  # "mysql" or "postgresql"
+    user_query: str
+
+@app.post("/remote")
+def analyze_remote(body: RemoteRequest):
+    """
+    Connect to a remote MySQL/PostgreSQL and run the natural language query.
+    """
+    db_type = body.db_type.lower()
+    if db_type not in ("mysql", "postgresql"):
+        return {"error": "db_type must be 'mysql' or 'postgresql'"}
+
+    # Extract schema
+    if db_type == "mysql":
+        schema = get_mysql_schema(
+            body.host, body.port, body.user, body.password, body.database
+        )
+    else:  # postgresql
+        schema = get_postgresql_schema(
+            body.host, body.port, body.user, body.password, body.database
+        )
+
+    db_info = {
+        "type": db_type,
+        "host": body.host,
+        "port": body.port,
+        "user": body.user,
+        "password": body.password,
+        "database": body.database,
+    }
+
+    sql, result = generate_and_execute_sql(body.user_query, schema, db_info)
+    return {"sql": sql, "result": result}

schema.py

@@ -0,0 +1,104 @@
+import sqlite3
+import sqliteschema
+from typing import Dict, Any
+
+# -------------------------
+# 1) Business SQLite Schema
+# -------------------------
+
+def get_sqlite_schema() -> Dict[str, Any]:
+    """
+    Extracts the schema of the hardcoded business SQLite database
+    and returns it as a dictionary.
+    """
+    db_path = "business_database.db"
+    extractor = sqliteschema.SQLiteSchemaExtractor(db_path)
+    schema_dict = extractor.fetch_database_schema_as_dict()
+    return schema_dict
+
+# -------------------------
+# 2) Uploaded SQLite Schema
+# -------------------------
+
+def get_uploaded_sqlite_schema(db_path: str) -> Dict[str, Any]:
+    """
+    Extracts the schema from a user uploaded SQLite file.
+    """
+    extractor = sqliteschema.SQLiteSchemaExtractor(db_path)
+    schema_dict = extractor.fetch_database_schema_as_dict()
+    return schema_dict
+
+# -------------------------
+# 3) MySQL Schema
+# -------------------------
+
+def get_mysql_schema(host: str, port: str, user: str, password: str, database: str) -> Dict[str, Any]:
+    """
+    Connects to a MySQL database and extracts its schema.
+    """
+    import mysql.connector
+
+    conn = mysql.connector.connect(
+        host=host,
+        port=int(port),
+        user=user,
+        password=password,
+        database=database
+    )
+    cursor = conn.cursor()
+    cursor.execute("SHOW TABLES")
+    tables = [t[0] for t in cursor.fetchall()]
+
+    schema: Dict[str, Any] = {}
+    for table in tables:
+        cursor.execute(f"DESCRIBE {table}")
+        cols = cursor.fetchall()
+        schema[table] = [
+            {"Field": col[0], "Type": col[1], "Null": col[2], "Key": col[3], "Default": col[4], "Extra": col[5]}
+            for col in cols
+        ]
+
+    conn.close()
+    return schema
+
+# -------------------------
+# 4) PostgreSQL Schema
+# -------------------------
+
+def get_postgresql_schema(host: str, port: str, user: str, password: str, database: str) -> Dict[str, Any]:
+    """
+    Connects to a PostgreSQL database and extracts its schema.
+    """
+    import psycopg2
+
+    conn = psycopg2.connect(
+        host=host,
+        port=int(port),
+        user=user,
+        password=password,
+        dbname=database
+    )
+    cursor = conn.cursor()
+
+    # list all public tables
+    cursor.execute("""
+        SELECT table_name
+        FROM information_schema.tables
+        WHERE table_schema = 'public'
+          AND table_type = 'BASE TABLE';
+    """)
+    tables = [t[0] for t in cursor.fetchall()]
+
+    schema: Dict[str, Any] = {}
+    for table in tables:
+        cursor.execute(f"""
+            SELECT column_name, data_type
+            FROM information_schema.columns
+            WHERE table_name = '{table}';
+        """)
+        cols = cursor.fetchall()
+        schema[table] = [{"column_name": col[0], "data_type": col[1]} for col in cols]
+
+    conn.close()
+    return schema
+

sql_analyst.py

@@ -0,0 +1,472 @@
+# import json
+# import sqlite3
+# from langchain_ollama import ChatOllama
+# from langchain_core.prompts import PromptTemplate
+# from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
+# from dotenv import load_dotenv
+
+# load_dotenv()
+
+# print("DEBUG: Initializing LLM")
+# llm = HuggingFaceEndpoint(
+#     model="Qwen/Qwen3-Coder-Next",
+#     task="text-generation"
+# )
+# model = ChatHuggingFace(llm=llm)
+
+# # ---------------------------------------------------------------------------- #
+# # Pilot Few-Shot SQL Examples (must be defined before use)
+# # ---------------------------------------------------------------------------- #
+# COMPLEX_EXAMPLES = """
+# -- Example A: Basic SELECT
+# User: "List all data from orders."
+# SQL:
+# SELECT * FROM orders;
+
+# -- Example B: Lookup by ID
+# User: "Get customer with ID 5."
+# SQL:
+# SELECT * FROM customers WHERE customer_id = 5;
+
+# -- Example C: Aggregation with Group
+# User: "Count number of products per category."
+# SQL:
+# SELECT category, COUNT(*) AS count_per_category
+# FROM products
+# GROUP BY category;
+
+# -- Example D: Update record
+# User: "Update customer ID 10's email."
+# SQL:
+# UPDATE customers SET email = 'new@example.com' WHERE customer_id = 10;
+
+# -- Example E: Delete specific records
+# User: "Delete cancelled orders before 2025."
+# SQL:
+# DELETE FROM orders WHERE status = 'cancelled' AND order_date < '2025-01-01';
+# """
+
+# # ---------------------------------------------------------------------------- #
+# # Prompt Template (stores DDL, examples, and user query)
+# # ---------------------------------------------------------------------------- #
+# prompt_template = PromptTemplate(
+#     template="""
+# -- You are an expert SQL generator.
+# -- Generate valid SQL (SELECT, INSERT, UPDATE, DELETE) to answer the user's request.
+# -- Avoid hallucinating tables/columns outside the schema.
+# -- Do not output CREATE, DROP, ALTER, or other schema changes.
+# -- If the question cannot be answered with the given schema, output:
+# -- Cannot answer with current schema.
+
+# -- SQL Dialect: {sql_dialect}
+
+# -- Schema DDL --
+# {schema_ddl}
+
+# -- Examples --
+# {few_shot_examples}
+
+# -- User Request --
+# {user_query}
+
+# -- OUTPUT ONLY the SQL:
+# """,
+#     input_variables=["schema_ddl", "few_shot_examples", "user_query", "sql_dialect"]
+# )
+
+# # ---------------------------------------------------------------------------- #
+# # Main: Generate + Execute SQL
+# # ---------------------------------------------------------------------------- #
+# def generate_and_execute_sql(user_query: str, schema: dict, db_config: dict = None):
+#     print("\nDEBUG: generate_and_execute_sql called")
+#     print("DEBUG: User query:", user_query)
+#     print("DEBUG: DB config:", db_config)
+
+#     # Determine SQL dialect
+#     if db_config is None or db_config.get("type") == "sqlite":
+#         dialect = "SQLite"
+#     elif db_config["type"] == "mysql":
+#         dialect = "MySQL"
+#     elif db_config["type"] == "postgresql":
+#         dialect = "PostgreSQL"
+#     else:
+#         dialect = "SQL"
+
+#     print("DEBUG: Using dialect:", dialect)
+
+#     ddl = convert_schema_to_ddl(schema)
+
+#     # Build final prompt with the schema and examples
+#     final_prompt = prompt_template.format(
+#         schema_ddl=ddl,
+#         few_shot_examples=COMPLEX_EXAMPLES.strip(),
+#         user_query=user_query,
+#         sql_dialect=dialect
+#     )
+
+#     print("DEBUG: Prompt sent to model:\n", final_prompt)
+
+#     # Ask LLM for SQL
+#     try:
+#         response = model.invoke([
+#             ("system", "Generate only SQL as instructed."),
+#             ("human", final_prompt),
+#         ])
+#     except Exception as e:
+#         print("ERROR: LLM invocation error:", e)
+#         return "", f"LLM invocation error: {str(e)}"
+
+#     generated_sql = response.content.strip()
+#     print("DEBUG: Model generated SQL:\n", generated_sql)
+
+#     # If the LLM says it cannot answer:
+#     if generated_sql.lower().startswith("cannot answer"):
+#         print("DEBUG: Model cannot answer with current schema")
+#         return generated_sql, "No execution (cannot answer)"
+
+#     # Execute the generated SQL
+#     print("DEBUG: Executing SQL...")
+#     result = execute_query(generated_sql, db_config)
+#     print("DEBUG: Execution result:", result)
+
+#     return generated_sql, result
+
+# # ---------------------------------------------------------------------------- #
+# # Unified SQL Executor (with column names)
+# # ---------------------------------------------------------------------------- #
+# def execute_query(query: str, db_config: dict):
+#     print("\nDEBUG: execute_query called")
+#     print("DEBUG: Query to execute:", query)
+#     print("DEBUG: DB Config:", db_config)
+
+#     db_type = db_config.get("type")
+
+#     try:
+#         if db_type == "mysql":
+#             print("DEBUG: Connecting to MySQL")
+#             import mysql.connector
+#             conn = mysql.connector.connect(
+#                 host=db_config["host"],
+#                 port=int(db_config["port"]),
+#                 user=db_config["user"],
+#                 password=db_config["password"],
+#                 database=db_config["database"],
+#             )
+
+#         elif db_type == "postgresql":
+#             print("DEBUG: Connecting to PostgreSQL")
+#             import psycopg2
+#             conn = psycopg2.connect(
+#                 host=db_config["host"],
+#                 port=int(db_config["port"]),
+#                 user=db_config["user"],
+#                 password=db_config["password"],
+#                 dbname=db_config["database"],
+#             )
+#             print("DEBUG: PostgreSQL connection successful")
+
+#         else:
+#             print("DEBUG: Connecting to SQLite")
+#             conn = sqlite3.connect(db_config.get("path"))
+
+#         cursor = conn.cursor()
+#         print("DEBUG: Cursor acquired, executing now...")
+
+#         cursor.execute(query)
+
+#         # SELECT → fetch results + columns
+#         if query.strip().lower().startswith("select"):
+#             print("DEBUG: SELECT query — fetching rows")
+#             rows = cursor.fetchall()
+
+#             # Get column names from cursor.description
+#             columns = [col[0] for col in cursor.description]  # first element is column name :contentReference[oaicite:1]{index=1}
+
+#             print("DEBUG: Column names:", columns)
+#             print("DEBUG: SELECT result rows:", rows)
+
+#             conn.close()
+#             return {"columns": columns, "rows": rows}
+
+#         # For INSERT/UPDATE/DELETE
+#         print("DEBUG: Non‑SELECT — commit and return affected count")
+#         conn.commit()
+#         affected = cursor.rowcount
+#         print("DEBUG: Affected rows:", affected)
+#         conn.close()
+#         return {"affected_rows": affected}
+
+#     except Exception as e:
+#         print("ERROR: Exception in execute_query:", e)
+#         try:
+#             conn.close()
+#         except:
+#             pass
+
+#         return f"SQL execution error: {str(e)}"
+
+# # ---------------------------------------------------------------------------- #
+# # Schema to DDL
+# # ---------------------------------------------------------------------------- #
+# def convert_schema_to_ddl(schema: dict) -> str:
+#     print("DEBUG: Converting schema to DDL")
+#     ddl_lines = []
+#     for table, cols in schema.items():
+#         col_defs = []
+#         for col in cols:
+#             col_name = (
+#                 col.get("name")
+#                 or col.get("Field")
+#                 or col.get("column_name")
+#                 or col.get("COLUMN_NAME")
+#                 or str(col)
+#             )
+#             col_type = (
+#                 col.get("type")
+#                 or col.get("Type")
+#                 or col.get("data_type")
+#                 or "TEXT"
+#             )
+#             col_defs.append(f"{col_name} {col_type}")
+#         ddl_lines.append(f"CREATE TABLE {table} ({', '.join(col_defs)});")
+#     ddl_text = "\n".join(ddl_lines)
+#     print("DEBUG: DDL:\n", ddl_text)
+#     return ddl_text
+import json
+import sqlite3
+from langchain_ollama import ChatOllama
+from langchain_core.prompts import PromptTemplate
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
+import os  # <-- for environment variables
+
+# Instead of load_dotenv(), get token from environment
+hf_token = os.environ.get("HUGGINGFACEHUB_ACCESS_TOKEN")
+if not hf_token:
+    raise ValueError("HUGGINGFACEHUB_ACCESS_TOKEN not found in environment variables!")
+
+print("DEBUG: Initializing LLM")
+llm = HuggingFaceEndpoint(
+    model="Qwen/Qwen3-Coder-Next",
+    task="text-generation",
+    huggingfacehub_api_token=hf_token  # <-- use environment token
+)
+model = ChatHuggingFace(llm=llm)
+
+# ---------------------------------------------------------------------------- #
+# Pilot Few-Shot SQL Examples (must be defined before use)
+# ---------------------------------------------------------------------------- #
+COMPLEX_EXAMPLES = """
+-- Example A: Basic SELECT
+User: "List all data from orders."
+SQL:
+SELECT * FROM orders;
+
+-- Example B: Lookup by ID
+User: "Get customer with ID 5."
+SQL:
+SELECT * FROM customers WHERE customer_id = 5;
+
+-- Example C: Aggregation with Group
+User: "Count number of products per category."
+SQL:
+SELECT category, COUNT(*) AS count_per_category
+FROM products
+GROUP BY category;
+
+-- Example D: Update record
+User: "Update customer ID 10's email."
+SQL:
+UPDATE customers SET email = 'new@example.com' WHERE customer_id = 10;
+
+-- Example E: Delete specific records
+User: "Delete cancelled orders before 2025."
+SQL:
+DELETE FROM orders WHERE status = 'cancelled' AND order_date < '2025-01-01';
+"""
+
+# ---------------------------------------------------------------------------- #
+# Prompt Template (stores DDL, examples, and user query)
+# ---------------------------------------------------------------------------- #
+prompt_template = PromptTemplate(
+    template="""
+-- You are an expert SQL generator.
+-- Generate valid SQL (SELECT, INSERT, UPDATE, DELETE) to answer the user's request.
+-- Avoid hallucinating tables/columns outside the schema.
+-- Do not output CREATE, DROP, ALTER, or other schema changes.
+-- If the question cannot be answered with the given schema, output:
+-- Cannot answer with current schema.
+
+-- SQL Dialect: {sql_dialect}
+
+-- Schema DDL --
+{schema_ddl}
+
+-- Examples --
+{few_shot_examples}
+
+-- User Request --
+{user_query}
+
+-- OUTPUT ONLY the SQL:
+""",
+    input_variables=["schema_ddl", "few_shot_examples", "user_query", "sql_dialect"]
+)
+
+# ---------------------------------------------------------------------------- #
+# Main: Generate + Execute SQL
+# ---------------------------------------------------------------------------- #
+def generate_and_execute_sql(user_query: str, schema: dict, db_config: dict = None):
+    print("\nDEBUG: generate_and_execute_sql called")
+    print("DEBUG: User query:", user_query)
+    print("DEBUG: DB config:", db_config)
+
+    # Determine SQL dialect
+    if db_config is None or db_config.get("type") == "sqlite":
+        dialect = "SQLite"
+    elif db_config["type"] == "mysql":
+        dialect = "MySQL"
+    elif db_config["type"] == "postgresql":
+        dialect = "PostgreSQL"
+    else:
+        dialect = "SQL"
+
+    print("DEBUG: Using dialect:", dialect)
+
+    ddl = convert_schema_to_ddl(schema)
+
+    # Build final prompt with the schema and examples
+    final_prompt = prompt_template.format(
+        schema_ddl=ddl,
+        few_shot_examples=COMPLEX_EXAMPLES.strip(),
+        user_query=user_query,
+        sql_dialect=dialect
+    )
+
+    print("DEBUG: Prompt sent to model:\n", final_prompt)
+
+    # Ask LLM for SQL
+    try:
+        response = model.invoke([
+            ("system", "Generate only SQL as instructed."),
+            ("human", final_prompt),
+        ])
+    except Exception as e:
+        print("ERROR: LLM invocation error:", e)
+        return "", f"LLM invocation error: {str(e)}"
+
+    generated_sql = response.content.strip()
+    print("DEBUG: Model generated SQL:\n", generated_sql)
+
+    # If the LLM says it cannot answer:
+    if generated_sql.lower().startswith("cannot answer"):
+        print("DEBUG: Model cannot answer with current schema")
+        return generated_sql, "No execution (cannot answer)"
+
+    # Execute the generated SQL
+    print("DEBUG: Executing SQL...")
+    result = execute_query(generated_sql, db_config)
+    print("DEBUG: Execution result:", result)
+
+    return generated_sql, result
+
+# ---------------------------------------------------------------------------- #
+# Unified SQL Executor (with column names)
+# ---------------------------------------------------------------------------- #
+def execute_query(query: str, db_config: dict):
+    print("\nDEBUG: execute_query called")
+    print("DEBUG: Query to execute:", query)
+    print("DEBUG: DB Config:", db_config)
+
+    db_type = db_config.get("type")
+
+    try:
+        if db_type == "mysql":
+            print("DEBUG: Connecting to MySQL")
+            import mysql.connector
+            conn = mysql.connector.connect(
+                host=db_config["host"],
+                port=int(db_config["port"]),
+                user=db_config["user"],
+                password=db_config["password"],
+                database=db_config["database"],
+            )
+
+        elif db_type == "postgresql":
+            print("DEBUG: Connecting to PostgreSQL")
+            import psycopg2
+            conn = psycopg2.connect(
+                host=db_config["host"],
+                port=int(db_config["port"]),
+                user=db_config["user"],
+                password=db_config["password"],
+                dbname=db_config["database"],
+            )
+            print("DEBUG: PostgreSQL connection successful")
+
+        else:
+            print("DEBUG: Connecting to SQLite")
+            conn = sqlite3.connect(db_config.get("path"))
+
+        cursor = conn.cursor()
+        print("DEBUG: Cursor acquired, executing now...")
+
+        cursor.execute(query)
+
+        # SELECT → fetch results + columns
+        if query.strip().lower().startswith("select"):
+            print("DEBUG: SELECT query — fetching rows")
+            rows = cursor.fetchall()
+
+            # Get column names from cursor.description
+            columns = [col[0] for col in cursor.description]
+
+            print("DEBUG: Column names:", columns)
+            print("DEBUG: SELECT result rows:", rows)
+
+            conn.close()
+            return {"columns": columns, "rows": rows}
+
+        # For INSERT/UPDATE/DELETE
+        print("DEBUG: Non‑SELECT — commit and return affected count")
+        conn.commit()
+        affected = cursor.rowcount
+        print("DEBUG: Affected rows:", affected)
+        conn.close()
+        return {"affected_rows": affected}
+
+    except Exception as e:
+        print("ERROR: Exception in execute_query:", e)
+        try:
+            conn.close()
+        except:
+            pass
+
+        return f"SQL execution error: {str(e)}"
+
+# ---------------------------------------------------------------------------- #
+# Schema to DDL
+# ---------------------------------------------------------------------------- #
+def convert_schema_to_ddl(schema: dict) -> str:
+    print("DEBUG: Converting schema to DDL")
+    ddl_lines = []
+    for table, cols in schema.items():
+        col_defs = []
+        for col in cols:
+            col_name = (
+                col.get("name")
+                or col.get("Field")
+                or col.get("column_name")
+                or col.get("COLUMN_NAME")
+                or str(col)
+            )
+            col_type = (
+                col.get("type")
+                or col.get("Type")
+                or col.get("data_type")
+                or "TEXT"
+            )
+            col_defs.append(f"{col_name} {col_type}")
+        ddl_lines.append(f"CREATE TABLE {table} ({', '.join(col_defs)});")
+    ddl_text = "\n".join(ddl_lines)
+    print("DEBUG: DDL:\n", ddl_text)
+    return ddl_text