Update engine.py

engine.py

@@ -1,282 +1,314 @@
-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)
-
-# =========================
-# 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 🙂"
-
-# =========================
-# 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
-
-# =========================
-# MAIN ENGINE
-# =========================
-
-def process_question(question):
-
-    question = correct_spelling(question)
-    question = normalize_time_question(question)
-
-    if any(x in question.lower() 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):
-        latest = 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 {latest}.",
-            "data": []
-        }
-
-    if not rows:
-        latest = get_latest_data_date()
-        return {
-            "status": "ok",
-            "message": friendly("No records found."),
-            "note": f"Available data is only up to {latest}.",
-            "data": []
-        }
-
-    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
+
+# =========================
+# 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"]
+
+# =========================
+# 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
+
+# =========================
+# MAIN ENGINE
+# =========================
+
+def process_question(question):
+    global LAST_PROMPT_TYPE
+
+    question = question.strip().lower()
+
+    # Handle confirmation replies like "yes"
+    if is_confirmation(question) 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\n\n"
+                "Just type one of these or ask your own question."
+            ),
+            "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"
+        latest = 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 {latest}.",
+            "data": []
+        }
+
+    if not rows:
+        LAST_PROMPT_TYPE = "NO_DATA"
+        latest = get_latest_data_date()
+        return {
+            "status": "ok",
+            "message": friendly("No records found."),
+            "note": f"Available data is only up to {latest}.",
+            "data": []
+        }
+
+    LAST_PROMPT_TYPE = None
+
+    return {
+        "status": "ok",
+        "sql": sql,
+        "columns": cols,
+        "data": rows
+    }