Upload engine.py

engine.py

@@ -1,391 +1,421 @@
-import os
-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"))
-conn = sqlite3.connect("hospital.db", check_same_thread=False)
-
-# =========================
-# CONVERSATION STATE
-# =========================
-
-LAST_PROMPT_TYPE = None
-LAST_SUGGESTED_DATE = None
-
-# =========================
-# HUMAN RESPONSE HELPERS
-# =========================
-
-def humanize(text):
-    return f"Sure 🙂\n\n{text}"
-
-def friendly(text):
-    return f"{text}\n\nIf you want, I can help you explore this further 🙂"
-
-def is_confirmation(text):
-    return text.strip().lower() in ["yes", "yep", "yeah", "ok", "okay", "sure"]
-
-def is_why_question(text):
-    return text.strip().lower().startswith("why")
-
-# =========================
-# SPELL CORRECTION
-# =========================
-
-KNOWN_TERMS = [
-    "patient", "patients", "condition", "conditions",
-    "encounter", "encounters", "visit", "visits",
-    "medication", "medications",
-    "admitted", "admission",
-    "year", "month", "last", "recent", "today"
-]
-
-def correct_spelling(q):
-    words = q.split()
-    fixed = []
-    for w in words:
-        clean = w.lower().strip(",.?")
-        match = get_close_matches(clean, KNOWN_TERMS, n=1, cutoff=0.8)
-        fixed.append(match[0] if match else w)
-    return " ".join(fixed)
-
-# =========================
-# SCHEMA
-# =========================
-
-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
-
-# =========================
-# HUMAN SCHEMA DESCRIPTION
-# =========================
-
-def describe_schema():
-    schema = load_ai_schema()
-
-    response = "Here’s the data I currently have access to:\n\n"
-
-    for table, meta in schema.items():
-        response += f"• **{table.capitalize()}** — {meta['description']}\n"
-        for col, desc in meta["columns"]:
-            response += f"  - {col}: {desc}\n"
-        response += "\n"
-
-    response += (
-        "You can ask things like:\n"
-        "• How many patients are there?\n"
-        "• Patient count by gender\n"
-        "• Admissions by year\n\n"
-        "Just tell me what you want to explore 🙂"
-    )
-
-    return response
-
-# =========================
-# TIME HANDLING
-# =========================
-
-def get_latest_data_date():
-    cur = conn.cursor()
-    r = cur.execute("SELECT MAX(start_date) FROM encounters").fetchone()
-    return r[0]
-
-def normalize_time_question(q):
-    latest = get_latest_data_date()
-    if not latest:
-        return q
-
-    if "today" in q:
-        return q.replace("today", f"on {latest[:10]}")
-
-    if "yesterday" in q:
-        return q.replace("yesterday", f"on {latest[:10]}")
-
-    return q
-
-# =========================
-# UNSUPPORTED QUESTIONS
-# =========================
-
-def get_unsupported_reason(q):
-    q = q.lower()
-
-    if any(w in q for w in ["consultant", "doctor"]):
-        return {
-            "reason": "Doctor or consultant-level data is not available.",
-            "suggestion": "Try asking about patients, visits, or admissions."
-        }
-
-    if any(w in q for w in ["department", "specialization"]):
-        return {
-            "reason": "Department-level data is not stored.",
-            "suggestion": "Try patient or visit related questions."
-        }
-
-    return None
-
-# =========================
-# SQL GENERATION
-# =========================
-
-def build_prompt(question):
-    schema = load_ai_schema()
-
-    prompt = """
-You are a hospital SQL assistant.
-
-Rules:
-- Use only SELECT
-- SQLite syntax
-- Use only listed tables/columns
-- Return ONLY SQL or NOT_ANSWERABLE
-"""
-
-    for table, meta in schema.items():
-        prompt += f"\nTable: {table}\n"
-        for col, desc in meta["columns"]:
-            prompt += f"- {col}: {desc}\n"
-
-    prompt += f"\nQuestion: {question}\n"
-    return prompt
-
-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()
-
-# =========================
-# SQL SAFETY
-# =========================
-
-def sanitize_sql(sql):
-    sql = sql.replace("```", "").replace("sql", "").strip()
-    sql = sql.split(";")[0]
-    return sql.replace("\n", " ").strip()
-
-def validate_sql(sql):
-    if not sql.lower().startswith("select"):
-        raise Exception("Only SELECT allowed")
-    return sql
-
-def run_query(sql):
-    cur = conn.cursor()
-    rows = cur.execute(sql).fetchall()
-    cols = [c[0] for c in cur.description]
-    return cols, rows
-
-# =========================
-# AGGREGATE SAFETY
-# =========================
-
-def is_aggregate_only_query(sql):
-    s = sql.lower()
-    return ("count(" in s or "sum(" in s or "avg(" in s) and "group by" not in s
-
-def has_underlying_data(sql):
-    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"
-
-    cur = conn.cursor()
-    return cur.execute(test_sql).fetchone() is not None
-
-def is_patient_summary_question(text):
-    t = text.lower()
-    return any(
-        phrase in t for phrase in [
-            "patient summary",
-            "patients summary",
-            "patient overview",
-            "summary of patients",
-            "give a patient summary",
-        ]
-    )
-
-# =========================
-# PATIENT SUMMARY
-# =========================
-
-def build_patient_summary():
-    cur = conn.cursor()
-
-    total = cur.execute(
-        "SELECT COUNT(*) FROM patients"
-    ).fetchone()[0]
-
-    genders = cur.execute(
-        "SELECT gender, COUNT(*) FROM patients GROUP BY gender"
-    ).fetchall()
-
-    msg = "Here’s a quick summary of patients:\n\n"
-    msg += f"• Total patients: {total}\n"
-
-    if genders:
-        msg += "• Gender distribution:\n"
-        for g, c in genders:
-            msg += f"  - {g}: {c}\n"
-
-    msg += (
-        "\nYou can also ask:\n"
-        "• Patients admitted by year\n"
-        "• Patient count by age\n"
-        "• Visit trends"
-    )
-
-    return msg
-
-
-# =========================
-# MAIN ENGINE
-# =========================
-
-def process_question(question):
-    global LAST_PROMPT_TYPE, LAST_SUGGESTED_DATE
-
-    q = question.strip().lower()
-    
-    # Patient summary intent
-    if is_patient_summary_question(q):
-        return {
-            "status": "ok",
-            "message": build_patient_summary(),
-            "data": [],
-            "sql": None,
-            "note": None
-        }
-
-    # Handle WHY questions
-    if is_why_question(q) and LAST_PROMPT_TYPE == "NO_DATA":
-        return {
-            "status": "ok",
-            "message": (
-                f"I suggested **{LAST_SUGGESTED_DATE[:4]}** because that’s the most recent year "
-                f"for which data exists in the system.\n\n"
-                "Your database doesn’t contain newer records, so questions like "
-                "“today” or “this year” can’t be answered yet.\n\n"
-                "You can still explore:\n"
-                "• Data from 2021\n"
-                "• Trends over time\n"
-                "• Patient summaries"
-            ),
-            "data": []
-        }
-
-    # Handle confirmations
-    if is_confirmation(q) and LAST_PROMPT_TYPE == "NO_DATA":
-        return {
-            "status": "ok",
-            "message": (
-                "Great 🙂\n\n"
-                "Here are some things you can ask:\n"
-                "• How many patients were admitted in 2021?\n"
-                "• Patient count by gender\n"
-                "• Total visits by month\n"
-                "• Most common conditions"
-            ),
-            "data": []
-        }
-
-    question = correct_spelling(question)
-    question = normalize_time_question(question)
-
-    if any(x in question for x in ["what data", "what tables", "which data"]):
-        return {
-            "status": "ok",
-            "message": humanize(describe_schema()),
-            "data": []
-        }
-
-    unsupported = get_unsupported_reason(question)
-    if unsupported:
-        return {
-            "status": "ok",
-            "message": (
-                f"{unsupported['reason']}\n\n"
-                f"{unsupported['suggestion']}\n\n"
-                "Example questions:\n"
-                "• How many patients were admitted last year?\n"
-                "• Total visits by month\n"
-                "• Patient count by gender"
-            ),
-            "data": []
-        }
-
-    sql = call_llm(build_prompt(question))
-
-    if sql == "NOT_ANSWERABLE":
-        return {
-            "status": "ok",
-            "message": "I don’t have enough data to answer that.",
-            "data": []
-        }
-
-    sql = validate_sql(sanitize_sql(sql))
-    cols, rows = run_query(sql)
-
-    if is_aggregate_only_query(sql) and not has_underlying_data(sql):
-        LAST_PROMPT_TYPE = "NO_DATA"
-        LAST_SUGGESTED_DATE = get_latest_data_date()
-        return {
-            "status": "ok",
-            "message": friendly("No data is available for that time period."),
-            "note": f"Available data is only up to {LAST_SUGGESTED_DATE}.",
-            "data": [],
-            "sql": None
-        }
-
-    if not rows:
-        LAST_PROMPT_TYPE = "NO_DATA"
-        LAST_SUGGESTED_DATE = get_latest_data_date()
-        return {
-            "status": "ok",
-            "message": friendly("No records found."),
-            "note": f"Available data is only up to {LAST_SUGGESTED_DATE}.",
-            "data": []
-        }
-
-    LAST_PROMPT_TYPE = None
-    LAST_SUGGESTED_DATE = None
-
-    return {
-        "status": "ok",
-        "sql": sql,
-        "columns": cols,
-        "data": rows
-    }
+import os
+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"))
+conn = sqlite3.connect("hospital.db", check_same_thread=False)
+
+# =========================
+# CONVERSATION STATE
+# =========================
+
+LAST_PROMPT_TYPE = None
+LAST_SUGGESTED_DATE = None
+
+# =========================
+# HUMAN RESPONSE HELPERS
+# =========================
+
+def humanize(text):
+    return f"Sure 🙂\n\n{text}"
+
+def friendly(text):
+    return f"{text}\n\nIf you want, I can help you explore this further 🙂"
+
+def is_confirmation(text):
+    return text.strip().lower() in ["yes", "yep", "yeah", "ok", "okay", "sure"]
+
+def is_why_question(text):
+    return text.strip().lower().startswith("why")
+
+# =========================
+# SPELL CORRECTION
+# =========================
+
+KNOWN_TERMS = [
+    "patient", "patients", "condition", "conditions",
+    "encounter", "encounters", "visit", "visits",
+    "medication", "medications",
+    "admitted", "admission",
+    "year", "month", "last", "recent", "today"
+]
+
+def correct_spelling(q):
+    words = q.split()
+    fixed = []
+    for w in words:
+        clean = w.lower().strip(",.?")
+        match = get_close_matches(clean, KNOWN_TERMS, n=1, cutoff=0.8)
+        fixed.append(match[0] if match else w)
+    return " ".join(fixed)
+
+# =========================
+# SCHEMA
+# =========================
+
+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
+
+# =========================
+# HUMAN SCHEMA DESCRIPTION
+# =========================
+
+def describe_schema():
+    schema = load_ai_schema()
+
+    response = "Here’s the data I currently have access to:\n\n"
+
+    for table, meta in schema.items():
+        response += f"• **{table.capitalize()}** — {meta['description']}\n"
+        for col, desc in meta["columns"]:
+            response += f"  - {col}: {desc}\n"
+        response += "\n"
+
+    response += (
+        "You can ask things like:\n"
+        "• How many patients are there?\n"
+        "• Patient count by gender\n"
+        "• Admissions by year\n\n"
+        "Just tell me what you want to explore 🙂"
+    )
+
+    return response
+
+# =========================
+# TIME HANDLING
+# =========================
+
+def get_latest_data_date():
+    cur = conn.cursor()
+    r = cur.execute("SELECT MAX(start_date) FROM encounters").fetchone()
+    return r[0]
+
+def normalize_time_question(q):
+    latest = get_latest_data_date()
+    if not latest:
+        return q
+
+    if "today" in q:
+        return q.replace("today", f"on {latest[:10]}")
+
+    if "yesterday" in q:
+        return q.replace("yesterday", f"on {latest[:10]}")
+
+    return q
+
+# =========================
+# UNSUPPORTED QUESTIONS
+# =========================
+
+def get_unsupported_reason(q):
+    q = q.lower()
+
+    if any(w in q for w in ["consultant", "doctor"]):
+        return {
+            "reason": "Doctor or consultant-level data is not available.",
+            "suggestion": "Try asking about patients, visits, or admissions."
+        }
+
+    if any(w in q for w in ["department", "specialization"]):
+        return {
+            "reason": "Department-level data is not stored.",
+            "suggestion": "Try patient or visit related questions."
+        }
+
+    return None
+
+# =========================
+# SQL GENERATION
+# =========================
+
+def build_prompt(question):
+    schema = load_ai_schema()
+
+    prompt = """
+You are a hospital SQL assistant.
+
+Rules:
+- Use only SELECT
+- SQLite syntax
+- Use only listed tables/columns
+- Return ONLY SQL or NOT_ANSWERABLE
+"""
+
+    for table, meta in schema.items():
+        prompt += f"\nTable: {table}\n"
+        for col, desc in meta["columns"]:
+            prompt += f"- {col}: {desc}\n"
+
+    prompt += f"\nQuestion: {question}\n"
+    return prompt
+
+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()
+
+# =========================
+# SQL SAFETY
+# =========================
+
+def sanitize_sql(sql):
+    sql = sql.replace("```", "").replace("sql", "").strip()
+    sql = sql.split(";")[0]
+    return sql.replace("\n", " ").strip()
+
+def validate_sql(sql):
+    if not sql.lower().startswith("select"):
+        raise Exception("Only SELECT allowed")
+    return sql
+
+def run_query(sql):
+    cur = conn.cursor()
+    rows = cur.execute(sql).fetchall()
+    cols = [c[0] for c in cur.description]
+    return cols, rows
+
+# =========================
+# AGGREGATE SAFETY
+# =========================
+
+def is_aggregate_only_query(sql):
+    s = sql.lower()
+    return ("count(" in s or "sum(" in s or "avg(" in s) and "group by" not in s
+
+def has_underlying_data(sql):
+    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"
+
+    cur = conn.cursor()
+    return cur.execute(test_sql).fetchone() is not None
+
+def is_patient_summary_question(text):
+    t = text.lower()
+    return any(
+        phrase in t for phrase in [
+            "patient summary",
+            "patients summary",
+            "patient overview",
+            "summary of patients",
+            "give a patient summary",
+        ]
+    )
+
+# =========================
+# PATIENT SUMMARY
+# =========================
+
+def build_patient_summary():
+    cur = conn.cursor()
+
+    total = cur.execute(
+        "SELECT COUNT(*) FROM patients"
+    ).fetchone()[0]
+
+    genders = cur.execute(
+        "SELECT gender, COUNT(*) FROM patients GROUP BY gender"
+    ).fetchall()
+
+    msg = "Here’s a quick summary of patients:\n\n"
+    msg += f"• Total patients: {total}\n"
+
+    if genders:
+        msg += "• Gender distribution:\n"
+        for g, c in genders:
+            msg += f"  - {g}: {c}\n"
+
+    msg += (
+        "\nYou can also ask:\n"
+        "• Patients admitted by year\n"
+        "• Patient count by age\n"
+        "• Visit trends"
+    )
+
+    return msg
+
+
+# =========================
+# MAIN ENGINE
+# =========================
+
+def process_question(question):
+    global LAST_PROMPT_TYPE, LAST_SUGGESTED_DATE
+
+    q = question.strip().lower()
+
+    # -------------------------------
+    # Patient summary intent
+    # -------------------------------
+    if is_patient_summary_question(q):
+        return {
+            "status": "ok",
+            "message": build_patient_summary(),
+            "data": [],
+            "sql": None,
+            "note": None
+        }
+
+    # -------------------------------
+    # WHY follow-up handling
+    # -------------------------------
+    if is_why_question(q) and LAST_PROMPT_TYPE == "NO_DATA":
+        year = LAST_SUGGESTED_DATE[:4] if LAST_SUGGESTED_DATE else "the latest available year"
+        return {
+            "status": "ok",
+            "message": (
+                f"I suggested **{year}** because that’s the most recent year "
+                f"for which data exists in the system.\n\n"
+                "Your database doesn’t contain newer records yet.\n\n"
+                "You can explore:\n"
+                "• Data from 2021\n"
+                "• Trends over time\n"
+                "• Patient summaries"
+            ),
+            "data": []
+        }
+
+    # -------------------------------
+    # YES / confirmation handling
+    # -------------------------------
+    if is_confirmation(q) and LAST_PROMPT_TYPE == "NO_DATA":
+        return {
+            "status": "ok",
+            "message": (
+                "Great 🙂\n\n"
+                "Here are some things you can ask:\n"
+                "• How many patients were admitted in 2021?\n"
+                "• Patient count by gender\n"
+                "• Total visits by month\n"
+                "• Most common conditions"
+            ),
+            "data": []
+        }
+
+    # -------------------------------
+    # Normalize question
+    # -------------------------------
+    question = correct_spelling(question)
+    question = normalize_time_question(question)
+
+    # Reset state once user asks a fresh question
+    LAST_PROMPT_TYPE = None
+    LAST_SUGGESTED_DATE = None
+
+    # -------------------------------
+    # Metadata queries
+    # -------------------------------
+    if any(x in question for x in ["what data", "what tables", "which data"]):
+        return {
+            "status": "ok",
+            "message": humanize(describe_schema()),
+            "data": []
+        }
+
+    # -------------------------------
+    # Unsupported questions
+    # -------------------------------
+    unsupported = get_unsupported_reason(question)
+    if unsupported:
+        return {
+            "status": "ok",
+            "message": (
+                f"{unsupported['reason']}\n\n"
+                f"{unsupported['suggestion']}\n\n"
+                "Example questions:\n"
+                "• How many patients were admitted last year?\n"
+                "• Total visits by month\n"
+                "• Patient count by gender"
+            ),
+            "data": []
+        }
+
+    # -------------------------------
+    # LLM → SQL
+    # -------------------------------
+    sql = call_llm(build_prompt(question))
+
+    if sql == "NOT_ANSWERABLE":
+        return {
+            "status": "ok",
+            "message": "I don’t have enough data to answer that.",
+            "data": []
+        }
+
+    sql = validate_sql(sanitize_sql(sql))
+    cols, rows = run_query(sql)
+
+    # -------------------------------
+    # No data (aggregate case)
+    # -------------------------------
+    if is_aggregate_only_query(sql) and not has_underlying_data(sql):
+        LAST_PROMPT_TYPE = "NO_DATA"
+        LAST_SUGGESTED_DATE = get_latest_data_date()
+
+        return {
+            "status": "ok",
+            "message": friendly("No data is available for that time period."),
+            "note": f"Available data is only up to {LAST_SUGGESTED_DATE}.",
+            "data": [],
+            "sql": None
+        }
+
+    # -------------------------------
+    # Empty result set
+    # -------------------------------
+    if not rows:
+        LAST_PROMPT_TYPE = "NO_DATA"
+        LAST_SUGGESTED_DATE = get_latest_data_date()
+
+        return {
+            "status": "ok",
+            "message": friendly("No records found."),
+            "note": f"Available data is only up to {LAST_SUGGESTED_DATE}.",
+            "data": []
+        }
+
+    # -------------------------------
+    # Successful response
+    # -------------------------------
+    return {
+        "status": "ok",
+        "sql": sql,
+        "columns": cols,
+        "data": rows
+    }