"""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}"]}