Upload 2 files

engine.py

@@ -13,7 +13,8 @@ api_key = os.getenv("OPENAI_API_KEY")
 if not api_key:
     raise ValueError("OPENAI_API_KEY environment variable is not set")
 client = OpenAI(api_key=api_key)
-conn = sqlite3.connect("hospital.db", check_same_thread=False)
+conn = sqlite3.connect("mimic_iv_demo.db", check_same_thread=False)
+
 
 # =========================
 # CONVERSATION STATE
@@ -74,32 +75,13 @@ def correct_spelling(q):
 # =========================
 # SCHEMA
 # =========================
+import json
 from functools import lru_cache
 
 @lru_cache(maxsize=1)
 def load_ai_schema():
-    cur = conn.cursor()
-    schema = {}
-
-    tables = cur.execute("""
-        SELECT table_name, description
-        FROM ai_tables
-        WHERE ai_enabled = 1
-    """).fetchall()
-
-    for table, desc in tables:
-        cols = cur.execute("""
-            SELECT column_name, description
-            FROM ai_columns
-            WHERE table_name = ? AND ai_allowed = 1
-        """, (table,)).fetchall()
-
-        schema[table] = {
-            "description": desc,
-            "columns": cols
-        }
-
-    return schema
+    with open("hospital_metadata.json", "r") as f:
+        return json.load(f)
 
 # =========================
 # TABLE MATCHING (CORE LOGIC)
@@ -216,14 +198,14 @@ def describe_schema(max_tables=10):
 # =========================
 
 def get_latest_data_date():
-    """Get the latest data date from encounters table."""
-    cur = conn.cursor()
     try:
-        r = cur.execute("SELECT MAX(start_date) FROM encounters").fetchone()
-        return r[0] if r and r[0] else None
-    except sqlite3.Error:
+        return conn.execute(
+            "SELECT MAX(admittime) FROM admissions"
+        ).fetchone()[0]
+    except:
         return None
 
+
 def normalize_time_question(q):
     latest = get_latest_data_date()
     if not latest:
@@ -319,7 +301,7 @@ If the question mentions "consultant" or "doctor", use the table name "encounter
 
     for table, meta in schema.items():
         prompt += f"\nTable: {table}\n"
-        for col, desc in meta["columns"]:
+        for col, desc in meta["columns"].items():
             prompt += f"- {col}: {desc}\n"
 
     prompt += f"\nQuestion: {question}\n"
@@ -471,7 +453,7 @@ def build_table_summary(table_name):
     summary += f"• Total records: {total}\n"
 
     # Try to summarize categorical columns using metadata
-    for col_name, col_desc in columns:
+    for col_name, col_desc in columns.items():
         # Validate column name
         if not validate_identifier(col_name):
             continue