db.py

"""
SQLite persistence for diagnosis history.

Table: diagnoses
  id               INTEGER PRIMARY KEY AUTOINCREMENT
  timestamp        TEXT    (ISO-8601 UTC)
  patient_age      TEXT
  patient_gender   TEXT
  symptoms_summary TEXT    (first 200 chars of symptoms)
  top_diagnosis    TEXT
  confidence       REAL    (0–100)
  risk_level       TEXT    (LOW / MEDIUM / HIGH / CRITICAL)
  full_result_json TEXT    (complete synthesis result as JSON)
"""

from __future__ import annotations

import json
import sqlite3
from datetime import datetime, timezone
from pathlib import Path

DB_PATH = Path(__file__).resolve().parent / "medagent.db"

_DDL = """
CREATE TABLE IF NOT EXISTS diagnoses (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp        TEXT    NOT NULL,
    patient_age      TEXT    NOT NULL DEFAULT '',
    patient_gender   TEXT    NOT NULL DEFAULT '',
    symptoms_summary TEXT    NOT NULL DEFAULT '',
    top_diagnosis    TEXT    NOT NULL DEFAULT '',
    confidence       REAL    NOT NULL DEFAULT 0,
    risk_level       TEXT    NOT NULL DEFAULT '',
    full_result_json TEXT    NOT NULL DEFAULT '{}'
);
"""


def _connect() -> sqlite3.Connection:
    conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    """Create the table if it doesn't exist. Safe to call on every startup."""
    with _connect() as conn:
        conn.executescript(_DDL)


def save_diagnosis(
    patient_data: dict,
    synthesis_result: dict,
    debate_result: dict,
) -> int:
    """
    Persist one completed diagnosis. Returns the new row id.

    patient_data     — the input dict (symptoms, medical_history, …)
    synthesis_result — output of SynthesisAgent.synthesize()
    debate_result    — output of DebateAgent.debate()
    """
    diff     = synthesis_result.get("differential", [])
    top      = diff[0] if diff else {}
    risk     = synthesis_result.get("risk_level", {})

    row = {
        "timestamp":        datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "patient_age":      str(patient_data.get("age", "")),
        "patient_gender":   str(patient_data.get("gender", "")),
        "symptoms_summary": str(patient_data.get("symptoms", ""))[:200],
        "top_diagnosis":    top.get("disease", ""),
        "confidence":       float(top.get("confidence_pct", 0)),
        "risk_level":       risk.get("label", ""),
        "full_result_json": json.dumps({
            "differential":      diff,
            "recommendations":   synthesis_result.get("recommendations", {}),
            "risk_level":        risk,
            "distinguishing_test": synthesis_result.get("distinguishing_test", ""),
            "debate_backend":    synthesis_result.get("debate_backend", "voting"),
            "debate_summary":    debate_result.get("summary", ""),
            "agent_a_argument":  debate_result.get("agent_a_argument", ""),
            "agent_b_argument":  debate_result.get("agent_b_argument", ""),
        }),
    }

    sql = """
        INSERT INTO diagnoses
            (timestamp, patient_age, patient_gender, symptoms_summary,
             top_diagnosis, confidence, risk_level, full_result_json)
        VALUES
            (:timestamp, :patient_age, :patient_gender, :symptoms_summary,
             :top_diagnosis, :confidence, :risk_level, :full_result_json)
    """
    with _connect() as conn:
        cur = conn.execute(sql, row)
        return cur.lastrowid


def get_recent(limit: int = 20) -> list[dict]:
    """Return the most recent *limit* diagnosis rows as plain dicts."""
    sql = """
        SELECT id, timestamp, patient_age, patient_gender,
               symptoms_summary, top_diagnosis, confidence, risk_level
        FROM diagnoses
        ORDER BY id DESC
        LIMIT ?
    """
    with _connect() as conn:
        rows = conn.execute(sql, (limit,)).fetchall()
    return [dict(r) for r in rows]


def get_full(row_id: int) -> dict | None:
    """Return full row including full_result_json, or None if not found."""
    sql = "SELECT * FROM diagnoses WHERE id = ?"
    with _connect() as conn:
        row = conn.execute(sql, (row_id,)).fetchone()
    if row is None:
        return None
    d = dict(row)
    try:
        d["full_result"] = json.loads(d["full_result_json"])
    except Exception:
        d["full_result"] = {}
    return d