agent/ui_common.py

"""Shared Streamlit helpers for the agent Space pages."""
from __future__ import annotations

import streamlit as st

from agent import agdb, config as C


def db_ready() -> tuple[bool, str]:
    """(ok, message) — is the Postgres config present and reachable?"""
    try:
        conn = agdb.connect()
        conn.close()
        return True, ""
    except Exception as exc:
        return False, str(exc)


@st.cache_resource
def boot() -> dict:
    """Once per server process: schema, seed queries, start scheduler.

    Delegates to agent.autostart.boot_once, which run_space.py already ran at
    container start; this is the fallback for a process started with plain
    `streamlit run app.py` (local development, AppTest)."""
    from agent import autostart
    return autostart.boot_once()


def is_admin() -> bool:
    return bool(st.session_state.get("agent_admin", False))


def admin_gate() -> bool:
    """Sidebar unlock for state-changing controls. Returns is_admin()."""
    with st.sidebar:
        if not C.ADMIN_PASSWORD:
            st.warning("AGENT_ADMIN_PASSWORD secret is not set — "
                       "controls are locked (read-only mode).")
            return False
        if is_admin():
            st.success("Admin unlocked")
            if st.button("Lock"):
                st.session_state["agent_admin"] = False
                st.rerun()
        else:
            pw = st.text_input("Admin password", type="password",
                               key="admin_pw_input")
            if st.button("Unlock") :
                if pw == C.ADMIN_PASSWORD:
                    st.session_state["agent_admin"] = True
                    st.rerun()
                else:
                    st.error("Wrong password")
    return is_admin()


def status_banner() -> None:
    info = boot()
    if not info["db"]:
        st.error(f"Database not reachable — set the DB_* secrets. ({info['error']})")
        st.stop()
    if not C.GEMINI_API_KEY:
        st.warning("GEMINI_API_KEY is not set — discovery/download will work, "
                   "but extraction is disabled until the secret is added.")


def df_or_empty(sql: str, params: tuple = ()):
    import pandas as pd
    try:
        conn = agdb.connect()
        try:
            return agdb.fetch_df(conn, sql, params)
        finally:
            conn.close()
    except Exception:
        return pd.DataFrame()