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")


# How far back the historian goes. Every WRPS history group on CI Server
# carries LIFE_TIME "1 weeks" (05-scada/modbus/export_his_group.qli), so there
# is nothing older than this to find — on imh or on the fixtures standing in
# for it.
#
# THIS IS NOT A TUNING KNOB. A question reaching past it returns zero rows, and
# zero rows outside retention means "the historian does not go back that far",
# NOT "nothing happened". Those are different answers and only one of them is
# true. `outside_retention` on the result below is what lets the answer say so;
# reporting the second when the first is the case would be an answer outside
# the evidence, which is the third line this system does not cross.
HISTORY_RETENTION_DAYS = 7


@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]
    # True when the window asked for reaches past what the historian keeps.
    # An empty result with this set must be reported as a retention limit.
    outside_retention: bool = False


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"
    note = f"rolling {days} days to {end.strftime('%Y-%m-%d %H:%M')} {end.tzname()}"
    if days > HISTORY_RETENTION_DAYS:
        # Deliberately NOT clamped. A silently shortened window would answer a
        # question nobody asked and read as though it had answered the one they
        # did. Let it run, return nothing, and let outside_retention say why.
        note += (
            f" - BEYOND RETENTION: the historian keeps "
            f"{HISTORY_RETENTION_DAYS} days, so part of this window does not exist"
        )
        log.warning("window of %d days exceeds the %d day historian retention",
                    days, HISTORY_RETENTION_DAYS)
    return start.strftime(fmt), end.strftime(fmt), note


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"]

    # Did the window reach past what the historian holds? Compared against the
    # window actually run, not the days argument, so a caller passing explicit
    # dates is checked the same way as one asking for a rolling window.
    outside_retention = False
    if isinstance(window, list) and window:
        try:
            asked_from = datetime.fromisoformat(window[0])
            if asked_from.tzinfo is None:
                asked_from = asked_from.replace(tzinfo=ZoneInfo(cfg.site_timezone))
            horizon = datetime.now(timezone.utc) - timedelta(days=HISTORY_RETENTION_DAYS)
            outside_retention = asked_from < horizon
        except ValueError:
            log.warning("could not parse window start %r for a retention check", window[0])

    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),
            "retention_days": str(HISTORY_RETENTION_DAYS),
        },
        annotation=body.get("annotation", {}),
        outside_retention=outside_retention,
    )

    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 = HISTORY_RETENTION_DAYS) -> 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": [
            # The CI Server ITEM name, which is what the historian is keyed on.
            # This used to be PS_STN_WET_WELL_LEVEL - a SCADA POINT name, one
            # layer up - and it matched nothing on a tag-level lookup. See
            # cube/model/process_values.yml for the four namespaces involved.
            {"member": "process_values.item_name", "operator": "equals",
             "values": ["AID.WRPS.STN.LEVEL"]}
        ],
    }