Ask AI: a read only question box in the top bar for administrators

An administrator with AI set up gets an Ask AI button in the top bar.  It
opens a chat panel for plain questions about the data, such as which
repositories carry a given CVE, or what a person owns.  Ask for a CSV or a
PDF and one is made to download.

The model is given a map of the tables it may read and how they join, and
writes SELECT statements.  None of them runs as written.  Each is parsed with
sqlglot, held to the listed tables and columns (a * only expands to those),
stripped of INTO, locking reads, session variables and file or sleep
functions, capped, and rebuilt.  Password hashes, TOTP secrets, tokens,
secret previews, snippets and site sign in details are not on the map, so
they cannot be read at all.

Nothing is written to the database.  Lookups run on a new read only session:
the database is told the connection is read only, any statement that is not
a SELECT is refused before it goes, the session is rolled back, and the
connection is thrown away rather than pooled.  These calls skip the AI usage
table and the audit log and go to the server log instead.  Queries stop
after 15 seconds.

Files are built from the rows, never from the model's text, with formula
cells neutralized, and held in memory for half an hour for the person who
asked.  Answers are shown as text, links are built on the server from id
columns, and each administrator gets 20 questions in 10 minutes.

Author: hackrange

app/core/ai.py

@@ -328,7 +328,8 @@ def low_effort(cfg: dict) -> dict:
 
 
 def complete(cfg: dict, system: str, prompt: str, *, max_tokens: int = None,
-             timeout: float = None, purpose: str = "other") -> dict:
+             timeout: float = None, purpose: str = "other",
+             record: bool = True) -> dict:
     """Ask once.  -> {ok, text, error, provider, model, input_tokens,
     output_tokens, cache_read_tokens, cache_write_tokens, finish, cut_off,
     duration_ms}
@@ -338,7 +339,8 @@ def complete(cfg: dict, system: str, prompt: str, *, max_tokens: int = None,
     the caller is a page and a stack trace is not something to show on one.
 
     Every call that actually went out is written to the usage table, under
-    purpose, good or bad.
+    purpose, good or bad, unless record is False.  The admin question box
+    passes False because it must never write to the database; it logs instead.
     """
     provider = cfg.get("ai.provider") or "disabled"
     out = {"ok": False, "text": "", "error": "", "provider": provider,
@@ -369,7 +371,12 @@ def complete(cfg: dict, system: str, prompt: str, *, max_tokens: int = None,
     if out["ok"] and not (out.get("text") or "").strip():
         out["ok"] = False
         out["error"] = "The model answered with nothing."
-    _record(out, purpose)
+    if record:
+        _record(out, purpose)
+    else:
+        log.info("AI %s: %s %s ok=%s in=%s out=%s %sms", purpose, provider,
+                 out.get("model") or "", out["ok"], out.get("input_tokens"),
+                 out.get("output_tokens"), out["duration_ms"])
     return out
 
 

app/core/ask.py

@@ -107,6 +107,61 @@ def session_scope():
         session.close()
 
 
+class WriteRefused(Exception):
+    """Something other than a SELECT was sent on a read only session."""
+
+
+def _select_only(conn, cursor, statement, parameters, context, executemany):
+    if not str(statement or "").lstrip().upper().startswith("SELECT"):
+        raise WriteRefused("Only SELECT is allowed here.")
+
+
+@contextlib.contextmanager
+def read_only_session(seconds: int = 0):
+    """A session that can only SELECT.  Nothing is ever committed.
+
+    Three locks, so no one of them has to be perfect.  The database itself is
+    told the connection is read only.  Every statement is checked before it
+    goes and anything that is not a SELECT raises.  And the session is rolled
+    back, never committed.  The connection is thrown away afterwards rather
+    than handed back to the pool, so the read only setting cannot leak into
+    a request that is allowed to write.
+    """
+    from sqlalchemy.orm import Session as OrmSession
+
+    conn = engine().connect()
+    try:
+        if settings.uses_sqlite:
+            conn.exec_driver_sql("PRAGMA query_only = ON")
+            if seconds:
+                # SQLite has no statement timeout, so a handler that is
+                # called every few thousand steps stops it instead.
+                deadline = time.monotonic() + seconds
+                raw = conn.connection.driver_connection
+                raw.set_progress_handler(
+                    lambda: 1 if time.monotonic() > deadline else 0, 10000)
+        else:
+            conn.exec_driver_sql("SET SESSION TRANSACTION READ ONLY")
+            if seconds:
+                try:        # MariaDB's name for it.  MySQL ignores this.
+                    conn.exec_driver_sql(
+                        f"SET SESSION max_statement_time = {int(seconds)}")
+                except Exception:  # noqa: BLE001
+                    conn.rollback()
+        conn.rollback()
+        event.listen(conn, "before_cursor_execute", _select_only)
+        session = OrmSession(bind=conn, autoflush=False, expire_on_commit=False,
+                             future=True)
+        try:
+            yield session
+        finally:
+            session.rollback()
+            session.close()
+    finally:
+        conn.invalidate()
+        conn.close()
+
+
 def wait_for_database(attempts: int = 30, delay: float = 2.0) -> None:
     """Give the database server time to come up.
 

app/core/db.py

@@ -9,11 +9,13 @@
 
 import logging
 import re
+import secrets
 from datetime import datetime, timedelta, timezone
 
 from fastapi import APIRouter, Depends, Form, HTTPException, Request
 from starlette.concurrency import run_in_threadpool
 
+from app.config import settings
 from app.core import live, rbac
 from app.core.audit import record
 from app.core.db import session_scope
@@ -5450,6 +5452,135 @@ async def test_ai(request: Request, user: User = Depends(current_user)):
         f"{result['duration_ms'] / 1000:.1f}s: {result['text'][:80]}"))
 
 
+def _ask_admin(user: User) -> None:
+    """The question box is for administrators and nobody else.
+
+    Checked here on every call, not by hiding the button.  An administrator
+    looking at the site as somebody else is that somebody else here, so they
+    are refused too.
+    """
+    if (user is None or user.disabled or user.role != UserRole.ADMIN
+            or getattr(user, "impersonated_by", "")):
+        raise HTTPException(status_code=403, detail="Not allowed")
+
+
+def ask_available(user: User) -> bool:
+    """Whether to offer the question box: an administrator, with AI set up."""
+    if (user is None or user.disabled or user.role != UserRole.ADMIN
+            or getattr(user, "impersonated_by", "")):
+        return False
+    import time
+
+    now = time.monotonic()
+    if now - _AI_READY["at"] < 15:
+        return _AI_READY["ready"]
+    from app.core import ai
+    from app.core.db import read_only_session
+
+    try:
+        with read_only_session() as db:
+            cfg = ai.config(db)
+        ready = ai.enabled(cfg) and not ai.missing(cfg)
+    except Exception:  # noqa: BLE001 - a missing button is not worth a 500
+        log.warning("could not tell whether AI is set up", exc_info=True)
+        ready = False
+    _AI_READY.update(at=now, ready=ready)
+    return ready
+
+
+# Whether AI is set up, for a few seconds, so every page does not ask.
+_AI_READY = {"at": -1000.0, "ready": False}
+
+
+def _ask_pdf(question: str, answer: str, results: list) -> bytes:
+    from weasyprint import HTML
+
+    from app.config import APP_NAME
+    from app.core import ask
+    from app.web.render import render_partial
+    from app.web.routes_pages import report_url_fetcher
+
+    tables = [{"title": r["title"], "columns": r["columns"],
+               "rows": r["rows"][:ask.PDF_ROWS],
+               "cut": len(r["rows"]) > ask.PDF_ROWS or r["truncated"]}
+              for r in results if not r["error"] and r["columns"]]
+    html = render_partial("ask_pdf.html", app_name=APP_NAME, question=question,
+                          answer=answer, tables=tables,
+                          now=datetime.now(timezone.utc))
+    return HTML(string=html, url_fetcher=report_url_fetcher("")).write_pdf()
+
+
+@router.post("/ai/ask")
+async def ai_ask(request: Request, user: User = Depends(current_user)):
+    """A plain question about the data.  Read only, administrators only.
+
+    Nothing here writes to the database: the lookups run on a session that
+    can only SELECT, and neither the usage table nor the audit log is
+    written.  The question and the queries go to the server log instead.
+    """
+    await check_csrf(request)
+    _ask_admin(user)
+
+    from fastapi.responses import JSONResponse
+
+    from app.core import ai, ask
+    from app.core.db import read_only_session
+
+    try:
+        body = await request.json()
+    except Exception:  # noqa: BLE001
+        body = None
+    if not isinstance(body, dict) or not isinstance(body.get("question"), str):
+        return JSONResponse({"error": "Send a question."}, status_code=400)
+    if not ask.allow(user.username):
+        return JSONResponse({"error": "That is a lot of questions.  Wait a few "
+                                      "minutes and ask again."}, status_code=429)
+
+    with read_only_session() as db:
+        cfg = ai.config(db)
+    if not ai.enabled(cfg) or ai.missing(cfg):
+        return JSONResponse({"error": "AI is not set up."}, status_code=400)
+
+    log.info("ask by %s from %s: %.200r", user.username, client_ip(request),
+             body["question"])
+    try:
+        out = await run_in_threadpool(ask.ask, cfg, user, body["question"],
+                                      body.get("history"), _ask_pdf)
+    except ask.AskError as exc:
+        return JSONResponse({"error": str(exc)}, status_code=400)
+    except Exception:  # noqa: BLE001 - never a stack trace on the page
+        ref = secrets.token_hex(4)
+        log.exception("ask failed, reference %s", ref)
+        return JSONResponse({"error": f"Something went wrong.  Reference {ref}."},
+                            status_code=500)
+    for q in out["queries"]:
+        log.info("ask by %s ran: %.500s %s", user.username, q["sql"],
+                 q["error"] or "ok")
+    for item in out["files"]:
+        item["href"] = f"{settings.base_path}/ai/ask/file/{item.pop('token')}"
+    out["links"] = [{"label": link["label"],
+                     "href": f"{settings.base_path}{link['href']}"}
+                    for link in out["links"]]
+    return JSONResponse(out, headers={"Cache-Control": "no-store"})
+
+
+@router.get("/ai/ask/file/{token}")
+def ai_ask_file(token: str, user: User = Depends(current_user)):
+    """A CSV or PDF made for an answer, to the person who asked, for half an hour."""
+    _ask_admin(user)
+
+    from fastapi.responses import Response
+
+    from app.core import ask
+
+    item = ask.take_file(user.username, token)
+    if item is None:
+        raise HTTPException(status_code=404, detail="That file is gone.  Ask again.")
+    return Response(content=item["body"], media_type=item["media"], headers={
+        "Content-Disposition": 'attachment; filename="%s"' % item["name"],
+        "Cache-Control": "no-store", "X-Content-Type-Options": "nosniff"})
+
+
 @router.post("/ai/models")
 async def ai_models(request: Request, user: User = Depends(current_user)):
     """The models a provider offers the saved key, for the picker."""

app/web/routes_api.py

@@ -237,9 +237,18 @@ def _shared(db, user: User) -> dict:
         "external_on": _external_on(db),
         "can_see_external": _external_on(db) and rbac.sees_external(db, user),
         "can_manage_external": _external_on(db) and rbac.can(user, "external.manage"),
+        # The question box in the top bar.  Administrators with AI set up only;
+        # the route checks again on every question.
+        "ask_ai": _ask_available(user),
     }
 
 
+def _ask_available(user) -> bool:
+    from app.web.routes_api import ask_available
+
+    return ask_available(user)
+
+
 def _external_on(db) -> bool:
     from app.core import external
 

app/web/routes_pages.py

@@ -1217,3 +1217,55 @@ tr.decided-out > td.decision-cell { opacity: 1; }
 .rec-actions { display: flex; gap: .8rem; align-items: center; flex-wrap: wrap; }
 .rec-actions form { display: flex; gap: .6rem; align-items: center; flex-wrap: wrap; margin: 0; }
 .rec-steps { margin: 0 0 .4rem 1.2rem; padding: 0; }
+
+/* ---- Ask AI, the question box in the top bar --------------------------- */
+.ask-toggle {
+    display: inline-flex; align-items: center; gap: .35rem;
+    background: rgba(255,255,255,.1); color: var(--text-invert);
+    border: 1px solid rgba(255,255,255,.25); border-radius: 999px;
+    padding: .25rem .7rem; font: inherit; font-size: .85rem; cursor: pointer;
+}
+.ask-toggle:hover, .ask-toggle[aria-expanded="true"] { background: rgba(255,255,255,.2); }
+.ask-toggle svg { width: 14px; height: 14px; }
+.ask-panel {
+    position: fixed; top: calc(var(--topbar-height) + 8px); right: 12px; z-index: 60;
+    width: min(30rem, calc(100vw - 24px)); height: min(36rem, calc(100vh - var(--topbar-height) - 24px));
+    display: flex; flex-direction: column;
+    background: var(--bg); color: var(--text);
+    border: 1px solid var(--border-strong); border-radius: var(--radius);
+    box-shadow: 0 10px 30px rgba(0,0,0,.25);
+}
+.ask-panel[hidden] { display: none; }
+.ask-head {
+    display: flex; align-items: center; gap: .6rem; padding: .6rem .8rem;
+    border-bottom: 1px solid var(--border);
+}
+.ask-head .faint { flex: 1; }
+.ask-clear, .ask-close {
+    background: none; border: 0; color: var(--text-soft); cursor: pointer; font: inherit;
+}
+.ask-close { font-size: 1.3rem; line-height: 1; }
+.ask-log { flex: 1; overflow-y: auto; padding: .8rem; display: flex; flex-direction: column; gap: .6rem; }
+.ask-msg {
+    max-width: 92%; padding: .55rem .75rem; border-radius: 10px;
+    white-space: pre-wrap; overflow-wrap: anywhere; line-height: 1.45; font-size: .9rem;
+}
+.ask-msg.user { align-self: flex-end; background: var(--header); color: var(--text-invert); }
+.ask-msg.bot { align-self: flex-start; background: var(--bg-soft); border: 1px solid var(--border); }
+.ask-msg.pending { color: var(--text-faint); font-style: italic; }
+.ask-msg.error { background: var(--bad-bg); color: var(--bad); border-color: transparent; }
+.ask-files, .ask-links { display: flex; flex-wrap: wrap; gap: .4rem .8rem; margin-top: .5rem; white-space: normal; }
+.ask-links a { font-size: .85rem; }
+.ask-table, .ask-sql { margin-top: .5rem; white-space: normal; }
+.ask-table summary, .ask-sql summary { cursor: pointer; font-size: .82rem; color: var(--text-soft); }
+.ask-table-scroll { max-height: 14rem; overflow: auto; margin-top: .3rem; }
+.ask-table table { font-size: .78rem; }
+.ask-table td, .ask-table th { white-space: nowrap; padding: .2rem .45rem; }
+.ask-sql pre { white-space: pre-wrap; font-size: .75rem; margin: .2rem 0 .5rem; }
+.ask-form { display: flex; gap: .5rem; padding: .6rem .8rem; border-top: 1px solid var(--border); align-items: flex-end; }
+.ask-form textarea {
+    flex: 1; resize: none; font: inherit; font-size: .9rem; padding: .4rem .5rem;
+    border: 1px solid var(--border-strong); border-radius: var(--radius);
+    background: var(--bg); color: var(--text);
+}
+@media (max-width: 640px) { .ask-toggle span { display: none; } }