|
| 1 | +"""Monitor forecast computation shared by the monitors dashboard and the MCP |
| 2 | +monitor tools. |
| 3 | +
|
| 4 | +The forecast a monitor displays depends on its threshold mode and type: |
| 5 | +
|
| 6 | +* Volume / Metric in Prediction Model mode: a value band extending forward. For |
| 7 | + monitors coupled to a Freshness monitor the band holds at the baseline until |
| 8 | + the next expected refresh, then steps to the forecast value (``gated_forecast_prediction``); |
| 9 | + otherwise it is the raw per-step prediction band. |
| 10 | +* Freshness in Prediction Model mode: a predicted next-update *time window* |
| 11 | + rather than a value band (``next_update_window``). |
| 12 | +
|
| 13 | +These functions are intentionally free of any UI dependency so both surfaces |
| 14 | +render the same forecast from the same source of truth. |
| 15 | +""" |
| 16 | + |
| 17 | +from datetime import UTC, date, datetime |
| 18 | + |
| 19 | +import pandas as pd |
| 20 | + |
| 21 | +from testgen.common.freshness_service import add_business_minutes, get_schedule_params, resolve_holiday_dates |
| 22 | +from testgen.common.models.test_definition import MonitorForecastPoint, TestDefinition, TestDefinitionSummary |
| 23 | +from testgen.common.models.test_suite import TestSuite |
| 24 | + |
| 25 | +# A monitor definition carrying the forecast-relevant fields — satisfied by both |
| 26 | +# the ORM row and the read-only summary dataclass. |
| 27 | +MonitorDefinition = TestDefinition | TestDefinitionSummary |
| 28 | + |
| 29 | + |
| 30 | +def resolve_suite_holiday_dates(test_suite: TestSuite) -> set[date] | None: |
| 31 | + """Holiday dates excluded from the suite's schedule, over a window spanning |
| 32 | + recent history through the forecast horizon. ``None`` when no calendars are set.""" |
| 33 | + if not test_suite.holiday_codes_list: |
| 34 | + return None |
| 35 | + now = pd.Timestamp.now("UTC") |
| 36 | + idx = pd.DatetimeIndex([now - pd.Timedelta(days=7), now + pd.Timedelta(days=30)]) |
| 37 | + return resolve_holiday_dates(test_suite.holiday_codes_list, idx) |
| 38 | + |
| 39 | + |
| 40 | +def next_update_window( |
| 41 | + freshness_definition: MonitorDefinition | None, |
| 42 | + last_detection_time: datetime | None, |
| 43 | + *, |
| 44 | + exclude_weekends: bool, |
| 45 | + holiday_dates: set[date] | None, |
| 46 | + cron_tz: str | None, |
| 47 | +) -> dict | None: |
| 48 | + """Predicted next-update window as ``{"start", "end"}`` epoch-ms, or ``None``. |
| 49 | +
|
| 50 | + The schedule-derived business-time interval from the last detected update out |
| 51 | + to the lower/upper staleness tolerance. ``start`` is ``None`` when only an |
| 52 | + upper tolerance is configured (the update is expected by ``end``). Returns |
| 53 | + ``None`` unless the Freshness monitor is in Prediction Model mode with a |
| 54 | + trained schedule and at least one observed update. |
| 55 | + """ |
| 56 | + if ( |
| 57 | + freshness_definition is None |
| 58 | + or freshness_definition.history_calculation != "PREDICT" |
| 59 | + or (freshness_definition.prediction and not freshness_definition.prediction.get("schedule_stage")) |
| 60 | + or freshness_definition.upper_tolerance is None |
| 61 | + or last_detection_time is None |
| 62 | + ): |
| 63 | + return None |
| 64 | + |
| 65 | + tz = cron_tz or "UTC" |
| 66 | + sched = get_schedule_params(freshness_definition.prediction) |
| 67 | + |
| 68 | + window_end = add_business_minutes( |
| 69 | + pd.Timestamp(last_detection_time), |
| 70 | + float(freshness_definition.upper_tolerance), |
| 71 | + exclude_weekends, |
| 72 | + holiday_dates, tz, |
| 73 | + excluded_days=sched.excluded_days, |
| 74 | + ) |
| 75 | + window_start = None |
| 76 | + if lower_minutes := (float(freshness_definition.lower_tolerance) if freshness_definition.lower_tolerance else None): |
| 77 | + window_start = add_business_minutes( |
| 78 | + pd.Timestamp(last_detection_time), |
| 79 | + lower_minutes, |
| 80 | + exclude_weekends, |
| 81 | + holiday_dates, tz, |
| 82 | + excluded_days=sched.excluded_days, |
| 83 | + ) |
| 84 | + |
| 85 | + return { |
| 86 | + "start": int(window_start.timestamp() * 1000) if window_start else None, |
| 87 | + "end": int(window_end.timestamp() * 1000), |
| 88 | + } |
| 89 | + |
| 90 | + |
| 91 | +def gated_forecast_prediction( |
| 92 | + definition: MonitorDefinition, |
| 93 | + freshness_window: dict | None, |
| 94 | + last_run_time: datetime | None, |
| 95 | +) -> dict | None: |
| 96 | + """Coupled forecast band for a Volume/Metric monitor whose value holds at its |
| 97 | + baseline between refreshes. |
| 98 | +
|
| 99 | + A flat baseline line from the latest run up to the predicted next-update |
| 100 | + window, then a step to the forecast's next-refresh value with the band |
| 101 | + opening to its tolerance. The anchor is never earlier than the latest run, so |
| 102 | + the forecast extends forward rather than back over history. Returns ``None`` |
| 103 | + when there is no usable forward window — no predicted window, the window has |
| 104 | + already elapsed, or no stored baseline — keyed as ``lower_tolerance`` / |
| 105 | + ``upper_tolerance`` (no sensitivity suffix); read with ``forecast_band_points``. |
| 106 | + """ |
| 107 | + baseline = definition.prediction.get("baseline_value") if definition.prediction else None |
| 108 | + now_ms = int(pd.Timestamp(last_run_time).timestamp() * 1000) if last_run_time is not None else None |
| 109 | + if freshness_window is None or now_ms is None or baseline is None: |
| 110 | + return None |
| 111 | + window_end = freshness_window.get("end") |
| 112 | + if window_end is None or window_end <= now_ms: |
| 113 | + return None |
| 114 | + |
| 115 | + forecast_means = (definition.prediction.get("mean") if definition.prediction else None) or {} |
| 116 | + next_refresh_mean = forecast_means[min(forecast_means, key=lambda k: int(k))] if forecast_means else baseline |
| 117 | + flat_anchor = max(freshness_window.get("start") or now_ms, now_ms) |
| 118 | + # lower/upper_tolerance are VARCHAR columns — coerce to float so the band dicts are numerically |
| 119 | + # typed throughout (baseline is already a float). |
| 120 | + lower_tol = float(definition.lower_tolerance) if definition.lower_tolerance is not None else None |
| 121 | + upper_tol = float(definition.upper_tolerance) if definition.upper_tolerance is not None else None |
| 122 | + return { |
| 123 | + "method": "predict", |
| 124 | + "mean": {flat_anchor: baseline, window_end: next_refresh_mean}, |
| 125 | + "lower_tolerance": {flat_anchor: baseline, window_end: lower_tol}, |
| 126 | + "upper_tolerance": {flat_anchor: baseline, window_end: upper_tol}, |
| 127 | + } |
| 128 | + |
| 129 | + |
| 130 | +def forecast_band_points(prediction: dict | None) -> list[MonitorForecastPoint]: |
| 131 | + """Convert a forecast band dict with plain ``lower_tolerance`` / ``upper_tolerance`` |
| 132 | + epoch-ms series (as built by ``gated_forecast_prediction``) into time-ordered |
| 133 | + forecast points. For the raw per-sensitivity prediction JSONB use |
| 134 | + ``forecast_points_from_prediction`` instead.""" |
| 135 | + if not prediction: |
| 136 | + return [] |
| 137 | + lower_series = prediction.get("lower_tolerance") or {} |
| 138 | + upper_series = prediction.get("upper_tolerance") or {} |
| 139 | + keys = sorted(set(lower_series) | set(upper_series), key=lambda k: int(k)) |
| 140 | + points: list[MonitorForecastPoint] = [] |
| 141 | + for k in keys: |
| 142 | + ts = datetime.fromtimestamp(int(k) / 1000.0, UTC) |
| 143 | + lower = lower_series.get(k) |
| 144 | + upper = upper_series.get(k) |
| 145 | + points.append(MonitorForecastPoint( |
| 146 | + test_time=ts, |
| 147 | + lower_bound=float(lower) if lower is not None else None, |
| 148 | + upper_bound=float(upper) if upper is not None else None, |
| 149 | + )) |
| 150 | + return points |
0 commit comments