yau-plant-assistant/api/tools/metrics.py

"""Cube client. The only path to plant history.

The agent never writes SQL against the historian. It builds a Cube query
object, guardrails caps it, Cube generates the SQL and answers from a
pre-aggregation where one exists. Three reasons, in order:

  * imh is a live system. Pre-aggregations in pg-ai keep "count alarms last
    week" off it entirely.
  * The definitions that make an answer right - what counts as an alarm, what
    "last week" means, what a pump-down is - live in the model files where a
    person can read and check them, not inside a generated string.
  * A query object can be validated. Generated SQL can only be inspected.

Timezone: storage is UTC and Cube converts once, using SITE_TIMEZONE. Never
convert here and never in a prompt.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any
from zoneinfo import ZoneInfo

import httpx
import jwt

from config import settings
from guardrails import check_cube_query

log = logging.getLogger("tools.metrics")


@dataclass
class MetricResult:
    query: dict[str, Any]
    rows: list[dict[str, Any]]
    row_count: int
    used_fixture_data: bool
    time_window: dict[str, str]
    annotation: dict[str, Any]


def _token() -> str:
    cfg = settings()
    return jwt.encode({"iss": "ai-api"}, cfg.cubejs_api_secret, algorithm="HS256")


def rolling_window(days: int) -> tuple[str, str, str]:
    """A rolling N x 24 h window in site local time, as Cube date strings.

    "Last week" means the rolling seven days, NOT the previous calendar week
    and NOT seven calendar days. Whatever it means, the answer states it - the
    third element is the description that goes into the response.
    """
    cfg = settings()
    tz = ZoneInfo(cfg.site_timezone)
    end = datetime.now(timezone.utc).astimezone(tz)
    start = end - timedelta(days=days)
    fmt = "%Y-%m-%dT%H:%M:%S"
    return (
        start.strftime(fmt),
        end.strftime(fmt),
        f"rolling {days} days to {end.strftime('%Y-%m-%d %H:%M')} {end.tzname()}",
    )


def run(query: dict[str, Any], *, trace=None) -> MetricResult:
    """Execute a Cube query. Raises GuardrailViolation, before it runs, on failure."""
    cfg = settings()
    capped = check_cube_query(
        query, max_rows=cfg.max_rows_returned, site_timezone=cfg.site_timezone
    )

    response = httpx.post(
        f"{cfg.cubejs_api_url}/load",
        json={"query": capped},
        headers={"Authorization": _token()},
        timeout=cfg.query_timeout_seconds,
    )
    response.raise_for_status()
    body = response.json()
    rows = body.get("data", [])

    window = capped["timeDimensions"][0]["dateRange"]
    result = MetricResult(
        query=capped,
        rows=rows,
        row_count=len(rows),
        # Anything sourced from the fixture schema is generated test data. The
        # flag rides all the way to the operator's screen.
        used_fixture_data=cfg.use_fixtures,
        time_window={
            "start": window[0] if isinstance(window, list) else str(window),
            "end": window[1] if isinstance(window, list) else str(window),
            # The timezone the query actually ran in, not the one we would
            # like it to have run in. check_cube_query pins it onto the query
            # itself, so these two can no longer disagree.
            "timezone": capped.get("timezone", cfg.site_timezone),
        },
        annotation=body.get("annotation", {}),
    )

    if trace is not None:
        try:
            trace.event(
                name="cube_query",
                metadata={
                    "query": capped,
                    "row_count": result.row_count,
                    "used_fixture_data": result.used_fixture_data,
                    # Whether a pre-aggregation served this. If it says false
                    # on imh, the Phase 5 gate has regressed and imh is being
                    # scanned - investigate before shipping the answer.
                    "pre_aggregation": body.get("usedPreAggregations", {}),
                },
            )
        except Exception:
            log.exception("failed to record Cube query in Langfuse")

    return result


# --- Query builders ---------------------------------------------------------
# Prebuilt shapes for the questions the demo actually asks. A builder is easier
# to check than a model-generated query object, and the ones below encode the
# definitions from the Cube models rather than restating them.


def alarm_count(
    *, equipment_id: str | None = None, alarm_type: str | None = None, days: int = 7
) -> dict[str, Any]:
    """Activations of an alarm over a rolling window.

    state = ACTIVE only, enforced inside the measure itself
    (cube/model/alarms.yml), not here - so a caller cannot forget it.
    """
    start, end, _ = rolling_window(days)
    filters = []
    if equipment_id:
        filters.append(
            {"member": "alarm_activity.equipment_equipment_id",
             "operator": "equals", "values": [equipment_id]}
        )
    if alarm_type:
        filters.append(
            {"member": "alarm_activity.alarm_type",
             "operator": "equals", "values": [alarm_type]}
        )
    return {
        "measures": ["alarm_activity.alarm_count"],
        "dimensions": ["alarm_activity.alarm_type"],
        "timeDimensions": [
            {"dimension": "alarm_activity.event_time", "dateRange": [start, end]}
        ],
        "filters": filters,
        "order": {"alarm_activity.alarm_count": "desc"},
    }


def alarm_detail(*, equipment_id: str | None = None, days: int = 7) -> dict[str, Any]:
    """The individual activations behind a count, so the answer can show them."""
    start, end, _ = rolling_window(days)
    filters = (
        [{"member": "alarm_activity.equipment_equipment_id",
          "operator": "equals", "values": [equipment_id]}]
        if equipment_id
        else []
    )
    return {
        "dimensions": [
            "alarm_activity.event_time",
            "alarm_activity.alarm_type",
            "alarm_activity.tag_id",
            "alarm_activity.value",
            "alarm_activity.priority",
        ],
        "timeDimensions": [
            {"dimension": "alarm_activity.event_time", "dateRange": [start, end]}
        ],
        "filters": filters + [
            {"member": "alarm_activity.state", "operator": "equals",
             "values": ["ACTIVE"]}
        ],
        "order": {"alarm_activity.event_time": "asc"},
        "limit": 200,
    }


def pump_down_evidence(*, days: int = 30) -> dict[str, Any]:
    """The evidence behind an advisory question about discharge rate.

    Rates actually used, how high the well got, how often it alarmed, how often
    it spilled - and the sample size, so a rate is never quoted without its
    denominator. This returns evidence. It does not return a recommendation,
    and AdvisoryAnswer rejects the response if one appears in the prose.
    """
    start, end, _ = rolling_window(days)
    return {
        "measures": [
            "operations.pump_down_count",
            "operations.avg_discharge_rate",
            "operations.min_discharge_rate",
            "operations.max_discharge_rate",
            "operations.avg_inflow_rate",
            "operations.max_level_reached",
            "operations.avg_max_level",
            "operations.high_alarm_count",
            "operations.high_alarm_rate",
            "operations.spill_count",
        ],
        "dimensions": ["operations.peak_pumps_running"],
        "timeDimensions": [
            {"dimension": "operations.start_time", "dateRange": [start, end]}
        ],
        "order": {"operations.peak_pumps_running": "asc"},
    }


def level_profile(*, days: int = 7, granularity: str = "hour") -> dict[str, Any]:
    """Wet well level over time. Percent of the weir crest, not millimetres."""
    start, end, _ = rolling_window(days)
    return {
        "measures": [
            "process_values.avg_value",
            "process_values.max_value",
            "process_values.sample_count",
        ],
        "timeDimensions": [
            {
                "dimension": "process_values.sample_time",
                "dateRange": [start, end],
                "granularity": granularity,
            }
        ],
        "filters": [
            {"member": "process_values.tag_id", "operator": "equals",
             "values": ["PS_STN_WET_WELL_LEVEL"]}
        ],
    }