Spaces:
Runtime error
Runtime error
| 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") | |
| import os | |
| # Use the secret set in HF Spaces | |
| hf_api_token = os.environ.get("HUGGINGFACEHUB_ACCESS_TOKEN") | |
| if hf_api_token is None: | |
| print("WARNING: HUGGINGFACEHUB_ACCESS_TOKEN not found in environment!") | |
| print("DEBUG: Initializing LLM with HF token") | |
| llm = HuggingFaceEndpoint( | |
| model="Qwen/Qwen3-Coder-Next", | |
| task="text-generation", | |
| huggingfacehub_api_token=hf_api_token # pass the env variable here | |
| ) | |
| 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 | |