Spaces:

decodingdatascience
/

sqlagent

Running

App Files Files Community

decodingdatascience commited on 6 days ago

Commit

b001e8e

verified ·

1 Parent(s): e57fb6c

Update app.py

Browse files

Files changed (1) hide show

app.py +231 -24

app.py CHANGED Viewed

@@ -14,6 +14,14 @@ from langchain.agents import create_agent
 from langchain.tools import tool
 from langgraph.checkpoint.memory import InMemorySaver
 # ------------------------------------------------------------
 # 1. Environment configuration
@@ -27,46 +35,225 @@ from langgraph.checkpoint.memory import InMemorySaver
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 MODEL_NAME = os.getenv("MODEL_NAME", "openai:gpt-5.4")
-DATABASE_PATH = Path(os.getenv("DATABASE_PATH", "data/Chinook_Db2.sql"))
 # ------------------------------------------------------------
-# 2. Database helpers
 # ------------------------------------------------------------
-def resolve_database_path() -> Path:
     """
-    Resolve the SQLite database path.
-    Default:
-    - data/Chinook_Sqlite.sqlite
-    You can override it in Hugging Face Spaces with:
-    DATABASE_PATH=/path/to/your/database.sqlite
     """
-    if DATABASE_PATH.exists():
-        return DATABASE_PATH
-    common_paths = [
-        Path("Chinook_Sqlite.sqlite"),
-        Path("chinook.db"),
-        Path("Chinook.db"),
-        Path("data/chinook.db"),
-        Path("data/Chinook.db"),
-    ]
-    for path in common_paths:
-        if path.exists():
-            return path
-    raise FileNotFoundError(
-        "SQLite database file was not found. "
-        "Upload your database file or set DATABASE_PATH in Hugging Face Variables."
     )
-DB_PATH = resolve_database_path()
 def get_database_schema(db_path: Path) -> str:
@@ -355,6 +542,25 @@ def normalize_history_to_messages(history):
 # 7. Gradio chat function
 # ------------------------------------------------------------
 def chat_with_sql_agent(message, history, thread_id):
     """
     Handles one user message from Gradio.
@@ -481,6 +687,7 @@ The agent can generate SQL, execute read-only queries, and remember follow-up qu
 **Model:** `{MODEL_NAME}`
 **Database:** `{DB_PATH}`
 """
         )

 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
 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"<html" in lower_sample or b"<!doctype html" in lower_sample:
+            return "html_file"
+        text_sample = sample.decode("utf-8", errors="ignore").lower()
+        sql_markers = [
+            "create table",
+            "insert into",
+            "begin transaction",
+            "pragma foreign_keys",
+            "drop table",
+        ]
+        if any(marker in text_sample for marker in sql_markers):
+            return "sql_script"
+        return "unknown"
+    except Exception:
+        return "unreadable"
+def resolve_runtime_db_path(path: Path) -> 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:
 # 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.
 **Model:** `{MODEL_NAME}`
 **Database:** `{DB_PATH}`
+**Hardware note:** CPU Basic is recommended. ZeroGPU compatibility is enabled for Spaces that require it.
 """
         )