Spaces:

JA26
/

ask-the-data

Running on Zero

Update agent.py

d231195 verified 3 days ago

20 kB

"""Ask the Data agent: question -> LLM writes SQL -> SQLite runs it -> LLM explains the result.
Numbers only ever come from the database. The LLM writes the query and the wording.
"""
import os
import re
import sqlite3
import time
 
import pandas as pd
 
DATA_FILE = os.environ.get("DATA_FILE", "data.xlsx")
MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-72B-Instruct")  # change in Space settings if unavailable
MAX_ROWS = 50
AI_FILE = os.environ.get("AI_FILE", "ai_outputs.xlsx")  # outputs from the batch AI notebook (optional)
 
# PRIVATE DATA: keep data.xlsx and ai_outputs.xlsx in a PRIVATE Hugging Face Dataset repo, not in this
# public Space. Set these in the Space settings:
#   DATA_REPO  (Variable) e.g. "your-username/ask-the-data-private"
#   DATA_TOKEN (Secret)   fine-grained token with READ access to that one repo only.
# Use the name DATA_TOKEN (not HF_TOKEN) so it does not clash with the token the LLM call uses.
# If DATA_REPO is empty, the files are read from the local folder (handy for testing).
DATA_REPO = os.environ.get("DATA_REPO", "").strip()
QUERY_TIMEOUT_SECONDS = 5

# LIVE DATABASE (Supabase / Postgres). Set SUPABASE_DB_URL as a Space SECRET to use it.
# If it is empty the agent falls back to the Excel files above (old behaviour).
DB_URL = os.environ.get("SUPABASE_DB_URL", "").strip()
PG_SCHEMA = os.environ.get("PG_SCHEMA", "agent")  # schema holding the read-only views the agent may query
 
 
def _data_path(filename):
    """Return a local path to the file: downloaded from the private repo, or the local file."""
    if not DATA_REPO:
        return filename
    try:
        from huggingface_hub import hf_hub_download
        return hf_hub_download(repo_id=DATA_REPO, filename=filename, repo_type="dataset",
                               token=os.environ.get("DATA_TOKEN"))
    except Exception as e:  # shows up in the Space Logs tab
        print(f"[data] could not download {filename} from {DATA_REPO}: {e}")
        return None
 
# sheet name in the Excel file -> table name the agent sees
SHEETS = {
    "Requests_Clean": "requests",
    "Complaints_Clean": "complaints",
    "Feedback_Clean": "feedback",
    "Lookup_Entity": "entities",
    "Lookup_Service": "services",
    "Lookup_Segment": "segments",
    "Dim_Channel": "channels",
    "SLA_Targets": "sla_targets",
}
DROP_COLS = ["Source_Sheet", "Original_Excel_Rows"]
 
NOTES = """
Notes about the data (follow these):
- requests: one row per request. Status values: Closed, Cancelled, Escalated, In Progress, Open. Response_Hours and Resolution_Hours are in hours.
- complaints.Request_ID and feedback.Request_ID join to requests.Request_ID. Join feedback via Request_ID, NOT via requests.Survey_ID (it is incomplete).
- services.Service_ID joins to requests/complaints/sla_targets; services.Entity_ID joins to entities. requests.Channel_ID joins to channels.Channel_ID, requests.Segment_ID to segments.Segment_ID.
- Overall_Satisfaction, Ease_Score, Communication_Score are 1-5 (5 best). Recommend_0_10 is 0-10.
- Rows where Text_Versions = 2 hold two combined comments; count each survey/complaint once.
- Quality_Issues is a text flag listing known data problems (NULL means no known issue).
- SVC009 and SVC021 have no First_Response_SLA (conflicting targets); do not compute first-response SLA compliance for them.
- Resolution SLA breach means Resolution_Hours > sla_targets.Resolution_SLA.
- Complaint rate = complaints divided by all requests: 100.0 * (SELECT COUNT(*) FROM complaints) / (SELECT COUNT(*) FROM requests). For a group (entity, service, channel) use 100.0 * COUNT(c.complaint_id) / COUNT(DISTINCT r.request_id) from requests r LEFT JOIN complaints c ON c.request_id = r.request_id, grouped by that column.
- For any rate or share, count the matching rows with COUNT(*) FILTER (WHERE condition) and divide by the total rows in the group. Never put the condition in the WHERE clause when those same rows are also the denominator.
- When answering from comments, begin with 'In a sample of comments' and do not say 'many' or 'most', because the agent only sees a sample.
- Low satisfaction means overall_satisfaction <= 2 (1 or 2 on the 1-5 scale; 3 is neutral). Low-satisfaction rate = 100.0 * COUNT(*) FILTER (WHERE overall_satisfaction <= 2) / COUNT(overall_satisfaction). Still rank 'happiest' or 'least happy' questions by AVG(overall_satisfaction).
"""
 
