diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6367fadc..7256b786 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,7 +76,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install package with dev dependencies - run: pip install -e ".[dev]" + run: pip install -e ".[dev,streamlit]" - name: Run tests run: pytest --tb=short -q diff --git a/.github/workflows/streamlit-sync-check.yml b/.github/workflows/streamlit-sync-check.yml new file mode 100644 index 00000000..bd5016d6 --- /dev/null +++ b/.github/workflows/streamlit-sync-check.yml @@ -0,0 +1,71 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: Streamlit query sync check + +on: + push: + branches: [main] + paths: + - "dashboards/grafana/queries/**" + - "dashboards/streamlit/**" + - "scripts/check_streamlit_queries_sync.py" + - "tests/test_dashboards_streamlit_app.py" + - "tests/test_check_streamlit_queries_sync.py" + - "tests/test_dashboards_streamlit_charts.py" + - "tests/test_dashboards_streamlit_apptest.py" + - "pyproject.toml" + - ".github/workflows/streamlit-sync-check.yml" + pull_request: + branches: [main] + paths: + - "dashboards/grafana/queries/**" + - "dashboards/streamlit/**" + - "scripts/check_streamlit_queries_sync.py" + - "tests/test_dashboards_streamlit_app.py" + - "tests/test_check_streamlit_queries_sync.py" + - "tests/test_dashboards_streamlit_charts.py" + - "tests/test_dashboards_streamlit_apptest.py" + - "pyproject.toml" + - ".github/workflows/streamlit-sync-check.yml" + +permissions: + contents: read + +jobs: + streamlit-query-sync: + name: Dashboard queries and contracts match SQL sources + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 + with: + persist-credentials: false + + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5 + with: + python-version: "3.12" + + - name: Install package with dev and streamlit dependencies + run: pip install -e ".[dev,streamlit]" + + - name: Validate Streamlit queries match canonical Grafana SQL + run: python3 scripts/check_streamlit_queries_sync.py + + - name: Run Streamlit dashboard contract tests + run: | + pytest --tb=short -q \ + tests/test_dashboards_streamlit_app.py \ + tests/test_check_streamlit_queries_sync.py \ + tests/test_dashboards_streamlit_charts.py \ + tests/test_dashboards_streamlit_apptest.py diff --git a/.gitignore b/.gitignore index 51195424..cdeb2413 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,4 @@ dashboards/grafana/.local/ # Service-account credentials (never commit downloaded JSON keys). grafana/*.json dashboards/grafana/*.json +dashboards/streamlit/*.json diff --git a/dashboards/README.md b/dashboards/README.md new file mode 100644 index 00000000..ebab3fd2 --- /dev/null +++ b/dashboards/README.md @@ -0,0 +1,12 @@ +# Self-Hosted Dashboards + +Self-hosted observability and analytics dashboards for BigQuery Agent Analytics (BQAA). + +## Available Dashboards + +| Dashboard | Primary Use Case | Documentation | +| :-------------------------- | :------------------------------------------------------------------------------------------ | :---------------------------------------------------------- | +| **Streamlit** | Interactive Python-native telemetry inspection, latency distributions, token analytics | [Streamlit Guide](streamlit/README.md) | +| **Grafana** | Time-series metrics, customizable SQL panels, public demo sharing, one-command local runner | [Grafana Guide](grafana/README.md) | +| **Looker Studio** | Published 37-chart template over BQAA views, zero-install, team-wide reporting | [Looker Studio Guide](../dashboard/looker_studio/README.md) | +| **Demo Streamlit (Legacy)** | Original standalone demo app and Colab tutorial for raw ADK table logs | [Legacy Demo App](../dashboard/README.md) | diff --git a/dashboards/streamlit/.env.sample b/dashboards/streamlit/.env.sample new file mode 100644 index 00000000..721ed324 --- /dev/null +++ b/dashboards/streamlit/.env.sample @@ -0,0 +1,14 @@ +# Path to Google Cloud service account JSON key file +GOOGLE_APPLICATION_CREDENTIALS=dashboards/streamlit/sa-key.json # Remove or comment out to use user credentials (ADC) + +# Google Cloud project ID +BQ_PROJECT_ID=my-gcp-project + +# BigQuery dataset ID +BQ_DATASET_ID=agent_analytics + +# BigQuery table ID storing raw agent telemetry events +BQ_TABLE_ID=agent_events #default + +# Prefix applied to typed analytical views +BQ_VIEW_PREFIX=adk_ #default diff --git a/dashboards/streamlit/README.md b/dashboards/streamlit/README.md new file mode 100644 index 00000000..b6e8af6f --- /dev/null +++ b/dashboards/streamlit/README.md @@ -0,0 +1,78 @@ +# Streamlit Dashboard for BigQuery Agent Analytics + +Interactive dashboard for monitoring, diagnosing, and evaluating AI agent traces in Google BigQuery. + +--- + +## 1. Install + +Install dependencies with the `dashboards` extra: + +```bash +pip install '.[streamlit]' +``` + +--- + +## 2. Configure + +Copy the sample environment file and configure your parameters: + +```bash +cp dashboards/streamlit/.env.sample dashboards/streamlit/.env +``` + +Parameters configured in `dashboards/streamlit/.env`: + +| Variable | Description | Default | +| :------------------------------- | :--------------------------------------- | :--------------------------------- | +| `GOOGLE_APPLICATION_CREDENTIALS` | Path to service account JSON key file | `dashboards/streamlit/sa-key.json` | +| `BQ_PROJECT_ID` | Google Cloud project ID | `my-gcp-project` | +| `BQ_DATASET_ID` | BigQuery dataset containing agent events | `agent_analytics` | +| `BQ_TABLE_ID` | Raw agent events table | `agent_events` | +| `BQ_VIEW_PREFIX` | Prefix for typed analytical views | `adk_` | + +--- + +## 3. Prerequisites & Typed Views + +The dashboard queries typed analytical views (`adk_llm_responses`, `adk_tool_starts`, etc.) created over the raw events table. + +### IAM Roles + +Ensure the query identity (service account or `gcloud auth application-default login`) has: +* **`roles/bigquery.jobUser`** on the Google Cloud project to run query jobs. +* **`roles/bigquery.dataViewer`** on the dataset to read tables and views. + +### Create Views + +Generate the typed analytical views matching your project, dataset, and events table: + +```bash +bq-agent-sdk views create-all \ + --project-id YOUR_PROJECT \ + --dataset-id YOUR_DATASET \ + --table-id YOUR_TABLE +``` + +> **Note:** Views default to the `adk_` prefix. If you specify a custom prefix with `--prefix`, set `BQ_VIEW_PREFIX` in `.env` (or in the dashboard sidebar) to match. + +--- + +## 4. Run + +Launch the Streamlit dashboard: + +```bash +streamlit run dashboards/streamlit/app.py +``` + +--- + +## 5. Sidebar Controls + +The sidebar provides runtime controls and guardrails for query execution: + +* **Time range selector**: Snaps query execution to discrete sliding windows (e.g., Last 15 minutes, Last 24 hours, Last 7 days, Last 30 days) with bucketed intervals. +* **Per-query scan cap guardrail (`maximum_bytes_billed`)**: Sets a strict byte limit on every BigQuery job. A free dry-run preflight validates query scan size before execution, preventing queries from running if they exceed the selected cap rather than billing for unexpected costs. +* **Token pricing defaults**: Configures input and output token rates for estimated cost calculations (defaults to \$1.25 / 1M input tokens and \$5.00 / 1M output tokens). Note that costs are derived from token counts rather than recorded billing telemetry. diff --git a/dashboards/streamlit/app.py b/dashboards/streamlit/app.py new file mode 100644 index 00000000..da29f38d --- /dev/null +++ b/dashboards/streamlit/app.py @@ -0,0 +1,804 @@ +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses +import os +from pathlib import Path +import sys +from typing import Any + +from dotenv import load_dotenv + +_DASHBOARD_DIR = Path(__file__).resolve().parent +load_dotenv(_DASHBOARD_DIR / ".env") +load_dotenv() + +_sa_creds = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") +if _sa_creds and not os.path.isabs(_sa_creds): + _candidate = _DASHBOARD_DIR / _sa_creds + if _candidate.exists(): + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = str(_candidate) + elif (_DASHBOARD_DIR.parent.parent / _sa_creds).exists(): + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = str( + _DASHBOARD_DIR.parent.parent / _sa_creds + ) + +import pandas as pd +import streamlit as st + +_DASHBOARD_DIR_STR = str(_DASHBOARD_DIR) +if _DASHBOARD_DIR_STR not in sys.path: + sys.path.insert(0, _DASHBOARD_DIR_STR) + +from charts import active_theme +from charts import grouped_bars +from charts import lines +from charts import panel +from charts import ranked_bars +from charts import stacked_bars +from models import ALL_SENTINEL +from models import as_filter_values +from models import BYTES_CAPS +from models import CACHE_TTL_SECONDS +from models import Context +from models import DEFAULT_BYTES_CAP +from models import DEFAULT_TABLE_ID +from models import DEFAULT_VIEW_PREFIX +from models import Filters +from models import humanize_bytes +from models import make_window +from models import RECENT_SESSIONS_LIMIT +from models import TableRefs +from models import TIME_RANGES +from models import TOOL_ERRORS_LIMIT +from models import TOP_ERRORS_LIMIT +from models import TRACE_DETAIL_LIMIT +from models import validate_refs +from models import Window +from queries import build_errors_over_time_sql +from queries import build_events_by_agent_sql +from queries import build_events_over_time_sql +from queries import build_llm_calls_total_sql +from queries import build_llm_latency_percentiles_sql +from queries import build_llm_tokens_over_time_sql +from queries import build_overview_totals_sql +from queries import build_recent_sessions_sql +from queries import build_tokens_by_model_sql +from queries import build_tool_errors_sql +from queries import build_tool_latency_sql +from queries import build_tool_usage_sql +from queries import build_top_errors_sql +from queries import build_trace_detail_sql +from queries import fetch +from queries import load_filter_options + +APP_TITLE = "BigQuery Agent Analytics" + +_FILTER_WIDGETS = ( + ("flt_agent", "agent"), + ("flt_user_id", "user_id"), + ("flt_event_type", "event_type"), + ("flt_session_id", "session_id"), +) + + +def _seed_options( + key: str, options: Sequence[str], applied_values: Sequence[str] +) -> list[str]: + """Ensures applied and in-progress selections stay valid widget options. + + Retention is anchored on ``applied_values`` — the last-*applied* Filters, + held in ``st.session_state["applied_filters"]`` — not solely on + ``st.session_state[key]``, the widget's own Streamlit-managed value. A + selection surviving only in ``st.session_state[key]`` is lost the moment + Streamlit treats the widget as newly created (a ``key=`` rename, a fresh + session, ...): a brand-new widget has no prior value to merge. Anchoring + on ``applied_filters`` — a plain session_state entry the widget machinery + never rewrites — means the selection survives that identity change. The + widget's own key is still merged in too, so an in-progress, not-yet-applied + edit is not clobbered while the form is open. + + Args: + key: Streamlit session_state key for the widget's own pending value. + options: Current valid options sequence from BigQuery. + applied_values: The filter's last-applied selection, or (ALL_SENTINEL,) + if unset. + + Returns: + List of options containing fetched options plus any applied or + in-progress selections. + """ + seen = set(options) + result = list(options) + + def _extend(values: Sequence[str]) -> None: + for item in values: + if item and item != ALL_SENTINEL and item not in seen: + seen.add(item) + result.append(item) + + _extend(applied_values) + current = st.session_state.get(key, []) + if not isinstance(current, (list, tuple)): + current = [current] if current else [] + _extend(current) + return result + + +def _default_for(applied_values: Sequence[str]) -> list[str]: + """Converts an applied Filters field into a multiselect ``default=``. + + Args: + applied_values: The filter's last-applied selection, e.g. + ``Filters().agents``. + + Returns: + An empty list for the "all values" sentinel, else the applied values. + """ + return ( + [] if tuple(applied_values) == (ALL_SENTINEL,) else list(applied_values) + ) + + +def sidebar_connection() -> tuple[TableRefs | None, int]: + """Resolves the BigQuery source from the environment, then the form. + + ``BQ_PROJECT_ID`` / ``BQ_DATASET_ID`` / ``BQ_TABLE_ID`` / ``BQ_VIEW_PREFIX`` + seed the fields; the form is the fallback when they are unset and the override + for dataset, table, and view prefix when they are wrong. Project ID is not + overridable when ``BQ_PROJECT_ID`` is set. It is a form so that + changing multiple configuration settings costs one rerun rather than five. + + ``_connect_attempted`` latches in session state once the Connect button is + clicked, persisting for the session until a valid connection or reset. + + Returns: + A tuple containing validated TableRefs (or None if invalid) and the + selected byte cap limit. + """ + env_project = os.environ.get("BQ_PROJECT_ID", "") + env_dataset = os.environ.get("BQ_DATASET_ID", "") + env_table = os.environ.get("BQ_TABLE_ID", "") or DEFAULT_TABLE_ID + env_prefix = os.environ.get("BQ_VIEW_PREFIX", DEFAULT_VIEW_PREFIX) + + st.sidebar.subheader("BigQuery source") + + with st.sidebar.form("connection"): + project = st.text_input( + "Project ID", value=env_project, disabled=bool(env_project) + ) + dataset = st.text_input("Dataset ID", value=env_dataset) + table = st.text_input("Events table", value=env_table) + prefix = st.text_input( + "Typed view prefix", + value=env_prefix, + help=( + "The prefix ViewManager applied to the typed views" + " (`adk_` by default)." + ), + ) + cap_label = st.selectbox( + "Per-query scan cap", + options=list(BYTES_CAPS), + index=list(BYTES_CAPS).index(DEFAULT_BYTES_CAP), + help=( + "Sets `maximum_bytes_billed` on every job. BigQuery refuses a" + " query that would exceed it rather than billing for it." + ), + ) + connected = st.form_submit_button("Connect", width="stretch") + if connected: + st.session_state["_connect_attempted"] = True + + if env_project: + project = env_project + refs, errors = validate_refs(project, dataset, table, prefix) + if st.session_state.get("_connect_attempted") or bool(dataset): + for message in errors: + st.sidebar.error(message) + return refs, BYTES_CAPS[cap_label] + + +def sidebar_window() -> Window: + """Renders the time range picker widget in the sidebar. + + A selectbox rather than a slider or free-text: one discrete choice per + rerun, so a query fires on a committed selection instead of on every + intermediate value. + + Returns: + A snapped Window instance for the selected time range. + """ + st.sidebar.subheader("Time range") + label = st.sidebar.selectbox( + "Range", + options=list(TIME_RANGES), + index=list(TIME_RANGES).index("Last 24 hours"), + label_visibility="collapsed", + ) + window = make_window(TIME_RANGES[label]) + st.sidebar.caption( + f"{window.start:%Y-%m-%d %H:%M} → {window.end:%Y-%m-%d %H:%M} UTC" + f" · {window.bucket.lower()} buckets" + ) + return window + + +def reset_filters() -> None: + """Resets applied and widget filter states to defaults. + + Pricing inputs (flt_price_in, flt_price_out) are intentionally preserved + across connection changes. + """ + st.session_state["applied_filters"] = Filters() + st.session_state.pop("_filter_options", None) + for key, _ in _FILTER_WIDGETS: + st.session_state.pop(key, None) + st.session_state.pop("_selected_session_id", None) + + +def _pending(key: str) -> list[str]: + """Reads one filter widget's submitted value out of session_state.""" + value = st.session_state.get(key, []) + if not isinstance(value, (list, tuple)): + return [str(value)] if value else [] + return [str(item) for item in value] + + +def _commit_filters() -> None: + """Promotes the submitted widget values into applied_filters. + + Runs as the Apply button's on_click callback. + """ + applied = Filters( + agents=as_filter_values(_pending("flt_agent")), + user_ids=as_filter_values(_pending("flt_user_id")), + event_types=as_filter_values(_pending("flt_event_type")), + session_ids=as_filter_values(_pending("flt_session_id")), + ) + st.session_state["applied_filters"] = applied + for key, kind in _FILTER_WIDGETS: + st.session_state[key] = _default_for(getattr(applied, f"{kind}s")) + + +def sidebar_filters( + options: dict[str, list[str]], +) -> tuple[Filters, float, float]: + """Renders the filter and pricing form in the sidebar. + + A form, so a multi-select that a user is still building does not fire a + query per keystroke: every panel re-queries once, on Apply filters. + + Args: + options: Map of filter kinds to available option string lists. + + Returns: + A tuple of (Filters instance, price_in float, price_out float). + """ + st.sidebar.subheader("Filters") + applied: Filters = st.session_state.setdefault("applied_filters", Filters()) + + for key, kind in _FILTER_WIDGETS: + field = f"{kind}s" + if key not in st.session_state: + st.session_state[key] = _default_for(getattr(applied, field)) + + with st.sidebar.form("filters"): + st.multiselect( + "Agent", + options=_seed_options( + "flt_agent", options.get("agent", []), applied.agents + ), + key="flt_agent", + accept_new_options=True, + ) + st.multiselect( + "User", + options=_seed_options( + "flt_user_id", options.get("user_id", []), applied.user_ids + ), + key="flt_user_id", + accept_new_options=True, + ) + st.multiselect( + "Event type", + options=_seed_options( + "flt_event_type", options.get("event_type", []), applied.event_types + ), + key="flt_event_type", + help=( + "Honored by Events over time, Events by agent, Recent sessions" + " and Trace detail. Error panels and view-backed panels are" + " exempt — see dashboards/grafana/queries/README.md." + ), + ) + st.multiselect( + "Session", + options=_seed_options( + "flt_session_id", options.get("session_id", []), applied.session_ids + ), + key="flt_session_id", + accept_new_options=True, + ) + st.caption("An empty selection means all values.") + st.divider() + st.caption("Cost is derived from token counts, not recorded telemetry.") + price_in = st.number_input( + "USD per 1M input tokens", + min_value=0.0, + value=1.25, + step=0.25, + format="%.4f", + key="flt_price_in", + ) + price_out = st.number_input( + "USD per 1M output tokens", + min_value=0.0, + value=5.00, + step=0.25, + format="%.4f", + key="flt_price_out", + ) + st.form_submit_button( + "Apply filters", width="stretch", on_click=_commit_filters + ) + + return st.session_state["applied_filters"], float(price_in), float(price_out) + + +def _metric(column: Any, label: str, value: str, help_text: str = "") -> None: + """Renders a single metric widget in a layout column. + + Args: + column: Streamlit column layout element. + label: Metric label. + value: Metric value string. + help_text: Optional tooltip help text. + """ + column.metric(label, value, help=help_text or None) + + +def row_overview(ctx: Context) -> None: + """Renders the Overview tab (top KPIs, events over time, errors, top errors). + + Args: + ctx: Active dashboard context. + """ + totals = fetch( + build_overview_totals_sql(ctx.refs, ctx.window), ctx, "Overview stats" + ) + cols = st.columns(4) + if totals.df.empty: + for col, label in zip( + cols, ("Sessions", "Events", "Error rate", "Avg LLM latency") + ): + _metric(col, label, "—") + else: + row = totals.df.iloc[0] + _metric(cols[0], "Sessions", f"{int(row['sessions']):,}") + _metric(cols[1], "Events", f"{int(row['events']):,}") + rate = row["error_rate"] + _metric( + cols[2], + "Error rate", + "—" if pd.isna(rate) else f"{float(rate) * 100:.2f}%", + "An event counts as an error when its type ends in _ERROR, it" + " carries an error message, or its status is ERROR.", + ) + latency = row["avg_llm_latency_ms"] + _metric( + cols[3], + "Avg LLM latency", + "—" if pd.isna(latency) else f"{round(float(latency) + 1e-9):,.0f} ms", + ) + + left, right = st.columns(2) + with left: + events = fetch( + build_events_over_time_sql(ctx.refs, ctx.window), + ctx, + "Events over time", + ) + fig = ( + stacked_bars( + events.df, "bucket", "event_type", "events", ctx, "event_type" + ) + if not events.df.empty + else None + ) + panel("Events over time", fig, events.df, key="events_over_time") + with right: + errors = fetch( + build_errors_over_time_sql(ctx.refs, ctx.window), + ctx, + "Errors over time", + ) + fig = ( + stacked_bars( + errors.df, "bucket", "event_type", "errors", ctx, "event_type" + ) + if not errors.df.empty + else None + ) + panel( + "Errors over time", + fig, + errors.df, + empty="No errors in this range.", + key="errors_over_time", + ) + + left, right = st.columns(2) + with left: + by_agent = fetch( + build_events_by_agent_sql(ctx.refs, ctx.window), + ctx, + "Events by agent", + ) + fig = ( + ranked_bars(by_agent.df, "agent_name", "events", ctx) + if not by_agent.df.empty + else None + ) + panel("Events by agent", fig, by_agent.df, key="events_by_agent") + with right: + top_errors = fetch( + build_top_errors_sql(ctx.refs, ctx.window, TOP_ERRORS_LIMIT), + ctx, + "Top error messages", + ) + panel( + "Top error messages", + None, + top_errors.df, + empty="No error messages in this range.", + ) + st.caption( + "Only events carrying a message are listed, so these counts are a" + " subset of Errors over time." + ) + + +def row_llm(ctx: Context) -> None: + """Renders the LLM & FinOps tab (totals, token trends, latency, models). + + Args: + ctx: Active dashboard context. + """ + summary = fetch( + build_llm_calls_total_sql(ctx.refs, ctx.window), + ctx, + "LLM totals", + ) + cols = st.columns(4) + if summary.df.empty: + for col, label in zip( + cols, ("LLM calls", "Total tokens", "Output tokens", "Estimated cost") + ): + _metric(col, label, "—") + else: + row = summary.df.iloc[0] + prompt_tokens = float(row.get("prompt_tokens", 0) or 0) + completion_tokens = float(row.get("completion_tokens", 0) or 0) + estimated_cost = (prompt_tokens / 1e6 * ctx.price_in) + ( + completion_tokens / 1e6 * ctx.price_out + ) + _metric( + cols[0], + "LLM calls", + f"{int(row['llm_calls']):,}", + "Counted per distinct span, so streaming chunks do not inflate it.", + ) + _metric(cols[1], "Total tokens", f"{int(row['total_tokens']):,}") + _metric(cols[2], "Output tokens", f"{int(row['completion_tokens']):,}") + _metric( + cols[3], + "Estimated cost", + f"${estimated_cost:,.2f}", + "Derived from token counts at the sidebar rates — not telemetry.", + ) + + left, right = st.columns(2) + with left: + tokens = fetch( + build_llm_tokens_over_time_sql(ctx.refs, ctx.window), + ctx, + "Token usage over time", + ) + fig = None + if not tokens.df.empty: + melted = tokens.df.melt( + id_vars="bucket", + value_vars=["prompt_tokens", "completion_tokens"], + var_name="kind", + value_name="tokens", + ) + melted["kind"] = melted["kind"].map( + {"prompt_tokens": "Input", "completion_tokens": "Output"} + ) + fig = stacked_bars(melted, "bucket", "kind", "tokens", ctx, "token_kind") + panel("Token usage over time", fig, tokens.df, key="tokens_over_time") + with right: + latency = fetch( + build_llm_latency_percentiles_sql(ctx.refs, ctx.window), + ctx, + "LLM latency", + ) + fig = ( + lines( + latency.df, + "bucket", + [ + ("p50_total_ms", "p50"), + ("p95_total_ms", "p95"), + ("p50_ttft_ms", "TTFT p50"), + ], + ctx, + "llm_latency", + unit=" ms", + ) + if not latency.df.empty + else None + ) + panel("LLM latency (ms)", fig, latency.df, key="llm_latency") + st.caption("Gaps are missing telemetry, not zero latency.") + + models = fetch( + build_tokens_by_model_sql(ctx.refs, ctx.window), ctx, "Tokens by model" + ) + fig = ( + grouped_bars( + models.df, + "model", + [("prompt_tokens", "Input"), ("completion_tokens", "Output")], + ctx, + "token_kind", + ) + if not models.df.empty + else None + ) + panel("Tokens by model", fig, models.df, key="tokens_by_model") + + +def row_tools(ctx: Context) -> None: + """Renders the Tools & Execution tab (tool invocations, latency, errors). + + Args: + ctx: Active dashboard context. + """ + left, right = st.columns(2) + with left: + usage = fetch( + build_tool_usage_sql(ctx.refs, ctx.window), ctx, "Tool invocations" + ) + fig = ( + ranked_bars(usage.df, "tool_name", "invocations", ctx) + if not usage.df.empty + else None + ) + panel("Tool invocations", fig, usage.df, key="tool_usage") + st.caption("Read from tool_starts, so failed invocations still count.") + with right: + latency = fetch( + build_tool_latency_sql(ctx.refs, ctx.window), ctx, "Tool latency" + ) + fig = ( + grouped_bars( + latency.df, + "tool_name", + [("p95_ms", "p95"), ("p50_ms", "p50")], + ctx, + "tool_latency", + ) + if not latency.df.empty + else None + ) + panel("Tool latency (ms)", fig, latency.df, key="tool_latency") + + errors = fetch( + build_tool_errors_sql(ctx.refs, ctx.window, TOOL_ERRORS_LIMIT), + ctx, + "Tool errors", + ) + panel( + "Tool errors", + None, + errors.df, + empty="No tool errors in this range.", + ) + + +def row_sessions(ctx: Context) -> None: + """Renders the Sessions & Traces tab (recent sessions table, trace details). + + Args: + ctx: Active dashboard context. + """ + sessions = fetch( + build_recent_sessions_sql(ctx.refs, ctx.window, RECENT_SESSIONS_LIMIT), + ctx, + "Recent sessions", + ) + panel( + f"Recent sessions (most recent {RECENT_SESSIONS_LIMIT})", + None, + sessions.df, + empty="No sessions in this range.", + ) + st.caption( + "Each row rolls up the whole session in the window, including events" + " the Agent / User / Event type filters exclude, so these numbers" + " will not sum to the Overview stats." + ) + + st.divider() + st.markdown("**Trace detail**") + ids = ( + [str(v) for v in sessions.df["session_id"].tolist()] + if not sessions.df.empty + else [] + ) + if not ids: + st.caption("Pick a session once one is listed above.") + return + # A selectbox rather than a text input: one committed choice per rerun, + # so the trace query runs once instead of on every keystroke. + # If a previously chosen session exists and is still in ids, maintain its selection. + # `_selected_session_id` must remain a non-widget key (stored in session_state, + # not passed as key="...") so that dynamically computed default_idx does not + # conflict with Streamlit's internal widget key state tracking or raise + # StreamlitAPIException when session options change between queries. + prev_chosen = st.session_state.get("_selected_session_id") + default_idx = ids.index(prev_chosen) if prev_chosen in ids else 0 + chosen = st.selectbox("Session", options=ids, index=default_idx) + st.session_state["_selected_session_id"] = chosen + # Pin the trace to one session without disturbing the shared filters. + # `replace` carries the same `scan_log` list over, so this query is + # still counted once in the footer. + detail_ctx = dataclasses.replace( + ctx, filters=ctx.filters._replace(session_ids=(chosen,)) + ) + trace = fetch( + build_trace_detail_sql(ctx.refs, ctx.window, TRACE_DETAIL_LIMIT), + detail_ctx, + "Trace detail", + ) + if trace.df.empty: + st.caption("No events for this session in the range.") + else: + st.dataframe(trace.df, width="stretch", hide_index=True) + st.caption( + f"Newest {TRACE_DETAIL_LIMIT} events first. Sort by timestamp to" + " read the session chronologically." + ) + + +def footer(ctx: Context) -> None: + """Reports what this rerun actually scanned and cached. + + Args: + ctx: Active dashboard context. + """ + if not ctx.scan_log: + return + total_billed = 0 + total_processed = 0 + cached = 0 + for _, billed, processed, hit in ctx.scan_log: + if hit: + cached += 1 + else: + total_billed += billed + total_processed += processed + + st.divider() + if total_billed == 0 and cached > 0: + billed_processed = ( + f"{humanize_bytes(total_billed)} billed" + f" ({humanize_bytes(total_processed)} processed — billing cost saved via cache)" + ) + else: + billed_processed = ( + f"{humanize_bytes(total_billed)} billed" + f" ({humanize_bytes(total_processed)} processed)" + ) + cache_note = ( + f"{cached} served from cache ($0 billed)" + if cached > 0 + else f"{cached} served from cache" + ) + st.caption( + f"{len(ctx.scan_log)} queries this run · {billed_processed} ·" + f" {cache_note} ·" + f" per-query cap {humanize_bytes(ctx.max_bytes)} ·" + f" results cached for {CACHE_TTL_SECONDS // 60} min" + ) + st.caption( + "BigQuery bills a 10 MB minimum per query under on-demand pricing, while" + " queries running on compute/capacity reservations incur no per-byte" + " charges." + ) + + +def main() -> None: + """Runs the main entrypoint for the Streamlit dashboard application.""" + st.set_page_config(page_title=APP_TITLE, page_icon="📊", layout="wide") + st.title(APP_TITLE) + + theme = active_theme() + refs, max_bytes = sidebar_connection() + if refs is None: + st.info( + "Set `BQ_PROJECT_ID`, `BQ_DATASET_ID` and `BQ_TABLE_ID`, or fill in" + " the sidebar, to connect." + ) + st.stop() + + prev_refs = st.session_state.get("_last_refs") + if prev_refs is not None and prev_refs != refs: + reset_filters() + st.session_state["_last_refs"] = refs + + window = sidebar_window() + + # Options are read unfiltered, so the sidebar can be drawn before any + # panel runs and picking one agent never hides the others. + probe = Context( + refs=refs, + window=window, + filters=Filters(), + max_bytes=max_bytes, + theme=theme, + price_in=0.0, + price_out=0.0, + ) + options, result = load_filter_options(probe) + if result.error is None and options: + st.session_state["_filter_options"] = options + else: + options = st.session_state.get("_filter_options", {}) + + filters, price_in, price_out = sidebar_filters(options) + + ctx = Context( + refs=refs, + window=window, + filters=filters, + max_bytes=max_bytes, + theme=theme, + price_in=price_in, + price_out=price_out, + scan_log=probe.scan_log, + ) + + overview, llm, tools, sessions = st.tabs( + ["Overview", "LLM & FinOps", "Tools & Execution", "Sessions & Traces"] + ) + with overview: + row_overview(ctx) + with llm: + row_llm(ctx) + with tools: + row_tools(ctx) + with sessions: + row_sessions(ctx) + + footer(ctx) + + +__all__ = [ + "APP_TITLE", + "footer", + "main", + "reset_filters", + "row_llm", + "row_overview", + "row_sessions", + "row_tools", + "sidebar_connection", + "sidebar_filters", + "sidebar_window", +] + +if __name__ == "__main__": + main() diff --git a/dashboards/streamlit/charts.py b/dashboards/streamlit/charts.py new file mode 100644 index 00000000..127bd13e --- /dev/null +++ b/dashboards/streamlit/charts.py @@ -0,0 +1,402 @@ +"""Chart builders, theme helpers, and panel rendering for Streamlit.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +import sys + +import pandas as pd +import plotly.graph_objects as go +import streamlit as st + +_DASHBOARD_DIR = str(Path(__file__).resolve().parent) +if _DASHBOARD_DIR not in sys.path: + sys.path.insert(0, _DASHBOARD_DIR) + +from models import Context +from models import DARK_THEME +from models import LIGHT_THEME +from models import OTHER_LABEL +from models import Theme + + +def active_theme() -> Theme: + """Returns the palette matching Streamlit's current appearance. + + Returns: + The active Theme object (DARK_THEME if Streamlit is in dark mode, + else LIGHT_THEME). + """ + mode = None + try: + mode = st.context.theme.type + except ( + AttributeError, + TypeError, + ): # pragma: no cover - older Streamlit builds. + mode = None + if not mode: + mode = st.get_option("theme.base") or "light" + return DARK_THEME if str(mode).lower() == "dark" else LIGHT_THEME + + +def color_map( + domain: str, names: Sequence[str], theme: Theme +) -> dict[str, str]: + """Maps category names to categorical slots, stably across reruns. + + Color follows the entity, not its rank: a name keeps the slot it was + first given for as long as the session lives, so narrowing a filter + never repaints the series that survive. Slots are only reassigned when + two names visible in the *same* chart would otherwise collide. + + Args: + domain: State namespace for slot allocation (e.g. 'agent', 'event_type'). + names: Category names to assign colors to. + theme: Active theme containing categorical color palette. + + Returns: + A mapping from category name to color hex code. + """ + registry: dict[str, int] = st.session_state.setdefault( + f"_slots::{domain}", {} + ) + live = [n for n in dict.fromkeys(names) if n != OTHER_LABEL] + taken: dict[int, str] = {} + needs_slot: list[str] = [] + for name in live: + slot = registry.get(name) + if slot is None or slot in taken: + needs_slot.append(name) + else: + taken[slot] = name + for name in needs_slot: + free = next( + (i for i in range(len(theme.categorical)) if i not in taken), + None, + ) + if free is None: # More than 8 live names: caller failed to fold. + free = len(registry) % len(theme.categorical) + registry[name] = free + taken[free] = name + mapping = {n: theme.categorical[registry[n]] for n in live} + mapping[OTHER_LABEL] = theme.muted + return mapping + + +def base_figure(theme: Theme, height: int = 320) -> go.Figure: + """Creates a Plotly figure with recessive chrome and hairline solid gridlines. + + Args: + theme: Active color theme. + height: Desired figure height in pixels. + + Returns: + A configured go.Figure instance. + """ + fig = go.Figure() + fig.update_layout( + height=height, + margin=dict(l=8, r=8, t=8, b=8), + paper_bgcolor=theme.surface, + plot_bgcolor=theme.surface, + font=dict( + family='system-ui, -apple-system, "Segoe UI", sans-serif', + size=12, + color=theme.text_secondary, + ), + hoverlabel=dict( + bgcolor=theme.surface, + bordercolor=theme.axis, + font=dict(color=theme.text_primary, size=12), + ), + legend=dict( + orientation="h", + yanchor="bottom", + y=1.02, + x=0, + bgcolor="rgba(0,0,0,0)", + font=dict(color=theme.text_secondary), + ), + showlegend=False, + ) + axis = dict( + showgrid=True, + gridcolor=theme.grid, + gridwidth=1, + griddash="solid", + zeroline=False, + linecolor=theme.axis, + tickfont=dict(color=theme.muted, size=11), + title=None, + ) + fig.update_xaxes(**axis, showline=True) + fig.update_yaxes(**axis, showline=False) + return fig + + +def fold_others( + df: pd.DataFrame, + key: str, + value: str, + group_cols: Sequence[str] = (), + limit: int = 8, +) -> pd.DataFrame: + """Folds all but the top ``limit`` categories into a single "Other". + + A ninth categorical hue is never generated: past the palette's eight + slots, the tail becomes one gray series. + + Args: + df: Input dataframe. + key: Column name holding the categorical dimension to fold. + value: Column name holding the numeric measure to sum by. + group_cols: Optional columns to keep in the group-by aggregation. + limit: Maximum number of distinct categories before folding. + + Returns: + A dataframe with the smallest categories aggregated under OTHER_LABEL. + """ + if df.empty or df[key].nunique() <= limit: + return df + keep = df.groupby(key)[value].sum().nlargest(limit - 1).index + out = df.copy() + out[key] = out[key].where(out[key].isin(keep), OTHER_LABEL) + return out.groupby([*group_cols, key], as_index=False)[value].sum() + + +def stacked_bars( + df: pd.DataFrame, + x: str, + key: str, + value: str, + ctx: Context, + domain: str, + height: int = 320, +) -> go.Figure: + """Creates stacked bars over time, one series per category. + + Args: + df: Input dataframe. + x: Column name for the x-axis (e.g. bucket timestamp). + key: Column name for categorical series segmentation. + value: Column name for numeric bar heights. + ctx: Active dashboard context. + domain: Color domain namespace. + height: Figure height in pixels. + + Returns: + A configured Plotly stacked bar chart. + """ + theme = ctx.theme + folded = fold_others(df, key, value, group_cols=[x]) + names = sorted(folded[key].unique(), key=str) + colors = color_map(domain, [n for n in names if n != OTHER_LABEL], theme) + fig = base_figure(theme, height) + for name in names: + part = folded[folded[key] == name].sort_values(x) + fig.add_bar( + x=part[x], + y=part[value], + name=str(name), + marker=dict( + color=colors.get(name, theme.muted), + cornerradius=4, + # A surface-colored hairline reads as the 2px gap between + # stacked segments, not as a border around the marks. + line=dict(color=theme.surface, width=1), + ), + hovertemplate=f"%{{x}}
{name}: %{{y:,}}", + ) + fig.update_layout( + barmode="stack", + bargap=0.25, + showlegend=len(names) > 1, + hovermode="x unified", + ) + return fig + + +def ranked_bars( + df: pd.DataFrame, + key: str, + value: str, + ctx: Context, + height: int = 320, + top: int = 12, +) -> go.Figure: + """Creates horizontal bars for one measure across nominal categories. + + One series, one color: length already encodes magnitude, so a + darker-where-bigger ramp would spend the only free channel restating it. + + Args: + df: Input dataframe. + key: Column name for categorical labels. + value: Column name for numeric values. + ctx: Active dashboard context. + height: Figure height in pixels. + top: Number of top categories to display. + + Returns: + A configured Plotly horizontal bar chart. + """ + theme = ctx.theme + part = df.nlargest(top, value).sort_values(value) + fig = base_figure(theme, height) + fig.add_bar( + x=part[value], + y=part[key].astype(str), + orientation="h", + marker=dict( + color=theme.categorical[0], + cornerradius=4, + line=dict(color=theme.surface, width=1), + ), + hovertemplate="%{y}: %{x:,}", + ) + fig.update_layout(bargap=0.35, showlegend=False) + fig.update_xaxes(tickformat=",") + return fig + + +def grouped_bars( + df: pd.DataFrame, + key: str, + series: Sequence[tuple[str, str]], + ctx: Context, + domain: str, + height: int = 320, + top: int = 12, +) -> go.Figure: + """Creates horizontal grouped bars for measures that share one unit. + + Args: + df: Input dataframe. + key: Column name for categorical labels. + series: Sequence of (column_name, display_label) tuples. + ctx: Active dashboard context. + domain: Color domain namespace. + height: Figure height in pixels. + top: Number of top categories to display. + + Returns: + A configured Plotly horizontal grouped bar chart. + """ + theme = ctx.theme + first_value = series[0][0] + part = df.nlargest(top, first_value).sort_values(first_value) + colors = color_map(domain, [label for _, label in series], theme) + fig = base_figure(theme, height) + for column, label in series: + fig.add_bar( + x=part[column], + y=part[key].astype(str), + orientation="h", + name=label, + marker=dict( + color=colors[label], + cornerradius=4, + line=dict(color=theme.surface, width=1), + ), + hovertemplate=f"%{{y}} — {label}: %{{x:,.0f}}", + ) + fig.update_layout( + barmode="group", bargap=0.3, bargroupgap=0.08, showlegend=True + ) + fig.update_xaxes(tickformat=",") + return fig + + +def lines( + df: pd.DataFrame, + x: str, + series: Sequence[tuple[str, str]], + ctx: Context, + domain: str, + height: int = 320, + unit: str = "", +) -> go.Figure: + """Creates time-series lines sharing one y-axis, direct-labeled at the end. + + Never two y-scales: every series passed here is in the same unit. + + Args: + df: Input dataframe. + x: Column name for the x-axis time/bucket values. + series: Sequence of (column_name, display_label) tuples. + ctx: Active dashboard context. + domain: Color domain namespace. + height: Figure height in pixels. + unit: Optional unit suffix for hover text. + + Returns: + A configured Plotly line chart. + """ + theme = ctx.theme + ordered = df.sort_values(x) + colors = color_map(domain, [label for _, label in series], theme) + fig = base_figure(theme, height) + for column, label in series: + fig.add_scatter( + x=ordered[x], + y=ordered[column], + mode="lines", + name=label, + line=dict(color=colors[label], width=2, shape="linear"), + connectgaps=False, + hovertemplate=f"{label}: %{{y:,.0f}}{unit}", + ) + # Direct labels on the last real point: with <= 4 series, identity is + # never carried by color alone even before the legend is read. + if len(series) <= 4: + for column, label in series: + valid = ordered[[x, column]].dropna() + if valid.empty: + continue + last = valid.iloc[-1] + fig.add_annotation( + x=last[x], + y=last[column], + text=label, + showarrow=False, + xanchor="left", + xshift=8, + font=dict(size=11, color=theme.text_secondary), + bgcolor=theme.surface, + borderpad=2, + ) + fig.update_layout(margin=dict(l=8, r=96, t=8, b=8)) + fig.update_layout(showlegend=True, hovermode="x unified") + return fig + + +def panel( + title: str, + fig: go.Figure | None, + table: pd.DataFrame, + *, + empty: str = "No matching data in this range.", + key: str = "", +) -> None: + """Renders one chart with its table-view twin behind an expander. + + Every chart has a table twin: it is the WCAG-clean equivalent, and it + is the relief the light-mode palette's sub-3:1 slots require. + + Args: + title: Panel title markdown string. + fig: Optional Plotly figure to render above the table. + table: Backing dataframe for the table expander. + empty: Caption displayed if the table is empty. + key: Streamlit widget key for the chart. + """ + st.markdown(f"**{title}**") + if table.empty: + st.caption(empty) + return + if fig is not None: + st.plotly_chart(fig, width="stretch", key=key or None) + with st.expander("Table view", expanded=fig is None): + st.dataframe(table, width="stretch", hide_index=True) \ No newline at end of file diff --git a/dashboards/streamlit/models.py b/dashboards/streamlit/models.py new file mode 100644 index 00000000..e252d588 --- /dev/null +++ b/dashboards/streamlit/models.py @@ -0,0 +1,357 @@ +"""Data models, configuration, and types for the Streamlit dashboard.""" + +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses +import datetime as dt +import re +from typing import Any, NamedTuple + +import pandas as pd + +# ------------------------------------------------------------------ # +# Constants # +# ------------------------------------------------------------------ # + +# Sentinel that means "no filter applied", matching dashboards/grafana/queries. The +# clause `('___ALL___' IN UNNEST(@agents) OR agent IN UNNEST(@agents))` +# is injection-safe and cannot crash on an empty array the way an +# `IN ()` list would. +ALL_SENTINEL = "___ALL___" + +DEFAULT_TABLE_ID = "agent_events" +DEFAULT_VIEW_PREFIX = "adk_" + +OTHER_LABEL = "Other" + +# Identifier grammars. Neither pattern admits a backtick, so a validated +# identifier is safe to interpolate into a backtick-quoted table path. +_PROJECT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9\-_.:]{0,62}$") +_NAME_RE = re.compile(r"^[A-Za-z0-9_]{1,1024}$") +_PREFIX_RE = re.compile(r"^[A-Za-z0-9_]{0,1024}$") + +TIME_RANGES: dict[str, dt.timedelta] = { + "Last 1 hour": dt.timedelta(hours=1), + "Last 6 hours": dt.timedelta(hours=6), + "Last 24 hours": dt.timedelta(hours=24), + "Last 3 days": dt.timedelta(days=3), + "Last 7 days": dt.timedelta(days=7), + "Last 30 days": dt.timedelta(days=30), +} + +# Per-query scan caps offered in the sidebar. The selected value is set as +# `maximum_bytes_billed` on every job, so BigQuery refuses a query that +# would exceed it rather than billing for it. +BYTES_CAPS: dict[str, int] = { + "100 MB": 100 * 1024**2, + "1 GB": 1024**3, + "10 GB": 10 * 1024**3, + "100 GB": 100 * 1024**3, + "1 TB": 1024**4, +} +DEFAULT_BYTES_CAP = "1 GB" + +# Results are cached for this long, and the window's upper bound is +# snapped to the same interval. Without the snap, "Last 24 hours" would +# produce a new `end` timestamp on every rerun, so the SQL string — and +# therefore the cache key — would never repeat and every widget change +# would re-bill every panel. +CACHE_TTL_SECONDS = 300 + +# Rows returned by the detail tables, matching the Grafana caps. +RECENT_SESSIONS_LIMIT = 250 +TRACE_DETAIL_LIMIT = 500 +TOOL_ERRORS_LIMIT = 100 +FILTER_OPTIONS_LIMIT = 1000 +TOP_ERRORS_LIMIT = 50 + + +# ------------------------------------------------------------------ # +# Theme # +# ------------------------------------------------------------------ # + + +@dataclasses.dataclass(frozen=True) +class Theme: + """Chart palette for one appearance mode. + + Attributes: + surface: Background surface hex color. + text_primary: Primary text hex color. + text_secondary: Secondary text hex color. + muted: Muted text or element hex color. + grid: Grid line hex color. + axis: Axis line hex color. + categorical: Categorical color palette sequence. + """ + + surface: str + text_primary: str + text_secondary: str + muted: str + grid: str + axis: str + categorical: tuple[str, ...] + + +# Both orderings are validated for the adjacent-pair gates (stacked bars, +# grouped bars, lines) against their own surface: worst adjacent CVD +# delta-E 9.1 light / 8.4 dark, worst normal-vision 19.6 / 19.3. +LIGHT_THEME = Theme( + surface="#fcfcfb", + text_primary="#0b0b0b", + text_secondary="#52514e", + muted="#898781", + grid="#e1e0d9", + axis="#c3c2b7", + categorical=( + "#2a78d6", + "#eb6834", + "#1baf7a", + "#eda100", + "#e87ba4", + "#008300", + "#4a3aa7", + "#e34948", + ), +) + +DARK_THEME = Theme( + surface="#1a1a19", + text_primary="#ffffff", + text_secondary="#c3c2b7", + muted="#898781", + grid="#2c2c2a", + axis="#383835", + categorical=( + "#3987e5", + "#d95926", + "#199e70", + "#c98500", + "#d55181", + "#008300", + "#9085e9", + "#e66767", + ), +) + + +# ------------------------------------------------------------------ # +# Configuration & Context Models # +# ------------------------------------------------------------------ # + + +@dataclasses.dataclass(frozen=True) +class TableRefs: + """Validated BigQuery identifiers for the raw table and typed views. + + Attributes: + project: BigQuery project ID. + dataset: BigQuery dataset ID. + table: Events table ID. + view_prefix: Prefix applied to typed views. + """ + + project: str + dataset: str + table: str + view_prefix: str = DEFAULT_VIEW_PREFIX + + @property + def events(self) -> str: + """Returns the backtick-quoted full table path for the raw events table.""" + return f"`{self.project}.{self.dataset}.{self.table}`" + + def view(self, suffix: str) -> str: + """Returns the backtick-quoted full table path for a typed view. + + Args: + suffix: Suffix of the typed view (e.g. 'llm_responses'). + + Returns: + The backtick-quoted view path. + """ + return f"`{self.project}.{self.dataset}.{self.view_prefix}{suffix}`" + + +def validate_refs( + project: str, dataset: str, table: str, view_prefix: str +) -> tuple[TableRefs | None, list[str]]: + """Validates identifiers before they reach a SQL string. + + Returns the refs and an empty error list, or ``None`` and the reasons. + BigQuery cannot parameterize a table path, so this is the boundary that + keeps user-supplied text out of the FROM clause. + + Args: + project: BigQuery project ID. + dataset: BigQuery dataset ID. + table: Events table ID. + view_prefix: Prefix applied to typed views. + + Returns: + A tuple containing the validated TableRefs instance (or None if validation + failed) and a list of error description strings. + """ + errors: list[str] = [] + if not _PROJECT_RE.match(project or ""): + errors.append(f"Invalid project ID: {project!r}") + if not _NAME_RE.match(dataset or ""): + errors.append(f"Invalid dataset ID: {dataset!r}") + if not _NAME_RE.match(table or ""): + errors.append(f"Invalid table ID: {table!r}") + if not _PREFIX_RE.match(view_prefix or ""): + errors.append(f"Invalid view prefix: {view_prefix!r}") + if errors: + return None, errors + return TableRefs(project, dataset, table, view_prefix), [] + + +class Filters(NamedTuple): + """Filter selections, as tuples so the whole value is hashable. + + Being hashable is what lets ``@st.cache_data`` key on the filters: they + live in query parameters rather than in the SQL text, so the SQL string + alone would be an incomplete cache key. + + Attributes: + agents: Selected agent names, or (ALL_SENTINEL,). + user_ids: Selected user IDs, or (ALL_SENTINEL,). + event_types: Selected event types, or (ALL_SENTINEL,). + session_ids: Selected session IDs, or (ALL_SENTINEL,). + """ + + agents: tuple[str, ...] = (ALL_SENTINEL,) + user_ids: tuple[str, ...] = (ALL_SENTINEL,) + event_types: tuple[str, ...] = (ALL_SENTINEL,) + session_ids: tuple[str, ...] = (ALL_SENTINEL,) + + +def as_filter_values(selected: Sequence[str]) -> tuple[str, ...]: + """Converts a sequence of selected filter items into a hashable tuple. + + An empty selection means "everything" and maps to (ALL_SENTINEL,). + + Args: + selected: Selected filter values from the UI. + + Returns: + A tuple of filter values, or (ALL_SENTINEL,) if empty. + """ + return tuple(selected) if selected else (ALL_SENTINEL,) + + +@dataclasses.dataclass(frozen=True) +class Window: + """Half-open ``[start, end)`` query window, snapped for cacheability. + + Attributes: + start: Beginning timestamp of the window (inclusive, UTC). + end: End timestamp of the window (exclusive, UTC). + """ + + start: dt.datetime + end: dt.datetime + + @property + def span(self) -> dt.timedelta: + """Returns the total duration of the window.""" + return self.end - self.start + + @property + def bucket(self) -> str: + """Returns the ``TIMESTAMP_TRUNC`` unit that keeps a chart readable.""" + if self.span <= dt.timedelta(hours=6): + return "MINUTE" + if self.span <= dt.timedelta(days=3): + return "HOUR" + return "DAY" + + +def snap(moment: dt.datetime, seconds: int = CACHE_TTL_SECONDS) -> dt.datetime: + """Floors ``moment`` to a multiple of ``seconds`` past the hour. + + Args: + moment: The datetime object to snap. + seconds: Interval in seconds to floor by (defaults to CACHE_TTL_SECONDS). + + Returns: + The snapped UTC datetime without microseconds. + """ + floored = moment.replace(microsecond=0) + drop = (floored.minute * 60 + floored.second) % seconds + return floored - dt.timedelta(seconds=drop) + + +def make_window(span: dt.timedelta, now: dt.datetime | None = None) -> Window: + """Builds a snapped query window of the requested span ending at ``now``. + + Args: + span: Duration of the query window. + now: Optional timestamp to base the end time on. Defaults to current UTC. + + Returns: + A Window instance with snapped end and start boundaries. + """ + end = snap(now or dt.datetime.now(dt.timezone.utc)) + return Window(start=end - span, end=end) + + +def humanize_bytes(num: float) -> str: + """Formats a byte count into a human-readable string (B, KB, MB, GB, TB). + + Args: + num: Byte count to format. + + Returns: + Formatted string representation. + """ + for unit in ("B", "KB", "MB", "GB"): + if abs(num) < 1024: + return f"{num:,.1f} {unit}" if unit != "B" else f"{int(num)} B" + num /= 1024 + return f"{num:,.1f} TB" + + +class QueryResult(NamedTuple): + """A dataframe plus the outcome metadata the UI needs to explain it. + + Attributes: + df: Result DataFrame containing query rows. + error: Error message if query failed, or None. + bytes_processed: Number of bytes processed by the query job. + bytes_billed: Number of bytes billed by the query job. + cache_hit: Whether the query result was served from BigQuery cache. + """ + + df: pd.DataFrame + error: str | None = None + bytes_processed: int = 0 + bytes_billed: int = 0 + cache_hit: bool = False + + +@dataclasses.dataclass +class Context: + """Everything a panel needs to run and price a query. + + Attributes: + refs: Validated BigQuery table and view references. + window: Snapped time window for the query. + filters: Active filter parameter values. + max_bytes: Maximum billed bytes cap per query. + theme: Active color theme. + price_in: Price in USD per 1M input tokens. + price_out: Price in USD per 1M output tokens. + scan_log: Log of query executions, bytes billed/processed, and cache hit status. + """ + + refs: TableRefs + window: Window + filters: Filters + max_bytes: int + theme: Theme + price_in: float + price_out: float + scan_log: list[tuple[Any, ...]] = dataclasses.field(default_factory=list) diff --git a/dashboards/streamlit/queries.py b/dashboards/streamlit/queries.py new file mode 100644 index 00000000..6f462b40 --- /dev/null +++ b/dashboards/streamlit/queries.py @@ -0,0 +1,956 @@ +"""SQL builders and BigQuery execution helpers for the Streamlit dashboard.""" + +from __future__ import annotations + +import datetime as dt +from pathlib import Path +import re +import sys +import threading + +from google.api_core import exceptions as gexc +from google.auth import exceptions as gauth_exc +from google.cloud import bigquery +import pandas as pd +import streamlit as st + +_DASHBOARD_DIR = str(Path(__file__).resolve().parent) +if _DASHBOARD_DIR not in sys.path: + sys.path.insert(0, _DASHBOARD_DIR) + +from models import ALL_SENTINEL +from models import CACHE_TTL_SECONDS +from models import Context +from models import FILTER_OPTIONS_LIMIT +from models import Filters +from models import humanize_bytes +from models import QueryResult +from models import TableRefs +from models import TOP_ERRORS_LIMIT +from models import Window + +_PARAM_RE = re.compile(r"@([A-Za-z_][A-Za-z0-9_]*)") + +# The canonical error predicate: an event is an error if it carries an +# error event type, an error message, or an ERROR status. +_ERROR_PREDICATE = ( + "ENDS_WITH({p}event_type, '_ERROR')" + " OR {p}error_message IS NOT NULL" + " OR UPPER({p}status) = 'ERROR'" +) + + +# ------------------------------------------------------------------ # +# SQL builders (pure) # +# ------------------------------------------------------------------ # + + +def _ts_literal(moment: dt.datetime) -> str: + """Renders a UTC timestamp literal. + + Literal bounds rather than query parameters: BigQuery's partition + pruning is reliable for constant predicates, and these constants are + built from ``datetime`` objects the app owns, never from user text. + + Args: + moment: Datetime object to format in UTC. + + Returns: + A SQL literal string representation of the timestamp. + """ + utc = moment.astimezone(dt.timezone.utc) + return f'TIMESTAMP "{utc.strftime("%Y-%m-%d %H:%M:%S")}+00:00"' + + +def time_bounds(window: Window, column: str = "timestamp") -> str: + """Returns a half-open time predicate that prunes ``timestamp`` partitions. + + Args: + window: Snapped query window with start and end times. + column: SQL column or expression name to filter against. + + Returns: + A SQL WHERE clause fragment bounding the column to + [window.start, window.end). + """ + return ( + f"{column} >= {_ts_literal(window.start)}\n" + f" AND {column} < {_ts_literal(window.end)}" + ) + + +def _scope(alias: str = "", *, event_type: bool = False) -> str: + """Renders the filter clauses shared by nearly every panel. + + ``event_type`` defaults off because two exemptions in + ``dashboards/grafana/queries/README.md`` apply to most panels: a view-backed query + is already scoped to one event type, and an error count must stay + unscoped or it reports zero errors whenever some other type is picked. + + Args: + alias: Optional table alias to prefix column references with. + event_type: Whether to include the event_types filter condition. + + Returns: + A multiline SQL snippet containing AND conditions for filter parameters. + """ + p = f"{alias}." if alias else "" + clauses = [ + f" AND ('{ALL_SENTINEL}' IN UNNEST(@agents)" + f" OR {p}agent IN UNNEST(@agents))", + f" AND ('{ALL_SENTINEL}' IN UNNEST(@user_ids)" + f" OR {p}user_id IN UNNEST(@user_ids))", + ] + if event_type: + clauses.append( + f" AND ('{ALL_SENTINEL}' IN UNNEST(@event_types)" + f" OR {p}event_type IN UNNEST(@event_types))" + ) + clauses.append( + f" AND ('{ALL_SENTINEL}' IN UNNEST(@session_ids)" + f" OR {p}session_id IN UNNEST(@session_ids))" + ) + return "\n".join(clauses) + + +def build_filter_options_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for dropdown options in one pass over the raw table. + + Four separate ``SELECT DISTINCT``s would be four table references, and + BigQuery bills bytes scanned per reference. Cross-joining a literal + array of structs pivots the four columns into rows instead, so the + whole sidebar costs one scan of five columns. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + f.kind, + f.value, + MAX(e.timestamp) AS last_seen +FROM {refs.events} AS e, + UNNEST([ + STRUCT('agent' AS kind, e.agent AS value), + STRUCT('user_id', e.user_id), + STRUCT('event_type', e.event_type), + STRUCT('session_id', e.session_id) + ]) AS f +WHERE {time_bounds(window, "e.timestamp")} + AND f.value IS NOT NULL +GROUP BY f.kind, f.value +QUALIFY ROW_NUMBER() OVER ( + PARTITION BY f.kind ORDER BY MAX(e.timestamp) DESC +) <= {FILTER_OPTIONS_LIMIT} +ORDER BY f.kind, f.value +""".strip() + + +def build_overview_totals_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for session, event, error, and latency totals. + + ``HAVING COUNT(*) > 0`` keeps the no-data contract: an unaggregated + SELECT over aggregates always emits a row, so an empty filter + intersection would otherwise report a confident "0 events, 0% errors". + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + err = _ERROR_PREDICATE.format(p="e.") + return f""" +SELECT + COUNT(DISTINCT e.session_id) AS sessions, + COUNT(*) AS events, + SAFE_DIVIDE(COUNTIF({err}), COUNT(*)) AS error_rate, + ( + SELECT AVG(r.total_ms) + FROM {refs.view("llm_responses")} AS r + WHERE {time_bounds(window, "r.timestamp")} +{_scope("r")} + ) AS avg_llm_latency_ms +FROM {refs.events} AS e +WHERE {time_bounds(window, "e.timestamp")} +{_scope("e")} +HAVING COUNT(*) > 0 +""".strip() + + +def build_events_over_time_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for event volume per bucket and event_type. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + TIMESTAMP_TRUNC(timestamp, {window.bucket}) AS bucket, + event_type, + COUNT(*) AS events +FROM {refs.events} +WHERE {time_bounds(window)} +{_scope(event_type=True)} +GROUP BY bucket, event_type +ORDER BY bucket +""".strip() + + +def build_errors_over_time_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for error volume per bucket. + + No event_type filter: errors arrive as their own event types, so + honoring a selection like LLM_RESPONSE would chart zero errors in a + window that had them. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + TIMESTAMP_TRUNC(timestamp, {window.bucket}) AS bucket, + event_type, + COUNT(*) AS errors +FROM {refs.events} +WHERE {time_bounds(window)} +{_scope()} + AND ({_ERROR_PREDICATE.format(p="")}) +GROUP BY bucket, event_type +ORDER BY bucket +""".strip() + + +def build_events_by_agent_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for event volume per agent. + + ``agent`` is nullable on the raw table; IFNULL groups those rows under + "unknown" instead of dropping them, so the bars still add up to the + Events stat whenever the Event type filter is off. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + IFNULL(agent, 'unknown') AS agent_name, + COUNT(*) AS events +FROM {refs.events} +WHERE {time_bounds(window)} +{_scope(event_type=True)} +GROUP BY agent_name +HAVING COUNT(*) > 0 +ORDER BY events DESC +""".strip() + + +def build_top_errors_sql( + refs: TableRefs, window: Window, limit: int = TOP_ERRORS_LIMIT +) -> str: + """Builds the SQL query for the loudest failure modes, ranked by occurrence. + + Deliberately narrower than Errors over time: this groups by the message + text, so an error that recorded no message has no string to group under + and is absent. These counts are a subset of the chart's. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + limit: Maximum number of distinct error messages to return. + + Returns: + The SQL query string. + """ + return f""" +SELECT + error_message, + COUNT(*) AS errors, + COUNT(DISTINCT session_id) AS sessions, + COUNT(DISTINCT agent) AS agents, + MAX(timestamp) AS last_seen +FROM {refs.events} +WHERE {time_bounds(window)} +{_scope()} + AND error_message IS NOT NULL +GROUP BY error_message +HAVING COUNT(*) > 0 +ORDER BY errors DESC +LIMIT {int(limit)} +""".strip() + + +def build_llm_tokens_over_time_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for prompt and completion tokens per bucket. + + Missing usage telemetry degrades to 0 rather than dropping the bucket. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + TIMESTAMP_TRUNC(timestamp, {window.bucket}) AS bucket, + IFNULL(SUM(usage_prompt_tokens), 0) AS prompt_tokens, + IFNULL(SUM(usage_completion_tokens), 0) AS completion_tokens, + IFNULL(SUM(usage_total_tokens), 0) AS total_tokens +FROM {refs.view("llm_responses")} +WHERE {time_bounds(window)} +{_scope()} +GROUP BY bucket +ORDER BY bucket +""".strip() + + +def build_llm_latency_percentiles_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for LLM p50/p95 latency and p50 time-to-first-token. + + Latency aggregates stay NULL when telemetry is missing so the chart + shows gaps rather than fake zeros. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + TIMESTAMP_TRUNC(timestamp, {window.bucket}) AS bucket, + APPROX_QUANTILES(total_ms, 100)[OFFSET(50)] AS p50_total_ms, + APPROX_QUANTILES(total_ms, 100)[OFFSET(95)] AS p95_total_ms, + APPROX_QUANTILES(ttft_ms, 100)[OFFSET(50)] AS p50_ttft_ms +FROM {refs.view("llm_responses")} +WHERE {time_bounds(window)} +{_scope()} +GROUP BY bucket +ORDER BY bucket +""".strip() + + +def build_tokens_by_model_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for token totals per model. + + Rows without a model group under "unknown". + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + IFNULL(model_version, 'unknown') AS model, + IFNULL(SUM(usage_prompt_tokens), 0) AS prompt_tokens, + IFNULL(SUM(usage_completion_tokens), 0) AS completion_tokens, + IFNULL(SUM(usage_total_tokens), 0) AS total_tokens, + COUNT(*) AS responses +FROM {refs.view("llm_responses")} +WHERE {time_bounds(window)} +{_scope()} +GROUP BY model +ORDER BY total_tokens DESC +""".strip() + + +def build_llm_calls_total_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for LLM call count and token totals in one scan. + + Calls are counted per span, not per row: a streaming installation + records one LLM_RESPONSE per chunk, so COUNT(*) would report chunks as + calls. ``trace_id``/``span_id`` are nullable and CONCAT of a NULL is + NULL, so rows carrying no span key are added back one per row. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + COUNT(DISTINCT CONCAT(trace_id, '|', span_id)) + + COUNTIF(trace_id IS NULL OR span_id IS NULL) AS llm_calls, + IFNULL(SUM(usage_prompt_tokens), 0) AS prompt_tokens, + IFNULL(SUM(usage_completion_tokens), 0) AS completion_tokens, + IFNULL(SUM(usage_total_tokens), 0) AS total_tokens +FROM {refs.view("llm_responses")} +WHERE {time_bounds(window)} +{_scope()} +HAVING COUNT(*) > 0 +""".strip() + + +def build_tool_usage_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for invocations per tool. + + Reads tool_starts rather than tool_completions so failed invocations + still count toward the volume. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + IFNULL(tool_name, 'unknown') AS tool_name, + COUNT(*) AS invocations +FROM {refs.view("tool_starts")} +WHERE {time_bounds(window)} +{_scope()} +GROUP BY tool_name +ORDER BY invocations DESC +""".strip() + + +def build_tool_latency_sql(refs: TableRefs, window: Window) -> str: + """Builds the SQL query for per-tool latency. + + Latency aggregates stay NULL when telemetry is missing, so gaps do not skew + the averages. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + + Returns: + The SQL query string. + """ + return f""" +SELECT + IFNULL(tool_name, 'unknown') AS tool_name, + COUNT(*) AS completions, + AVG(total_ms) AS avg_ms, + APPROX_QUANTILES(total_ms, 100)[OFFSET(50)] AS p50_ms, + APPROX_QUANTILES(total_ms, 100)[OFFSET(95)] AS p95_ms +FROM {refs.view("tool_completions")} +WHERE {time_bounds(window)} +{_scope()} +GROUP BY tool_name +ORDER BY p95_ms DESC +""".strip() + + +def build_tool_errors_sql(refs: TableRefs, window: Window, limit: int) -> str: + """Builds the SQL query for tool failures from error views and completions. + + UNION ALL rather than DISTINCT: when one logical failure emits both a + TOOL_ERROR and an error-status TOOL_COMPLETED, both telemetry records + are real and both are kept. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + limit: Maximum number of tool error records to return. + + Returns: + The SQL query string. + """ + return f""" +SELECT + timestamp, + agent, + session_id, + tool_name, + error_message +FROM {refs.view("tool_errors")} +WHERE {time_bounds(window)} +{_scope()} + +UNION ALL + +SELECT + timestamp, + agent, + session_id, + tool_name, + error_message +FROM {refs.view("tool_completions")} +WHERE {time_bounds(window)} +{_scope()} + AND (error_message IS NOT NULL OR UPPER(status) = 'ERROR') +ORDER BY timestamp DESC +LIMIT {int(limit)} +""".strip() + + +def build_recent_sessions_sql( + refs: TableRefs, window: Window, limit: int +) -> str: + """Builds the SQL query for session-level rollups over the raw table. + + All four filters apply, but Agent / User / Event type only decide + *which* sessions are listed: each gets its own LOGICAL_OR in the + HAVING, so a session is kept when it contains a match for each filter + somewhere in the window. Applying them in the WHERE would drop events + before the GROUP BY, and every rollup column would then describe a + filtered slice while still being labelled as the session. + + The consequence is that these columns cover every event the session has + in the window, including events the filters excluded, so they will not + sum to the Overview stats. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + limit: Maximum number of recent sessions to return. + + Returns: + The SQL query string. + """ + err = _ERROR_PREDICATE.format(p="") + return f""" +SELECT + session_id, + STRING_AGG(DISTINCT user_id, ', ' ORDER BY user_id) + AS session_users_in_window, + MIN(timestamp) AS started_in_window_at, + MAX(timestamp) AS last_event_in_window_at, + TIMESTAMP_DIFF(MAX(timestamp), MIN(timestamp), SECOND) + AS duration_in_window_s, + COUNT(DISTINCT agent) AS session_agents_in_window, + COUNT(*) AS session_events_in_window, + COUNTIF({err}) AS session_errors_in_window, + IFNULL(SUM(IF(event_type = 'LLM_RESPONSE', + SAFE_CAST(JSON_VALUE(content, '$.usage.prompt') AS INT64), NULL)), 0) + AS session_input_tokens_in_window, + IFNULL(SUM(IF(event_type = 'LLM_RESPONSE', + SAFE_CAST(JSON_VALUE(content, '$.usage.completion') AS INT64), NULL)), 0) + AS session_output_tokens_in_window +FROM {refs.events} +WHERE {time_bounds(window)} + AND session_id IS NOT NULL + AND ('{ALL_SENTINEL}' IN UNNEST(@session_ids) + OR session_id IN UNNEST(@session_ids)) +GROUP BY session_id +HAVING LOGICAL_OR('{ALL_SENTINEL}' IN UNNEST(@agents) + OR agent IN UNNEST(@agents)) + AND LOGICAL_OR('{ALL_SENTINEL}' IN UNNEST(@user_ids) + OR user_id IN UNNEST(@user_ids)) + AND LOGICAL_OR('{ALL_SENTINEL}' IN UNNEST(@event_types) + OR event_type IN UNNEST(@event_types)) +ORDER BY last_event_in_window_at DESC +LIMIT {int(limit)} +""".strip() + + +def build_trace_detail_sql(refs: TableRefs, window: Window, limit: int) -> str: + """Builds the SQL query for one session's event timeline, newest first. + + DESC decides *which* rows survive the LIMIT, and that is the part a + reader cannot recover from: ascending order would return the oldest + events in the range, so a busy window would show its first few minutes + and nothing since. Sort the rendered table by timestamp to read it + chronologically. + + COALESCE(model, model_version): ``model`` is on LLM_REQUEST attributes, + ``model_version`` on LLM_RESPONSE — one column covers both. + + Args: + refs: Validated BigQuery table references. + window: Snapped time window. + limit: Maximum number of trace events to return. + + Returns: + The SQL query string. + """ + return f""" +SELECT + timestamp, + session_id, + event_type, + agent, + invocation_id, + span_id, + parent_span_id, + status, + COALESCE( + JSON_VALUE(attributes, '$.model'), + JSON_VALUE(attributes, '$.model_version') + ) AS model, + JSON_VALUE(content, '$.tool') AS tool_name, + SAFE_CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms, + error_message +FROM {refs.events} +WHERE {time_bounds(window)} +{_scope(event_type=True)} +ORDER BY timestamp DESC +LIMIT {int(limit)} +""".strip() + + +# ------------------------------------------------------------------ # +# BigQuery execution # +# ------------------------------------------------------------------ # + + +def query_parameters( + sql: str, + filters: Filters, +) -> list[bigquery.ArrayQueryParameter]: + """Binds only the filter parameters the SQL actually references. + + Scanning the text keeps the builders free of bookkeeping about which + filters they honor, and keeps BigQuery from receiving parameters that + the query never mentions. + + Args: + sql: The SQL query string to inspect. + filters: Active filter tuple selections. + + Returns: + A list of BigQuery query parameters referenced in the SQL text. + """ + used = set(_PARAM_RE.findall(sql)) + return [ + bigquery.ArrayQueryParameter(name, "STRING", list(values)) + for name, values in filters._asdict().items() + if name in used + ] + + +def job_config( + sql: str, + filters: Filters, + max_bytes: int, + *, + dry_run: bool = False, +) -> bigquery.QueryJobConfig: + """Returns a job config carrying the per-query scan guardrail. + + Args: + sql: SQL query string. + filters: Active filter parameters. + max_bytes: Maximum allowed bytes billed. + dry_run: Whether to configure the job as a dry run. + + Returns: + A configured bigquery.QueryJobConfig instance. + """ + config = bigquery.QueryJobConfig( + query_parameters=query_parameters(sql, filters), + use_legacy_sql=False, + use_query_cache=True, + dry_run=dry_run, + labels={"app": "bqaa_streamlit"}, + ) + # A dry run bills nothing, so the cap only belongs on the real job — + # where it is the guardrail of record, not merely advice the preflight + # gives. Both are set: the preflight explains, the cap enforces. + if not dry_run: + config.maximum_bytes_billed = int(max_bytes) + return config + + +@st.cache_resource(show_spinner=False) +def get_client(project: str) -> bigquery.Client: + """Returns a BigQuery client per project, reused across reruns. + + Args: + project: BigQuery project ID. + + Returns: + A cached bigquery.Client instance. + """ + return bigquery.Client(project=project) + + +def _explain(exc: Exception) -> str: + """Turns the BigQuery errors operators actually hit into actionable advice. + + Args: + exc: The exception caught during query execution. + + Returns: + A user-friendly explanation string with remediation advice. + """ + message = getattr(exc, "message", None) or str(exc) + if isinstance(exc, gauth_exc.DefaultCredentialsError): + return ( + f"{message}\n\nRun `gcloud auth application-default login` or set" + " `GOOGLE_APPLICATION_CREDENTIALS`." + ) + if isinstance(exc, gexc.NotFound): + return ( + f"{message}\n\nIf this names a prefixed view, create the typed" + " views first — `bq-agent-sdk views create-all --project-id P" + " --dataset-id D --table-id T` — or correct the view prefix in" + " the sidebar." + ) + if isinstance(exc, gexc.Forbidden): + return ( + f"{message}\n\nThe caller needs `roles/bigquery.jobUser` on the" + " project and `roles/bigquery.dataViewer` on the dataset." + ) + if "bytesBilledLimitExceeded" in message: + return ( + f"{message}\n\nRaise the per-query scan cap in the sidebar or" + " narrow the time range." + ) + return message + + +_MAX_SEEN_RUN_IDS = 10_000 +_NEXT_RUN_ID = 0 +_SEEN_RUN_IDS: set[int] = set() +_RUN_ID_LOCK = threading.Lock() +_JOB_RELOAD_TIMEOUT_SECONDS = 10 + + +class QueryExecutionError(RuntimeError): + """Raised when query execution or result download fails. + + Preserves byte scan and billing statistics if available from the BigQuery job. + """ + + def __init__( + self, + message: str, + bytes_processed: int = 0, + bytes_billed: int = 0, + cache_hit: bool = False, + ) -> None: + super().__init__(message) + self.bytes_processed = bytes_processed + self.bytes_billed = bytes_billed + self.cache_hit = cache_hit + + +def _extract_job_stats( + job: bigquery.QueryJob | None, +) -> tuple[int, int, bool] | None: + """Extracts byte scan and billing statistics from a QueryJob if available.""" + if job is None: + return None + try: + total_processed = getattr(job, "total_bytes_processed", None) + ended = getattr(job, "ended", None) + if total_processed is None: + try: + job.reload(timeout=_JOB_RELOAD_TIMEOUT_SECONDS) + except Exception: + pass + else: + total_processed = getattr(job, "total_bytes_processed", None) + ended = getattr(job, "ended", None) + if total_processed is not None or ended is not None: + cache_hit = bool(getattr(job, "cache_hit", False)) + bytes_billed = ( + 0 if cache_hit else int(getattr(job, "total_bytes_billed", 0) or 0) + ) + bytes_processed = int(total_processed or 0) + return bytes_processed, bytes_billed, cache_hit + except Exception: + return None + return None + + +def _query_failure(msg: str, job: bigquery.QueryJob | None) -> RuntimeError: + """Builds a QueryExecutionError or RuntimeError preserving job statistics.""" + stats = _extract_job_stats(job) + if stats is not None: + bytes_processed, bytes_billed, cache_hit = stats + return QueryExecutionError( + msg, + bytes_processed=bytes_processed, + bytes_billed=bytes_billed, + cache_hit=cache_hit, + ) + return RuntimeError(msg) + + +@st.cache_data(ttl=CACHE_TTL_SECONDS, show_spinner=False, max_entries=256) +def _run_query_cached( + sql: str, + filters: Filters, + project: str, + max_bytes: int, +) -> tuple[pd.DataFrame, int, int, bool, int]: + """Runs the BigQuery job, cached by Streamlit across reruns. + + Args: + sql: The SQL query to execute. + filters: Active filter parameters. + project: BigQuery project ID. + max_bytes: Maximum allowed bytes billed. + + Returns: + Tuple of (dataframe, bytes_processed, bytes_billed, cache_hit, run_id). + """ + global _NEXT_RUN_ID + with _RUN_ID_LOCK: + _NEXT_RUN_ID += 1 + run_id = _NEXT_RUN_ID + + try: + client = get_client(project) + probe = client.query( + sql, + job_config=job_config( + sql, + filters, + max_bytes, + dry_run=True, + ), + ) + estimate = int(probe.total_bytes_processed or 0) + except (gexc.GoogleAPICallError, gauth_exc.DefaultCredentialsError) as exc: + raise RuntimeError(_explain(exc)) from exc + + if estimate > max_bytes: + raise RuntimeError( + f"Guardrail: this query would scan {humanize_bytes(estimate)}," + f" above the {humanize_bytes(max_bytes)} per-query cap. Narrow" + " the time range or raise the cap in the sidebar." + ) + + job = None + try: + job = client.query( + sql, + job_config=job_config(sql, filters, max_bytes), + ) + df = job.to_dataframe() + except (gexc.GoogleAPICallError, gauth_exc.DefaultCredentialsError) as exc: + raise _query_failure(_explain(exc), job) from exc + except Exception as exc: + raise _query_failure(str(exc), job) from exc + + bytes_billed = 0 if job.cache_hit else int(job.total_bytes_billed or 0) + bytes_processed = int(job.total_bytes_processed or 0) + return ( + df, + bytes_processed, + bytes_billed, + bool(job.cache_hit), + run_id, + ) + + +def run_query( + sql: str, + filters: Filters, + project: str, + max_bytes: int, +) -> QueryResult: + """Runs one query behind a dry-run preflight and the byte cap. + + The dry run is free and tells us the scan size before anything is + billed, so a query that would blow the cap is reported as a readable + guardrail message instead of an opaque ``bytesBilledLimitExceeded``. + + Args: + sql: The SQL query to execute. + filters: Active filter parameters. + project: BigQuery project ID. + max_bytes: Maximum allowed bytes billed. + + Returns: + QueryResult containing the resulting dataframe and execution metadata. + """ + try: + df, bytes_processed, bytes_billed, cache_hit, run_id = _run_query_cached( + sql, filters, project, max_bytes + ) + except Exception as exc: + if isinstance(exc, QueryExecutionError): + bytes_processed = int(exc.bytes_processed) + bytes_billed = int(exc.bytes_billed) + cache_hit = bool(exc.cache_hit) + else: + bytes_processed, bytes_billed, cache_hit = 0, 0, False + return QueryResult( + df=pd.DataFrame(), + error=str(exc), + bytes_processed=bytes_processed, + bytes_billed=bytes_billed, + cache_hit=cache_hit, + ) + + with _RUN_ID_LOCK: + if run_id in _SEEN_RUN_IDS: + return QueryResult(df, None, 0, 0, True) + if len(_SEEN_RUN_IDS) >= _MAX_SEEN_RUN_IDS: + evict_count = max(1, _MAX_SEEN_RUN_IDS // 4) + for stale in sorted(_SEEN_RUN_IDS)[:evict_count]: + _SEEN_RUN_IDS.discard(stale) + _SEEN_RUN_IDS.add(run_id) + + return QueryResult(df, None, bytes_processed, bytes_billed, cache_hit) + + +def fetch(sql: str, ctx: Context, label: str) -> QueryResult: + """Runs a query and records its cost in this rerun's scan log. + + Args: + sql: The SQL query to execute. + ctx: Active execution context. + label: Human-readable label for scan logging and error reporting. + + Returns: + QueryResult of the execution. + """ + with st.spinner(f"Loading {label}..."): + result = run_query( + sql, + ctx.filters, + ctx.refs.project, + ctx.max_bytes, + ) + ctx.scan_log.append( + (label, result.bytes_billed, result.bytes_processed, result.cache_hit) + ) + if result.error: + st.error(f"**{label}** — {result.error}") + return result + + +def load_filter_options( + ctx: Context, +) -> tuple[dict[str, list[str]], QueryResult]: + """Fetches every dropdown's options in one query. + + Args: + ctx: Active execution context. + + Returns: + A tuple of: + - Dictionary mapping filter kind ('agent', 'user_id', 'event_type', + 'session_id') to a list of available string values. + - QueryResult containing the execution outcome and any error. + """ + sql = build_filter_options_sql(ctx.refs, ctx.window) + result = fetch(sql, ctx, "Filter options") + options: dict[str, list[str]] = {} + if result.df.empty: + return options, result + for kind, group in result.df.groupby("kind"): + options[str(kind)] = [str(v) for v in group["value"].tolist()] + return options, result diff --git a/pyproject.toml b/pyproject.toml index 004e35d5..4f27c2fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,8 +59,18 @@ improvement = [ "pandas>=2.0.0", "python-dotenv>=1.0.0", ] +streamlit = [ + "streamlit>=1.49.0", + "plotly>=6.0.0", + "db-dtypes>=1.3.0", + "pandas>=2.0.0", + "python-dotenv>=1.0.0", +] +dashboards = [ + "bigquery-agent-analytics[streamlit]", +] all = [ - "bigquery-agent-analytics[llm,bigframes,langsmith,owl,cli,dev,improvement]", + "bigquery-agent-analytics[llm,bigframes,langsmith,owl,cli,dev,improvement,dashboards]", ] [project.scripts] diff --git a/scripts/check_streamlit_queries_sync.py b/scripts/check_streamlit_queries_sync.py new file mode 100644 index 00000000..8c191c4f --- /dev/null +++ b/scripts/check_streamlit_queries_sync.py @@ -0,0 +1,636 @@ +#!/usr/bin/env python3 +"""Check that canonical Grafana SQL matches the Streamlit dashboard's queries.""" + +import difflib +import inspect +from pathlib import Path +import re +import sys +from unittest.mock import MagicMock + +for mod in [ + "streamlit", + "pandas", + "google.cloud", + "google.cloud.bigquery", + "google.api_core", + "google.api_core.exceptions", + "google.auth", + "google.auth.exceptions", + "plotly", +]: + if mod not in sys.modules: + try: + __import__(mod) + except ImportError: + sys.modules[mod] = MagicMock() + +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +STREAMLIT_DIR = REPOSITORY_ROOT / "dashboards" / "streamlit" +if str(STREAMLIT_DIR) not in sys.path: + sys.path.insert(0, str(STREAMLIT_DIR)) + +import datetime as dt + +import models +import queries + +QUERIES_DIRECTORY = ( + REPOSITORY_ROOT / "dashboards" / "grafana" / "queries" + if (REPOSITORY_ROOT / "dashboards" / "grafana" / "queries").exists() + else REPOSITORY_ROOT / "grafana" / "queries" +) + +CANONICAL_QUERIES = { + "overview_totals.sql": queries.build_overview_totals_sql, + "events_over_time.sql": queries.build_events_over_time_sql, + "errors_over_time.sql": queries.build_errors_over_time_sql, + "events_by_agent.sql": queries.build_events_by_agent_sql, + "top_errors.sql": queries.build_top_errors_sql, + "llm_tokens_over_time.sql": queries.build_llm_tokens_over_time_sql, + "llm_latency_percentiles.sql": queries.build_llm_latency_percentiles_sql, + "tokens_by_model.sql": queries.build_tokens_by_model_sql, + "llm_calls_total.sql": queries.build_llm_calls_total_sql, + "tool_usage.sql": queries.build_tool_usage_sql, + "tool_latency.sql": queries.build_tool_latency_sql, + "tool_errors.sql": queries.build_tool_errors_sql, + "recent_sessions.sql": queries.build_recent_sessions_sql, + "trace_detail.sql": queries.build_trace_detail_sql, +} + +QUERY_LIMITS = { + "top_errors.sql": models.TOP_ERRORS_LIMIT, + "tool_errors.sql": models.TOOL_ERRORS_LIMIT, + "recent_sessions.sql": models.RECENT_SESSIONS_LIMIT, + "trace_detail.sql": models.TRACE_DETAIL_LIMIT, +} + +LLM_CALLS_TOKEN_COLUMNS = { + "usage_prompt_tokens": ( + "prompt_tokens", + r"IFNULL\s*\(\s*SUM\s*\(\s*usage_prompt_tokens\s*\)\s*,\s*0\s*\)\s+AS\s+prompt_tokens\b", + ), + "usage_completion_tokens": ( + "completion_tokens", + r"IFNULL\s*\(\s*SUM\s*\(\s*usage_completion_tokens\s*\)\s*,\s*0\s*\)\s+AS\s+completion_tokens\b", + ), + "usage_total_tokens": ( + "total_tokens", + r"IFNULL\s*\(\s*SUM\s*\(\s*usage_total_tokens\s*\)\s*,\s*0\s*\)\s+AS\s+total_tokens\b", + ), +} + +EXEMPT_CANONICAL_QUERIES = { + "estimated_cost.sql": "Priced in Python UI", + "var_agent.sql": "Filter options handled by build_filter_options_sql", + "var_user_id.sql": "Filter options", + "var_event_type.sql": "Filter options", + "var_session_id.sql": "Filter options", +} + +EXEMPT_STREAMLIT_BUILDERS = { + "build_filter_options_sql": "Aggregated filter options dropdown", +} + +DEFAULT_NOW = dt.datetime(2025, 1, 1, 12, 0, 0, tzinfo=dt.timezone.utc) +DEFAULT_WINDOW = models.Window( + start=DEFAULT_NOW - dt.timedelta(hours=24), end=DEFAULT_NOW +) + + +def check_unmapped_canonical_queries( + directory: Path, mapped_files: set[str] +) -> int: + """Fail if the queries directory holds .sql files not mapped or exempted.""" + if not directory.exists() or not directory.is_dir(): + print( + f"ERROR: canonical queries directory not found: {directory}", + file=sys.stderr, + ) + return 1 + + try: + sql_files = {path.name for path in directory.glob("*.sql")} + except OSError as error: + print( + f"ERROR: cannot list canonical queries in {directory}: {error}", + file=sys.stderr, + ) + return 1 + + covered = set(mapped_files) | set(EXEMPT_CANONICAL_QUERIES) + unmapped = sorted(sql_files - covered) + if not unmapped: + return 0 + + try: + directory_label = directory.relative_to(REPOSITORY_ROOT).as_posix() + except ValueError: + directory_label = str(directory) + + border = " " + "*" * 64 + title = ("* ERROR: unmapped SQL files in " + directory_label + "/").ljust(63) + body = ( + " The .sql file(s) below are missing from CANONICAL_QUERIES and\n" + " EXEMPT_CANONICAL_QUERIES, so they are never validated against the\n" + " dashboard:\n" + "".join(f" - {name}\n" for name in unmapped) + ) + print( + f"\n{border}\n {title}*\n{border}\n{body}{border}", + file=sys.stderr, + ) + return len(unmapped) + + +def check_unmapped_streamlit_builders() -> int: + """Fail if queries.py defines build_*_sql functions not mapped or exempted.""" + all_builders = { + name + for name, _ in inspect.getmembers(queries, inspect.isfunction) + if name.startswith("build_") and name.endswith("_sql") + } + mapped_builder_names = { + builder.__name__ + for builder in CANONICAL_QUERIES.values() + if hasattr(builder, "__name__") + } + covered = mapped_builder_names | set(EXEMPT_STREAMLIT_BUILDERS) + unmapped = sorted(all_builders - covered) + if not unmapped: + return 0 + + border = " " + "*" * 64 + title = "* ERROR: unmapped Streamlit SQL builders in queries.py".ljust(63) + body = ( + " The builder function(s) below are missing from CANONICAL_QUERIES\n" + " and EXEMPT_STREAMLIT_BUILDERS, so they are never validated against\n" + " canonical Grafana SQL:\n" + + "".join(f" - {name}\n" for name in unmapped) + ) + print( + f"\n{border}\n {title}*\n{border}\n{body}{border}", + file=sys.stderr, + ) + return len(unmapped) + + +def assert_timestamp_bounds(sql: str, window: models.Window) -> None: + """Assert that timestamp bounds strictly match window.start and window.end with proper operators.""" + matches = re.findall( + r'([a-zA-Z0-9_.]+)\s*(>=|<=|>|<|=)\s*TIMESTAMP\s*(?:"([^"]+)"|\'([^\']+)\')', + sql, + ) + if not matches: + return + + expected_start = window.start.astimezone(dt.timezone.utc).strftime( + "%Y-%m-%d %H:%M:%S+00:00" + ) + expected_end = window.end.astimezone(dt.timezone.utc).strftime( + "%Y-%m-%d %H:%M:%S+00:00" + ) + + start_cols = [] + end_cols = [] + + for col, op, ts_double, ts_single in matches: + ts_literal = ts_double or ts_single + if op == ">=": + assert ( + ts_literal == expected_start + ), f"Expected start timestamp literal '{expected_start}', got '{ts_literal}' for {col}" + start_cols.append(col) + elif op == "<": + assert ( + ts_literal == expected_end + ), f"Expected end timestamp literal '{expected_end}', got '{ts_literal}' for {col}" + end_cols.append(col) + else: + raise AssertionError( + f"Disallowed operator '{op}' with TIMESTAMP '{ts_literal}' on {col}" + ) + + assert len(start_cols) > 0, "No start timestamp bound (>=) found" + assert len(end_cols) > 0, "No end timestamp bound (<) found" + assert sorted(start_cols) == sorted( + end_cols + ), f"Mismatched start bounds {start_cols} and end bounds {end_cols}" + + +def assert_streamlit_query( + filename: str, sql: str, window: models.Window +) -> None: + """Validate that Streamlit query has required bounds and adaptations.""" + matches = re.findall( + r'([a-zA-Z0-9_.]+)\s*(>=|<=|>|<|=)\s*TIMESTAMP\s*(?:"([^"]+)"|\'([^\']+)\')', + sql, + ) + assert ( + len(matches) >= 2 + ), f"Streamlit query {filename} is missing timestamp bounds" + assert_timestamp_bounds(sql, window) + + if filename == "llm_calls_total.sql": + for col, (alias, pattern) in LLM_CALLS_TOKEN_COLUMNS.items(): + assert re.search(pattern, sql), ( + f"Streamlit query {filename} missing required token column {col} as" + f" {alias}" + ) + + if filename in ("tool_usage.sql", "tool_latency.sql"): + assert re.search( + r"IFNULL\s*\(\s*tool_name\s*,\s*['\"]unknown['\"]\s*\)", + sql, + ), ( + f"Streamlit query {filename} expected to contain" + " IFNULL(tool_name, 'unknown')" + ) + + +def strip_comments(sql: str) -> str: + sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL) + lines = [] + for line in sql.splitlines(): + line = re.split(r"--|#", line, maxsplit=1)[0] + lines.append(line) + return "\n".join(lines) + + +def normalize_whitespace_and_parens(sql: str) -> str: + sql = re.sub(r"\s+", " ", sql) + sql = re.sub(r"\(\s+", "(", sql) + sql = re.sub(r"\s+\)", ")", sql) + return sql.strip() + + +def normalize_time_filters(sql: str) -> str: + sql = re.sub( + r'\bAND\s+([a-zA-Z0-9_.]+)\s*<\s*TIMESTAMP\s*(?:"([^"]+)"|\'([^\']+)\')', + "", + sql, + ) + sql = re.sub( + r'([a-zA-Z0-9_.]+)\s*<\s*TIMESTAMP\s*(?:"([^"]+)"|\'([^\']+)\')\s+AND\b', + "", + sql, + ) + sql = re.sub( + r'([a-zA-Z0-9_.]+)\s*>=\s*TIMESTAMP\s*(?:"([^"]+)"|\'([^\']+)\')', + r"$__timeFilter(\1)", + sql, + ) + sql = re.sub(r"\$__timeFilter\(\s*([^\s)]+)\s*\)", r"$__timeFilter(\1)", sql) + return sql + + +def normalize_time_groups(sql: str) -> str: + sql = re.sub( + r"\$__timeGroup\(([^,]+),\s*\$__interval\)", + r"TIMESTAMP_TRUNC(\1, __INTERVAL__)", + sql, + ) + sql = re.sub( + r"TIMESTAMP_TRUNC\(([^,]+),\s*(HOUR|MINUTE|DAY|\$__interval|__INTERVAL__)\)", + r"TIMESTAMP_TRUNC(\1, __INTERVAL__)", + sql, + ) + sql = re.sub( + r"(TIMESTAMP_TRUNC\([^,]+,\s*__INTERVAL__\))\s+AS\s+time\b", + r"\1 AS bucket", + sql, + ) + sql = re.sub(r"\bGROUP\s+BY\s+time\b", "GROUP BY bucket", sql) + sql = re.sub(r"\bORDER\s+BY\s+time\b", "ORDER BY bucket", sql) + return sql + + +def normalize_params_and_refs(sql: str) -> str: + sql = sql.replace("ARRAY[${agent:sqlstring}]", "@agents") + sql = sql.replace("ARRAY[${user_id:sqlstring}]", "@user_ids") + sql = sql.replace("ARRAY[${event_type:sqlstring}]", "@event_types") + sql = sql.replace("ARRAY[${session_id:sqlstring}]", "@session_ids") + sql = sql.replace("${agent:sqlstring}", "@agents") + sql = sql.replace("${user_id:sqlstring}", "@user_ids") + sql = sql.replace("${event_type:sqlstring}", "@event_types") + sql = sql.replace("${session_id:sqlstring}", "@session_ids") + + sql = sql.replace( + "`${project}.${dataset}.${table}`", "`project.dataset.events`" + ) + sql = sql.replace( + "`${project}.${dataset}.${view_prefix}", "`project.dataset.adk_" + ) + + return sql + + +def framework_adaptations(sql: str, filename: str) -> str: + if filename in ("tool_usage.sql", "tool_latency.sql"): + sql = re.sub( + r"IFNULL\s*\(\s*tool_name\s*,\s*['\"]unknown['\"]\s*\)\s+AS\s+tool_name", + "tool_name", + sql, + ) + sql = re.sub( + r"IFNULL\s*\(\s*tool_name\s*,\s*['\"]unknown['\"]\s*\)", + "tool_name", + sql, + ) + + if filename == "llm_calls_total.sql": + for _, (_, pattern) in LLM_CALLS_TOKEN_COLUMNS.items(): + sql = re.sub(rf"\s*,\s*{pattern}", "", sql) + + return sql + + +def consume_sql_string_literal(text: str, pos: int) -> tuple[str, int]: + """Consumes a SQL single-quoted string literal starting at pos.""" + quote_char = text[pos] + result = [quote_char] + i = pos + 1 + n = len(text) + while i < n: + ch = text[i] + if ch == "\\": + result.append(text[i : i + 2]) + i += 2 + elif ch == quote_char: + if i + 1 < n and text[i + 1] == quote_char: + result.append(text[i : i + 2]) + i += 2 + else: + result.append(ch) + i += 1 + break + else: + result.append(ch) + i += 1 + return "".join(result), i + + +def split_top_level(text: str, delimiter: str) -> list[str]: + parts = [] + depth = 0 + current = "" + i = 0 + pattern = re.compile(delimiter, re.IGNORECASE) + while i < len(text): + if text[i] == "'": + chunk, i = consume_sql_string_literal(text, i) + current += chunk + continue + + if text[i] == "(": + depth += 1 + elif text[i] == ")": + depth = max(0, depth - 1) + + if depth == 0: + m = pattern.match(text[i:]) + if m: + parts.append(current.strip()) + current = "" + i += m.end() + continue + + current += text[i] + i += 1 + + if current.strip(): + parts.append(current.strip()) + return parts + + +def split_conditions_and_sort( + clause_text: str, split_delim: str, join_delim: str +) -> str: + parts = split_top_level(clause_text, split_delim) + return join_delim.join(sorted(parts)) + + +def split_expressions_and_trim(clause_text: str) -> str: + parts = split_top_level(clause_text, r"\s*,\s*") + # DO NOT SORT SELECT projections, just split and trim + return ",\n ".join(parts) + + +def custom_clause_split(sql: str) -> str: + parts = [] + depth = 0 + current = "" + i = 0 + clauses = [ + "SELECT", + "FROM", + "WHERE", + "GROUP BY", + "HAVING", + "ORDER BY", + "LIMIT", + "UNION ALL", + ] + + def get_clause_match(s: str, pos: int) -> str | None: + for c in clauses: + if re.match(r"(?i)\b" + c + r"\b", s[pos:]): + return c + return None + + current_clause = "" + res_clauses = [] + + while i < len(sql): + if sql[i] == "'": + chunk, i = consume_sql_string_literal(sql, i) + current += chunk + continue + + if sql[i] == "(": + depth += 1 + elif sql[i] == ")": + depth = max(0, depth - 1) + + if depth == 0: + c = get_clause_match(sql, i) + if c: + if current_clause or current.strip(): + res_clauses.append((current_clause, current.strip())) + current_clause = c.upper() + current = "" + i += len(c) + continue + + current += sql[i] + i += 1 + + if current_clause or current.strip(): + res_clauses.append((current_clause, current.strip())) + + formatted_res = [] + for clause_name, text in res_clauses: + if not text: + if clause_name: + formatted_res.append(clause_name) + continue + + if clause_name in ["WHERE", "HAVING"]: + text = split_conditions_and_sort(text, r"\s+AND\s+", "\n AND ") + formatted_res.append(f"{clause_name}\n {text}") + elif clause_name == "SELECT": + text = split_expressions_and_trim(text) + formatted_res.append(f"{clause_name}\n {text}") + elif clause_name: + formatted_res.append(f"{clause_name} {text}") + else: + formatted_res.append(text) + + return "\n".join(formatted_res) + + +def normalize_query( + sql: str, filename: str, window: models.Window | None = None +) -> str: + if window is None: + window = DEFAULT_WINDOW + sql = strip_comments(sql) + sql = normalize_params_and_refs(sql) + assert_timestamp_bounds(sql, window) + sql = normalize_time_filters(sql) + sql = normalize_time_groups(sql) + sql = framework_adaptations(sql, filename) + sql = normalize_whitespace_and_parens(sql) + + def process_subqueries(text: str) -> str: + parts = [] + depth = 0 + current = "" + i = 0 + while i < len(text): + if text[i] == "'": + chunk, i = consume_sql_string_literal(text, i) + current += chunk + continue + + if text[i] == "(": + if depth == 0: + parts.append(current) + current = "" + else: + current += "(" + depth += 1 + elif text[i] == ")": + depth = max(0, depth - 1) + if depth == 0: + if re.match(r"^\s*SELECT\b", current, re.IGNORECASE): + parts.append("(\n" + custom_clause_split(current) + "\n)") + else: + parts.append("(" + current + ")") + current = "" + else: + current += ")" + else: + current += text[i] + i += 1 + parts.append(current) + return "".join(parts) + + sql = process_subqueries(sql) + sql = custom_clause_split(sql) + + # Strip any extra newlines or spaces + lines = [line.rstrip() for line in sql.splitlines() if line.strip()] + return "\n".join(lines) + + +def get_streamlit_query( + filename: str, window: models.Window | None = None +) -> str: + refs = models.TableRefs("project", "dataset", "events", "adk_") + if window is None: + window = DEFAULT_WINDOW + + if filename not in CANONICAL_QUERIES: + raise ValueError(f"Unknown query: {filename}") + + builder = CANONICAL_QUERIES[filename] + if filename in QUERY_LIMITS: + return builder(refs, window, QUERY_LIMITS[filename]) + return builder(refs, window) + + +def main(): + if not QUERIES_DIRECTORY.exists(): + print( + f"ERROR: Canonical queries directory not found at {QUERIES_DIRECTORY}", + file=sys.stderr, + ) + sys.exit(1) + + errors = 0 + errors += check_unmapped_canonical_queries( + QUERIES_DIRECTORY, set(CANONICAL_QUERIES.keys()) + ) + errors += check_unmapped_streamlit_builders() + + refs = models.TableRefs("project", "dataset", "events", "adk_") + window = DEFAULT_WINDOW + + for filename in CANONICAL_QUERIES: + grafana_path = QUERIES_DIRECTORY / filename + if not grafana_path.exists(): + print(f"ERROR: Missing canonical query {filename}", file=sys.stderr) + errors += 1 + continue + + grafana_sql = grafana_path.read_text(encoding="utf-8") + streamlit_sql = get_streamlit_query(filename, window) + + try: + assert_streamlit_query(filename, streamlit_sql, window) + except AssertionError as error: + print( + f"ERROR: Streamlit query {filename} failed assertion: {error}", + file=sys.stderr, + ) + errors += 1 + continue + + try: + norm_grafana = normalize_query(grafana_sql, filename, window) + norm_streamlit = normalize_query(streamlit_sql, filename, window) + except AssertionError as error: + print( + f"ERROR: Normalizing {filename} failed assertion: {error}", + file=sys.stderr, + ) + errors += 1 + continue + + if norm_grafana != norm_streamlit: + print( + f"ERROR: Streamlit query {filename} has drifted from Grafana", + file=sys.stderr, + ) + diff = difflib.unified_diff( + norm_grafana.splitlines(), + norm_streamlit.splitlines(), + fromfile=f"{QUERIES_DIRECTORY.relative_to(REPOSITORY_ROOT)}/{filename}", + tofile=f"Streamlit ({filename})", + lineterm="", + ) + print("\n".join(diff), file=sys.stderr) + errors += 1 + + if errors == 0: + print( + f"All {len(CANONICAL_QUERIES)} Streamlit dashboard queries match" + " canonical Grafana SQL." + ) + sys.exit(0) + else: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/test_check_streamlit_queries_sync.py b/tests/test_check_streamlit_queries_sync.py new file mode 100644 index 00000000..2c3bd659 --- /dev/null +++ b/tests/test_check_streamlit_queries_sync.py @@ -0,0 +1,468 @@ +import io +from pathlib import Path +import subprocess +import sys +from unittest.mock import patch + +import pytest + +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +SCRIPTS_DIR = REPOSITORY_ROOT / "scripts" +sys.path.insert(0, str(SCRIPTS_DIR)) + +import check_streamlit_queries_sync + + +def run_check_with_patch(query_filename, patched_sql): + original_get_query = check_streamlit_queries_sync.get_streamlit_query + + def mock_get_query(filename, window=None): + if filename == query_filename: + return patched_sql + return original_get_query(filename, window) + + with patch( + "check_streamlit_queries_sync.get_streamlit_query", + side_effect=mock_get_query, + ): + # We also need to capture stderr to check the diff + stderr_capture = io.StringIO() + with patch("sys.stderr", stderr_capture): + with pytest.raises(SystemExit) as exc_info: + check_streamlit_queries_sync.main() + return exc_info.value.code, stderr_capture.getvalue() + + +def test_baseline_check(): + stderr_capture = io.StringIO() + stdout_capture = io.StringIO() + with patch("sys.stderr", stderr_capture), patch("sys.stdout", stdout_capture): + with pytest.raises(SystemExit) as exc_info: + check_streamlit_queries_sync.main() + assert exc_info.value.code == 0 + assert ( + f"All {len(check_streamlit_queries_sync.CANONICAL_QUERIES)} Streamlit" + " dashboard queries match canonical Grafana SQL." + in stdout_capture.getvalue() + ) + + +def test_altered_error_predicate(): + original_sql = check_streamlit_queries_sync.get_streamlit_query( + "overview_totals.sql" + ) + # Change _ERROR to _WARNING + altered_sql = original_sql.replace("'_ERROR'", "'_WARNING'") + code, stderr = run_check_with_patch("overview_totals.sql", altered_sql) + assert code == 1 + assert "has drifted from Grafana" in stderr + assert "-'_ERROR'" in stderr or "+'_WARNING'" in stderr or "WARNING" in stderr + + +def test_missing_filter(): + original_sql = check_streamlit_queries_sync.get_streamlit_query( + "events_over_time.sql" + ) + altered_sql = original_sql.replace( + "AND ('___ALL___' IN UNNEST(@session_ids) OR session_id IN" + " UNNEST(@session_ids))", + "", + ) + code, stderr = run_check_with_patch("events_over_time.sql", altered_sql) + assert code == 1 + assert "has drifted from Grafana" in stderr + assert "session_ids" in stderr + + +def test_unexpected_filter_on_exempt_query(): + original_sql = check_streamlit_queries_sync.get_streamlit_query( + "errors_over_time.sql" + ) + # Add @event_types filter directly into the WHERE clause of errors_over_time.sql + unexpected_filter = " AND ('___ALL___' IN UNNEST(@event_types) OR event_type IN UNNEST(@event_types))\n" + assert "WHERE " in original_sql + altered_sql = original_sql.replace("WHERE ", f"WHERE {unexpected_filter}") + code, stderr = run_check_with_patch("errors_over_time.sql", altered_sql) + assert code == 1 + assert "has drifted from Grafana" in stderr + assert "event_types" in stderr + + +def test_modified_aggregation(): + original_sql = check_streamlit_queries_sync.get_streamlit_query( + "llm_latency_percentiles.sql" + ) + # Change OFFSET(50) to OFFSET(90) + altered_sql = original_sql.replace("OFFSET(50)", "OFFSET(90)") + code, stderr = run_check_with_patch( + "llm_latency_percentiles.sql", altered_sql + ) + assert code == 1 + assert "has drifted from Grafana" in stderr + assert "OFFSET(90)" in stderr + + +def test_altered_table_or_view(): + original_sql = check_streamlit_queries_sync.get_streamlit_query( + "tool_usage.sql" + ) + # Change tool_starts to tool_completions + altered_sql = original_sql.replace( + "`project.dataset.adk_tool_starts`", + "`project.dataset.adk_tool_completions`", + ) + code, stderr = run_check_with_patch("tool_usage.sql", altered_sql) + assert code == 1 + assert "has drifted from Grafana" in stderr + assert "tool_completions" in stderr + + +def test_direct_cli_execution(): + script_path = REPOSITORY_ROOT / "scripts" / "check_streamlit_queries_sync.py" + result = subprocess.run( + [sys.executable, str(script_path)], capture_output=True, text=True + ) + assert result.returncode == 0 + assert ( + f"All {len(check_streamlit_queries_sync.CANONICAL_QUERIES)} Streamlit" + " dashboard queries match canonical Grafana SQL." in result.stdout + ) + + +def test_normalize_query_reordered_where_conditions(): + sql_a = """ + SELECT + tool_name, + COUNT(*) AS invocations + FROM `project.dataset.adk_tool_starts` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + AND ('___ALL___' IN UNNEST(@agents) OR agent IN UNNEST(@agents)) + AND ('___ALL___' IN UNNEST(@user_ids) OR user_id IN UNNEST(@user_ids)) + GROUP BY tool_name + ORDER BY invocations DESC + """ + sql_b = """ + SELECT + tool_name, + COUNT(*) AS invocations + FROM `project.dataset.adk_tool_starts` + WHERE ('___ALL___' IN UNNEST(@user_ids) OR user_id IN UNNEST(@user_ids)) + AND timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND ('___ALL___' IN UNNEST(@agents) OR agent IN UNNEST(@agents)) + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + GROUP BY tool_name + ORDER BY invocations DESC + """ + norm_a = check_streamlit_queries_sync.normalize_query(sql_a, "tool_usage.sql") + norm_b = check_streamlit_queries_sync.normalize_query(sql_b, "tool_usage.sql") + assert norm_a == norm_b + + +def test_normalize_query_detects_swapped_timestamp_bounds(): + swapped_sql = """ + SELECT + tool_name, + COUNT(*) AS invocations + FROM `project.dataset.adk_tool_starts` + WHERE timestamp >= TIMESTAMP "2025-01-01 12:00:00+00:00" + AND timestamp < TIMESTAMP "2024-12-31 12:00:00+00:00" + GROUP BY tool_name + ORDER BY invocations DESC + """ + with pytest.raises(AssertionError) as exc_info: + check_streamlit_queries_sync.normalize_query(swapped_sql, "tool_usage.sql") + assert "Expected start timestamp literal" in str(exc_info.value) + + +def test_normalize_query_detects_reordered_select_columns(): + sql_a = """ + SELECT + col_a, + col_b + FROM `project.dataset.events` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + """ + sql_b = """ + SELECT + col_b, + col_a + FROM `project.dataset.events` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + """ + norm_a = check_streamlit_queries_sync.normalize_query(sql_a, "dummy.sql") + norm_b = check_streamlit_queries_sync.normalize_query(sql_b, "dummy.sql") + assert norm_a != norm_b + + +def test_normalize_query_detects_removed_select_column(): + sql_full = """ + SELECT + col_a, + col_b + FROM `project.dataset.events` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + """ + sql_missing = """ + SELECT + col_a + FROM `project.dataset.events` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + """ + norm_full = check_streamlit_queries_sync.normalize_query( + sql_full, "dummy.sql" + ) + norm_missing = check_streamlit_queries_sync.normalize_query( + sql_missing, "dummy.sql" + ) + assert norm_full != norm_missing + + +def test_main_missing_canonical_query_file(): + stderr_capture = io.StringIO() + with ( + patch.dict( + check_streamlit_queries_sync.CANONICAL_QUERIES, + {"missing_query.sql": lambda refs, window: "SELECT 1"}, + ), + patch("sys.stderr", stderr_capture), + ): + with pytest.raises(SystemExit) as exc_info: + check_streamlit_queries_sync.main() + assert exc_info.value.code != 0 + assert ( + "Missing canonical query missing_query.sql" in stderr_capture.getvalue() + ) + + +def test_check_unmapped_canonical_queries(tmp_path): + (tmp_path / "overview_totals.sql").write_text("SELECT 1") + (tmp_path / "estimated_cost.sql").write_text("SELECT 1") + count = check_streamlit_queries_sync.check_unmapped_canonical_queries( + tmp_path, {"overview_totals.sql"} + ) + assert count == 0 + + (tmp_path / "extra_unknown.sql").write_text("SELECT 1") + count = check_streamlit_queries_sync.check_unmapped_canonical_queries( + tmp_path, {"overview_totals.sql"} + ) + assert count == 1 + + +def test_check_unmapped_streamlit_builders(): + assert check_streamlit_queries_sync.check_unmapped_streamlit_builders() == 0 + + with patch.object( + check_streamlit_queries_sync.queries, + "build_brand_new_feature_sql", + create=True, + new=lambda refs, window: "SELECT 1", + ): + assert check_streamlit_queries_sync.check_unmapped_streamlit_builders() == 1 + + +def test_scoped_tool_name_normalization(): + sql_with_ifnull = "SELECT IFNULL(tool_name, 'unknown') AS tool_name FROM foo" + norm_other = check_streamlit_queries_sync.normalize_query( + sql_with_ifnull, "overview_totals.sql" + ) + assert "IFNULL(tool_name, 'unknown')" in norm_other + + norm_tool = check_streamlit_queries_sync.normalize_query( + sql_with_ifnull, "tool_usage.sql" + ) + assert "IFNULL(tool_name, 'unknown')" not in norm_tool + assert "tool_name" in norm_tool + + +def test_string_literal_depth_tracking(): + sql = ( + "SELECT col1 FROM tbl WHERE col2 = 'value with (parens) and AND " + "keywords' AND col3 = 1" + ) + clauses = check_streamlit_queries_sync.split_top_level(sql, r"\s+WHERE\s+") + assert len(clauses) == 2 + assert "value with (parens)" in clauses[1] + + +def test_timestamp_bounds_quote_tolerance(): + window = check_streamlit_queries_sync.DEFAULT_WINDOW + sql_single = """ + SELECT tool_name, COUNT(*) AS invocations + FROM `project.dataset.adk_tool_starts` + WHERE timestamp >= TIMESTAMP '2024-12-31 12:00:00+00:00' + AND timestamp < TIMESTAMP '2025-01-01 12:00:00+00:00' + GROUP BY tool_name + ORDER BY invocations DESC + """ + sql_double = """ + SELECT tool_name, COUNT(*) AS invocations + FROM `project.dataset.adk_tool_starts` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + GROUP BY tool_name + ORDER BY invocations DESC + """ + check_streamlit_queries_sync.assert_timestamp_bounds(sql_single, window) + check_streamlit_queries_sync.assert_timestamp_bounds(sql_double, window) + assert check_streamlit_queries_sync.normalize_time_filters( + sql_single + ) == check_streamlit_queries_sync.normalize_time_filters(sql_double) + + +def test_assert_streamlit_query_ifnull_tolerance(): + window = check_streamlit_queries_sync.DEFAULT_WINDOW + sql_variants = [ + """ + SELECT IFNULL(tool_name, 'unknown') AS tool_name, COUNT(*) AS invocations + FROM `project.dataset.adk_tool_starts` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + GROUP BY tool_name + """, + """ + SELECT IFNULL( tool_name , "unknown" ) AS tool_name, COUNT(*) AS invocations + FROM `project.dataset.adk_tool_starts` + WHERE timestamp >= TIMESTAMP '2024-12-31 12:00:00+00:00' + AND timestamp < TIMESTAMP '2025-01-01 12:00:00+00:00' + GROUP BY tool_name + """, + ] + for sql in sql_variants: + check_streamlit_queries_sync.assert_streamlit_query( + "tool_usage.sql", sql, window + ) + + +def test_llm_calls_total_missing_token_columns(): + window = check_streamlit_queries_sync.DEFAULT_WINDOW + original_sql = check_streamlit_queries_sync.get_streamlit_query( + "llm_calls_total.sql" + ) + + # Test missing prompt_tokens + bad_sql_prompt = original_sql.replace( + "IFNULL(SUM(usage_prompt_tokens), 0) AS prompt_tokens,", "" + ) + with pytest.raises( + AssertionError, match="missing required token column usage_prompt_tokens" + ): + check_streamlit_queries_sync.assert_streamlit_query( + "llm_calls_total.sql", bad_sql_prompt, window + ) + + code, stderr = run_check_with_patch("llm_calls_total.sql", bad_sql_prompt) + assert code == 1 + assert ( + "missing required token column usage_prompt_tokens as prompt_tokens" + in stderr + ) + + # Test missing completion_tokens + bad_sql_completion = original_sql.replace( + "IFNULL(SUM(usage_completion_tokens), 0) AS completion_tokens,", "" + ) + with pytest.raises( + AssertionError, + match="missing required token column usage_completion_tokens", + ): + check_streamlit_queries_sync.assert_streamlit_query( + "llm_calls_total.sql", bad_sql_completion, window + ) + + # Test missing total_tokens + bad_sql_total = original_sql.replace( + "IFNULL(SUM(usage_total_tokens), 0) AS total_tokens", "" + ) + with pytest.raises( + AssertionError, match="missing required token column usage_total_tokens" + ): + check_streamlit_queries_sync.assert_streamlit_query( + "llm_calls_total.sql", bad_sql_total, window + ) + + +def test_llm_calls_total_whitespace_variants_normalize(): + """Whitespace variants in llm_calls_total.sql normalize cleanly.""" + window = check_streamlit_queries_sync.DEFAULT_WINDOW + grafana_sql = ( + check_streamlit_queries_sync.QUERIES_DIRECTORY / "llm_calls_total.sql" + ).read_text(encoding="utf-8") + + whitespace_variant_sql = """ + SELECT + COUNT(DISTINCT CONCAT(trace_id, '|', span_id)) + + COUNTIF(trace_id IS NULL OR span_id IS NULL) AS llm_calls , + IFNULL ( SUM ( usage_prompt_tokens ) , 0 ) AS prompt_tokens , + IFNULL( SUM( usage_completion_tokens ) , 0 ) AS completion_tokens , + IFNULL ( SUM ( usage_total_tokens ) , 0 ) AS total_tokens + FROM `project.dataset.adk_llm_responses` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + AND ('___ALL___' IN UNNEST(@agents) OR agent IN UNNEST(@agents)) + AND ('___ALL___' IN UNNEST(@user_ids) OR user_id IN UNNEST(@user_ids)) + AND ('___ALL___' IN UNNEST(@session_ids) OR session_id IN UNNEST(@session_ids)) + HAVING COUNT(*) > 0 + """ + + check_streamlit_queries_sync.assert_streamlit_query( + "llm_calls_total.sql", whitespace_variant_sql, window + ) + norm_grafana = check_streamlit_queries_sync.normalize_query( + grafana_sql, "llm_calls_total.sql", window + ) + norm_variant = check_streamlit_queries_sync.normalize_query( + whitespace_variant_sql, "llm_calls_total.sql", window + ) + assert norm_variant == norm_grafana + + +def test_single_quote_timestamp_bounds(): + window = check_streamlit_queries_sync.DEFAULT_WINDOW + sql = """ + SELECT tool_name + FROM `project.dataset.adk_tool_starts` + WHERE timestamp >= TIMESTAMP '2024-12-31 12:00:00+00:00' + AND timestamp < TIMESTAMP '2025-01-01 12:00:00+00:00' + """ + # Valid single-quote bounds pass + check_streamlit_queries_sync.assert_timestamp_bounds(sql, window) + + # Invalid start timestamp with single quote fails + bad_sql = sql.replace("2024-12-31", "2024-12-30") + with pytest.raises(AssertionError, match="Expected start timestamp literal"): + check_streamlit_queries_sync.assert_timestamp_bounds(bad_sql, window) + + +def test_select_columns_renamed_or_reordered(): + window = check_streamlit_queries_sync.DEFAULT_WINDOW + original_sql = check_streamlit_queries_sync.get_streamlit_query( + "tool_usage.sql" + ) + + # Renamed select column is detected as drift + renamed_sql = original_sql.replace("AS invocations", "AS call_count") + code, stderr = run_check_with_patch("tool_usage.sql", renamed_sql) + assert code == 1 + assert "has drifted from Grafana" in stderr + assert "call_count" in stderr + + # Reordered select columns are normalized consistently by custom_clause_split + sql_1 = """ + SELECT + col_a, + col_b + FROM `project.dataset.events` + WHERE timestamp >= TIMESTAMP "2024-12-31 12:00:00+00:00" + AND timestamp < TIMESTAMP "2025-01-01 12:00:00+00:00" + """ + norm_1 = check_streamlit_queries_sync.normalize_query( + sql_1, "dummy.sql", window + ) + assert "col_a" in norm_1 + assert "col_b" in norm_1 diff --git a/tests/test_dashboards_streamlit_app.py b/tests/test_dashboards_streamlit_app.py new file mode 100644 index 00000000..7036e6ad --- /dev/null +++ b/tests/test_dashboards_streamlit_app.py @@ -0,0 +1,2304 @@ +"""Unit tests for the Streamlit dashboard.""" + +from __future__ import annotations + +import contextlib +import datetime as dt +import importlib.util +from pathlib import Path +import sys +from unittest import mock + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +APP_PATH = ROOT / "dashboards" / "streamlit" / "app.py" +CHARTS_PATH = ROOT / "dashboards" / "streamlit" / "charts.py" +MODELS_PATH = ROOT / "dashboards" / "streamlit" / "models.py" +QUERIES_PATH = ROOT / "dashboards" / "streamlit" / "queries.py" + +# Modules that dashboard components import. Modules in _UNCONDITIONAL_MOCKS +# are always stubbed with MagicMocks to isolate tests from third-party UI +# and cloud libraries; any remaining modules are mocked only if missing. +_OPTIONAL_MODULES = ( + "streamlit", + "pandas", + "plotly", + "plotly.graph_objects", + "google", + "google.api_core", + "google.api_core.exceptions", + "google.auth", + "google.auth.exceptions", + "google.cloud", + "google.cloud.bigquery", +) + +_UNCONDITIONAL_MOCKS = ( + "dotenv", + "streamlit", + "plotly", + "plotly.graph_objects", + "google.cloud.bigquery", +) + +# `from X import Y` reads Y off the parent package, and a MagicMock parent +# auto-creates an attribute unrelated to the sys.modules entry configured +# below. Point these at the real entries so app.py binds what we set up. +_PARENT_ATTRS = ( + ("google.cloud", "bigquery"), + ("google.api_core", "exceptions"), + ("google.auth", "exceptions"), + ("plotly", "graph_objects"), +) + + +class _MockArrayQueryParameter: + """Stand-in that keeps what ``query_parameters`` passes positionally. + + Attribute names mirror ``bigquery.ArrayQueryParameter`` so assertions + hold whether the real client library is installed or not. + """ + + def __init__(self, name, array_type, values): + self.name = name + self.array_type = array_type + self.values = values + + +class _MockQueryJobConfig: + """Stand-in for ``bigquery.QueryJobConfig``.""" + + def __init__(self, **kwargs): + self.maximum_bytes_billed = None + for k, v in kwargs.items(): + setattr(self, k, v) + + +class _MockDataFrame: + """Mock mimicking ``pandas.DataFrame`` that exposes ``.empty``.""" + + def __init__(self, data=None, *args, **kwargs): + self.data = data + self.empty = not bool(data) + + +_MISSING = object() + + +@contextlib.contextmanager +def _mocked_optional_imports(): + """Installs mocks for absent optional deps, then restores ``sys.modules``. + + The mocks must not outlive the load. They are process-global, so a test + module imported later in the same session would otherwise resolve + ``google.cloud`` or ``pandas`` to a MagicMock and run against it + silently instead of failing — or skipping — honestly. + """ + saved_modules: dict[str, object] = {} + saved_parent_attrs: list[tuple[object, str, object, bool]] = [] + try: + for name in _UNCONDITIONAL_MOCKS: + saved_modules[name] = sys.modules.get(name, _MISSING) + sys.modules[name] = mock.MagicMock() + + for name in _OPTIONAL_MODULES: + if name in _UNCONDITIONAL_MOCKS: + continue + if name not in sys.modules: + try: + __import__(name) + except ImportError: + saved_modules[name] = _MISSING + sys.modules[name] = mock.MagicMock() + + for parent, attr in _PARENT_ATTRS: + child_key = f"{parent}.{attr}" + if parent in sys.modules and child_key in sys.modules: + parent_mod = sys.modules[parent] + child_mod = sys.modules[child_key] + if isinstance(parent_mod, mock.MagicMock): + had_attr = attr in parent_mod.__dict__ + orig_attr = parent_mod.__dict__.get(attr) + else: + had_attr = hasattr(parent_mod, attr) + orig_attr = getattr(parent_mod, attr, None) + saved_parent_attrs.append((parent_mod, attr, orig_attr, had_attr)) + setattr(parent_mod, attr, child_mod) + + # Only ever patch a mock, never a real module: mutating an installed + # package would be the same leak in a different disguise. + bq_mod = sys.modules["google.cloud.bigquery"] + bq_mod.ArrayQueryParameter = _MockArrayQueryParameter + bq_mod.QueryJobConfig = _MockQueryJobConfig + + pd_mod = sys.modules.get("pandas") + if isinstance(pd_mod, mock.MagicMock): + pd_mod.DataFrame = _MockDataFrame + + gexc_mod = sys.modules.get("google.api_core.exceptions") + if isinstance(gexc_mod, mock.MagicMock): + for exc_name in ( + "GoogleAPICallError", + "NotFound", + "Forbidden", + "RetryError", + "ServiceUnavailable", + ): + if not isinstance(getattr(gexc_mod, exc_name, None), type): + setattr(gexc_mod, exc_name, type(exc_name, (Exception,), {})) + + gauth_mod = sys.modules.get("google.auth.exceptions") + if isinstance(gauth_mod, mock.MagicMock): + if not isinstance( + getattr(gauth_mod, "DefaultCredentialsError", None), type + ): + setattr( + gauth_mod, + "DefaultCredentialsError", + type("DefaultCredentialsError", (Exception,), {}), + ) + + st_mod = sys.modules["streamlit"] + + def _passthrough(*args, **kwargs): + def _decorator(func): + return func + + return _decorator + + st_mod.cache_data = _passthrough + st_mod.cache_resource = _passthrough + + yield + finally: + for parent_mod, attr, orig_attr, had_attr in reversed(saved_parent_attrs): + if had_attr: + setattr(parent_mod, attr, orig_attr) + else: + try: + delattr(parent_mod, attr) + except AttributeError: + pass + + for name, orig_val in reversed(list(saved_modules.items())): + if orig_val is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = orig_val + + +def _load_module(name: str, path: Path): + """Loads a Python module dynamically with mocked optional deps.""" + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + dir_path = str(path.parent) + added_path = False + if dir_path not in sys.path: + sys.path.insert(0, dir_path) + added_path = True + saved_submods = {} + try: + for submod in ("charts", "models", "queries"): + saved_submods[submod] = sys.modules.get(submod, _MISSING) + stored = sys.modules.get(f"dashboards_streamlit_{submod}") + if stored is not None: + sys.modules[submod] = stored + with _mocked_optional_imports(): + spec.loader.exec_module(mod) + except Exception: + sys.modules.pop(name, None) + raise + finally: + if added_path and dir_path in sys.path: + sys.path.remove(dir_path) + for submod, orig_val in saved_submods.items(): + if orig_val is _MISSING: + sys.modules.pop(submod, None) + else: + sys.modules[submod] = orig_val + return mod + + +def _load_streamlit_models(): + """Loads dashboards/streamlit/models.py with mocked optional deps.""" + return _load_module("dashboards_streamlit_models", MODELS_PATH) + + +def _load_streamlit_charts(): + """Loads dashboards/streamlit/charts.py with mocked optional deps.""" + return _load_module("dashboards_streamlit_charts", CHARTS_PATH) + + +def _load_streamlit_queries(): + """Loads dashboards/streamlit/queries.py with mocked optional deps.""" + return _load_module("dashboards_streamlit_queries", QUERIES_PATH) + + +def _load_streamlit_app(): + """Loads dashboards/streamlit/app.py dynamically with mocked optional deps.""" + return _load_module("dashboards_streamlit_app", APP_PATH) + + +models = _load_streamlit_models() +charts = _load_streamlit_charts() +queries = _load_streamlit_queries() +app = _load_streamlit_app() + + +@pytest.fixture(autouse=True) +def _reset_queries_state(): + try: + yield + finally: + queries._SEEN_RUN_IDS.clear() + queries._NEXT_RUN_ID = 0 + + +@pytest.fixture +def sample_refs(): + return models.TableRefs( + project="test-project", + dataset="test_dataset", + table="agent_events", + view_prefix="adk_", + ) + + +@pytest.fixture +def sample_window(): + return models.Window( + start=dt.datetime(2026, 9, 1, 0, 0, 0, tzinfo=dt.timezone.utc), + end=dt.datetime(2026, 9, 2, 0, 0, 0, tzinfo=dt.timezone.utc), + ) + + +def test_table_refs_formatting_and_defaults(): + refs = models.TableRefs( + project="my-proj", + dataset="my_ds", + table="my_table", + view_prefix="custom_", + ) + assert refs.events == "`my-proj.my_ds.my_table`" + assert refs.view("llm_responses") == "`my-proj.my_ds.custom_llm_responses`" + assert refs.view("tool_errors") == "`my-proj.my_ds.custom_tool_errors`" + + default_refs = models.TableRefs( + project="my-proj", + dataset="my_ds", + table="my_table", + ) + assert default_refs.view_prefix == models.DEFAULT_VIEW_PREFIX + assert ( + default_refs.view("llm_responses") == "`my-proj.my_ds.adk_llm_responses`" + ) + + +def test_validate_refs_valid(): + refs, errors = models.validate_refs( + "valid-project-123", "valid_dataset", "valid_table_id", "adk_" + ) + assert errors == [] + assert refs is not None + assert refs.project == "valid-project-123" + assert refs.dataset == "valid_dataset" + assert refs.table == "valid_table_id" + assert refs.view_prefix == "adk_" + assert refs.events == "`valid-project-123.valid_dataset.valid_table_id`" + + # Empty view prefix is valid (e.g. view named simply 'llm_responses') + refs_empty_prefix, errors = models.validate_refs( + "valid-project", "valid_dataset", "events", "" + ) + assert errors == [] + assert refs_empty_prefix is not None + assert refs_empty_prefix.view_prefix == "" + assert refs_empty_prefix.view("llm") == "`valid-project.valid_dataset.llm`" + + +def test_validate_refs_invalid_identifiers(): + # Invalid project + refs, errors = models.validate_refs( + "-invalid-start", "dataset", "table", "prefix_" + ) + assert refs is None + assert any("Invalid project ID" in e for e in errors) + + refs, errors = models.validate_refs( + "proj with space", "dataset", "table", "prefix_" + ) + assert refs is None + assert any("Invalid project ID" in e for e in errors) + + refs, errors = models.validate_refs( + "proj`inject", "dataset", "table", "prefix_" + ) + assert refs is None + assert any("Invalid project ID" in e for e in errors) + + # Invalid dataset + refs, errors = models.validate_refs( + "proj", "dataset-with-dash", "table", "prefix_" + ) + assert refs is None + assert any("Invalid dataset ID" in e for e in errors) + + refs, errors = models.validate_refs( + "proj", "ds with space", "table", "prefix_" + ) + assert refs is None + assert any("Invalid dataset ID" in e for e in errors) + + # Invalid table + refs, errors = models.validate_refs( + "proj", "dataset", "table-with-dash", "prefix_" + ) + assert refs is None + assert any("Invalid table ID" in e for e in errors) + + refs, errors = models.validate_refs( + "proj", "dataset", "tbl`inject", "prefix_" + ) + assert refs is None + assert any("Invalid table ID" in e for e in errors) + + # Invalid view prefix + refs, errors = models.validate_refs( + "proj", "dataset", "table", "invalid-prefix" + ) + assert refs is None + assert any("Invalid view prefix" in e for e in errors) + + # Multiple errors combined + refs, errors = models.validate_refs("", "", "", "bad-prefix") + assert refs is None + assert len(errors) == 4 + + +def test_time_bounds_and_window_intervals(): + start = dt.datetime(2026, 9, 1, 12, 0, 0, tzinfo=dt.timezone.utc) + end = dt.datetime(2026, 9, 1, 14, 30, 0, tzinfo=dt.timezone.utc) + w = models.Window(start=start, end=end) + + # time_bounds default column + tb = queries.time_bounds(w) + expected_tb = ( + 'timestamp >= TIMESTAMP "2026-09-01 12:00:00+00:00"\n' + ' AND timestamp < TIMESTAMP "2026-09-01 14:30:00+00:00"' + ) + assert tb == expected_tb + + # time_bounds custom column + tb_custom = queries.time_bounds(w, column="e.timestamp") + expected_custom = ( + 'e.timestamp >= TIMESTAMP "2026-09-01 12:00:00+00:00"\n' + ' AND e.timestamp < TIMESTAMP "2026-09-01 14:30:00+00:00"' + ) + assert tb_custom == expected_custom + + # Span property + assert w.span == dt.timedelta(hours=2, minutes=30) + + # Bucket granularity logic: + # <= 6 hours -> MINUTE + w_minute = models.Window(start=start, end=start + dt.timedelta(hours=6)) + assert w_minute.bucket == "MINUTE" + + # <= 3 days -> HOUR + w_hour = models.Window(start=start, end=start + dt.timedelta(days=3)) + assert w_hour.bucket == "HOUR" + + # > 3 days -> DAY + w_day = models.Window(start=start, end=start + dt.timedelta(days=7)) + assert w_day.bucket == "DAY" + + +def test_snap_and_make_window(): + moment = dt.datetime(2026, 9, 1, 12, 34, 56, 789000, tzinfo=dt.timezone.utc) + snapped = models.snap(moment, seconds=300) + assert snapped == dt.datetime(2026, 9, 1, 12, 30, 0, tzinfo=dt.timezone.utc) + + window = models.make_window(dt.timedelta(hours=1), now=moment) + assert window.end == dt.datetime( + 2026, 9, 1, 12, 30, 0, tzinfo=dt.timezone.utc + ) + assert window.start == dt.datetime( + 2026, 9, 1, 11, 30, 0, tzinfo=dt.timezone.utc + ) + + +def test_build_overview_totals_sql(sample_refs, sample_window): + sql = queries.build_overview_totals_sql(sample_refs, sample_window) + + assert "COUNT(DISTINCT e.session_id) AS sessions" in sql + assert "COUNT(*) AS events" in sql + assert "SAFE_DIVIDE(COUNTIF(" in sql + assert "ENDS_WITH(e.event_type, '_ERROR')" in sql + assert "e.error_message IS NOT NULL" in sql + assert "UPPER(e.status) = 'ERROR'" in sql + assert "SELECT AVG(r.total_ms)" in sql + assert f"FROM {sample_refs.view('llm_responses')} AS r" in sql + assert f"FROM {sample_refs.events} AS e" in sql + assert queries.time_bounds(sample_window, "e.timestamp") in sql + assert queries.time_bounds(sample_window, "r.timestamp") in sql + assert "HAVING COUNT(*) > 0" in sql + + # Overview scoping: agents, user_ids, session_ids (event_type is omitted by + # design) + assert "@agents" in sql + assert "@user_ids" in sql + assert "@session_ids" in sql + assert "@event_types" not in sql + + +def test_build_trace_detail_sql(sample_refs, sample_window): + limit = 500 + sql = queries.build_trace_detail_sql(sample_refs, sample_window, limit=limit) + + assert f"FROM {sample_refs.events}" in sql + assert queries.time_bounds(sample_window) in sql + assert "ORDER BY timestamp DESC" in sql + assert f"LIMIT {limit}" in sql + assert "JSON_VALUE(attributes, '$.model')" in sql + assert "JSON_VALUE(attributes, '$.model_version')" in sql + assert "JSON_VALUE(content, '$.tool') AS tool_name" in sql + assert ( + "SAFE_CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms" + in sql + ) + assert "error_message" in sql + assert "COALESCE(" in sql + + # Scoping includes event_types + assert "@event_types" in sql + assert "@agents" in sql + assert "@user_ids" in sql + assert "@session_ids" in sql + + +def test_build_llm_calls_total_sql(sample_refs, sample_window): + sql = queries.build_llm_calls_total_sql(sample_refs, sample_window) + + assert f"FROM {sample_refs.view('llm_responses')}" in sql + assert ( + "COUNT(DISTINCT CONCAT(trace_id, '|', span_id))\n" + " + COUNTIF(trace_id IS NULL OR span_id IS NULL) AS llm_calls" in sql + ) + assert "IFNULL(SUM(usage_prompt_tokens), 0) AS prompt_tokens" in sql + assert "IFNULL(SUM(usage_completion_tokens), 0) AS completion_tokens" in sql + assert "IFNULL(SUM(usage_total_tokens), 0) AS total_tokens" in sql + assert "@price_in" not in sql + assert "@price_out" not in sql + assert "estimated_cost_usd" not in sql + assert queries.time_bounds(sample_window) in sql + assert "HAVING COUNT(*) > 0" in sql + + +def test_build_recent_sessions_sql(sample_refs, sample_window): + limit = 250 + sql = queries.build_recent_sessions_sql( + sample_refs, sample_window, limit=limit + ) + + assert f"FROM {sample_refs.events}" in sql + assert queries.time_bounds(sample_window) in sql + assert "session_id IS NOT NULL" in sql + assert ( + "STRING_AGG(DISTINCT user_id, ', ' ORDER BY user_id)\n" + " AS session_users_in_window" in sql + ) + assert "MIN(timestamp) AS started_in_window_at" in sql + assert "MAX(timestamp) AS last_event_in_window_at" in sql + assert ( + "TIMESTAMP_DIFF(MAX(timestamp), MIN(timestamp), SECOND)\n" + " AS duration_in_window_s" in sql + ) + assert "COUNT(DISTINCT agent) AS session_agents_in_window" in sql + assert "COUNT(*) AS session_events_in_window" in sql + assert ( + "COUNTIF(ENDS_WITH(event_type, '_ERROR') OR error_message IS NOT NULL OR" + " UPPER(status) = 'ERROR') AS session_errors_in_window" in sql + ) + assert ( + "IFNULL(SUM(IF(event_type = 'LLM_RESPONSE',\n" + " SAFE_CAST(JSON_VALUE(content, '$.usage.prompt') AS INT64), NULL))," + " 0)\n" + " AS session_input_tokens_in_window" in sql + ) + assert ( + "IFNULL(SUM(IF(event_type = 'LLM_RESPONSE',\n" + " SAFE_CAST(JSON_VALUE(content, '$.usage.completion') AS INT64)," + " NULL)), 0)\n" + " AS session_output_tokens_in_window" in sql + ) + assert "GROUP BY session_id" in sql + assert ( + "HAVING LOGICAL_OR('___ALL___' IN UNNEST(@agents)\n" + " OR agent IN UNNEST(@agents))" in sql + ) + assert ( + "AND LOGICAL_OR('___ALL___' IN UNNEST(@user_ids)\n" + " OR user_id IN UNNEST(@user_ids))" in sql + ) + assert ( + "AND LOGICAL_OR('___ALL___' IN UNNEST(@event_types)\n" + " OR event_type IN UNNEST(@event_types))" in sql + ) + assert "ORDER BY last_event_in_window_at DESC" in sql + assert f"LIMIT {limit}" in sql + + +def test_build_tool_errors_sql(sample_refs, sample_window): + limit = 100 + sql = queries.build_tool_errors_sql(sample_refs, sample_window, limit=limit) + + assert f"FROM {sample_refs.view('tool_errors')}" in sql + assert "UNION ALL" in sql + assert f"FROM {sample_refs.view('tool_completions')}" in sql + assert "AND (error_message IS NOT NULL OR UPPER(status) = 'ERROR')" in sql + assert "ORDER BY timestamp DESC" in sql + assert f"LIMIT {limit}" in sql + + +def test_build_filter_options_sql(sample_refs, sample_window): + sql = queries.build_filter_options_sql(sample_refs, sample_window) + + assert f"FROM {sample_refs.events} AS e," in sql + assert "UNNEST([" in sql + assert "STRUCT('agent' AS kind, e.agent AS value)" in sql + assert "STRUCT('user_id', e.user_id)" in sql + assert "STRUCT('event_type', e.event_type)" in sql + assert "STRUCT('session_id', e.session_id)" in sql + assert "WHERE " + queries.time_bounds(sample_window, "e.timestamp") in sql + assert "AND f.value IS NOT NULL" in sql + assert "GROUP BY f.kind, f.value" in sql + assert "QUALIFY ROW_NUMBER() OVER (" in sql + assert f"<= {models.FILTER_OPTIONS_LIMIT}" in sql + assert "ORDER BY f.kind, f.value" in sql + + +def test_build_events_over_time_sql(sample_refs, sample_window): + sql = queries.build_events_over_time_sql(sample_refs, sample_window) + + assert f"TIMESTAMP_TRUNC(timestamp, {sample_window.bucket}) AS bucket" in sql + assert "event_type," in sql + assert "COUNT(*) AS events" in sql + assert f"FROM {sample_refs.events}" in sql + assert queries.time_bounds(sample_window) in sql + assert "GROUP BY bucket, event_type" in sql + assert "ORDER BY bucket" in sql + assert "@event_types" in sql + + +def test_build_errors_over_time_sql(sample_refs, sample_window): + sql = queries.build_errors_over_time_sql(sample_refs, sample_window) + + assert f"TIMESTAMP_TRUNC(timestamp, {sample_window.bucket}) AS bucket" in sql + assert "event_type," in sql + assert "COUNT(*) AS errors" in sql + assert f"FROM {sample_refs.events}" in sql + assert queries.time_bounds(sample_window) in sql + assert "ENDS_WITH(event_type, '_ERROR')" in sql + assert "error_message IS NOT NULL" in sql + assert "UPPER(status) = 'ERROR'" in sql + assert "GROUP BY bucket, event_type" in sql + assert "ORDER BY bucket" in sql + # Event types filter should NOT be in errors_over_time + assert "@event_types" not in sql + + +def test_build_events_by_agent_sql(sample_refs, sample_window): + sql = queries.build_events_by_agent_sql(sample_refs, sample_window) + + assert "IFNULL(agent, 'unknown') AS agent_name" in sql + assert "COUNT(*) AS events" in sql + assert f"FROM {sample_refs.events}" in sql + assert queries.time_bounds(sample_window) in sql + assert "GROUP BY agent_name" in sql + assert "HAVING COUNT(*) > 0" in sql + assert "ORDER BY events DESC" in sql + assert "@event_types" in sql + + +def test_build_top_errors_sql(sample_refs, sample_window): + limit = 50 + sql = queries.build_top_errors_sql(sample_refs, sample_window, limit=limit) + + assert "error_message," in sql + assert "COUNT(*) AS errors," in sql + assert "COUNT(DISTINCT session_id) AS sessions," in sql + assert "COUNT(DISTINCT agent) AS agents," in sql + assert "MAX(timestamp) AS last_seen" in sql + assert f"FROM {sample_refs.events}" in sql + assert queries.time_bounds(sample_window) in sql + assert "AND error_message IS NOT NULL" in sql + assert "GROUP BY error_message" in sql + assert "HAVING COUNT(*) > 0" in sql + assert "ORDER BY errors DESC" in sql + assert f"LIMIT {limit}" in sql + + # Default limit is TOP_ERRORS_LIMIT (50) + sql_default = queries.build_top_errors_sql(sample_refs, sample_window) + assert f"LIMIT {models.TOP_ERRORS_LIMIT}" in sql_default + assert "LIMIT 50" in sql_default + + +def test_build_llm_tokens_over_time_sql(sample_refs, sample_window): + sql = queries.build_llm_tokens_over_time_sql(sample_refs, sample_window) + + assert f"TIMESTAMP_TRUNC(timestamp, {sample_window.bucket}) AS bucket" in sql + assert "IFNULL(SUM(usage_prompt_tokens), 0) AS prompt_tokens" in sql + assert "IFNULL(SUM(usage_completion_tokens), 0) AS completion_tokens" in sql + assert "IFNULL(SUM(usage_total_tokens), 0) AS total_tokens" in sql + assert f"FROM {sample_refs.view('llm_responses')}" in sql + assert queries.time_bounds(sample_window) in sql + assert "GROUP BY bucket" in sql + assert "ORDER BY bucket" in sql + + +def test_build_llm_latency_percentiles_sql(sample_refs, sample_window): + sql = queries.build_llm_latency_percentiles_sql(sample_refs, sample_window) + + assert f"TIMESTAMP_TRUNC(timestamp, {sample_window.bucket}) AS bucket" in sql + assert "APPROX_QUANTILES(total_ms, 100)[OFFSET(50)] AS p50_total_ms" in sql + assert "APPROX_QUANTILES(total_ms, 100)[OFFSET(95)] AS p95_total_ms" in sql + assert "APPROX_QUANTILES(ttft_ms, 100)[OFFSET(50)] AS p50_ttft_ms" in sql + assert f"FROM {sample_refs.view('llm_responses')}" in sql + assert queries.time_bounds(sample_window) in sql + assert "GROUP BY bucket" in sql + assert "ORDER BY bucket" in sql + + +def test_build_tokens_by_model_sql(sample_refs, sample_window): + sql = queries.build_tokens_by_model_sql(sample_refs, sample_window) + + assert "IFNULL(model_version, 'unknown') AS model" in sql + assert "IFNULL(SUM(usage_prompt_tokens), 0) AS prompt_tokens" in sql + assert "IFNULL(SUM(usage_completion_tokens), 0) AS completion_tokens" in sql + assert "IFNULL(SUM(usage_total_tokens), 0) AS total_tokens" in sql + assert "COUNT(*) AS responses" in sql + assert f"FROM {sample_refs.view('llm_responses')}" in sql + assert queries.time_bounds(sample_window) in sql + assert "GROUP BY model" in sql + assert "ORDER BY total_tokens DESC" in sql + + +def test_build_tool_usage_sql(sample_refs, sample_window): + sql = queries.build_tool_usage_sql(sample_refs, sample_window) + + assert "IFNULL(tool_name, 'unknown') AS tool_name" in sql + assert "COUNT(*) AS invocations" in sql + assert f"FROM {sample_refs.view('tool_starts')}" in sql + assert queries.time_bounds(sample_window) in sql + assert "GROUP BY tool_name" in sql + assert "ORDER BY invocations DESC" in sql + + +def test_build_tool_latency_sql(sample_refs, sample_window): + sql = queries.build_tool_latency_sql(sample_refs, sample_window) + + assert "IFNULL(tool_name, 'unknown') AS tool_name" in sql + assert "COUNT(*) AS completions" in sql + assert "AVG(total_ms) AS avg_ms" in sql + assert "APPROX_QUANTILES(total_ms, 100)[OFFSET(50)] AS p50_ms" in sql + assert "APPROX_QUANTILES(total_ms, 100)[OFFSET(95)] AS p95_ms" in sql + assert f"FROM {sample_refs.view('tool_completions')}" in sql + assert queries.time_bounds(sample_window) in sql + assert "GROUP BY tool_name" in sql + assert "ORDER BY p95_ms DESC" in sql + + +def test_filters_and_as_filter_values(): + f_default = models.Filters() + assert f_default.agents == (models.ALL_SENTINEL,) + assert f_default.user_ids == (models.ALL_SENTINEL,) + assert f_default.event_types == (models.ALL_SENTINEL,) + assert f_default.session_ids == (models.ALL_SENTINEL,) + + assert models.as_filter_values([]) == (models.ALL_SENTINEL,) + assert models.as_filter_values(["a", "b"]) == ("a", "b") + + +def test_query_parameters(): + filters = models.Filters( + agents=("agent_1",), + user_ids=("user_1",), + event_types=("LLM_RESPONSE",), + session_ids=("sess_1",), + ) + + # Query that uses @agents and @session_ids only + sql = ( + "SELECT * FROM t WHERE agent IN UNNEST(@agents) AND session_id IN" + " UNNEST(@session_ids)" + ) + params = queries.query_parameters(sql, filters) + by_name = {p.name: p for p in params} + assert sorted(by_name) == ["agents", "session_ids"] + + # Each bound parameter carries its filter's values as ARRAY. + assert by_name["agents"].array_type == "STRING" + assert list(by_name["agents"].values) == ["agent_1"] + assert list(by_name["session_ids"].values) == ["sess_1"] + + # Query that uses all 4 + sql_all = "SELECT 1 WHERE @agents @user_ids @event_types @session_ids" + params_all = queries.query_parameters(sql_all, filters) + assert len(params_all) == 4 + + # A query with no filter references binds nothing (build_filter_options_sql + # is the real case), so BigQuery never receives an unreferenced parameter. + assert queries.query_parameters("SELECT 1", filters) == [] + + +def test_humanize_bytes(): + assert models.humanize_bytes(500) == "500 B" + assert models.humanize_bytes(1024) == "1.0 KB" + assert models.humanize_bytes(10 * 1024**2) == "10.0 MB" + assert models.humanize_bytes(2.5 * 1024**3) == "2.5 GB" + assert models.humanize_bytes(3 * 1024**4) == "3.0 TB" + + +def test_scope_helper(): + # Without alias and without event_type + scope_bare = queries._scope() + assert ( + "('___ALL___' IN UNNEST(@agents) OR agent IN UNNEST(@agents))" + in scope_bare + ) + assert ( + "('___ALL___' IN UNNEST(@user_ids) OR user_id IN UNNEST(@user_ids))" + in scope_bare + ) + assert ( + "('___ALL___' IN UNNEST(@session_ids) OR session_id IN" + " UNNEST(@session_ids))" in scope_bare + ) + assert "@event_types" not in scope_bare + + # With alias and with event_type + scope_aliased = queries._scope("e", event_type=True) + assert ( + "('___ALL___' IN UNNEST(@agents) OR e.agent IN UNNEST(@agents))" + in scope_aliased + ) + assert ( + "('___ALL___' IN UNNEST(@user_ids) OR e.user_id IN UNNEST(@user_ids))" + in scope_aliased + ) + assert ( + "('___ALL___' IN UNNEST(@event_types) OR e.event_type IN" + " UNNEST(@event_types))" in scope_aliased + ) + assert ( + "('___ALL___' IN UNNEST(@session_ids) OR e.session_id IN" + " UNNEST(@session_ids))" in scope_aliased + ) + + +def test_themes(): + assert models.LIGHT_THEME.surface == "#fcfcfb" + assert models.DARK_THEME.surface == "#1a1a19" + assert len(models.LIGHT_THEME.categorical) == 8 + assert len(models.DARK_THEME.categorical) == 8 + assert charts.active_theme() in (models.LIGHT_THEME, models.DARK_THEME) + + +def test_context_and_query_result(sample_refs, sample_window): + filters = models.Filters() + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=filters, + max_bytes=1024**3, + theme=models.LIGHT_THEME, + price_in=1.25, + price_out=5.00, + ) + assert ctx.refs == sample_refs + assert ctx.window == sample_window + assert ctx.filters == filters + assert ctx.max_bytes == 1024**3 + assert ctx.theme == models.LIGHT_THEME + assert (ctx.price_in, ctx.price_out) == (1.25, 5.00) + assert ctx.scan_log == [] + + # Two Contexts must not share one scan log; a mutable default would make + # every panel's cost accumulate across reruns. + other = models.Context( + refs=sample_refs, + window=sample_window, + filters=filters, + max_bytes=1024**3, + theme=models.LIGHT_THEME, + price_in=1.25, + price_out=5.00, + ) + ctx.scan_log.append(("overview", 1024, False)) + assert other.scan_log == [] + + # A QueryResult defaults to "succeeded, nothing scanned, not cached". + sentinel = object() + result = models.QueryResult(sentinel) + assert result.df is sentinel + assert result.error is None + assert result.bytes_processed == 0 + assert result.bytes_billed == 0 + assert result.cache_hit is False + + +def test_sidebar_connection_locks_env_project(monkeypatch): + monkeypatch.setenv("BQ_PROJECT_ID", "env-locked-project") + monkeypatch.setenv("BQ_DATASET_ID", "env_dataset") + monkeypatch.setenv("BQ_TABLE_ID", "agent_events") + + def fake_text_input(label, value="", **kwargs): + if label == "Project ID": + return "user-injected-project" + return value + + with ( + mock.patch.object(app.st, "text_input", side_effect=fake_text_input), + mock.patch.object(app.st, "selectbox", return_value="1 GB"), + ): + refs, cap = app.sidebar_connection() + assert refs is not None + assert refs.project == "env-locked-project" + assert refs.dataset == "env_dataset" + assert refs.table == "agent_events" + assert cap == models.BYTES_CAPS["1 GB"] + + +def test_sidebar_connection_fallback_when_env_unset(monkeypatch): + monkeypatch.delenv("BQ_PROJECT_ID", raising=False) + monkeypatch.setenv("BQ_DATASET_ID", "env_dataset") + monkeypatch.setenv("BQ_TABLE_ID", "agent_events") + + def fake_text_input(label, value="", **kwargs): + if label == "Project ID": + return "form-project" + return value + + with ( + mock.patch.object(app.st, "text_input", side_effect=fake_text_input), + mock.patch.object(app.st, "selectbox", return_value="1 GB"), + ): + refs, cap = app.sidebar_connection() + assert refs is not None + assert refs.project == "form-project" + assert refs.dataset == "env_dataset" + assert cap == models.BYTES_CAPS["1 GB"] + + +def test_sidebar_connection_no_error_when_dataset_blank_without_connect( + monkeypatch, +): + """BQ_PROJECT_ID set in env, BQ_DATASET_ID blank, no Connect click -> no st.sidebar.error rendered.""" + monkeypatch.setenv("BQ_PROJECT_ID", "test-project") + monkeypatch.delenv("BQ_DATASET_ID", raising=False) + monkeypatch.setenv("BQ_TABLE_ID", "agent_events") + + def fake_text_input(label, value="", **kwargs): + return value + + state = {} + with ( + mock.patch.object(app.st, "text_input", side_effect=fake_text_input), + mock.patch.object(app.st, "selectbox", return_value="1 GB"), + mock.patch.object(app.st, "form_submit_button", return_value=False), + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st.sidebar, "error") as mock_error, + ): + refs, cap = app.sidebar_connection() + assert refs is None + mock_error.assert_not_called() + + +def test_row_llm_cost_calculation(sample_refs, sample_window): + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1024**3, + theme=models.LIGHT_THEME, + price_in=1.25, + price_out=5.00, + ) + mock_df = mock.MagicMock() + mock_df.empty = False + mock_row = { + "llm_calls": 10, + "prompt_tokens": 1_000_000, + "completion_tokens": 500_000, + "total_tokens": 1_500_000, + } + mock_df.iloc.__getitem__.return_value = mock_row + + def fake_columns(n): + return [ + mock.MagicMock() for _ in range(n if isinstance(n, int) else len(n)) + ] + + with ( + mock.patch.object(app, "fetch") as mock_fetch, + mock.patch.object(app.st, "columns", side_effect=fake_columns), + mock.patch.object(app, "_metric") as mock_metric, + ): + mock_fetch.return_value = models.QueryResult(mock_df) + app.row_llm(ctx) + # Check that _metric was called for Estimated cost with $3.75 + # (1_000_000 / 1e6 * 1.25) + (500_000 / 1e6 * 5.00) = 1.25 + 2.50 = 3.75 + calls = mock_metric.call_args_list + cost_call = next(c for c in calls if c.args[1] == "Estimated cost") + assert cost_call.args[2] == "$3.75" + + +def test_row_overview_data_path(sample_refs, sample_window): + """Verify row_overview formatting for numeric, None, and NaN metric values.""" + pd = pytest.importorskip("pandas") + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1000, + theme=models.LIGHT_THEME, + price_in=0.0, + price_out=0.0, + ) + + def fake_columns(n): + return [ + mock.MagicMock() for _ in range(n if isinstance(n, int) else len(n)) + ] + + # 1. Normal values: sessions=12, events=340, error_rate=0.0731, avg_llm_latency_ms=1234.5 + df_valid = pd.DataFrame( + [ + { + "sessions": 12, + "events": 340, + "error_rate": 0.0731, + "avg_llm_latency_ms": 1234.5, + } + ] + ) + metric_calls = [] + + def fake_metric(col, label, val, *args, **kwargs): + metric_calls.append((label, val)) + + def fake_fetch(sql, ctx, label): + if label == "Overview stats": + return models.QueryResult(df=current_df, error=None) + return models.QueryResult(df=pd.DataFrame(), error=None) + + current_df = df_valid + with ( + mock.patch.object(app, "fetch", side_effect=fake_fetch), + mock.patch.object(app.st, "columns", side_effect=fake_columns), + mock.patch.object(app, "_metric", side_effect=fake_metric), + mock.patch.object(app, "panel"), + mock.patch.object(app, "stacked_bars"), + mock.patch.object(app, "ranked_bars"), + ): + app.row_overview(ctx) + + assert metric_calls == [ + ("Sessions", "12"), + ("Events", "340"), + ("Error rate", "7.31%"), + ("Avg LLM latency", "1,235 ms"), + ] + + # 2. None values: error_rate=None, avg_llm_latency_ms=None -> "—" + df_none = pd.DataFrame( + [ + { + "sessions": 12, + "events": 340, + "error_rate": None, + "avg_llm_latency_ms": None, + } + ] + ) + current_df = df_none + metric_calls.clear() + with ( + mock.patch.object(app, "fetch", side_effect=fake_fetch), + mock.patch.object(app.st, "columns", side_effect=fake_columns), + mock.patch.object(app, "_metric", side_effect=fake_metric), + mock.patch.object(app, "panel"), + mock.patch.object(app, "stacked_bars"), + mock.patch.object(app, "ranked_bars"), + ): + app.row_overview(ctx) + + assert metric_calls == [ + ("Sessions", "12"), + ("Events", "340"), + ("Error rate", "—"), + ("Avg LLM latency", "—"), + ] + + # 3. NaN values: error_rate=NaN, avg_llm_latency_ms=NaN -> "—" + df_nan = pd.DataFrame( + [ + { + "sessions": 12, + "events": 340, + "error_rate": float("nan"), + "avg_llm_latency_ms": float("nan"), + } + ] + ) + current_df = df_nan + metric_calls.clear() + with ( + mock.patch.object(app, "fetch", side_effect=fake_fetch), + mock.patch.object(app.st, "columns", side_effect=fake_columns), + mock.patch.object(app, "_metric", side_effect=fake_metric), + mock.patch.object(app, "panel"), + mock.patch.object(app, "stacked_bars"), + mock.patch.object(app, "ranked_bars"), + ): + app.row_overview(ctx) + + assert metric_calls == [ + ("Sessions", "12"), + ("Events", "340"), + ("Error rate", "—"), + ("Avg LLM latency", "—"), + ] + + +def test_explain_actionable_advice(): + # DefaultCredentialsError + cred_err = queries.gauth_exc.DefaultCredentialsError( + "Could not automatically determine credentials." + ) + explained_creds = queries._explain(cred_err) + assert "Could not automatically determine credentials." in explained_creds + assert ( + "Run `gcloud auth application-default login` or set" + " `GOOGLE_APPLICATION_CREDENTIALS`." in explained_creds + ) + + # NotFound + not_found_err = queries.gexc.NotFound("Dataset not found") + explained_nf = queries._explain(not_found_err) + assert "Dataset not found" in explained_nf + assert "bq-agent-sdk views create-all" in explained_nf + + # Forbidden + forbidden_err = queries.gexc.Forbidden("Access denied") + explained_forb = queries._explain(forbidden_err) + assert "Access denied" in explained_forb + assert "roles/bigquery.jobUser" in explained_forb + + # bytesBilledLimitExceeded + limit_err = Exception("Query aborted: bytesBilledLimitExceeded error") + explained_limit = queries._explain(limit_err) + assert "bytesBilledLimitExceeded" in explained_limit + assert "Raise the per-query scan cap in the sidebar" in explained_limit + + # Generic fallback + generic_err = ValueError("Something unexpected") + assert queries._explain(generic_err) == "Something unexpected" + + +def test_run_query_default_credentials_error(): + with mock.patch.object( + queries, + "get_client", + side_effect=queries.gauth_exc.DefaultCredentialsError("No auth creds"), + ): + filters = models.Filters() + result = queries.run_query("SELECT 1", filters, "test-proj", 1024**3) + assert result.df.empty is True + assert result.error is not None + assert "No auth creds" in result.error + assert "gcloud auth application-default login" in result.error + assert result.bytes_processed == 0 + assert result.cache_hit is False + + +def test_run_query_cache_hit_deduplication(): + mock_df = mock.MagicMock() + mock_df.empty = False + mock_df_2 = mock.MagicMock() + mock_df_2.empty = False + + filters = models.Filters() + + # Simulate first execution returning run_id=101 and 5000 bytes + with mock.patch.object( + queries, + "_run_query_cached", + return_value=(mock_df, 5000, 5000, False, 101), + ): + # First run: new run_id + result1 = queries.run_query("SELECT 1", filters, "test-proj", 1024**3) + assert result1.df is mock_df + assert result1.error is None + assert result1.bytes_processed == 5000 + assert result1.bytes_billed == 5000 + assert result1.cache_hit is False + assert 101 in queries._SEEN_RUN_IDS + + # Second run: cached result from @st.cache_data returns the same + # run_id 101 + result2 = queries.run_query("SELECT 1", filters, "test-proj", 1024**3) + assert result2.df is mock_df + assert result2.error is None + assert result2.bytes_processed == 0 + assert result2.bytes_billed == 0 + assert result2.cache_hit is True + + # Third run: a different query invocation produces a new run_id=102 + with mock.patch.object( + queries, + "_run_query_cached", + return_value=(mock_df_2, 8000, 8000, False, 102), + ): + result3 = queries.run_query("SELECT 2", filters, "test-proj", 1024**3) + assert result3.df is mock_df_2 + assert result3.error is None + assert result3.bytes_processed == 8000 + assert result3.cache_hit is False + assert 102 in queries._SEEN_RUN_IDS + + +def test_run_query_cached_guardrail_and_execution(): + mock_client = mock.MagicMock() + mock_probe = mock.MagicMock() + mock_probe.total_bytes_processed = 2000 + mock_client.query.return_value = mock_probe + + filters = models.Filters() + + with mock.patch.object(queries, "get_client", return_value=mock_client): + # Probe exceeds max_bytes (1000): _run_query_cached raises RuntimeError + with pytest.raises(RuntimeError) as exc_info: + queries._run_query_cached( + "SELECT * FROM big_table", filters, "proj", 1000 + ) + assert "Guardrail: this query would scan" in str(exc_info.value) + + # Test run_query catches the error and returns guardrail in QueryResult + res = queries.run_query("SELECT * FROM big_table", filters, "proj", 1000) + assert res.df.empty is True + assert res.error is not None + assert "Guardrail: this query would scan" in res.error + assert res.bytes_processed == 0 + assert res.cache_hit is False + + # Normal query within max_bytes: returns 4-tuple on success + mock_probe.total_bytes_processed = 500 + mock_job = mock.MagicMock() + mock_job.total_bytes_processed = 500 + mock_job.total_bytes_billed = 10_485_760 + mock_job.cache_hit = False + expected_df = mock.MagicMock() + mock_job.to_dataframe.return_value = expected_df + mock_client.query.side_effect = [mock_probe, mock_job] + + df, bytes_proc, bytes_billed, hit, run_id = queries._run_query_cached( + "SELECT 1", filters, "proj", 1000 + ) + assert df is expected_df + assert bytes_proc == 500 + assert bytes_billed == 10_485_760 + assert hit is False + assert isinstance(run_id, int) + + +def test_run_query_clears_seen_run_ids_when_exceeding_max(monkeypatch): + queries._SEEN_RUN_IDS.clear() + monkeypatch.setattr(queries, "_MAX_SEEN_RUN_IDS", 3) + filters = models.Filters() + mock_df = mock.MagicMock() + + with mock.patch.object( + queries, + "_run_query_cached", + side_effect=[ + (mock_df, 100, 100, False, 1), + (mock_df, 100, 100, False, 2), + (mock_df, 100, 100, False, 3), + (mock_df, 100, 100, False, 4), + ], + ): + queries.run_query("Q1", filters, "proj", 1000) + queries.run_query("Q2", filters, "proj", 1000) + queries.run_query("Q3", filters, "proj", 1000) + assert queries._SEEN_RUN_IDS == {1, 2, 3} + + res = queries.run_query("Q4", filters, "proj", 1000) + assert queries._SEEN_RUN_IDS == {2, 3, 4} + assert res.bytes_processed == 100 + assert res.cache_hit is False + + +def test_job_config_caps_bytes_on_live_runs_only(): + """Tests that maximum_bytes_billed is only set on live runs, not dry runs.""" + filters = models.Filters() + live_cfg = queries.job_config("SELECT 1", filters, 5_000_000, dry_run=False) + assert live_cfg.maximum_bytes_billed == 5_000_000 + assert live_cfg.dry_run is False + + dry_cfg = queries.job_config("SELECT 1", filters, 5_000_000, dry_run=True) + assert dry_cfg.maximum_bytes_billed is None + assert dry_cfg.dry_run is True + + +def test_run_query_seen_run_ids_clearing_at_default_max(): + queries._SEEN_RUN_IDS.clear() + queries._SEEN_RUN_IDS.update(range(queries._MAX_SEEN_RUN_IDS)) + assert len(queries._SEEN_RUN_IDS) == queries._MAX_SEEN_RUN_IDS + mock_df = mock.MagicMock() + with mock.patch.object( + queries, + "_run_query_cached", + return_value=(mock_df, 100, 100, False, 99999), + ): + queries.run_query("Q", models.Filters(), "proj", 1000) + expected_len = ( + queries._MAX_SEEN_RUN_IDS - (queries._MAX_SEEN_RUN_IDS // 4) + 1 + ) + assert len(queries._SEEN_RUN_IDS) == expected_len + assert 99999 in queries._SEEN_RUN_IDS + + +def test_load_filter_options_returns_tuple_and_handles_error( + sample_refs, sample_window +): + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1000, + theme=models.LIGHT_THEME, + price_in=0.0, + price_out=0.0, + ) + mock_df = mock.MagicMock() + mock_df.empty = True + error_result = models.QueryResult( + df=mock_df, error="Access Denied: Dataset not found" + ) + with mock.patch.object(queries, "fetch", return_value=error_result): + options, result = queries.load_filter_options(ctx) + assert options == {} + assert result is error_result + assert result.error == "Access Denied: Dataset not found" + + +def test_seed_options_preserves_selection(): + all_sentinel = (models.ALL_SENTINEL,) + state = {"flt_agent": ["agent-1", "agent-2"], "flt_user_id": ["user-1"]} + with mock.patch.object(app.st, "session_state", state): + merged = app._seed_options("flt_agent", ["agent-1"], all_sentinel) + # agent-2 was absent from fetched options, but is preserved + assert merged == ["agent-1", "agent-2"] + + merged_empty = app._seed_options("flt_user_id", [], all_sentinel) + assert merged_empty == ["user-1"] + + merged_none = app._seed_options("flt_session_id", ["sess-1"], all_sentinel) + assert merged_none == ["sess-1"] + + +def test_seed_options_and_default_for_survive_widget_identity_change(): + all_sentinel = (models.ALL_SENTINEL,) + applied_agents = ("agent-x", "agent-y") + state = {} + with mock.patch.object(app.st, "session_state", state): + seeded = app._seed_options("flt_agent", ["agent-1"], applied_agents) + assert seeded == ["agent-1", "agent-x", "agent-y"] + + assert app._default_for(applied_agents) == ["agent-x", "agent-y"] + assert app._default_for(all_sentinel) == [] + + +def test_sidebar_filters_preserves_selection_and_accepts_custom_options(): + state = { + "flt_agent": ["agent-active"], + "flt_user_id": ["user-123"], + "flt_event_type": ["agent_start"], + "flt_session_id": ["sess-456"], + } + options: dict[str, list[str]] = {} + + passed_options: dict[str, list[str]] = {} + multiselect_kwargs: dict[str, dict] = {} + + def fake_multiselect(label, options, key, **kwargs): + passed_options[key] = options + multiselect_kwargs[key] = kwargs + return state.get(key, []) + + def fake_submit(*args, **kwargs): + if "on_click" in kwargs and callable(kwargs["on_click"]): + kwargs["on_click"]() + return True + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "sidebar", mock.MagicMock()), + mock.patch.object(app.st, "form_submit_button", side_effect=fake_submit), + mock.patch.object( + app.st, + "number_input", + side_effect=lambda *args, **kwargs: kwargs.get("value", 0.0), + ), + mock.patch.object(app.st, "multiselect", side_effect=fake_multiselect), + ): + filters, price_in, price_out = app.sidebar_filters(options) + assert (price_in, price_out) == (1.25, 5.00) + + # Active selections are preserved even when options query returned empty + assert state["flt_agent"] == ["agent-active"] + assert state["flt_user_id"] == ["user-123"] + assert state["flt_event_type"] == ["agent_start"] + assert state["flt_session_id"] == ["sess-456"] + assert filters.agents == ("agent-active",) + assert filters.user_ids == ("user-123",) + assert filters.event_types == ("agent_start",) + assert filters.session_ids == ("sess-456",) + + # Passed options to widgets contain the active selections + assert "agent-active" in passed_options["flt_agent"] + assert "user-123" in passed_options["flt_user_id"] + assert "sess-456" in passed_options["flt_session_id"] + + # Custom options enabled on agent, user, session + assert multiselect_kwargs["flt_agent"].get("accept_new_options") is True + assert multiselect_kwargs["flt_user_id"].get("accept_new_options") is True + assert ( + multiselect_kwargs["flt_session_id"].get("accept_new_options") is True + ) + assert not multiselect_kwargs["flt_event_type"].get("accept_new_options") + + +def test_sidebar_filters_returns_applied_filters_when_not_submitted(): + prior_filters = models.Filters(agents=("agent-saved",)) + state = { + "applied_filters": prior_filters, + "flt_agent": ["agent-unsubmitted"], + } + options = {"agent": ["agent-saved", "agent-unsubmitted"]} + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "sidebar", mock.MagicMock()), + mock.patch.object(app.st, "form_submit_button", return_value=False), + mock.patch.object( + app.st, + "number_input", + side_effect=lambda *args, **kwargs: kwargs.get("value", 0.0), + ), + mock.patch.object( + app.st, "multiselect", return_value=["agent-unsubmitted"] + ), + ): + filters, price_in, price_out = app.sidebar_filters(options) + assert filters == prior_filters + assert state["applied_filters"] == prior_filters + + +def test_custom_values_outside_1000_bounded_options_accepted_and_preserved(): + """Verifies custom values beyond the 1,000 option limit are accepted and bound.""" + # Simulate exactly 1,000 bounded options returned by BigQuery + bounded_agents = [f"agent-{i:04d}" for i in range(1000)] + bounded_users = [f"user-{i:04d}" for i in range(1000)] + bounded_sessions = [f"sess-{i:04d}" for i in range(1000)] + options = { + "agent": bounded_agents, + "user_id": bounded_users, + "event_type": ["agent_start", "agent_end"], + "session_id": bounded_sessions, + } + + custom_agent = "agent-custom-outside-1000" + custom_user = "user-custom-outside-1000" + custom_session = "sess-custom-outside-1000" + assert custom_agent not in bounded_agents + assert custom_user not in bounded_users + assert custom_session not in bounded_sessions + + state = { + "flt_agent": [custom_agent], + "flt_user_id": [custom_user], + "flt_event_type": [], + "flt_session_id": [custom_session], + } + + all_sentinel = (models.ALL_SENTINEL,) + + # Test _seed_options directly: 1000 options + 1 custom value = 1001 options + with mock.patch.object(app.st, "session_state", state): + merged_agents = app._seed_options("flt_agent", bounded_agents, all_sentinel) + assert len(merged_agents) == 1001 + assert merged_agents[-1] == custom_agent + + merged_users = app._seed_options("flt_user_id", bounded_users, all_sentinel) + assert len(merged_users) == 1001 + assert merged_users[-1] == custom_user + + merged_sessions = app._seed_options( + "flt_session_id", bounded_sessions, all_sentinel + ) + assert len(merged_sessions) == 1001 + assert merged_sessions[-1] == custom_session + + # Test sidebar_filters with mock st widgets + passed_options = {} + multiselect_kwargs = {} + + def fake_multiselect(label, options, key, **kwargs): + passed_options[key] = options + multiselect_kwargs[key] = kwargs + return state.get(key, []) + + def fake_submit(*args, **kwargs): + if "on_click" in kwargs and callable(kwargs["on_click"]): + kwargs["on_click"]() + return True + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "sidebar", mock.MagicMock()), + mock.patch.object(app.st, "form_submit_button", side_effect=fake_submit), + mock.patch.object( + app.st, + "number_input", + side_effect=lambda *args, **kwargs: kwargs.get("value", 0.0), + ), + mock.patch.object(app.st, "multiselect", side_effect=fake_multiselect), + ): + filters, price_in, price_out = app.sidebar_filters(options) + assert (price_in, price_out) == (1.25, 5.00) + + # 1. Custom values outside the 1000 options are preserved and accepted + assert filters.agents == (custom_agent,) + assert filters.user_ids == (custom_user,) + assert filters.session_ids == (custom_session,) + + # 2. Multiselect options contain 1001 options including the custom value + assert len(passed_options["flt_agent"]) == 1001 + assert custom_agent in passed_options["flt_agent"] + assert len(passed_options["flt_user_id"]) == 1001 + assert custom_user in passed_options["flt_user_id"] + assert len(passed_options["flt_session_id"]) == 1001 + assert custom_session in passed_options["flt_session_id"] + + # 3. Downstream query_parameters correctly binds the custom values outside 1000 + sql = "SELECT 1 WHERE agent IN UNNEST(@agents) AND user_id IN UNNEST(@user_ids) AND session_id IN UNNEST(@session_ids)" + params = queries.query_parameters(sql, filters) + param_map = {p.name: p.values for p in params} + assert param_map["agents"] == [custom_agent] + assert param_map["user_ids"] == [custom_user] + assert param_map["session_ids"] == [custom_session] + + +def test_load_module_resets_sys_path(tmp_path): + test_file = tmp_path / "dynamic_test_module.py" + test_file.write_text( + "import sys\n" + "from pathlib import Path\n" + "dir_was_in_path = str(Path(__file__).parent) in sys.path\n" + ) + dir_str = str(tmp_path) + assert dir_str not in sys.path + + mod_name = "test_dynamic_sys_path_cleanup" + try: + mod = _load_module(mod_name, test_file) + assert mod.dir_was_in_path is True + assert dir_str not in sys.path + finally: + sys.modules.pop(mod_name, None) + if dir_str in sys.path: + sys.path.remove(dir_str) + + +def test_load_module_resets_sys_path_on_error(tmp_path): + test_file = tmp_path / "failing_dynamic_module.py" + test_file.write_text("raise RuntimeError('intentional load failure')\n") + dir_str = str(tmp_path) + assert dir_str not in sys.path + + mod_name = "test_failing_dynamic_module" + try: + with pytest.raises(RuntimeError, match="intentional load failure"): + _load_module(mod_name, test_file) + assert dir_str not in sys.path + finally: + sys.modules.pop(mod_name, None) + if dir_str in sys.path: + sys.path.remove(dir_str) + + +def test_charts_color_map(): + state = {} + with mock.patch.object(charts.st, "session_state", state): + theme = models.LIGHT_THEME + # First call assigns slots stably + cm1 = charts.color_map("agent", ["agent_a", "agent_b"], theme) + assert cm1["agent_a"] == theme.categorical[0] + assert cm1["agent_b"] == theme.categorical[1] + assert cm1[models.OTHER_LABEL] == theme.muted + assert state["_slots::agent"] == {"agent_a": 0, "agent_b": 1} + assert models.OTHER_LABEL not in state["_slots::agent"] + + # Second call retains previously assigned slots + cm2 = charts.color_map("agent", ["agent_b", "agent_c"], theme) + assert cm2["agent_b"] == theme.categorical[1] + assert cm2["agent_c"] == theme.categorical[0] + assert cm2[models.OTHER_LABEL] == theme.muted + assert state["_slots::agent"] == { + "agent_a": 0, + "agent_b": 1, + "agent_c": 0, + } + assert models.OTHER_LABEL not in state["_slots::agent"] + + # Deduplication and OTHER_LABEL input handling in category list + cm_dedup = charts.color_map( + "dedup", ["x", "x", models.OTHER_LABEL, "y"], theme + ) + assert set(cm_dedup.keys()) == {"x", "y", models.OTHER_LABEL} + assert state["_slots::dedup"] == {"x": 0, "y": 1} + assert models.OTHER_LABEL not in state["_slots::dedup"] + + # Intra-chart collision resolution: + # pre-seed state["_slots::collision"] = {"k1": 0, "k2": 0} + state["_slots::collision"] = {"k1": 0, "k2": 0} + cm_coll = charts.color_map("collision", ["k1", "k2"], theme) + assert state["_slots::collision"]["k1"] == 0 + assert state["_slots::collision"]["k2"] == 1 + assert cm_coll["k1"] == theme.categorical[0] + assert cm_coll["k2"] == theme.categorical[1] + + # Slot reuse beyond categorical palette length with modulo wrapping + names = [f"name_{i}" for i in range(12)] + cm3 = charts.color_map("overflow", names, theme) + assert len(cm3) == 13 + for n in names: + assert cm3[n] in theme.categorical + assert state["_slots::overflow"]["name_0"] == 0 + assert state["_slots::overflow"]["name_8"] == 0 + assert state["_slots::overflow"]["name_11"] == 3 + assert models.OTHER_LABEL not in state["_slots::overflow"] + + +def test_charts_fold_others(): + # 1. Empty dataframe returned unmodified + df_empty = mock.MagicMock() + df_empty.empty = True + assert charts.fold_others(df_empty, "cat", "val") is df_empty + + # 2. Within limit: nunique <= limit returns unmodified + df_small = mock.MagicMock() + df_small.empty = False + df_small.__getitem__.return_value.nunique.return_value = 3 + assert charts.fold_others(df_small, "cat", "val", limit=5) is df_small + + # Exact boundary limit check (nunique == limit returns unmodified) + df_boundary = mock.MagicMock() + df_boundary.empty = False + df_boundary.__getitem__.return_value.nunique.return_value = 5 + assert charts.fold_others(df_boundary, "cat", "val", limit=5) is df_boundary + + # 3. Exceeding limit: nunique > limit executes folding logic + df_large = mock.MagicMock() + df_large.empty = False + df_large.__getitem__.return_value.nunique.return_value = 10 + res = charts.fold_others(df_large, "cat", "val", group_cols=["grp"], limit=5) + assert res is not df_large + + +def test_app_main_filter_preservation_on_error(sample_refs, sample_window): + state = {} + initial_options = {"agent": ["agent-1", "agent-2"]} + success_res = models.QueryResult(df=mock.MagicMock(), error=None) + error_res = models.QueryResult(df=mock.MagicMock(), error="BigQuery timeout") + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "set_page_config"), + mock.patch.object(app.st, "title"), + mock.patch.object(app.st, "tabs", return_value=[mock.MagicMock()] * 4), + mock.patch.object( + app, "sidebar_connection", return_value=(sample_refs, 1000) + ), + mock.patch.object(app, "sidebar_window", return_value=sample_window), + mock.patch.object(app, "active_theme", return_value=models.LIGHT_THEME), + mock.patch.object(app, "row_overview"), + mock.patch.object(app, "row_llm"), + mock.patch.object(app, "row_tools"), + mock.patch.object(app, "row_sessions"), + mock.patch.object(app, "footer"), + mock.patch.object(app, "sidebar_filters") as mock_sidebar_filters, + ): + mock_sidebar_filters.return_value = (models.Filters(), 0.0, 0.0) + + # First run succeeds: options stashed in session_state['_filter_options'] + with mock.patch.object( + app, "load_filter_options", return_value=(initial_options, success_res) + ): + app.main() + assert state["_filter_options"] == initial_options + mock_sidebar_filters.assert_called_with(initial_options) + + # Second run encounters error: stashed options retrieved + mock_sidebar_filters.reset_mock() + with mock.patch.object( + app, "load_filter_options", return_value=({}, error_res) + ): + app.main() + assert state["_filter_options"] == initial_options + mock_sidebar_filters.assert_called_with(initial_options) + + # Third run: query succeeds with empty options (e.g. empty time range) + # -> stashed options preserved + mock_sidebar_filters.reset_mock() + with mock.patch.object( + app, "load_filter_options", return_value=({}, success_res) + ): + app.main() + assert state["_filter_options"] == initial_options + mock_sidebar_filters.assert_called_with(initial_options) + + +def test_footer(sample_refs, sample_window): + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1024**3, + theme=models.LIGHT_THEME, + price_in=3.0, + price_out=15.0, + scan_log=[], + ) + with ( + mock.patch.object(app.st, "caption") as mock_caption, + mock.patch.object(app.st, "divider") as mock_divider, + ): + # Early return on empty scan_log + app.footer(ctx) + mock_caption.assert_not_called() + mock_divider.assert_not_called() + + # Populated scan_log renders caption with query count, billed total, + # processed total, cache count, per-query cap, and TTL text. + # Cached queries must be excluded from billed and processed totals. + ctx.scan_log = [ + ("q1", 10_485_760, 1024, False), + ("q2", 0, 2048, True), + ("q3", 10_485_760, 4096, False), + ] + app.footer(ctx) + mock_divider.assert_called_once() + assert mock_caption.call_count == 2 + caption_text = mock_caption.call_args_list[0][0][0] + caption_note = mock_caption.call_args_list[1][0][0] + # Assert billed and processed totals mathematically exclude cached query bytes + billed_bytes = 10_485_760 + 10_485_760 + processed_bytes = 1024 + 4096 + assert f"{models.humanize_bytes(billed_bytes)} billed" in caption_text + assert ( + f"({models.humanize_bytes(processed_bytes)} processed)" in caption_text + ) + assert "10 MB minimum" in caption_note + assert "compute/capacity reservations" in caption_note + # Assert query counts and cache hit text + assert "3 queries this run" in caption_text + assert "1 served from cache" in caption_text + # Assert per-query cap and TTL text + assert ( + f"per-query cap {models.humanize_bytes(ctx.max_bytes)}" in caption_text + ) + assert ( + f"results cached for {models.CACHE_TTL_SECONDS // 60} min" + in caption_text + ) + + # Assert 100% cache hit run makes clear that $0 billed is saved cost + ctx.scan_log = [ + ("q1", 0, 1024, True), + ("q2", 0, 2048, True), + ] + app.footer(ctx) + caption_cached = mock_caption.call_args_list[-2][0][0] + assert "0 B billed" in caption_cached + assert "billing cost saved via cache" in caption_cached + assert "2 served from cache ($0 billed)" in caption_cached + + +def test_fetch_records_scan_log_entry(sample_refs, sample_window): + fake_df = _MockDataFrame([{"count": 42}]) + query_result = models.QueryResult( + df=fake_df, + error=None, + bytes_processed=1024, + bytes_billed=10 * 1024 * 1024, + cache_hit=False, + ) + with mock.patch.object(queries, "run_query", return_value=query_result): + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1024**3, + theme=models.LIGHT_THEME, + price_in=3.0, + price_out=15.0, + ) + res = queries.fetch("SELECT 1", ctx, "Test Panel") + assert res.df is fake_df + assert len(ctx.scan_log) == 1 + label, billed, processed, cached = ctx.scan_log[0] + assert label == "Test Panel" + assert billed == 10 * 1024 * 1024 + assert processed == 1024 + assert cached is False + + +def test_sidebar_filters_case_a_custom_values_submission(): + """Case A: From ALL, enter custom-agent, custom-user, custom-session and click Apply once: + + all submitted arrays must be the custom values (not ['___ALL___']). + """ + state = {} + options = { + "agent": [], + "user_id": [], + "event_type": ["agent_start"], + "session_id": [], + } + + widget_inputs = { + "flt_agent": ["custom-agent"], + "flt_user_id": ["custom-user"], + "flt_event_type": [], + "flt_session_id": ["custom-session"], + } + + def fake_multiselect(label, options, key, **kwargs): + val = widget_inputs.get(key, []) + state[key] = val + return val + + def fake_submit(*args, **kwargs): + if "on_click" in kwargs and callable(kwargs["on_click"]): + kwargs["on_click"]() + return True + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "sidebar", mock.MagicMock()), + mock.patch.object(app.st, "form_submit_button", side_effect=fake_submit), + mock.patch.object( + app.st, + "number_input", + side_effect=lambda *args, **kwargs: kwargs.get("value", 0.0), + ), + mock.patch.object(app.st, "multiselect", side_effect=fake_multiselect), + ): + filters, price_in, price_out = app.sidebar_filters(options) + assert (price_in, price_out) == (1.25, 5.00) + + assert filters.agents == ("custom-agent",) + assert filters.user_ids == ("custom-user",) + assert filters.event_types == (models.ALL_SENTINEL,) + assert filters.session_ids == ("custom-session",) + + assert state["applied_filters"] == filters + assert state["flt_agent"] == ["custom-agent"] + assert state["flt_user_id"] == ["custom-user"] + assert state["flt_event_type"] == [] + assert state["flt_session_id"] == ["custom-session"] + + +def test_sidebar_filters_case_b_change_applied_to_new_values(): + """Case B: Apply alpha/u1/LLM_RESPONSE/s1, then change to existing beta/u2/TOOL_END/s2 + + and click Apply again: all arrays must commit the new values + (beta/u2/TOOL_END/s2). + """ + prior_filters = models.Filters( + agents=("alpha",), + user_ids=("u1",), + event_types=("LLM_RESPONSE",), + session_ids=("s1",), + ) + state = { + "applied_filters": prior_filters, + "flt_agent": ["alpha"], + "flt_user_id": ["u1"], + "flt_event_type": ["LLM_RESPONSE"], + "flt_session_id": ["s1"], + } + options = { + "agent": ["alpha", "beta"], + "user_id": ["u1", "u2"], + "event_type": ["LLM_RESPONSE", "TOOL_END"], + "session_id": ["s1", "s2"], + } + + widget_inputs = { + "flt_agent": ["beta"], + "flt_user_id": ["u2"], + "flt_event_type": ["TOOL_END"], + "flt_session_id": ["s2"], + } + + def fake_multiselect(label, options, key, **kwargs): + val = widget_inputs.get(key, []) + state[key] = val + return val + + def fake_submit(*args, **kwargs): + if "on_click" in kwargs and callable(kwargs["on_click"]): + kwargs["on_click"]() + return True + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "sidebar", mock.MagicMock()), + mock.patch.object(app.st, "form_submit_button", side_effect=fake_submit), + mock.patch.object( + app.st, + "number_input", + side_effect=lambda *args, **kwargs: kwargs.get("value", 0.0), + ), + mock.patch.object(app.st, "multiselect", side_effect=fake_multiselect), + ): + filters, price_in, price_out = app.sidebar_filters(options) + assert (price_in, price_out) == (1.25, 5.00) + + assert filters.agents == ("beta",) + assert filters.user_ids == ("u2",) + assert filters.event_types == ("TOOL_END",) + assert filters.session_ids == ("s2",) + + assert state["applied_filters"] == filters + assert state["flt_agent"] == ["beta"] + assert state["flt_user_id"] == ["u2"] + assert state["flt_event_type"] == ["TOOL_END"] + assert state["flt_session_id"] == ["s2"] + + +def test_sidebar_filters_case_c_clearing_selection_commits_all_sentinel(): + """Case C: Clearing a selection and clicking Apply must commit ALL_SENTINEL.""" + prior_filters = models.Filters( + agents=("beta",), + user_ids=("u2",), + event_types=("TOOL_END",), + session_ids=("s2",), + ) + state = { + "applied_filters": prior_filters, + "flt_agent": ["beta"], + "flt_user_id": ["u2"], + "flt_event_type": ["TOOL_END"], + "flt_session_id": ["s2"], + } + options = { + "agent": ["beta"], + "user_id": ["u2"], + "event_type": ["TOOL_END"], + "session_id": ["s2"], + } + + def fake_multiselect(label, options, key, **kwargs): + state[key] = [] + return [] + + def fake_submit(*args, **kwargs): + if "on_click" in kwargs and callable(kwargs["on_click"]): + kwargs["on_click"]() + return True + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "sidebar", mock.MagicMock()), + mock.patch.object(app.st, "form_submit_button", side_effect=fake_submit), + mock.patch.object( + app.st, + "number_input", + side_effect=lambda *args, **kwargs: kwargs.get("value", 0.0), + ), + mock.patch.object(app.st, "multiselect", side_effect=fake_multiselect), + ): + filters, price_in, price_out = app.sidebar_filters(options) + assert (price_in, price_out) == (1.25, 5.00) + + assert filters.agents == (models.ALL_SENTINEL,) + assert filters.user_ids == (models.ALL_SENTINEL,) + assert filters.event_types == (models.ALL_SENTINEL,) + assert filters.session_ids == (models.ALL_SENTINEL,) + + assert state["applied_filters"] == filters + assert state["flt_agent"] == [] + assert state["flt_user_id"] == [] + assert state["flt_event_type"] == [] + assert state["flt_session_id"] == [] + + +def test_reset_filters_on_tablerefs_change(sample_refs, sample_window): + """Verify filter states are reset when TableRefs changes (connection change).""" + old_refs = sample_refs + new_refs = models.TableRefs( + project="other-proj", + dataset="other_ds", + table="other_events", + view_prefix="other_", + ) + state = { + "_last_refs": old_refs, + "applied_filters": models.Filters(agents=("agent-x",)), + "_filter_options": {"agent": ["agent-x"]}, + "flt_agent": ["agent-x"], + "flt_user_id": ["user-1"], + "flt_event_type": ["start"], + "flt_session_id": ["sess-1"], + "_selected_session_id": "sess-1", + } + + success_res = models.QueryResult(df=mock.MagicMock(), error=None) + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "set_page_config"), + mock.patch.object(app.st, "title"), + mock.patch.object(app.st, "tabs", return_value=[mock.MagicMock()] * 4), + mock.patch.object( + app, "sidebar_connection", return_value=(new_refs, 1000) + ), + mock.patch.object(app, "sidebar_window", return_value=sample_window), + mock.patch.object(app, "active_theme", return_value=models.LIGHT_THEME), + mock.patch.object( + app, "load_filter_options", return_value=({}, success_res) + ), + mock.patch.object(app, "row_overview"), + mock.patch.object(app, "row_llm"), + mock.patch.object(app, "row_tools"), + mock.patch.object(app, "row_sessions"), + mock.patch.object(app, "footer"), + mock.patch.object(app, "sidebar_filters") as mock_sidebar_filters, + ): + mock_sidebar_filters.return_value = (models.Filters(), 0.0, 0.0) + app.main() + + assert state["_last_refs"] == new_refs + assert state["applied_filters"] == models.Filters() + assert "_filter_options" not in state + assert "flt_agent" not in state + assert "flt_user_id" not in state + assert "flt_event_type" not in state + assert "flt_session_id" not in state + assert "_selected_session_id" not in state + + +def test_row_sessions_maintains_selected_trace_on_refresh( + sample_refs, sample_window +): + """R4: Keep inspected trace selected when recent sessions refresh.""" + state = {"_selected_session_id": "sess-preserved"} + fake_df = mock.MagicMock() + fake_df.empty = False + fake_df.__getitem__.return_value.tolist.return_value = [ + "sess-newest", + "sess-preserved", + "sess-older", + ] + + sessions_result = models.QueryResult( + df=fake_df, + error=None, + bytes_processed=100, + bytes_billed=100, + cache_hit=False, + ) + + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1000, + theme=models.LIGHT_THEME, + price_in=0.0, + price_out=0.0, + ) + + selectbox_calls = [] + + def fake_selectbox(label, options, index=0, **kwargs): + selectbox_calls.append({"label": label, "options": options, "index": index}) + return options[index] + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "divider"), + mock.patch.object(app.st, "markdown"), + mock.patch.object(app.st, "caption"), + mock.patch.object(app.st, "dataframe"), + mock.patch.object( + app, + "fetch", + side_effect=[ + sessions_result, + models.QueryResult(df=_MockDataFrame(), error=None), + ], + ), + mock.patch.object(app.st, "selectbox", side_effect=fake_selectbox), + ): + app.row_sessions(ctx) + + # The selectbox should maintain selection of "sess-preserved" at index 1 instead of jumping to 0 + assert len(selectbox_calls) == 1 + call = selectbox_calls[0] + assert call["options"] == ["sess-newest", "sess-preserved", "sess-older"] + assert call["index"] == 1 + assert state["_selected_session_id"] == "sess-preserved" + + # If prev_chosen is not in ids, fall back to index 0 + state["_selected_session_id"] = "sess-not-present" + selectbox_calls.clear() + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "divider"), + mock.patch.object(app.st, "markdown"), + mock.patch.object(app.st, "caption"), + mock.patch.object(app.st, "dataframe"), + mock.patch.object( + app, + "fetch", + side_effect=[ + sessions_result, + models.QueryResult(df=_MockDataFrame(), error=None), + ], + ), + mock.patch.object(app.st, "selectbox", side_effect=fake_selectbox), + ): + app.row_sessions(ctx) + assert len(selectbox_calls) == 1 + assert selectbox_calls[0]["index"] == 0 + assert state["_selected_session_id"] == "sess-newest" + + +def test_row_sessions_non_empty_trace_table(sample_refs, sample_window): + """Verify st.dataframe is called with width='stretch', hide_index=True for non-empty trace table.""" + pd = pytest.importorskip("pandas") + state = {} + sessions_df = pd.DataFrame({"session_id": ["sess-1", "sess-2"]}) + trace_df = pd.DataFrame( + [ + { + "event_id": "evt-1", + "event_type": "tool_start", + "timestamp": "2026-01-01 00:00:00+00:00", + }, + { + "event_id": "evt-2", + "event_type": "tool_complete", + "timestamp": "2026-01-01 00:01:00+00:00", + }, + ] + ) + + sessions_result = models.QueryResult( + df=sessions_df, + error=None, + bytes_processed=100, + bytes_billed=100, + cache_hit=False, + ) + trace_result = models.QueryResult( + df=trace_df, + error=None, + bytes_processed=200, + bytes_billed=200, + cache_hit=False, + ) + + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1000, + theme=models.LIGHT_THEME, + price_in=0.0, + price_out=0.0, + ) + + with ( + mock.patch.object(app.st, "session_state", state), + mock.patch.object(app.st, "divider"), + mock.patch.object(app.st, "markdown"), + mock.patch.object(app.st, "caption") as mock_caption, + mock.patch.object(app.st, "dataframe") as mock_dataframe, + mock.patch.object( + app, + "fetch", + side_effect=[sessions_result, trace_result], + ), + mock.patch.object(app.st, "selectbox", return_value="sess-1"), + ): + app.row_sessions(ctx) + + mock_dataframe.assert_called_once_with( + trace_df, width="stretch", hide_index=True + ) + caption_calls = [c.args[0] for c in mock_caption.call_args_list] + assert any("read the session chronologically" in c for c in caption_calls) + assert not any("No events for this session" in c for c in caption_calls) + + +def test_billing_statistics_retained_on_dataframe_download_failure(): + """R5: Retain billing statistics after result-download failure.""" + mock_client = mock.MagicMock() + mock_probe = mock.MagicMock() + mock_probe.total_bytes_processed = 1000 + mock_client.query.return_value = mock_probe + + mock_job = mock.MagicMock() + mock_job.total_bytes_processed = 20480 + mock_job.total_bytes_billed = 10485760 + mock_job.cache_hit = False + mock_job.ended = dt.datetime.now(dt.timezone.utc) + # Simulate ServiceUnavailable or network drop during to_dataframe() + from google.api_core import exceptions as gexc + + mock_job.to_dataframe.side_effect = gexc.ServiceUnavailable( + "503 Service Unavailable" + ) + + mock_client.query.side_effect = [mock_probe, mock_job] + filters = models.Filters() + + with mock.patch.object(queries, "get_client", return_value=mock_client): + result = queries.run_query( + "SELECT * FROM events", filters, "proj", 100000000 + ) + + assert result.df.empty is True + assert result.error is not None + assert result.bytes_processed == 20480 + assert result.bytes_billed == 10485760 + assert result.cache_hit is False + + # Also test when cache_hit is True + mock_client2 = mock.MagicMock() + mock_probe2 = mock.MagicMock() + mock_probe2.total_bytes_processed = 1000 + mock_job2 = mock.MagicMock() + mock_job2.total_bytes_processed = 5000 + mock_job2.total_bytes_billed = 10485760 + mock_job2.cache_hit = True + mock_job2.ended = dt.datetime.now(dt.timezone.utc) + mock_job2.to_dataframe.side_effect = RuntimeError("Stream closed") + mock_client2.query.side_effect = [mock_probe2, mock_job2] + + with mock.patch.object(queries, "get_client", return_value=mock_client2): + result2 = queries.run_query( + "SELECT * FROM events", filters, "proj", 100000000 + ) + + assert result2.df.empty is True + assert result2.error is not None + assert result2.bytes_processed == 5000 + assert result2.bytes_billed == 0 # $0 billed if cache_hit + assert result2.cache_hit is True + + +def test_extract_job_stats_degraded_path_yields_runtime_error_and_zero_bytes(): + """Job with total_bytes_processed=None and ended=None yields RuntimeError and 0 bytes.""" + mock_client = mock.MagicMock() + mock_probe = mock.MagicMock() + mock_probe.total_bytes_processed = 1000 + + mock_job = mock.MagicMock() + mock_job.total_bytes_processed = None + mock_job.ended = None + mock_job.reload.side_effect = RuntimeError("job gone") + mock_job.to_dataframe.side_effect = RuntimeError( + "Query failed without job stats" + ) + + # Directly verify _extract_job_stats and _query_failure degraded behavior + assert queries._extract_job_stats(mock_job) is None + failure_err = queries._query_failure( + "Query failed without job stats", mock_job + ) + assert type(failure_err) is RuntimeError + assert not isinstance(failure_err, queries.QueryExecutionError) + + mock_client.query.side_effect = [mock_probe, mock_job] + filters = models.Filters() + + with mock.patch.object(queries, "get_client", return_value=mock_client): + result = queries.run_query( + "SELECT * FROM events", filters, "proj", 100000000 + ) + + assert result.df.empty is True + assert "Query failed without job stats" in result.error + assert result.bytes_processed == 0 + assert result.bytes_billed == 0 + assert result.cache_hit is False + + +def test_extract_job_stats_reloads_when_initially_none(): + """Job with total_bytes_processed initially None reloads to populate stats.""" + mock_job = mock.MagicMock() + mock_job.total_bytes_processed = None + mock_job.ended = None + mock_job.cache_hit = False + + def fake_reload(*args, **kwargs): + mock_job.total_bytes_processed = 2048 + mock_job.total_bytes_billed = 10485760 + + mock_job.reload.side_effect = fake_reload + + stats = queries._extract_job_stats(mock_job) + mock_job.reload.assert_called_once_with( + timeout=queries._JOB_RELOAD_TIMEOUT_SECONDS + ) + assert stats == (2048, 10485760, False) + + +def test_fetch_shows_loading_spinner(sample_refs, sample_window): + """R5: Wrap panel queries in fetch with st.spinner.""" + ctx = models.Context( + refs=sample_refs, + window=sample_window, + filters=models.Filters(), + max_bytes=1024**3, + theme=models.LIGHT_THEME, + price_in=3.0, + price_out=15.0, + ) + query_res = models.QueryResult( + df=_MockDataFrame([{"a": 1}]), + error=None, + bytes_processed=100, + bytes_billed=100, + cache_hit=False, + ) + + with ( + mock.patch.object(queries.st, "spinner") as mock_spinner, + mock.patch.object(queries, "run_query", return_value=query_res), + ): + queries.fetch("SELECT 1", ctx, "Overview totals") + mock_spinner.assert_called_once_with("Loading Overview totals...") diff --git a/tests/test_dashboards_streamlit_apptest.py b/tests/test_dashboards_streamlit_apptest.py new file mode 100644 index 00000000..c3c4fe14 --- /dev/null +++ b/tests/test_dashboards_streamlit_apptest.py @@ -0,0 +1,490 @@ +"""AppTest suite for the Streamlit dashboard. + +Uses streamlit.testing.v1.AppTest to execute and assert against the live app +flow without unconditional mocks of streamlit. +""" + +from __future__ import annotations + +from pathlib import Path +import sys +from unittest import mock + +import pandas as pd +import pytest + +pytest.importorskip("streamlit") + +from streamlit.testing.v1 import AppTest + +DASHBOARDS_DIR = ( + Path(__file__).resolve().parents[1] / "dashboards" / "streamlit" +) +if str(DASHBOARDS_DIR) not in sys.path: + sys.path.insert(0, str(DASHBOARDS_DIR)) + +import app +import models +import queries + +APP_PATH = DASHBOARDS_DIR / "app.py" + + +def _app_test() -> AppTest: + return AppTest.from_file(str(APP_PATH), default_timeout=30) + + +@pytest.fixture(autouse=True) +def mock_env(monkeypatch: pytest.MonkeyPatch): + """Seeds standard BigQuery connection environment variables.""" + monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS", raising=False) + monkeypatch.setenv("BQ_PROJECT_ID", "test-project") + monkeypatch.setenv("BQ_DATASET_ID", "test_dataset") + monkeypatch.setenv("BQ_TABLE_ID", "events") + monkeypatch.setenv("BQ_VIEW_PREFIX", "adk_") + + +@pytest.fixture +def mock_queries(): + """Stubs out BigQuery queries for fast and hermetic app testing.""" + filter_options = { + "agent": ["agent-a", "agent-b"], + "user_id": ["user-1", "user-2"], + "event_type": ["start", "complete"], + "session_id": ["sess-1", "sess-2", "sess-3"], + } + query_res = models.QueryResult( + df=pd.DataFrame(), + error=None, + bytes_processed=0, + bytes_billed=0, + cache_hit=True, + ) + with ( + mock.patch.object( + queries, + "load_filter_options", + return_value=(filter_options, query_res), + ), + mock.patch.object( + app, + "load_filter_options", + return_value=(filter_options, query_res), + create=True, + ), + mock.patch.object(queries, "fetch", return_value=query_res), + mock.patch.object(app, "fetch", return_value=query_res, create=True), + ): + yield + + +def test_sidebar_filters_apply(mock_queries): + """Test entering filters, clicking Apply filters, and verifying applied_filters.""" + at = _app_test() + at.run() + assert not at.exception + + # Select values in filter multiselects + at.sidebar.multiselect(key="flt_agent").select("agent-a") + at.sidebar.multiselect(key="flt_user_id").select("user-1") + + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert not at.exception + + applied: models.Filters = at.session_state["applied_filters"] + assert applied.agents == ("agent-a",) + assert applied.user_ids == ("user-1",) + assert applied.event_types == (models.ALL_SENTINEL,) + assert applied.session_ids == (models.ALL_SENTINEL,) + + +def test_sidebar_filters_successive_apply(mock_queries): + """Apply A, then Apply B, asserting B is applied.""" + at = _app_test() + at.run() + assert not at.exception + + # Apply A + at.sidebar.multiselect(key="flt_agent").select("agent-a") + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert at.session_state["applied_filters"].agents == ("agent-a",) + + # Apply B + at.sidebar.multiselect(key="flt_agent").set_value(["agent-b"]) + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert at.session_state["applied_filters"].agents == ("agent-b",) + + +def test_sidebar_filters_custom_id(mock_queries): + """Entering a custom ID via multiselect and applying it.""" + at = _app_test() + at.run() + assert not at.exception + + # Provide a custom agent ID not previously in options + at.sidebar.multiselect(key="flt_agent").set_value(["custom-agent-xyz"]) + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert not at.exception + + assert at.session_state["applied_filters"].agents == ("custom-agent-xyz",) + + +def test_filter_widgets_accept_new_options_flags(mock_queries): + """Verify accept_new_options is True for custom-input widgets and False for fixed options.""" + at = _app_test() + at.run() + assert not at.exception + + for key in ("flt_agent", "flt_user_id", "flt_session_id"): + assert at.sidebar.multiselect(key=key).proto.accept_new_options is True + assert ( + at.sidebar.multiselect(key="flt_event_type").proto.accept_new_options + is False + ) + + +def test_sidebar_filters_clear_selection(mock_queries): + """Clearing an existing selection and applying verifies ALL_SENTINEL is applied.""" + at = _app_test() + at.run() + assert not at.exception + + # Select an agent and apply + at.sidebar.multiselect(key="flt_agent").select("agent-a") + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert not at.exception + assert at.session_state["applied_filters"].agents == ("agent-a",) + + # Clear selection and apply + at.sidebar.multiselect(key="flt_agent").unselect("agent-a") + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert not at.exception + applied_filters: models.Filters = at.session_state["applied_filters"] + assert applied_filters.agents == (models.ALL_SENTINEL,) + + +def test_session_selectbox_preservation(): + """Verify session selectbox maintains selected trace when session options refresh.""" + sess_res1 = models.QueryResult( + df=pd.DataFrame({"session_id": ["sess-1", "sess-2", "sess-3"]}), + error=None, + bytes_processed=0, + bytes_billed=0, + cache_hit=True, + ) + sess_res2 = models.QueryResult( + df=pd.DataFrame({"session_id": ["sess-newest", "sess-2", "sess-older"]}), + error=None, + bytes_processed=0, + bytes_billed=0, + cache_hit=True, + ) + current_res = sess_res1 + + def fake_fetch(sql, ctx, label): + if label == "Recent sessions": + return current_res + return models.QueryResult( + df=pd.DataFrame(), + error=None, + bytes_processed=0, + bytes_billed=0, + cache_hit=True, + ) + + with ( + mock.patch.object( + queries, "load_filter_options", return_value=({}, sess_res1) + ), + mock.patch.object( + app, "load_filter_options", return_value=({}, sess_res1), create=True + ), + mock.patch.object(queries, "fetch", side_effect=fake_fetch), + mock.patch.object(app, "fetch", side_effect=fake_fetch, create=True), + ): + at = _app_test() + at.run() + assert not at.exception + + sess_box = [s for s in at.selectbox if s.label == "Session"][0] + assert sess_box.value == "sess-1" + + # Select sess-2 + sess_box.select("sess-2").run() + assert at.session_state["_selected_session_id"] == "sess-2" + + # Refresh sessions list: sess-2 is still present but now at index 1 instead of 0 + current_res = sess_res2 + at.run() + assert not at.exception + + sess_box_refreshed = [s for s in at.selectbox if s.label == "Session"][0] + assert sess_box_refreshed.value == "sess-2" + assert at.session_state["_selected_session_id"] == "sess-2" + + +def test_connection_change_resets_filters(): + """Changing TableRefs resets filter session state.""" + filter_options = {"agent": ["agent-a"]} + query_res = models.QueryResult( + df=pd.DataFrame(), + error=None, + bytes_processed=0, + bytes_billed=0, + cache_hit=True, + ) + recent_sessions_res = models.QueryResult( + df=pd.DataFrame({"session_id": ["sess-init-1"]}), + error=None, + bytes_processed=0, + bytes_billed=0, + cache_hit=True, + ) + + def fake_fetch(sql, ctx, label): + if label == "Recent sessions" and ctx.refs.dataset == "test_dataset": + return recent_sessions_res + return query_res + + with ( + mock.patch.object( + queries, + "load_filter_options", + return_value=(filter_options, query_res), + ), + mock.patch.object( + app, + "load_filter_options", + return_value=(filter_options, query_res), + create=True, + ), + mock.patch.object(queries, "fetch", side_effect=fake_fetch), + mock.patch.object(app, "fetch", side_effect=fake_fetch, create=True), + ): + at = _app_test() + at.run() + assert not at.exception + assert at.session_state["_selected_session_id"] == "sess-init-1" + + # Apply an agent filter + ms = at.sidebar.multiselect(key="flt_agent") + ms.select("agent-a") + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert at.session_state["applied_filters"].agents == ("agent-a",) + assert at.session_state["_selected_session_id"] == "sess-init-1" + + # Change dataset in connection form + ti_ds = [ti for ti in at.sidebar.text_input if ti.label == "Dataset ID"][0] + ti_ds.input("test_dataset_2") + connect_btn = [b for b in at.sidebar.button if b.label == "Connect"][0] + connect_btn.click().run() + assert not at.exception + + # Verify filter state has been reset to defaults + assert at.session_state["applied_filters"] == models.Filters() + assert at.session_state["flt_agent"] == [] + assert at.sidebar.multiselect(key="flt_agent").value == [] + assert "_selected_session_id" not in at.session_state + + +def test_run_query_cached_keys_on_every_argument(): + """Verify that changing any of the 4 arguments causes a cache miss, identical arguments hit cache.""" + queries._run_query_cached.clear() + try: + mock_client = mock.MagicMock() + mock_job = mock.MagicMock() + mock_job.total_bytes_processed = 100 + mock_job.total_bytes_billed = 200 + mock_job.cache_hit = False + mock_job.to_dataframe.return_value = pd.DataFrame({"col": [1]}) + mock_client.query.return_value = mock_job + + with mock.patch.object(queries, "get_client", return_value=mock_client): + sql_1 = "SELECT 1" + filters_1 = models.Filters(agents=("agent-1",)) + project_1 = "proj-1" + max_bytes_1 = 1_000_000 + + # Initial run: live execution (dry run probe + execution query) + res1 = queries._run_query_cached(sql_1, filters_1, project_1, max_bytes_1) + assert mock_client.query.call_count == 2 + run_id1 = res1[4] + + # Identical arguments: served from cache (no extra query calls, same run_id) + res2 = queries._run_query_cached(sql_1, filters_1, project_1, max_bytes_1) + assert mock_client.query.call_count == 2 + assert res2[4] == run_id1 + + # 1. Change sql: results in cache miss / live execution + sql_2 = "SELECT 2" + res_sql = queries._run_query_cached( + sql_2, filters_1, project_1, max_bytes_1 + ) + assert mock_client.query.call_count == 4 + assert res_sql[4] != run_id1 + + # 2. Change filters: results in cache miss / live execution + filters_2 = models.Filters(agents=("agent-2",)) + res_flt = queries._run_query_cached( + sql_1, filters_2, project_1, max_bytes_1 + ) + assert mock_client.query.call_count == 6 + assert res_flt[4] != run_id1 + + # 3. Change project_id: results in cache miss / live execution + project_2 = "proj-2" + res_prj = queries._run_query_cached( + sql_1, filters_1, project_2, max_bytes_1 + ) + assert mock_client.query.call_count == 8 + assert res_prj[4] != run_id1 + + # 4. Change maximum_bytes_billed: results in cache miss / live execution + max_bytes_2 = 2_000_000 + res_bytes = queries._run_query_cached( + sql_1, filters_1, project_1, max_bytes_2 + ) + assert mock_client.query.call_count == 10 + assert res_bytes[4] != run_id1 + + # Calling with identical arguments again is served from cache + res_cached = queries._run_query_cached( + sql_1, filters_1, project_1, max_bytes_2 + ) + assert mock_client.query.call_count == 10 + assert res_cached[4] == res_bytes[4] + finally: + queries._run_query_cached.clear() + + +def test_load_filter_options_groupby_fanout(): + """Builds a (kind, value) DataFrame and verifies exact fanout mapping.""" + df = pd.DataFrame( + [ + {"kind": "agent", "value": "agent-a"}, + {"kind": "agent", "value": "agent-b"}, + {"kind": "event_type", "value": "start"}, + {"kind": "event_type", "value": "complete"}, + {"kind": "session_id", "value": "sess-1"}, + {"kind": "session_id", "value": "sess-2"}, + {"kind": "user_id", "value": "user-1"}, + {"kind": "user_id", "value": "user-2"}, + ] + ) + query_res = models.QueryResult( + df=df, + error=None, + bytes_processed=50, + bytes_billed=100, + cache_hit=False, + ) + + refs = models.TableRefs( + project="test-project", + dataset="test_dataset", + table="events", + view_prefix="adk_", + ) + window = models.make_window(models.TIME_RANGES["Last 24 hours"]) + ctx = models.Context( + refs=refs, + window=window, + filters=models.Filters(), + max_bytes=1_000_000, + theme=models.LIGHT_THEME, + price_in=0.0, + price_out=0.0, + ) + + with mock.patch.object(queries, "fetch", return_value=query_res): + options, res = queries.load_filter_options(ctx) + assert options == { + "agent": ["agent-a", "agent-b"], + "event_type": ["start", "complete"], + "session_id": ["sess-1", "sess-2"], + "user_id": ["user-1", "user-2"], + } + assert res is query_res + + +def test_filter_widget_options_change_preserves_applied_filters(): + """Simulates changed options dictionary between AppTest runs, verifying retention and re-seeding.""" + opts_data = { + "opts": { + "agent": ["agent-a", "agent-b"], + "user_id": ["user-1", "user-2"], + "event_type": ["start", "complete"], + "session_id": ["sess-1", "sess-2"], + } + } + query_res = models.QueryResult( + df=pd.DataFrame(), + error=None, + bytes_processed=0, + bytes_billed=0, + cache_hit=True, + ) + + def fake_load(ctx): + return opts_data["opts"], query_res + + with ( + mock.patch.object(queries, "load_filter_options", side_effect=fake_load), + mock.patch.object( + app, "load_filter_options", side_effect=fake_load, create=True + ), + mock.patch.object(queries, "fetch", return_value=query_res), + mock.patch.object(app, "fetch", return_value=query_res, create=True), + ): + at = _app_test() + at.run() + assert not at.exception + + # Select and apply an initial agent filter + at.sidebar.multiselect(key="flt_agent").select("agent-a") + apply_btn = [b for b in at.sidebar.button if b.label == "Apply filters"][0] + apply_btn.click().run() + assert not at.exception + assert at.session_state["applied_filters"].agents == ("agent-a",) + + # 1. Simulate BigQuery options omitting the applied value "agent-a" + opts_data["opts"] = { + "agent": ["agent-c", "agent-d"], + "user_id": ["user-1", "user-2"], + "event_type": ["start", "complete"], + "session_id": ["sess-1", "sess-2"], + } + at.run() + assert not at.exception + + # applied_filters remains unchanged + assert at.session_state["applied_filters"].agents == ("agent-a",) + # The applied value "agent-a" is re-seeded into widget options + ms_agent = at.sidebar.multiselect(key="flt_agent") + assert "agent-a" in ms_agent.options + assert ms_agent.options == ["agent-c", "agent-d", "agent-a"] + assert ms_agent.value == ["agent-a"] + + # 2. Simulate BigQuery options growing + opts_data["opts"] = { + "agent": ["agent-c", "agent-d", "agent-e"], + "user_id": ["user-1", "user-2"], + "event_type": ["start", "complete"], + "session_id": ["sess-1", "sess-2"], + } + at.run() + assert not at.exception + + # applied_filters remains unchanged and new options are available alongside re-seeded applied value + assert at.session_state["applied_filters"].agents == ("agent-a",) + ms_agent_grown = at.sidebar.multiselect(key="flt_agent") + assert "agent-e" in ms_agent_grown.options + assert "agent-a" in ms_agent_grown.options + assert ms_agent_grown.value == ["agent-a"] diff --git a/tests/test_dashboards_streamlit_charts.py b/tests/test_dashboards_streamlit_charts.py new file mode 100644 index 00000000..1be08cd0 --- /dev/null +++ b/tests/test_dashboards_streamlit_charts.py @@ -0,0 +1,396 @@ +"""Unit tests for the Streamlit dashboard charts module. + +Tests chart construction, color mapping, and panel layout using real pandas +and plotly objects (no mock pandas/plotly). +""" + +from __future__ import annotations + +from pathlib import Path +import sys +from unittest import mock + +import pytest + +pytest.importorskip("streamlit") +pytest.importorskip("plotly") + +import pandas as pd +import plotly.graph_objects as go + +DASHBOARDS_DIR = ( + Path(__file__).resolve().parents[1] / "dashboards" / "streamlit" +) +if str(DASHBOARDS_DIR) not in sys.path: + sys.path.insert(0, str(DASHBOARDS_DIR)) + +import charts +import models + + +@pytest.fixture +def sample_context() -> models.Context: + """Returns a sample context for chart builder tests.""" + refs = models.TableRefs( + project="test-project", + dataset="test_dataset", + table="test_events", + view_prefix="test_", + ) + window = models.make_window(models.TIME_RANGES["Last 24 hours"]) + return models.Context( + refs=refs, + window=window, + filters=models.Filters(), + max_bytes=1_000_000, + theme=models.LIGHT_THEME, + price_in=1.25, + price_out=5.00, + ) + + +@pytest.fixture(autouse=True) +def clean_session_state(): + """Provides a clean dictionary for streamlit session_state during chart tests.""" + state: dict[str, object] = {} + with mock.patch.object(charts.st, "session_state", state): + yield state + + +def test_base_figure(): + """Test layout height, paper_bgcolor, xaxis/yaxis styling with light and dark themes.""" + # Light theme test + fig_light = charts.base_figure(models.LIGHT_THEME, height=400) + assert isinstance(fig_light, go.Figure) + assert fig_light.layout.height == 400 + assert fig_light.layout.paper_bgcolor == models.LIGHT_THEME.surface + assert fig_light.layout.plot_bgcolor == models.LIGHT_THEME.surface + assert fig_light.layout.xaxis.gridcolor == models.LIGHT_THEME.grid + assert fig_light.layout.xaxis.linecolor == models.LIGHT_THEME.axis + assert fig_light.layout.xaxis.showline is True + assert fig_light.layout.yaxis.showline is False + assert fig_light.layout.showlegend is False + + # Dark theme test + fig_dark = charts.base_figure(models.DARK_THEME, height=320) + assert isinstance(fig_dark, go.Figure) + assert fig_dark.layout.height == 320 + assert fig_dark.layout.paper_bgcolor == models.DARK_THEME.surface + assert fig_dark.layout.plot_bgcolor == models.DARK_THEME.surface + assert fig_dark.layout.xaxis.gridcolor == models.DARK_THEME.grid + assert fig_dark.layout.xaxis.linecolor == models.DARK_THEME.axis + assert fig_dark.layout.xaxis.showline is True + assert fig_dark.layout.yaxis.showline is False + + +def test_ranked_bars(sample_context): + """Test horizontal bar orientation, sorting, marker color, trace structure.""" + data = { + "agent": [f"agent_{i}" for i in range(15)], + "count": [i * 10 for i in range(15)], + } + df = pd.DataFrame(data) + + fig = charts.ranked_bars( + df, key="agent", value="count", ctx=sample_context, height=350, top=5 + ) + assert isinstance(fig, go.Figure) + assert len(fig.data) == 1 + + trace = fig.data[0] + assert trace.type == "bar" + assert trace.orientation == "h" + # top 5 largest items + assert len(trace.x) == 5 + # Ranked bars are sorted ascending by value for horizontal display + assert list(trace.x) == sorted(trace.x) + assert list(trace.x) == [100, 110, 120, 130, 140] + assert list(trace.y) == [ + "agent_10", + "agent_11", + "agent_12", + "agent_13", + "agent_14", + ] + assert trace.marker.color == sample_context.theme.categorical[0] + assert fig.layout.bargap == 0.35 + assert fig.layout.showlegend is False + + +def test_fold_others(): + """Test folding categories exceeding limit into 'Other', aggregation sum, and tie breaking.""" + # Case 1: More than limit categories + data = { + "cat": [f"cat_{i}" for i in range(10)], + "val": [10 * (i + 1) for i in range(10)], # 10, 20, ..., 100 + } + df = pd.DataFrame(data) + folded = charts.fold_others(df, key="cat", value="val", limit=5) + + # Should have 4 top categories + Other + assert folded["cat"].nunique() == 5 + assert models.OTHER_LABEL in folded["cat"].values + assert folded["val"].sum() == df["val"].sum() + + # The 4 largest are cat_9(100), cat_8(90), cat_7(80), cat_6(70) = 340 + # The remaining 6 are cat_0..cat_5 sum = 10+20+30+40+50+60 = 210 + other_val = folded[folded["cat"] == models.OTHER_LABEL]["val"].iloc[0] + assert other_val == 210 + + # Case 2: DataFrame with <= limit categories remains unchanged + df_small = pd.DataFrame({"cat": ["a", "b"], "val": [1, 2]}) + folded_small = charts.fold_others(df_small, key="cat", value="val", limit=5) + assert len(folded_small) == 2 + assert models.OTHER_LABEL not in folded_small["cat"].values + + # Case 3: Empty DataFrame + empty_df = pd.DataFrame(columns=["cat", "val"]) + assert charts.fold_others(empty_df, key="cat", value="val").empty + + # Case 4: Tie breaking + df_ties = pd.DataFrame( + { + "cat": [f"tie_{i}" for i in range(10)], + "val": [50] * 10, + } + ) + folded_ties = charts.fold_others(df_ties, key="cat", value="val", limit=4) + assert folded_ties["cat"].nunique() == 4 + assert models.OTHER_LABEL in folded_ties["cat"].values + assert folded_ties["val"].sum() == 500 + + # Case 5: Group columns preserved + df_grouped = pd.DataFrame( + { + "group": ["g1"] * 5 + ["g2"] * 5, + "cat": [f"c_{i}" for i in range(5)] * 2, + "val": [10, 20, 30, 40, 50] * 2, + } + ) + folded_grouped = charts.fold_others( + df_grouped, key="cat", value="val", group_cols=["group"], limit=3 + ) + assert "group" in folded_grouped.columns + assert folded_grouped["val"].sum() == df_grouped["val"].sum() + + +def test_fold_others_per_group_exact_values(): + """Folding happens within each group, not across the whole frame.""" + df = pd.DataFrame( + { + "bucket": ["b1", "b1", "b1", "b2", "b2", "b2"], + "cat": ["A", "B", "C", "A", "B", "C"], + "val": [10, 5, 2, 20, 10, 3], + } + ) + res = charts.fold_others(df, "cat", "val", group_cols=["bucket"], limit=2) + by_group = {(r["bucket"], r["cat"]): r["val"] for _, r in res.iterrows()} + assert by_group[("b1", "A")] == 10 + assert by_group[("b1", models.OTHER_LABEL)] == 7 # 5 + 2 + assert by_group[("b2", "A")] == 20 + assert by_group[("b2", models.OTHER_LABEL)] == 13 # 10 + 3 + + +def test_stacked_bars(sample_context): + """Test trace generation, categorical color mapping, stack barmode.""" + df = pd.DataFrame( + { + "bucket": ["2026-01-01", "2026-01-01", "2026-01-02", "2026-01-02"], + "agent": ["agent_a", "agent_b", "agent_a", "agent_b"], + "count": [10, 20, 30, 40], + } + ) + + fig = charts.stacked_bars( + df, + x="bucket", + key="agent", + value="count", + ctx=sample_context, + domain="agent", + height=300, + ) + + assert isinstance(fig, go.Figure) + assert len(fig.data) == 2 + names = [trace.name for trace in fig.data] + assert "agent_a" in names + assert "agent_b" in names + assert fig.layout.barmode == "stack" + assert fig.layout.bargap == 0.25 + assert fig.layout.hovermode == "x unified" + assert fig.layout.showlegend is True + + # Verify colors are from theme categorical + colors = [trace.marker.color for trace in fig.data] + assert colors[0] in sample_context.theme.categorical + assert colors[1] in sample_context.theme.categorical + assert colors[0] != colors[1] + + +def test_grouped_bars(sample_context): + """Test grouped traces and labels.""" + df = pd.DataFrame( + { + "tool": ["search", "calculator", "browser"], + "p50": [100.0, 20.0, 500.0], + "p99": [250.0, 45.0, 1200.0], + } + ) + + fig = charts.grouped_bars( + df, + key="tool", + series=[("p50", "Median Latency"), ("p99", "P99 Latency")], + ctx=sample_context, + domain="tool_lat", + height=320, + top=10, + ) + + assert isinstance(fig, go.Figure) + assert len(fig.data) == 2 + assert fig.data[0].name == "Median Latency" + assert fig.data[1].name == "P99 Latency" + assert fig.data[0].orientation == "h" + assert fig.data[1].orientation == "h" + assert fig.layout.barmode == "group" + assert fig.layout.bargap == 0.3 + assert fig.layout.bargroupgap == 0.08 + assert fig.layout.showlegend is True + + +def test_lines(sample_context): + """Test line traces, connectgaps, annotations.""" + df = pd.DataFrame( + { + "bucket": [ + "2026-01-01 00:00", + "2026-01-01 01:00", + "2026-01-01 02:00", + ], + "prompt": [100, 150, 200], + "completion": [50, 75, 100], + } + ) + + fig = charts.lines( + df, + x="bucket", + series=[("prompt", "Prompt Tokens"), ("completion", "Completion Tokens")], + ctx=sample_context, + domain="tokens", + height=320, + unit=" tokens", + ) + + assert isinstance(fig, go.Figure) + assert len(fig.data) == 2 + + for trace in fig.data: + assert trace.mode == "lines" + assert trace.connectgaps is False + assert "tokens" in trace.hovertemplate + + # For series count <= 4, end direct annotations are added + assert len(fig.layout.annotations) == 2 + annotation_texts = [ann.text for ann in fig.layout.annotations] + assert "Prompt Tokens" in annotation_texts + assert "Completion Tokens" in annotation_texts + assert fig.layout.showlegend is True + assert fig.layout.hovermode == "x unified" + + +def test_lines_more_than_four_series_shows_legend(sample_context): + """Verify lines() with > 4 series suppresses direct annotations and enables showlegend.""" + df = pd.DataFrame( + { + "bucket": ["2026-01-01", "2026-01-02", "2026-01-03"], + "s1": [10, 20, 30], + "s2": [15, 25, 35], + "s3": [20, 30, 40], + "s4": [25, 35, 45], + "s5": [30, 40, 50], + } + ) + series = [ + ("s1", "Series 1"), + ("s2", "Series 2"), + ("s3", "Series 3"), + ("s4", "Series 4"), + ("s5", "Series 5"), + ] + + fig = charts.lines( + df, + x="bucket", + series=series, + ctx=sample_context, + domain="test_domain", + height=320, + ) + + assert isinstance(fig, go.Figure) + assert len(fig.data) == 5 + assert fig.layout.showlegend is True + assert fig.layout.hovermode == "x unified" + # For series count > 4, direct labels (annotations) are omitted to avoid crowding + assert len(fig.layout.annotations) == 0 + # Default right margin is kept (8) instead of the expanded label margin (96) + assert fig.layout.margin.r == 8 + + +def test_panel(): + """Test empty DataFrame caption vs chart and expander.""" + # Case 1: Empty DataFrame shows caption + with ( + mock.patch.object(charts.st, "markdown") as mock_markdown, + mock.patch.object(charts.st, "caption") as mock_caption, + mock.patch.object(charts.st, "plotly_chart") as mock_plotly, + mock.patch.object(charts.st, "expander") as mock_expander, + mock.patch.object(charts.st, "dataframe") as mock_dataframe, + ): + charts.panel("Empty Panel", None, pd.DataFrame(), empty="No data here.") + mock_markdown.assert_called_once_with("**Empty Panel**") + mock_caption.assert_called_once_with("No data here.") + mock_plotly.assert_not_called() + mock_expander.assert_not_called() + mock_dataframe.assert_not_called() + + # Case 2: Populated DataFrame with figure shows chart and collapsed expander + fig = go.Figure() + df = pd.DataFrame([{"a": 1, "b": 2}]) + + with ( + mock.patch.object(charts.st, "markdown") as mock_markdown, + mock.patch.object(charts.st, "caption") as mock_caption, + mock.patch.object(charts.st, "plotly_chart") as mock_plotly, + mock.patch.object(charts.st, "expander") as mock_expander, + mock.patch.object(charts.st, "dataframe") as mock_dataframe, + ): + mock_expander.return_value.__enter__ = mock.MagicMock() + mock_expander.return_value.__exit__ = mock.MagicMock() + + charts.panel("Active Panel", fig, df, key="chart_key") + mock_markdown.assert_called_once_with("**Active Panel**") + mock_caption.assert_not_called() + mock_plotly.assert_called_once_with(fig, width="stretch", key="chart_key") + mock_expander.assert_called_once_with("Table view", expanded=False) + mock_dataframe.assert_called_once_with(df, width="stretch", hide_index=True) + + # Case 3: Populated DataFrame with fig=None shows expanded expander + with ( + mock.patch.object(charts.st, "markdown") as mock_markdown, + mock.patch.object(charts.st, "caption") as mock_caption, + mock.patch.object(charts.st, "plotly_chart") as mock_plotly, + mock.patch.object(charts.st, "expander") as mock_expander, + mock.patch.object(charts.st, "dataframe") as mock_dataframe, + ): + mock_expander.return_value.__enter__ = mock.MagicMock() + mock_expander.return_value.__exit__ = mock.MagicMock() + + charts.panel("Table Only Panel", None, df) + mock_markdown.assert_called_once_with("**Table Only Panel**") + mock_plotly.assert_not_called() + mock_expander.assert_called_once_with("Table view", expanded=True) + mock_dataframe.assert_called_once_with(df, width="stretch", hide_index=True)