Update app.py

app.py

@@ -22,10 +22,8 @@ from typing import List, Dict, Any, Tuple, Optional
 
 import gradio as gr
 import pandas as pd
-import numpy as np
 
 # -------------------- OpenAI (optional) --------------------
-USE_RESPONSES_API = True
 OPENAI_AVAILABLE = True
 DEFAULT_MODEL = os.getenv("OPENAI_MODEL")  # optional override
 try:
@@ -37,11 +35,10 @@ except Exception:
 
 def _candidate_models():
     base = [
-        DEFAULT_MODEL,
+        (DEFAULT_MODEL or "").strip() or None,
         "gpt-4o-mini",
         "gpt-4o",
         "gpt-4.1-mini",
-        "o3-mini",
     ]
     seen = set()
     return [m for m in base if m and (m not in seen and not seen.add(m))]
@@ -50,7 +47,6 @@ def _candidate_models():
 DB_DIR = "/data" if os.path.exists("/data") else "."
 DB_PATH = os.path.join(DB_DIR, "sql_trainer_dynamic.db")
 EXPORT_DIR = "."
-ADMIN_KEY = os.getenv("ADMIN_KEY", "demo")
 RANDOM_SEED = int(os.getenv("RANDOM_SEED", "7"))
 random.seed(RANDOM_SEED)
 SYS_RAND = random.SystemRandom()
@@ -60,8 +56,8 @@ DB_LOCK = threading.RLock()
 
 def connect_db():
     """
-    Single shared connection that can be used across threads.
-    All operations (reads + writes) are serialized via DB_LOCK.
+    Single shared connection usable across threads.
+    All operations (reads + writes) serialized via DB_LOCK.
     WAL mode enables concurrent reads.
     """
     con = sqlite3.connect(DB_PATH, check_same_thread=False)
@@ -221,137 +217,91 @@ FALLBACK_QUESTIONS = [
      "requires_aliases":False,"required_aliases":[]},
 ]
 
-# -------------------- OpenAI JSON schema --------------------
+# -------------------- OpenAI JSON schema (validated after parse) --------------------
 DOMAIN_AND_QUESTIONS_SCHEMA = {
-    "name": "DomainSQLPack",
-    "schema": {
-        "type": "object",
-        "additionalProperties": False,
-        "properties": {
-            "domain": {"type":"string"},
-            "tables": {
-                "type":"array",
-                "items": {
-                    "type":"object",
-                    "additionalProperties": False,
-                    "properties": {
-                        "name": {"type":"string"},
-                        "pk": {"type":"array","items":{"type":"string"}},
-                        "columns": {
-                            "type":"array",
-                            "items": {
-                                "type":"object",
-                                "additionalProperties": False,
-                                "properties": {"name":{"type":"string"}, "type":{"type":"string"}},
-                                "required":["name","type"]
-                            }
-                        },
-                        "fks": {
-                            "type":"array",
-                            "items": {
-                                "type":"object",
-                                "additionalProperties": False,
-                                "properties": {
-                                    "columns":{"type":"array","items":{"type":"string"}},
-                                    "ref_table":{"type":"string"},
-                                    "ref_columns":{"type":"array","items":{"type":"string"}}
-                                },
-                                "required":["columns","ref_table","ref_columns"]
-                            }
-                        },
-                        "rows": {"type":"array","items":{"type":["object","array"]}}
-                    },
-                    "required":["name","pk","columns","fks","rows"]
-                },
-                "minItems":3,"maxItems":4
-            },
-            "questions": {
-                "type":"array",
-                "items": {
-                    "type":"object",
-                    "additionalProperties": False,
-                    "properties": {
-                        "id":{"type":"string"},
-                        "category":{"type":"string"},
-                        "difficulty":{"type":"integer"},
-                        "prompt_md":{"type":"string"},
-                        "answer_sql":{"type":"array","items":{"type":"string"}},
-                        "requires_aliases":{"type":"boolean"},
-                        "required_aliases":{"type":"array","items":{"type":"string"}}
-                    },
-                    "required":["id","category","difficulty","prompt_md","answer_sql"]
-                },
-                "minItems":8,"maxItems":12
-            }
-        },
-        "required":["domain","tables","questions"]
-    },
-    "strict": True
+    "required": ["domain", "tables", "questions"]
 }
 
 def _domain_prompt(prev_domain: Optional[str]) -> str:
     extra = f" Avoid using the previous domain '{prev_domain}' if possible." if prev_domain else ""
     return f"""
-You are designing a small relational dataset and training questions for SQL basics.{extra}
-
-1) Choose ONE domain at random from:
-   - bookstore, retail sales, wholesaler, sales tax, oil and gas wells, marketing.
-
-2) Produce exactly 3–4 tables that fit together (SQLite-friendly):
-   - Use snake_case, avoid reserved words.
-   - Types: INTEGER, REAL, TEXT, NUMERIC, DATE (no advanced features).
-   - Primary keys (pk) and foreign keys (fks) must align.
-   - Provide 8–15 small, realistic seed rows per table (not huge).
-
-3) Generate 8–12 SQL questions covering basics with varied, natural language:
-   - Categories from: "SELECT *", "SELECT columns", "WHERE", "Aliases",
-     "JOIN (INNER)", "JOIN (LEFT)", "Aggregation", "VIEW", "CTAS / SELECT INTO".
-   - Include a few joins and at least one LEFT JOIN.
-   - Include one view creation.
-   - Include one table creation from SELECT (either CTAS or SELECT INTO).
-   - Prefer SQLite-compatible SQL. DO NOT use RIGHT/FULL OUTER JOIN.
-   - Offer 1–3 acceptable answer_sql variants per question.
-   - For 1–2 questions, require table aliases (set requires_aliases=true and list required_aliases).
-
-Return JSON only.
+Return ONLY a valid JSON object (no markdown, no prose).
+The JSON must have: domain (string), tables (3–4 table objects), and questions (8–12 question objects).{extra}
+
+Rules:
+- One domain chosen from: bookstore, retail sales, wholesaler, sales tax, oil and gas wells, marketing.
+- Tables: SQLite-friendly. Use snake_case. Each table has: name, pk (list of column names),
+  columns (list of {{name,type}}), fks (list of {{columns,ref_table,ref_columns}}), rows (8–15 small seed rows).
+- Questions: diverse natural language. Categories: "SELECT *", "SELECT columns", "WHERE", "Aliases",
+  "JOIN (INNER)", "JOIN (LEFT)", "Aggregation", "VIEW", "CTAS / SELECT INTO".
+  Include at least one LEFT JOIN, one VIEW creation, one CTAS or SELECT INTO.
+  Provide 1–3 'answer_sql' strings per question. Prefer SQLite-compatible SQL. Do NOT use RIGHT/FULL OUTER JOIN.
+  For 1–2 questions, set requires_aliases=true and list required_aliases.
+
+Example top-level keys (do not include comments in output):
+{{
+  "domain": "retail sales",
+  "tables": [...],
+  "questions": [...]
+}}
 """
 
