Upload engine.py

engine.py

@@ -3,13 +3,16 @@ import re
 import sqlite3
 from openai import OpenAI
 from difflib import get_close_matches
-from datetime import datetime
 
 # =========================
 # SETUP
 # =========================
 
-client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+# Validate API key
+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)
 
 # =========================
@@ -56,12 +59,6 @@ KNOWN_TERMS = [
     "admitted", "admission",
     "year", "month", "last", "recent", "today"
 ]
-DOMAIN_ALIASES = {
-    "consultant": ["provider", "encounter"],
-    "doctor": ["provider"],
-    "appointment": ["encounter"],
-    "visit": ["encounter"],
-}
 
 def correct_spelling(q):
     words = q.split()
@@ -219,9 +216,13 @@ def describe_schema(max_tables=10):
 # =========================
 
 def get_latest_data_date():
+    """Get the latest data date from encounters table."""
     cur = conn.cursor()
-    r = cur.execute("SELECT MAX(start_date) FROM encounters").fetchone()
-    return r[0]
+    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 None
 
 def normalize_time_question(q):
     latest = get_latest_data_date()
@@ -327,22 +328,32 @@ If the question mentions "consultant" or "doctor", use the table name "encounter
 
 
 def call_llm(prompt):
-    res = client.chat.completions.create(
-        model="gpt-4.1-mini",
-        messages=[
-            {"role": "system", "content": "Return only SQL or NOT_ANSWERABLE"},
-            {"role": "user", "content": prompt}
-        ],
-        temperature=0
-    )
-    return res.choices[0].message.content.strip()
+    """Call OpenAI API with error handling."""
+    try:
+        res = client.chat.completions.create(
+            model="gpt-4.1-mini",
+            messages=[
+                {"role": "system", "content": "Return only SQL or NOT_ANSWERABLE"},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0
+        )
+        if not res.choices or not res.choices[0].message.content:
+            raise ValueError("Empty response from OpenAI API")
+        return res.choices[0].message.content.strip()
+    except Exception as e:
+        raise ValueError(f"OpenAI API error: {str(e)}")
 
 # =========================
 # SQL SAFETY
 # =========================
 
 def sanitize_sql(sql):
-    sql = sql.replace("```", "").replace("sql", "").strip()
+    # Remove code fence markers but preserve legitimate SQL
+    sql = sql.replace("```sql", "").replace("```", "").strip()
+    # Remove leading/trailing markdown code markers
+    if sql.startswith("sql"):
+        sql = sql[3:].strip()
     sql = sql.split(";")[0]
     return sql.replace("\n", " ").strip()
 
@@ -373,14 +384,21 @@ def correct_table_names(sql):
 
 def validate_sql(sql):
     if not sql.lower().startswith("select"):
-        raise Exception("Only SELECT allowed")
+        raise ValueError("Only SELECT allowed")
     return sql
 
 def run_query(sql):
+    """Execute SQL query with proper error handling."""
     cur = conn.cursor()
-    rows = cur.execute(sql).fetchall()
-    cols = [c[0] for c in cur.description]
-    return cols, rows
+    try:
+        rows = cur.execute(sql).fetchall()
+        if cur.description:
+            cols = [c[0] for c in cur.description]
+        else:
+            cols = []
+        return cols, rows
+    except sqlite3.Error as e:
+        raise ValueError(f"Database query error: {str(e)}")
 
 # =========================
 # AGGREGATE SAFETY
@@ -391,32 +409,61 @@ def is_aggregate_only_query(sql):
     return ("count(" in s or "sum(" in s or "avg(" in s) and "group by" not in s
 
 def has_underlying_data(sql):
+    """Check if underlying data exists for the SQL query."""
     base = sql.lower()
     if "from" not in base:
         return False
 
     base = base.split("from", 1)[1]
-    test_sql = "SELECT 1 FROM " + base.split("group by")[0] + " LIMIT 1"
+    # Split at GROUP BY, ORDER BY, LIMIT, etc. to get just the FROM clause
+    for clause in ["group by", "order by", "limit", "having"]:
+        base = base.split(clause)[0]
+    
+    test_sql = "SELECT 1 FROM " + base.strip() + " LIMIT 1"
 
     cur = conn.cursor()
-    return cur.execute(test_sql).fetchone() is not None
+    try:
+        return cur.execute(test_sql).fetchone() is not None
+    except sqlite3.Error:
+        return False
 
 # =========================
 # PATIENT SUMMARY
 # =========================
 
-def build_table_summary(table_name):
-    cur = conn.cursor()
-
-    # Total rows (still need to query actual data for count)
-    total = cur.execute(
-        f"SELECT COUNT(*) FROM {table_name}"
-    ).fetchone()[0]
+def validate_identifier(name):
+    """Validate that identifier is safe (only alphanumeric and underscores)."""
+    if not name or not isinstance(name, str):
+        return False
+    # Check for SQL injection attempts
+    forbidden = [";", "--", "/*", "*/", "'", '"', "`", "(", ")", " ", "\n", "\t"]
+    if any(char in name for char in forbidden):
+        return False
+    # Must start with letter or underscore, rest alphanumeric/underscore
+    return bool(re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', name))
 
-    # Get column info from METADATA (ai_columns) not database structure
+def build_table_summary(table_name):
+    """Build summary for a table using metadata."""
+    # Validate table name against metadata first
     schema = load_ai_schema()
     if table_name not in schema:
         return f"Table {table_name} not found in metadata."
+    
+    # Additional safety check
+    if not validate_identifier(table_name):
+        return f"Invalid table name: {table_name}"
+
+    cur = conn.cursor()
+    
+    # Total rows (still need to query actual data for count)
+    # Note: SQLite doesn't support parameterized table names
+    # Since we validated table_name against metadata, it's safe
+    try:
+        total = cur.execute(
+            f"SELECT COUNT(*) FROM {table_name}"
+        ).fetchone()[0]
+    except sqlite3.Error as e:
+        return f"Error querying table {table_name}: {str(e)}"
 
     columns = schema[table_name]["columns"]  # [(col_name, description), ...]
 
@@ -425,6 +472,10 @@ def build_table_summary(table_name):
 
     # Try to summarize categorical columns using metadata
     for col_name, col_desc in columns:
+        # Validate column name
+        if not validate_identifier(col_name):
+            continue
+            
         # Try to determine if it's a categorical column based on name/description
         # Skip likely numeric/date columns
         col_lower = col_name.lower()
@@ -433,6 +484,7 @@ def build_table_summary(table_name):
             
         # Try to get breakdown for text-like columns
         try:
+            # Note: SQLite doesn't support parameterized identifiers, so we validate
             rows = cur.execute(
                 f"""
                 SELECT {col_name}, COUNT(*)
@@ -447,8 +499,9 @@ def build_table_summary(table_name):
                 summary += f"\n• {col_name.capitalize()} breakdown:\n"
                 for val, count in rows:
                     summary += f"  - {val}: {count}\n"
-        except:
-            pass  # ignore columns that can't be grouped
+        except (sqlite3.Error, sqlite3.OperationalError) as e:
+            # Ignore columns that can't be grouped (likely not categorical)
+            pass
 
     summary += "\nYou can ask more detailed questions about this data."