NOTES_SENTIMENT = """- feedback.Sentiment is exactly 'Positive', 'Neutral' or 'Negative' (AI-generated from Feedback_Text ONLY); Sentiment_Confidence is 0-1. Complaints are never sentiment-scored; do not compute sentiment for complaints.
"""
NOTES_CATEGORY = """- complaints.Complaint_Category is AI-generated (zero-shot) from Complaint_Text; treat categories as estimates. Low-confidence rows are marked by Category_Confidence (0-1, lower is less certain) and by a review flag column (e.g. Low_Confidence or Needs_Review) ONLY if that column appears in the schema above; never use a column that is not in the schema.
"""
NOTES_SUMMARIES = """- summaries: precomputed AI text summaries per scope. Scope_Type is Overall, Entity or Service; Scope_ID is ALL, an Entity_ID or a Service_ID; Feedback_Count and Complaint_Count are the records covered. Use it for "summarize ..." questions.
"""
NOTES_TEXT = """- For questions about what citizens say or why they are unhappy, return a SAMPLE of matching comments (LIMIT 15) using LIKE on feedback.Feedback_Text or complaints.Complaint_Text, joined to services for context. Do not answer these with counts alone.
"""
 
_conn = None
_schema = None
_notes = NOTES
_ai = None
_table_names = list(SHEETS.values())
 
# ---------------------------------------------------------------------------
# DONE for Supabase: set the Space secret SUPABASE_DB_URL and the agent reads the live database
# (views in the 'agent' schema). The notes below describe the original plan.
#
#   Database type (SQLite / Postgres / other):  ______
#   Host / connection string:                   ______  (store it as a Space secret, e.g. DB_URL)
#   Read-only user:                             ______  (agent must only ever read)
#   Table names (if different from ours):       ______
#
# What changes when we switch:
#   1. load_db() connects to the real database instead of loading data.xlsx
#      (keep the "read-only" rule; no writes from the agent).
#   2. get_schema() reads the real table/column list instead of the SQLite one.
#   3. If it is Postgres (not SQLite), update SQL_SYSTEM below to say "PostgreSQL"
#      and check date/text functions in the prompt.
#   Everything else (LLM calls, SQL safety checks, explanation) stays the same.
# ---------------------------------------------------------------------------
 
 
PG_NOTES = """
PostgreSQL notes (follow these):
- Table and column names are lowercase. Never put double quotes around table or column names. The notes above may write names with capitals; write them unquoted.
- Round with ROUND(value::numeric, 2). Divide with a decimal, e.g. 100.0 * a / b.
- Date columns are real timestamps: use date_trunc('month', col), EXTRACT(year FROM col) or to_char(col, 'YYYY-MM').
- Use ILIKE for case-insensitive text matching.
"""


def _pg_query(sql, max_rows=None):
    """Run one query on Supabase inside a READ ONLY transaction with a timeout.
    The database role is also read-only; this is a second layer."""
    import psycopg2
    conn = psycopg2.connect(DB_URL, connect_timeout=10, sslmode="require")
    try:
        conn.set_session(readonly=True)  # every statement runs in BEGIN READ ONLY
        with conn.cursor() as cur:
            cur.execute(f"SET LOCAL statement_timeout = {int(QUERY_TIMEOUT_SECONDS * 1000)}")
            cur.execute(f"SET LOCAL search_path = {PG_SCHEMA}")
            cur.execute(sql)
            cols = [d[0] for d in cur.description]
            rows = cur.fetchmany(max_rows) if max_rows else cur.fetchall()
        return cols, rows
    finally:
        conn.close()


def _load_pg():
    """Read the table/column list from the database (once) and build the prompt schema."""
    global _schema, _notes, _table_names
    if _schema is not None:
        return True
    try:
        _, rows = _pg_query(
            "select table_name, column_name, data_type from information_schema.columns "
            f"where table_schema = '{PG_SCHEMA}' order by table_name, ordinal_position")
    except Exception as e:  # shows up in the Space Logs tab; retried on the next question
        print(f"[db] could not read the schema from the database: {type(e).__name__}: {e}", flush=True)
        return None
    if not rows:
        print(f"[db] no tables found in schema {PG_SCHEMA!r}", flush=True)
        return None
    tables = {}
    for t, c, d in rows:
        tables.setdefault(t, []).append((c, d))
    _table_names = list(tables)
    names = lambda t: {c.lower() for c, _ in tables.get(t, [])}
    _notes = NOTES + NOTES_TEXT
    if "sentiment" in names("feedback"):
        _notes += NOTES_SENTIMENT
    if "complaint_category" in names("complaints"):
        _notes += NOTES_CATEGORY
    if "summaries" in tables:
        _notes += NOTES_SUMMARIES
    _notes += PG_NOTES
    parts = [f"{t}: " + ", ".join(f"{c} ({d})" for c, d in cols) for t, cols in tables.items()]
    _schema = "PostgreSQL tables:\n" + "\n".join(parts) + "\n" + _notes
    return True


def load_ai():
    """Read the batch AI outputs (optional file). Sheets: feedback_ai, complaints_ai, summaries."""
    global _ai
    if _ai is None and DB_URL:
        try:
            cols, rows = _pg_query("select * from summaries")
            want = ["Scope_Type", "Scope_ID", "Scope_Name", "Feedback_Summary", "Complaint_Summary",
                    "Feedback_Count", "Complaint_Count", "Pct_Negative", "Top_Complaint_Category"]
            df = pd.DataFrame(rows, columns=cols).rename(columns={w.lower(): w for w in want})
            _ai = {"summaries": df}
        except Exception as e:
            print(f"[db] could not read summaries: {type(e).__name__}: {e}", flush=True)
            return {}  # not cached, so the next call retries
    if _ai is None:
        _ai = {}
        ai_path = _data_path(AI_FILE)
        if ai_path and os.path.exists(ai_path):
            _ai = {k.strip().lower(): v for k, v in pd.read_excel(ai_path, sheet_name=None).items()}
    return _ai
 
 
