# ============================================================ # DDS SQL Agent with Modern LangChain Memory + Gradio UI # Hugging Face Spaces version # ============================================================ import os import re import sqlite3 from pathlib import Path from uuid import uuid4 import gradio as gr from langchain.agents import create_agent from langchain.tools import tool from langgraph.checkpoint.memory import InMemorySaver # Optional Hugging Face ZeroGPU support. # This is useful only if you select ZeroGPU hardware in Space settings. # For this OpenAI API app, CPU Basic is recommended. try: import spaces except Exception: spaces = None # ------------------------------------------------------------ # 1. Environment configuration # ------------------------------------------------------------ # Add this in Hugging Face Space Settings -> Variables and Secrets: # Secret name: OPENAI_API_KEY # # Optional Space variables: # MODEL_NAME = openai:gpt-5.4 # DATABASE_PATH = data/Chinook_Sqlite.sqlite OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") MODEL_NAME = os.getenv("MODEL_NAME", "openai:gpt-5.4") # Download the real Chinook SQLite DB directly from GitHub. # This avoids manually uploading the DB file to Hugging Face Spaces. CHINOOK_URL = os.getenv( "CHINOOK_URL", "https://github.com/lerocha/chinook-database/raw/master/ChinookDatabase/DataSources/Chinook_Sqlite.sqlite", ) # Local runtime path inside the Space. # You can override this with DATABASE_PATH if needed. DB_PATH = Path(os.getenv("DATABASE_PATH", "Chinook.db")) # ------------------------------------------------------------ # Optional ZeroGPU mode # ------------------------------------------------------------ # Recommended for this app: CPU Basic or CPU Upgrade. # # Why? # - The LLM is called through the OpenAI API. # - The model is not loaded locally on Hugging Face. # - SQLite and Gradio do not require GPU. # # If you selected ZeroGPU hardware and see "no GPU function", # keep USE_ZEROGPU=true. The @spaces.GPU decorator tells HF # that this function is allowed to request ZeroGPU. # # If you run on CPU hardware, this can stay true; HF says the # decorator is effect-free in non-ZeroGPU environments. USE_ZEROGPU = os.getenv("USE_ZEROGPU", "true").strip().lower() in { "1", "true", "yes", "y", } # ------------------------------------------------------------ # 2. Database download + validation helpers # ------------------------------------------------------------ APP_DIR = Path(__file__).resolve().parent def is_sqlite_database_file(path: Path) -> bool: """ A valid SQLite database starts with: SQLite format 3\x00 """ if not path.exists() or not path.is_file(): return False try: with open(path, "rb") as file: header = file.read(16) return header == b"SQLite format 3\x00" except Exception: return False def inspect_file_type(path: Path) -> str: """ Diagnose common file issues. """ if not path.exists(): return "missing" if path.is_dir(): return "directory" try: with open(path, "rb") as file: sample = file.read(4096) if sample.startswith(b"SQLite format 3\x00"): return "sqlite" if sample.startswith(b"PK"): return "zip_file" lower_sample = sample.lower() if b"version https://git-lfs.github.com/spec" in lower_sample: return "git_lfs_pointer" if b" Path: """ Resolve DB path inside Hugging Face Spaces. If DATABASE_PATH is relative, place it relative to the app directory. Example: DATABASE_PATH=Chinook.db becomes: /home/user/app/Chinook.db """ if path.is_absolute(): return path return APP_DIR / path def download_file(url: str, output_path: Path) -> None: """ Download file using Python standard library. requests is intentionally avoided to keep requirements simpler. """ import urllib.request output_path.parent.mkdir(parents=True, exist_ok=True) request = urllib.request.Request( url, headers={ "User-Agent": "Mozilla/5.0 HuggingFaceSpace SQLite Downloader", }, ) with urllib.request.urlopen(request, timeout=60) as response: content = response.read() output_path.write_bytes(content) def download_chinook_database_if_needed() -> Path: """ Download the real Chinook SQLite database from GitHub if needed. This function fixes: - missing DB files - corrupted files - HTML files saved as DB - Git LFS pointer files - SQL scripts renamed as .db/.sqlite """ runtime_db_path = resolve_runtime_db_path(DB_PATH) if is_sqlite_database_file(runtime_db_path): print(f"Using existing valid SQLite database: {runtime_db_path}") return runtime_db_path if runtime_db_path.exists(): existing_type = inspect_file_type(runtime_db_path) print( f"Existing database path is not valid SQLite: {runtime_db_path}. " f"Detected type: {existing_type}. Re-downloading..." ) try: runtime_db_path.unlink() except Exception: pass else: print(f"Database not found at {runtime_db_path}. Downloading...") print(f"Downloading Chinook database from: {CHINOOK_URL}") download_file(CHINOOK_URL, runtime_db_path) if not is_sqlite_database_file(runtime_db_path): detected_type = inspect_file_type(runtime_db_path) raise sqlite3.DatabaseError( f""" Downloaded file is not a valid SQLite database. Download URL: {CHINOOK_URL} Saved path: {runtime_db_path} Detected file type: {detected_type} Possible fixes: 1. Check that CHINOOK_URL points to a raw SQLite file. 2. Use this default URL: https://github.com/lerocha/chinook-database/raw/master/ChinookDatabase/DataSources/Chinook_Sqlite.sqlite 3. Do not use a normal GitHub webpage URL. 4. Do not use a .sql dump URL unless you add SQL conversion logic. """ ) print(f"Successfully downloaded valid SQLite database: {runtime_db_path}") print(f"Database size: {runtime_db_path.stat().st_size:,} bytes") return runtime_db_path DB_PATH = download_chinook_database_if_needed() def get_database_schema(db_path: Path) -> str: """ Extract table and column information from the SQLite database. This schema is injected into the system prompt so the agent knows the DB structure. """ conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute( """ SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name; """ ) tables = [row[0] for row in cursor.fetchall()] schema_lines = [] for table in tables: schema_lines.append(f"\nTable: {table}") cursor.execute(f"PRAGMA table_info({table});") columns = cursor.fetchall() for column in columns: # PRAGMA table_info columns: # cid, name, type, notnull, dflt_value, pk _, name, col_type, notnull, _, pk = column flags = [] if pk: flags.append("PRIMARY KEY") if notnull: flags.append("NOT NULL") flag_text = f" ({', '.join(flags)})" if flags else "" schema_lines.append(f"- {name}: {col_type}{flag_text}") conn.close() return "\n".join(schema_lines) DATABASE_SCHEMA = get_database_schema(DB_PATH) def strip_sql_code_fences(query: str) -> str: """ Removes markdown code fences if the model returns SQL inside ```sql ... ```. """ query = query.strip() if query.startswith("```"): query = re.sub(r"^```(?:sql)?", "", query, flags=re.IGNORECASE).strip() query = re.sub(r"```$", "", query).strip() return query def is_read_only_sql(query: str) -> bool: """ Basic read-only protection. Allows SELECT, WITH, PRAGMA, and EXPLAIN. Blocks INSERT, UPDATE, DELETE, DROP, ALTER, CREATE, etc. """ cleaned = strip_sql_code_fences(query) cleaned = re.sub(r"/\*.*?\*/", "", cleaned, flags=re.DOTALL) cleaned = re.sub(r"--.*?$", "", cleaned, flags=re.MULTILINE) cleaned = cleaned.strip().lower() allowed_starts = ("select", "with", "pragma", "explain") if not cleaned.startswith(allowed_starts): return False blocked_keywords = [ "insert ", "update ", "delete ", "drop ", "alter ", "create ", "replace ", "truncate ", "attach ", "detach ", "vacuum", "reindex", ] return not any(keyword in cleaned for keyword in blocked_keywords) def rows_to_markdown(columns, rows, max_rows: int = 50) -> str: """ Convert SQL rows to a Markdown table for readable chatbot output. """ if not rows: return "Query executed successfully, but returned no rows." rows = rows[:max_rows] def clean_cell(value): if value is None: return "" text = str(value) text = text.replace("\n", " ").replace("|", "\\|") return text header = "| " + " | ".join(columns) + " |" separator = "| " + " | ".join(["---"] * len(columns)) + " |" body_lines = [] for row in rows: body_lines.append("| " + " | ".join(clean_cell(value) for value in row) + " |") return "\n".join([header, separator] + body_lines) # ------------------------------------------------------------ # 3. SQL tool # ------------------------------------------------------------ @tool def execute_sql(query: str) -> str: """ Execute a read-only SQLite SQL query against the Chinook database. Use this tool when the user asks analytical questions that require database access. Only SELECT, WITH, PRAGMA, and EXPLAIN queries are allowed. """ query = strip_sql_code_fences(query) if not is_read_only_sql(query): return ( "Blocked for safety. Only read-only SQL is allowed. " "Please use SELECT, WITH, PRAGMA, or EXPLAIN queries." ) try: conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() cursor.execute(query) rows = cursor.fetchall() columns = [description[0] for description in cursor.description] if cursor.description else [] conn.close() if not columns: return "Query executed successfully." result_table = rows_to_markdown(columns, rows) if len(rows) > 50: result_table += f"\n\nShowing first 50 rows out of {len(rows)} rows." return result_table except Exception as e: return f"SQL execution error: {str(e)}" # ------------------------------------------------------------ # 4. System prompt # ------------------------------------------------------------ SYSTEM_PROMPT = f""" You are a helpful SQL data analyst for the Chinook SQLite database. Your job: - Understand the user's business/data question. - Write correct SQLite queries. - Use the execute_sql tool to query the database. - Explain the result clearly and concisely. - For follow-up questions, use the conversation memory. Important rules: - Use only read-only SQL. - Never modify the database. - Prefer clear SQL with explicit table joins. - When useful, explain the SQL logic briefly. - If the user asks a vague question, make a reasonable interpretation and proceed. - If the database does not contain enough information, say that clearly. Available database schema: {DATABASE_SCHEMA} """ # ------------------------------------------------------------ # 5. Create LangChain agent with short-term memory # ------------------------------------------------------------ # InMemorySaver gives thread-level memory during the live Space session. # For production-grade persistent memory, replace this with a database-backed checkpointer. checkpointer = InMemorySaver() sql_agent_with_memory = create_agent( model=MODEL_NAME, tools=[execute_sql], system_prompt=SYSTEM_PROMPT, checkpointer=checkpointer, ) # ------------------------------------------------------------ # 6. Gradio helpers # ------------------------------------------------------------ def content_to_text(content): """ Convert LangChain message content into displayable text. """ if isinstance(content, str): return content if isinstance(content, list): text_parts = [] for item in content: if isinstance(item, dict): if "text" in item: text_parts.append(item["text"]) elif "content" in item: text_parts.append(str(item["content"])) else: text_parts.append(str(item)) else: text_parts.append(str(item)) return "\n".join(text_parts) return str(content) def create_thread_id(): """ Same thread_id = same LangGraph memory. New thread_id = fresh conversation. """ return f"dds-sql-agent-{uuid4()}" def normalize_history_to_messages(history): """ Gradio expects messages format: [ {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."} ] """ if history is None: return [] normalized = [] for item in history: if isinstance(item, dict) and "role" in item and "content" in item: role = item.get("role") if role in ["user", "assistant"]: normalized.append( { "role": role, "content": content_to_text(item.get("content", "")), } ) return normalized # ------------------------------------------------------------ # 7. Gradio chat function # ------------------------------------------------------------ def zerogpu_compatible(fn): """ Optional Hugging Face ZeroGPU wrapper. If ZeroGPU hardware is selected, Hugging Face expects at least one function to be decorated with @spaces.GPU. For this app, GPU is not technically required because the LLM runs through the OpenAI API. CPU Basic is recommended. This wrapper exists only to make the Space compatible with ZeroGPU if selected. """ if USE_ZEROGPU and spaces is not None: return spaces.GPU(duration=120)(fn) return fn @zerogpu_compatible def chat_with_sql_agent(message, history, thread_id): """ Handles one user message from Gradio. This returns messages format without passing type="messages" to gr.Chatbot, because some Gradio 6 runtimes expect messages but do not accept the type argument. """ history = normalize_history_to_messages(history) if not OPENAI_API_KEY: assistant_message = ( "OPENAI_API_KEY is missing. In Hugging Face Spaces, go to " "Settings → Variables and Secrets → New Secret, then add:\n\n" "`OPENAI_API_KEY = your_openai_api_key`" ) return history + [ {"role": "user", "content": message or ""}, {"role": "assistant", "content": assistant_message}, ], "", thread_id or create_thread_id() if not thread_id: thread_id = create_thread_id() if not message or not message.strip(): return history, "", thread_id user_message = message.strip() try: result = sql_agent_with_memory.invoke( { "messages": [ { "role": "user", "content": user_message, } ] }, config={ "configurable": { "thread_id": thread_id } }, ) assistant_message = content_to_text(result["messages"][-1].content) except Exception as e: assistant_message = f""" Something went wrong while running the SQL agent. Error: ```text {str(e)} ``` Check: 1. OPENAI_API_KEY is set in Hugging Face Secrets. 2. MODEL_NAME is available in your OpenAI account. 3. The SQLite database file exists at: `{DB_PATH}` """ updated_history = history + [ { "role": "user", "content": user_message, }, { "role": "assistant", "content": assistant_message, }, ] return updated_history, "", thread_id def reset_chat(): """ Clears UI history and starts a fresh memory thread. """ return [], create_thread_id() def example_question(question): """ Puts an example question into the textbox. """ return question # ------------------------------------------------------------ # 8. Build Gradio app # ------------------------------------------------------------ custom_css = """ #main-container { max-width: 1100px; margin: 0 auto; } .dds-note { font-size: 0.95rem; opacity: 0.85; } """ with gr.Blocks(title="DDS SQL Agent", css=custom_css) as demo: thread_id_state = gr.State(value=create_thread_id()) with gr.Column(elem_id="main-container"): gr.Markdown( f""" # DDS SQL Agent with Memory Ask questions about the Chinook SQLite database. The agent can generate SQL, execute read-only queries, and remember follow-up questions in the same session. **Model:** `{MODEL_NAME}` **Database:** `{DB_PATH}` **Hardware note:** CPU Basic is recommended. ZeroGPU compatibility is enabled for Spaces that require it. """ ) if not OPENAI_API_KEY: gr.Markdown( """ > **Setup needed:** `OPENAI_API_KEY` is not set. > Add it in Hugging Face Spaces under **Settings → Variables and Secrets → New Secret**. """ ) chatbot = gr.Chatbot( value=[], height=560, label="SQL Agent Chat", placeholder="Ask a question about the database...", ) with gr.Row(): user_input = gr.Textbox( placeholder="Example: Which customer spent the most money?", label="Your question", scale=8, ) submit_btn = gr.Button( "Ask", scale=1, variant="primary", ) with gr.Row(): clear_btn = gr.Button("New Chat / Reset Memory") gr.Markdown("### Example questions") with gr.Row(): ex1 = gr.Button("Which customer spent the most money?") ex2 = gr.Button("Show total sales by country.") ex3 = gr.Button("Which genre has the most tracks?") ex4 = gr.Button("What are the top-selling tracks?") ex1.click(example_question, inputs=[gr.State("Which customer spent the most money?")], outputs=[user_input]) ex2.click(example_question, inputs=[gr.State("Show total sales by country.")], outputs=[user_input]) ex3.click(example_question, inputs=[gr.State("Which genre has the most tracks?")], outputs=[user_input]) ex4.click(example_question, inputs=[gr.State("What are the top-selling tracks?")], outputs=[user_input]) submit_btn.click( fn=chat_with_sql_agent, inputs=[user_input, chatbot, thread_id_state], outputs=[chatbot, user_input, thread_id_state], ) user_input.submit( fn=chat_with_sql_agent, inputs=[user_input, chatbot, thread_id_state], outputs=[chatbot, user_input, thread_id_state], ) clear_btn.click( fn=reset_chat, inputs=[], outputs=[chatbot, thread_id_state], ) # ------------------------------------------------------------ # 9. Launch for Hugging Face Spaces # ------------------------------------------------------------ if __name__ == "__main__": demo.queue().launch()