Spaces:

Emannn2004
/

ai_Dataanalyst

Running

File size: 16,439 Bytes

0d3d494

# 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