+def _loose_json_parse(s: str) -> Optional[dict]:
+    """Extract the first JSON object from a possibly-wrapped string."""
+    try:
+        return json.loads(s)
+    except Exception:
+        pass
+    # Try to find the first {...} block
+    start = s.find("{")
+    end = s.rfind("}")
+    if start != -1 and end != -1 and end > start:
+        try:
+            return json.loads(s[start:end+1])
+        except Exception:
+            return None
+    return None
+
 def llm_generate_domain_and_questions(prev_domain: Optional[str]) -> Tuple[Optional[Dict[str,Any]], Optional[str], Optional[str]]:
     """
     Returns (obj, error_message, model_used).
+    Uses Chat Completions JSON mode if available; otherwise falls back to strict-instruction parsing.
     """
     if not OPENAI_AVAILABLE or not os.getenv("OPENAI_API_KEY"):
         return None, "OpenAI client not available or OPENAI_API_KEY missing.", None
 
     errors = []
+    prompt = _domain_prompt(prev_domain)
+
     for model in _candidate_models():
+        # Try JSON mode first (if supported)
         try:
-            prompt = _domain_prompt(prev_domain)
-            if USE_RESPONSES_API:
-                resp = _client.responses.create(
+            try:
+                chat = _client.chat.completions.create(
                     model=model,
-                    response_format={"type":"json_schema","json_schema":DOMAIN_AND_QUESTIONS_SCHEMA},
-                    input=[{"role":"user","content": prompt}],
+                    messages=[{"role":"user","content": prompt}],
                     temperature=0.6,
+                    response_format={"type":"json_object"}  # newer SDKs
                 )
-                data_text = getattr(resp, "output_text", None)
-                if not data_text:
-                    try:
-                        data_text = resp.output[0].content[0].text  # older SDK layout
-                    except Exception:
-                        data_text = None
-            else:
+                data_text = chat.choices[0].message.content
+            except TypeError:
+                # Older SDKs: no response_format argument → plain completion with strict instructions
                 chat = _client.chat.completions.create(
                     model=model,
-                    messages=[{"role":"user","content": prompt}],
+                    messages=[{"role":"system","content":"Return ONLY a JSON object. No markdown."},
+                              {"role":"user","content": prompt}],
                     temperature=0.6
                 )
                 data_text = chat.choices[0].message.content
 
-            if not data_text:
-                raise RuntimeError("Empty response from model.")
+            obj = _loose_json_parse(data_text or "")
+            if not obj:
+                raise RuntimeError("Could not parse JSON from model output.")
 
-            obj = json.loads(data_text)
+            # Minimal validation
+            for k in DOMAIN_AND_QUESTIONS_SCHEMA["required"]:
+                if k not in obj:
+                    raise RuntimeError(f"Missing key '{k}'")
             # Guardrails: strip RIGHT/FULL joins from answers
             clean_qs = []
             for q in obj.get("questions", []):