def _merge_ai(base, extra, key):
    """Add the AI columns to a base table, matched on its ID column."""
    if extra is None or key not in extra.columns:
        return base
    extra = extra.drop_duplicates(key)
    new_cols = [c for c in extra.columns if c == key or c not in base.columns]
    return base.merge(extra[new_cols], on=key, how="left")
 
 
def get_summaries(scope=""):
    """Precomputed summaries for the dashboard. Works even without the main data."""
    ai = load_ai()
    sm = ai.get("summaries")
    if sm is None or sm.empty:
        return {"scope": scope or "Overall", "feedback_summary": None, "complaint_summary": None,
                "available_scopes": [], "note": "Summaries are not available yet."}
    scopes = sm["Scope_ID"].astype(str).tolist()
    want = (scope or "").strip().lower() or "all"
    hit = sm[(sm["Scope_ID"].astype(str).str.lower() == want) |
             (sm.get("Scope_Name", sm["Scope_ID"]).astype(str).str.lower() == want)]
    if hit.empty and want in ("overall", "all"):
        hit = sm[sm["Scope_Type"].astype(str).str.lower() == "overall"]
    if hit.empty:
        return {"scope": scope, "feedback_summary": None, "complaint_summary": None,
                "available_scopes": scopes, "note": "Unknown scope. Pick one of available_scopes."}
    r = hit.iloc[0]
    def clean(v):
        return None if pd.isna(v) else str(v)
    return {"scope": str(r["Scope_ID"]), "scope_name": clean(r.get("Scope_Name")),
            "feedback_summary": clean(r.get("Feedback_Summary")),
            "complaint_summary": clean(r.get("Complaint_Summary")),
            "feedback_count": None if pd.isna(r.get("Feedback_Count", float("nan"))) else int(r["Feedback_Count"]),
            "complaint_count": None if pd.isna(r.get("Complaint_Count", float("nan"))) else int(r["Complaint_Count"]),
            "stats": {k: (None if pd.isna(r[k]) else (r[k].item() if hasattr(r[k], "item") else r[k]))
                      for k in ("Pct_Negative", "Top_Complaint_Category") if k in r.index},
            "available_scopes": scopes}
 
 
def load_db():
    """Load the Excel file into an in-memory SQLite database (once).
    TEMPORARY: replace with a live database connection (see TODO above)."""
    global _conn, _schema
    if DB_URL:
        return _load_pg()
    if _conn is not None:
        return _conn
    data_path = _data_path(DATA_FILE)
    if not data_path or not os.path.exists(data_path):
        return None  # no data connected yet: the agent replies "not connected"
    conn = sqlite3.connect(":memory:", check_same_thread=False)
    frames = {}
    for sheet, table in SHEETS.items():
        df = pd.read_excel(data_path, sheet_name=sheet, header=4)
        df = df.drop(columns=[c for c in DROP_COLS if c in df.columns])
        frames[table] = df
    # light fixes for known gaps, so answers are not skewed
    req, cmp_ = frames["requests"], frames["complaints"]
    cmp_["Service_ID"] = cmp_["Service_ID"].str.upper()
    fill = cmp_["Request_ID"].map(req.set_index("Request_ID")["Service_ID"])
    cmp_["Service_ID"] = cmp_["Service_ID"].fillna(fill)
    svc_entity = frames["services"].set_index("Service_ID")["Entity_ID"]
    req["Entity_ID"] = req["Entity_ID"].fillna(req["Service_ID"].map(svc_entity))
    ai = load_ai()
    frames["feedback"] = _merge_ai(frames["feedback"], ai.get("feedback_ai"), "Survey_ID")
    frames["complaints"] = _merge_ai(frames["complaints"], ai.get("complaints_ai"), "Complaint_ID")
    if ai.get("summaries") is not None:
        frames["summaries"] = ai["summaries"]
    global _notes, _table_names
    _notes = NOTES + NOTES_TEXT
    if "Sentiment" in frames["feedback"].columns:
        _notes += NOTES_SENTIMENT
    if "Complaint_Category" in frames["complaints"].columns:
        _notes += NOTES_CATEGORY
    if "summaries" in frames:
        _notes += NOTES_SUMMARIES
    _table_names = list(frames)
    for table, df in frames.items():
        for col in df.columns:
            if pd.api.types.is_datetime64_any_dtype(df[col]):
                df[col] = df[col].dt.strftime("%Y-%m-%d %H:%M:%S")
        df.to_sql(table, conn, index=False)
    conn.execute("PRAGMA query_only = ON")
    parts = []
    for table in frames:
        cols = ", ".join(f"{c} ({t})" for c, t in
                         [(r[1], r[2]) for r in conn.execute(f"PRAGMA table_info({table})")])
        parts.append(f"{table}: {cols}")
    _schema = "SQLite tables:\n" + "\n".join(parts) + "\n" + _notes
    # Hard read-only rule, enforced by SQLite itself (set AFTER the tables are loaded):
    # only reading is allowed. Any write, PRAGMA, ATTACH, etc. is refused.
    allowed = {20, 21, 31, 33}  # SQLITE_READ, SQLITE_SELECT, SQLITE_FUNCTION, SQLITE_RECURSIVE
    conn.set_authorizer(lambda action, *_: sqlite3.SQLITE_OK if action in allowed else sqlite3.SQLITE_DENY)
    _conn = conn
    return conn
 
 
