app.py

"""
Network text-to-SQL chat demo.

Run locally:
    python db/seed_from_csv.py          # one-time: create db/network.db
    python app.py

On Hugging Face Spaces: set the Space SDK to Gradio, push this repo, and it
will build automatically (app.py is the default entrypoint).
"""

import sqlite3
from pathlib import Path

import gradio as gr

import simulator
from guardrail import validate_and_sanitize, SQLGuardrailError
from infer import generate_sql, synthesize_answer
from prompt_template import build_prompt

DB_PATH = Path(__file__).parent / "network.db"


def get_connection():
    # read-only connection: SQLite URI mode blocks writes at the driver level,
    # a second line of defense behind the guardrail.
    return sqlite3.connect(f"file:{DB_PATH}?mode=ro", uri=True, check_same_thread=False)


def answer_question(question: str, history):
    if not DB_PATH.exists():
        return "Database not found. Run `python seed_from_csv.py` first.", ""

    prompt = build_prompt(question)
    raw_sql = generate_sql(prompt)

    try:
        safe_sql = validate_and_sanitize(raw_sql)
    except SQLGuardrailError as e:
        return f"I couldn't safely run that query ({e}). Try rephrasing.", raw_sql

    try:
        conn = get_connection()
        cur = conn.cursor()
        cur.execute(safe_sql)
        rows = cur.fetchall()
        columns = [d[0] for d in cur.description]
        conn.close()
    except sqlite3.Error as e:
        return f"The query failed to execute ({e}).", safe_sql

    nl_answer = synthesize_answer(question, columns, rows)
    return nl_answer, safe_sql


def chat_fn(message, history):
    answer, sql = answer_question(message, history)
    shown = answer
    if sql:
        shown += f"\n\n<details><summary>Show SQL</summary>\n\n```sql\n{sql}\n```\n</details>"
    return shown


EXAMPLE_QUESTIONS = [
    "Which devices are currently down?",
    "What are the top 5 talkers by bandwidth in the last hour?",
    "Any unresolved critical alerts right now?",
    "Which interfaces have packet loss above 2% right now?",
    "How many devices are in a degraded state right now?",
]

with gr.Blocks(title="Network State Assistant") as demo:
    gr.Markdown(
        "# Network State Assistant (demo)\n"
        "Ask about device status, traffic, alerts, and performance in plain English."
    )
    chat = gr.ChatInterface(
        fn=chat_fn,
        examples=EXAMPLE_QUESTIONS,
        title=None,
    )

if __name__ == "__main__":
    simulator.start(interval_seconds=5)
    demo.launch()