def get_schema():
    load_db()
    return _schema
 
 
# ---------- LLM ----------
def llm(messages):
    """Call a hosted model through Hugging Face. Replaced in tests."""
    from huggingface_hub import InferenceClient
    token = os.environ.get("HFR_TOKEN") or os.environ.get("HF_TOKEN")
    client = InferenceClient(model=MODEL_ID, token=token)
    out = client.chat_completion(messages=messages, max_tokens=700, temperature=0.0)
    return out.choices[0].message.content
 
 
# ---------- SQL safety ----------
FORBIDDEN = re.compile(
    r"\b(insert|update|delete|drop|alter|create|replace|attach|detach|pragma|vacuum|reindex|truncate|grant|revoke|copy|call|execute|listen|notify)\b",
    re.I)
 
 
def extract_sql(text):
    m = re.search(r"```(?:sql)?\s*(.*?)```", text, re.S | re.I)
    sql = (m.group(1) if m else text).strip()
    return sql.rstrip(";").strip()
 
 
def _without_literals(sql):
    """Blank out 'quoted text' so words inside it (e.g. 'Property Address Update') are not flagged."""
    return re.sub(r"'(?:[^']|'')*'", "''", sql)
 
 
def validate_sql(sql):
    if not sql:
        raise ValueError("empty query")
    check = _without_literals(sql)
    if ";" in check:
        raise ValueError("only one statement is allowed")
    if not re.match(r"^(select|with)\b", check, re.I):
        raise ValueError("only SELECT queries are allowed")
    if FORBIDDEN.search(check):
        raise ValueError("query contains a forbidden keyword")
    return sql
 
 
def run_sql(sql):
    if DB_URL:
        return _pg_query(sql, MAX_ROWS)
    conn = load_db()
    deadline = time.time() + QUERY_TIMEOUT_SECONDS  # stop runaway queries
    conn.set_progress_handler(lambda: 1 if time.time() > deadline else 0, 100000)
    try:
        cur = conn.execute(sql)
        cols = [d[0] for d in cur.description]
        rows = cur.fetchmany(MAX_ROWS)
    finally:
        conn.set_progress_handler(None, 0)
    return cols, rows
 
 
def tables_used(sql):
    names = set(_table_names)
    return sorted(t for t in names if re.search(rf"\b{t}\b", sql, re.I))
 
 
# ---------- main entry ----------
SQL_SYSTEM = (
    "You write a single read-only SQLite SELECT query that answers the user's question about "
    "a citizen service dataset. Use only the tables and columns listed. Return only the query "
    "in a ```sql block, nothing else. If the question cannot be answered from this data, return "
    "exactly: CANNOT_ANSWER\n\n"
)
EXPLAIN_SYSTEM = (
    "You explain database results to a government service manager in plain English. Use ONLY the "
    "numbers in the result table; never invent or estimate figures. Be concise (2-4 sentences). "
    "If the result contains comments, summarize the common themes and say they are a sample, not all comments. If the result is empty, say no matching records were found. Mention relevant data caveats "
    "from the notes if they affect the answer."
)
 
 
def _sql_system():
    return SQL_SYSTEM.replace("SQLite", "PostgreSQL") if DB_URL else SQL_SYSTEM


def answer_question(question):
    question = (question or "").strip()
    if not question:
        return {"answer": "Please type a question.", "sources": []}
    if load_db() is None:
        return {"answer": "The database isn't connected yet, so I can't answer data questions right now.",
                "sources": []}
    schema = get_schema()
    messages = [{"role": "system", "content": _sql_system() + schema},
                {"role": "user", "content": question}]
    sql, error = None, None
    for _ in range(2):  # one retry if the query fails
        reply = llm(messages)
        if "CANNOT_ANSWER" in reply:
            return {"answer": "I can't answer that from the current dataset.", "sources": []}
        try:
            sql = validate_sql(extract_sql(reply))
            cols, rows = run_sql(sql)
            error = None
            break
        except Exception as e:  # validation or SQL error
            error = str(e)
            messages += [{"role": "assistant", "content": reply},
                         {"role": "user", "content": f"That failed: {error}. Fix the query."}]
    if error:
        return {"answer": "I couldn't build a valid query for that question. Try rephrasing it.",
                "sources": []}
    table = " | ".join(cols) + "\n" + "\n".join(" | ".join(str(v) for v in r) for r in rows)
    text = llm([
        {"role": "system", "content": EXPLAIN_SYSTEM + "\n" + _notes},
        {"role": "user", "content": f"Question: {question}\n\nSQL used:\n{sql}\n\nResult "
                                     f"({len(rows)} rows, max {MAX_ROWS}):\n{table}"},
    ])
    return {"answer": text.strip(), "sources": tables_used(sql) + [f"SQL: {sql}"]}

	"""Ask the Data agent: question -> LLM writes SQL -> SQLite runs it -> LLM explains the result.
	Numbers only ever come from the database. The LLM writes the query and the wording.
	"""
	import os
	import re
	import sqlite3
	import time

	import pandas as pd

	DATA_FILE = os.environ.get("DATA_FILE", "data.xlsx")
	MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-72B-Instruct") # change in Space settings if unavailable
	MAX_ROWS = 50
	AI_FILE = os.environ.get("AI_FILE", "ai_outputs.xlsx") # outputs from the batch AI notebook (optional)

	# PRIVATE DATA: keep data.xlsx and ai_outputs.xlsx in a PRIVATE Hugging Face Dataset repo, not in this
	# public Space. Set these in the Space settings:
	# DATA_REPO (Variable) e.g. "your-username/ask-the-data-private"
	# DATA_TOKEN (Secret) fine-grained token with READ access to that one repo only.
	# Use the name DATA_TOKEN (not HF_TOKEN) so it does not clash with the token the LLM call uses.
	# If DATA_REPO is empty, the files are read from the local folder (handy for testing).
	DATA_REPO = os.environ.get("DATA_REPO", "").strip()
	QUERY_TIMEOUT_SECONDS = 5

	# LIVE DATABASE (Supabase / Postgres). Set SUPABASE_DB_URL as a Space SECRET to use it.
	# If it is empty the agent falls back to the Excel files above (old behaviour).
	DB_URL = os.environ.get("SUPABASE_DB_URL", "").strip()
	PG_SCHEMA = os.environ.get("PG_SCHEMA", "agent") # schema holding the read-only views the agent may query


	def _data_path(filename):
	"""Return a local path to the file: downloaded from the private repo, or the local file."""
	if not DATA_REPO:
	return filename
	try:
	from huggingface_hub import hf_hub_download
	return hf_hub_download(repo_id=DATA_REPO, filename=filename, repo_type="dataset",
	token=os.environ.get("DATA_TOKEN"))
	except Exception as e: # shows up in the Space Logs tab
	print(f"[data] could not download {filename} from {DATA_REPO}: {e}")
	return None

	# sheet name in the Excel file -> table name the agent sees
	SHEETS = {
	"Requests_Clean": "requests",
	"Complaints_Clean": "complaints",
	"Feedback_Clean": "feedback",
	"Lookup_Entity": "entities",
	"Lookup_Service": "services",
	"Lookup_Segment": "segments",
	"Dim_Channel": "channels",
	"SLA_Targets": "sla_targets",
	}
	DROP_COLS = ["Source_Sheet", "Original_Excel_Rows"]

	NOTES = """
	Notes about the data (follow these):
	- requests: one row per request. Status values: Closed, Cancelled, Escalated, In Progress, Open. Response_Hours and Resolution_Hours are in hours.
	- complaints.Request_ID and feedback.Request_ID join to requests.Request_ID. Join feedback via Request_ID, NOT via requests.Survey_ID (it is incomplete).
	- services.Service_ID joins to requests/complaints/sla_targets; services.Entity_ID joins to entities. requests.Channel_ID joins to channels.Channel_ID, requests.Segment_ID to segments.Segment_ID.
	- Overall_Satisfaction, Ease_Score, Communication_Score are 1-5 (5 best). Recommend_0_10 is 0-10.
	- Rows where Text_Versions = 2 hold two combined comments; count each survey/complaint once.
	- Quality_Issues is a text flag listing known data problems (NULL means no known issue).
	- SVC009 and SVC021 have no First_Response_SLA (conflicting targets); do not compute first-response SLA compliance for them.
	- Resolution SLA breach means Resolution_Hours > sla_targets.Resolution_SLA.
	- Complaint rate = complaints divided by all requests: 100.0 * (SELECT COUNT() FROM complaints) / (SELECT COUNT() FROM requests). For a group (entity, service, channel) use 100.0 * COUNT(c.complaint_id) / COUNT(DISTINCT r.request_id) from requests r LEFT JOIN complaints c ON c.request_id = r.request_id, grouped by that column.
	- For any rate or share, count the matching rows with COUNT(*) FILTER (WHERE condition) and divide by the total rows in the group. Never put the condition in the WHERE clause when those same rows are also the denominator.
	- When answering from comments, begin with 'In a sample of comments' and do not say 'many' or 'most', because the agent only sees a sample.
	- Low satisfaction means overall_satisfaction <= 2 (1 or 2 on the 1-5 scale; 3 is neutral). Low-satisfaction rate = 100.0 * COUNT(*) FILTER (WHERE overall_satisfaction <= 2) / COUNT(overall_satisfaction). Still rank 'happiest' or 'least happy' questions by AVG(overall_satisfaction).
	"""

	NOTES_SENTIMENT = """- feedback.Sentiment is exactly 'Positive', 'Neutral' or 'Negative' (AI-generated from Feedback_Text ONLY); Sentiment_Confidence is 0-1. Complaints are never sentiment-scored; do not compute sentiment for complaints.
	"""
	NOTES_CATEGORY = """- complaints.Complaint_Category is AI-generated (zero-shot) from Complaint_Text; treat categories as estimates. Low-confidence rows are marked by Category_Confidence (0-1, lower is less certain) and by a review flag column (e.g. Low_Confidence or Needs_Review) ONLY if that column appears in the schema above; never use a column that is not in the schema.
	"""
	NOTES_SUMMARIES = """- summaries: precomputed AI text summaries per scope. Scope_Type is Overall, Entity or Service; Scope_ID is ALL, an Entity_ID or a Service_ID; Feedback_Count and Complaint_Count are the records covered. Use it for "summarize ..." questions.
	"""
	NOTES_TEXT = """- For questions about what citizens say or why they are unhappy, return a SAMPLE of matching comments (LIMIT 15) using LIKE on feedback.Feedback_Text or complaints.Complaint_Text, joined to services for context. Do not answer these with counts alone.
	"""

	_conn = None
	_schema = None
	_notes = NOTES
	_ai = None
	_table_names = list(SHEETS.values())

	# ---------------------------------------------------------------------------
	# DONE for Supabase: set the Space secret SUPABASE_DB_URL and the agent reads the live database
	# (views in the 'agent' schema). The notes below describe the original plan.
	#
	# Database type (SQLite / Postgres / other): ______
	# Host / connection string: ______ (store it as a Space secret, e.g. DB_URL)
	# Read-only user: ______ (agent must only ever read)
	# Table names (if different from ours): ______
	#
	# What changes when we switch:
	# 1. load_db() connects to the real database instead of loading data.xlsx
	# (keep the "read-only" rule; no writes from the agent).
	# 2. get_schema() reads the real table/column list instead of the SQLite one.
	# 3. If it is Postgres (not SQLite), update SQL_SYSTEM below to say "PostgreSQL"
	# and check date/text functions in the prompt.
	# Everything else (LLM calls, SQL safety checks, explanation) stays the same.
	# ---------------------------------------------------------------------------


	PG_NOTES = """
	PostgreSQL notes (follow these):
	- Table and column names are lowercase. Never put double quotes around table or column names. The notes above may write names with capitals; write them unquoted.
	- Round with ROUND(value::numeric, 2). Divide with a decimal, e.g. 100.0 * a / b.
	- Date columns are real timestamps: use date_trunc('month', col), EXTRACT(year FROM col) or to_char(col, 'YYYY-MM').
	- Use ILIKE for case-insensitive text matching.
	"""


	def _pg_query(sql, max_rows=None):
	"""Run one query on Supabase inside a READ ONLY transaction with a timeout.
	The database role is also read-only; this is a second layer."""
	import psycopg2
	conn = psycopg2.connect(DB_URL, connect_timeout=10, sslmode="require")
	try:
	conn.set_session(readonly=True) # every statement runs in BEGIN READ ONLY
	with conn.cursor() as cur:
	cur.execute(f"SET LOCAL statement_timeout = {int(QUERY_TIMEOUT_SECONDS * 1000)}")
	cur.execute(f"SET LOCAL search_path = {PG_SCHEMA}")
	cur.execute(sql)
	cols = [d[0] for d in cur.description]
	rows = cur.fetchmany(max_rows) if max_rows else cur.fetchall()
	return cols, rows
	finally:
	conn.close()


	def _load_pg():
	"""Read the table/column list from the database (once) and build the prompt schema."""
	global _schema, _notes, _table_names
	if _schema is not None:
	return True
	try:
	_, rows = _pg_query(
	"select table_name, column_name, data_type from information_schema.columns "
	f"where table_schema = '{PG_SCHEMA}' order by table_name, ordinal_position")
	except Exception as e: # shows up in the Space Logs tab; retried on the next question
	print(f"[db] could not read the schema from the database: {type(e).__name__}: {e}", flush=True)
	return None
	if not rows:
	print(f"[db] no tables found in schema {PG_SCHEMA!r}", flush=True)
	return None
	tables = {}
	for t, c, d in rows:
	tables.setdefault(t, []).append((c, d))
	_table_names = list(tables)
	names = lambda t: {c.lower() for c, _ in tables.get(t, [])}
	_notes = NOTES + NOTES_TEXT
	if "sentiment" in names("feedback"):
	_notes += NOTES_SENTIMENT
	if "complaint_category" in names("complaints"):
	_notes += NOTES_CATEGORY
	if "summaries" in tables:
	_notes += NOTES_SUMMARIES
	_notes += PG_NOTES
	parts = [f"{t}: " + ", ".join(f"{c} ({d})" for c, d in cols) for t, cols in tables.items()]
	_schema = "PostgreSQL tables:\n" + "\n".join(parts) + "\n" + _notes
	return True


	def load_ai():
	"""Read the batch AI outputs (optional file). Sheets: feedback_ai, complaints_ai, summaries."""
	global _ai
	if _ai is None and DB_URL:
	try:
	cols, rows = _pg_query("select * from summaries")
	want = ["Scope_Type", "Scope_ID", "Scope_Name", "Feedback_Summary", "Complaint_Summary",
	"Feedback_Count", "Complaint_Count", "Pct_Negative", "Top_Complaint_Category"]
	df = pd.DataFrame(rows, columns=cols).rename(columns={w.lower(): w for w in want})
	_ai = {"summaries": df}
	except Exception as e:
	print(f"[db] could not read summaries: {type(e).__name__}: {e}", flush=True)
	return {} # not cached, so the next call retries
	if _ai is None:
	_ai = {}
	ai_path = _data_path(AI_FILE)
	if ai_path and os.path.exists(ai_path):
	_ai = {k.strip().lower(): v for k, v in pd.read_excel(ai_path, sheet_name=None).items()}
	return _ai


	def _merge_ai(base, extra, key):
	"""Add the AI columns to a base table, matched on its ID column."""
	if extra is None or key not in extra.columns:
	return base
	extra = extra.drop_duplicates(key)
	new_cols = [c for c in extra.columns if c == key or c not in base.columns]
	return base.merge(extra[new_cols], on=key, how="left")


	def get_summaries(scope=""):
	"""Precomputed summaries for the dashboard. Works even without the main data."""
	ai = load_ai()
	sm = ai.get("summaries")
	if sm is None or sm.empty:
	return {"scope": scope or "Overall", "feedback_summary": None, "complaint_summary": None,
	"available_scopes": [], "note": "Summaries are not available yet."}
	scopes = sm["Scope_ID"].astype(str).tolist()
	want = (scope or "").strip().lower() or "all"
	hit = sm[(sm["Scope_ID"].astype(str).str.lower() == want) \|
	(sm.get("Scope_Name", sm["Scope_ID"]).astype(str).str.lower() == want)]
	if hit.empty and want in ("overall", "all"):
	hit = sm[sm["Scope_Type"].astype(str).str.lower() == "overall"]
	if hit.empty:
	return {"scope": scope, "feedback_summary": None, "complaint_summary": None,
	"available_scopes": scopes, "note": "Unknown scope. Pick one of available_scopes."}
	r = hit.iloc[0]
	def clean(v):
	return None if pd.isna(v) else str(v)
	return {"scope": str(r["Scope_ID"]), "scope_name": clean(r.get("Scope_Name")),
	"feedback_summary": clean(r.get("Feedback_Summary")),
	"complaint_summary": clean(r.get("Complaint_Summary")),
	"feedback_count": None if pd.isna(r.get("Feedback_Count", float("nan"))) else int(r["Feedback_Count"]),
	"complaint_count": None if pd.isna(r.get("Complaint_Count", float("nan"))) else int(r["Complaint_Count"]),
	"stats": {k: (None if pd.isna(r[k]) else (r[k].item() if hasattr(r[k], "item") else r[k]))
	for k in ("Pct_Negative", "Top_Complaint_Category") if k in r.index},
	"available_scopes": scopes}


	def load_db():
	"""Load the Excel file into an in-memory SQLite database (once).
	TEMPORARY: replace with a live database connection (see TODO above)."""
	global _conn, _schema
	if DB_URL:
	return _load_pg()
	if _conn is not None:
	return _conn
	data_path = _data_path(DATA_FILE)
	if not data_path or not os.path.exists(data_path):
	return None # no data connected yet: the agent replies "not connected"
	conn = sqlite3.connect(":memory:", check_same_thread=False)
	frames = {}
	for sheet, table in SHEETS.items():
	df = pd.read_excel(data_path, sheet_name=sheet, header=4)
	df = df.drop(columns=[c for c in DROP_COLS if c in df.columns])
	frames[table] = df
	# light fixes for known gaps, so answers are not skewed
	req, cmp_ = frames["requests"], frames["complaints"]
	cmp_["Service_ID"] = cmp_["Service_ID"].str.upper()
	fill = cmp_["Request_ID"].map(req.set_index("Request_ID")["Service_ID"])
	cmp_["Service_ID"] = cmp_["Service_ID"].fillna(fill)
	svc_entity = frames["services"].set_index("Service_ID")["Entity_ID"]
	req["Entity_ID"] = req["Entity_ID"].fillna(req["Service_ID"].map(svc_entity))
	ai = load_ai()
	frames["feedback"] = _merge_ai(frames["feedback"], ai.get("feedback_ai"), "Survey_ID")
	frames["complaints"] = _merge_ai(frames["complaints"], ai.get("complaints_ai"), "Complaint_ID")
	if ai.get("summaries") is not None:
	frames["summaries"] = ai["summaries"]
	global _notes, _table_names
	_notes = NOTES + NOTES_TEXT
	if "Sentiment" in frames["feedback"].columns:
	_notes += NOTES_SENTIMENT
	if "Complaint_Category" in frames["complaints"].columns:
	_notes += NOTES_CATEGORY
	if "summaries" in frames:
	_notes += NOTES_SUMMARIES
	_table_names = list(frames)
	for table, df in frames.items():
	for col in df.columns:
	if pd.api.types.is_datetime64_any_dtype(df[col]):
	df[col] = df[col].dt.strftime("%Y-%m-%d %H:%M:%S")
	df.to_sql(table, conn, index=False)
	conn.execute("PRAGMA query_only = ON")
	parts = []
	for table in frames:
	cols = ", ".join(f"{c} ({t})" for c, t in
	[(r[1], r[2]) for r in conn.execute(f"PRAGMA table_info({table})")])
	parts.append(f"{table}: {cols}")
	_schema = "SQLite tables:\n" + "\n".join(parts) + "\n" + _notes
	# Hard read-only rule, enforced by SQLite itself (set AFTER the tables are loaded):
	# only reading is allowed. Any write, PRAGMA, ATTACH, etc. is refused.
	allowed = {20, 21, 31, 33} # SQLITE_READ, SQLITE_SELECT, SQLITE_FUNCTION, SQLITE_RECURSIVE
	conn.set_authorizer(lambda action, *_: sqlite3.SQLITE_OK if action in allowed else sqlite3.SQLITE_DENY)
	_conn = conn
	return conn


	def get_schema():
	load_db()
	return _schema


	# ---------- LLM ----------
	def llm(messages):
	"""Call a hosted model through Hugging Face. Replaced in tests."""
	from huggingface_hub import InferenceClient
	token = os.environ.get("HFR_TOKEN") or os.environ.get("HF_TOKEN")
	client = InferenceClient(model=MODEL_ID, token=token)
	out = client.chat_completion(messages=messages, max_tokens=700, temperature=0.0)
	return out.choices[0].message.content


	# ---------- SQL safety ----------
	FORBIDDEN = re.compile(
	r"\b(insert\|update\|delete\|drop\|alter\|create\|replace\|attach\|detach\|pragma\|vacuum\|reindex\|truncate\|grant\|revoke\|copy\|call\|execute\|listen\|notify)\b",
	re.I)


	def extract_sql(text):
	m = re.search(r"```(?:sql)?\s(.?)```", text, re.S \| re.I)
	sql = (m.group(1) if m else text).strip()
	return sql.rstrip(";").strip()


	def _without_literals(sql):
	"""Blank out 'quoted text' so words inside it (e.g. 'Property Address Update') are not flagged."""
	return re.sub(r"'(?:[^']\|'')*'", "''", sql)


	def validate_sql(sql):
	if not sql:
	raise ValueError("empty query")
	check = _without_literals(sql)
	if ";" in check:
	raise ValueError("only one statement is allowed")
	if not re.match(r"^(select\|with)\b", check, re.I):
	raise ValueError("only SELECT queries are allowed")
	if FORBIDDEN.search(check):
	raise ValueError("query contains a forbidden keyword")
	return sql


	def run_sql(sql):
	if DB_URL:
	return _pg_query(sql, MAX_ROWS)
	conn = load_db()
	deadline = time.time() + QUERY_TIMEOUT_SECONDS # stop runaway queries
	conn.set_progress_handler(lambda: 1 if time.time() > deadline else 0, 100000)
	try:
	cur = conn.execute(sql)
	cols = [d[0] for d in cur.description]
	rows = cur.fetchmany(MAX_ROWS)
	finally:
	conn.set_progress_handler(None, 0)
	return cols, rows


	def tables_used(sql):
	names = set(_table_names)
	return sorted(t for t in names if re.search(rf"\b{t}\b", sql, re.I))


	# ---------- main entry ----------
	SQL_SYSTEM = (
	"You write a single read-only SQLite SELECT query that answers the user's question about "
	"a citizen service dataset. Use only the tables and columns listed. Return only the query "
	"in a ```sql block, nothing else. If the question cannot be answered from this data, return "
	"exactly: CANNOT_ANSWER\n\n"
	)
	EXPLAIN_SYSTEM = (
	"You explain database results to a government service manager in plain English. Use ONLY the "
	"numbers in the result table; never invent or estimate figures. Be concise (2-4 sentences). "
	"If the result contains comments, summarize the common themes and say they are a sample, not all comments. If the result is empty, say no matching records were found. Mention relevant data caveats "
	"from the notes if they affect the answer."
	)


	def _sql_system():
	return SQL_SYSTEM.replace("SQLite", "PostgreSQL") if DB_URL else SQL_SYSTEM


	def answer_question(question):
	question = (question or "").strip()
	if not question:
	return {"answer": "Please type a question.", "sources": []}
	if load_db() is None:
	return {"answer": "The database isn't connected yet, so I can't answer data questions right now.",
	"sources": []}
	schema = get_schema()
	messages = [{"role": "system", "content": _sql_system() + schema},
	{"role": "user", "content": question}]
	sql, error = None, None
	for _ in range(2): # one retry if the query fails
	reply = llm(messages)
	if "CANNOT_ANSWER" in reply:
	return {"answer": "I can't answer that from the current dataset.", "sources": []}
	try:
	sql = validate_sql(extract_sql(reply))
	cols, rows = run_sql(sql)
	error = None
	break
	except Exception as e: # validation or SQL error
	error = str(e)
	messages += [{"role": "assistant", "content": reply},
	{"role": "user", "content": f"That failed: {error}. Fix the query."}]
	if error:
	return {"answer": "I couldn't build a valid query for that question. Try rephrasing it.",
	"sources": []}
	table = " \| ".join(cols) + "\n" + "\n".join(" \| ".join(str(v) for v in r) for r in rows)
	text = llm([
	{"role": "system", "content": EXPLAIN_SYSTEM + "\n" + _notes},
	{"role": "user", "content": f"Question: {question}\n\nSQL used:\n{sql}\n\nResult "
	f"({len(rows)} rows, max {MAX_ROWS}):\n{table}"},
	])
	return {"answer": text.strip(), "sources": tables_used(sql) + [f"SQL: {sql}"]}