vpp-ai-platform/packages/runtime/src/workflows/day-ahead-situation.ts

import { createStep, createWorkflow } from '@mastra/core/workflows'
import { z } from 'zod'
import type { ForecastBundle, ForecastKind, SituationReport, ToolCallRef } from '@vpp/domain'
import { ForecastBundle as ForecastBundleSchema, MarketDate, SituationReport as SituationReportSchema, ToolCallRef as ToolCallRefSchema } from '@vpp/domain'
import type { RuntimeContext } from '../context.js'
import { LlmUnavailable } from '../llm.js'
import { registerSkill } from '../skills.js'

/**
 * D-1 06:00 日前态势 (docs/07): periodic trigger → forecasts (tools) →
 * analysis agent writes findings around the numbers → SituationReport on
 * the bus. Runs with the LLM down (I6): findings fall back to a template.
 */
export const SituationInput = z.object({ market_date: MarketDate })

export const SituationOutput = z.object({
  report: SituationReportSchema,
  report_ref: z.string(),
  forecasts: z.object({ LOAD: ForecastBundleSchema, PV: ForecastBundleSchema, PRICE: ForecastBundleSchema }),
  tool_calls: z.array(ToolCallRefSchema),
  llm_backend: z.string(),
  llm_used: z.boolean(),
})

const SERIES: Record<ForecastKind, { series: string; unit: 'mw' | 'yuan_per_mwh' }> = {
  LOAD: { series: 'load:aggregate', unit: 'mw' },
  PV: { series: 'pv:aggregate', unit: 'mw' },
  PRICE: { series: 'price:da', unit: 'yuan_per_mwh' },
}

const Findings = z.object({
  findings: z.array(z.object({ kind: z.enum(['TREND', 'ANOMALY', 'RISK', 'ATTRIBUTION']), summary: z.string().min(1) })).min(1).max(6),
})

export async function runForecasts(ctx: RuntimeContext, marketDate: string): Promise<{ forecasts: Record<ForecastKind, ForecastBundle>; calls: ToolCallRef[] }> {
  const d = new Date(`${marketDate}T00:00:00Z`)
  const from = new Date(d.getTime() - ctx.config.forecastWindowDays * 86_400_000).toISOString().slice(0, 10)
  const to = new Date(d.getTime() - 86_400_000).toISOString().slice(0, 10)
  const forecasts = {} as Record<ForecastKind, ForecastBundle>
  const calls: ToolCallRef[] = []
  for (const kind of ['LOAD', 'PV', 'PRICE'] as const) {
    const history = ctx.timeseries.range(SERIES[kind].series, from, to).map((r) => r.curve)
    if (history.length === 0) throw new Error(`no history for ${SERIES[kind].series} in [${from}, ${to}] — ingest data first`)
    const featuresRef = ctx.snapshots.put({ series: SERIES[kind].series, from, to, revisions: history.length })
    ctx.lineage.addDataRef(featuresRef)
    const skill = registerSkill(
      { lineage: ctx.lineage, snapshots: ctx.snapshots, idGen: () => ctx.newId('tc') },
      { id: `${kind.toLowerCase()}-forecast`, version: '1.0.0', invoke: (req: Parameters<typeof ctx.skills.forecast>[1]) => ctx.skills.forecast(kind, req) },
    )
    const { output, ref } = await skill.call({ kind, market_date: marketDate, unit: SERIES[kind].unit, history, exogenous: {}, features_snapshot_ref: featuresRef })
    forecasts[kind] = output
    calls.push(ref)
  }
  return { forecasts, calls }
}

/** Risk level is deterministic (it gates the abnormal-day protocol, docs/13 §1); the LLM only narrates. */
export function riskLevel(ctx: RuntimeContext, price: ForecastBundle): SituationReport['risk_level'] {
  const ratio = Math.max(...price.quantiles.p90.values.map((v, i) => Number(v) / Math.max(Number(price.quantiles.p50.values[i]), 1e-9)))
  const extreme = Number(ctx.config.extremeDayPriceRatio)
  if (ratio >= extreme) return 'EXTREME'
  if (ratio >= 1 + (extreme - 1) * 0.66) return 'HIGH'
  if (ratio >= 1 + (extreme - 1) * 0.33) return 'MEDIUM'
  return 'LOW'
}

export function createDayAheadSituation(ctx: RuntimeContext) {
  const forecastStep = createStep({
    id: 'forecast',
    inputSchema: SituationInput,
    outputSchema: z.object({ market_date: MarketDate, forecasts: SituationOutput.shape.forecasts, tool_calls: z.array(ToolCallRefSchema) }),
    execute: async ({ inputData }) => {
      return ctx.lineage.run(async () => {
        const { forecasts, calls } = await runForecasts(ctx, inputData.market_date)
        return { market_date: inputData.market_date, forecasts, tool_calls: calls }
      })
    },
  })

  const analyzeStep = createStep({
    id: 'analyze',
    inputSchema: forecastStep.outputSchema,
    outputSchema: SituationOutput,
    execute: async ({ inputData }) => {
      const { forecasts, tool_calls } = inputData
      const refs = tool_calls.map((t) => t.outputs_ref)
      const level = riskLevel(ctx, forecasts.PRICE)
      const peakPrice = Math.max(...forecasts.PRICE.quantiles.p50.values.map(Number))
      const peakLoad = Math.max(...forecasts.LOAD.quantiles.p50.values.map(Number))

      let findings: SituationReport['findings']
      let llmUsed = false
      const fallback: SituationReport['findings'] = [
        { kind: 'TREND', summary: `P50 peak load ${peakLoad.toFixed(1)} MW; P50 peak price ${peakPrice.toFixed(2)} yuan/MWh (template — LLM unavailable)`, refs },
        { kind: 'RISK', summary: `risk level ${level} from price band width`, refs },
      ]
      try {
        const prompt =
          `Market date ${inputData.market_date}. Tool outputs (snapshot refs): ` +
          tool_calls.map((t) => `${t.tool}→${t.outputs_ref.slice(0, 12)}`).join(', ') +
          `. Deterministic risk level: ${level}. P50 peak load ${peakLoad.toFixed(1)} MW, P50 peak price ${peakPrice.toFixed(2)} yuan/MWh. ` +
          'Write 2–4 findings (TREND/ANOMALY/RISK) that only restate these numbers.'
        const out = await ctx.llm.structured('analysis-agent', prompt, Findings)
        findings = out.findings.map((f) => ({ ...f, refs }))
        llmUsed = true
      } catch (e) {
        if (!(e instanceof LlmUnavailable)) throw e
        findings = fallback
      }

      const report: SituationReport = {
        id: ctx.newId('sit'),
        market_date: inputData.market_date,
        risk_level: level,
        findings,
        forecast_refs: [forecasts.LOAD.id, forecasts.PV.id, forecasts.PRICE.id],
        generated_at: ctx.clock(),
      }
      const report_ref = ctx.snapshots.put(report)
      ctx.events.append({
        event_type: 'SituationPublished',
        payload: { report, report_ref, forecast_tool_calls: tool_calls },
        correlation_id: `situation-${inputData.market_date}`,
      })
      return { report, report_ref, forecasts, tool_calls, llm_backend: ctx.llm.backend, llm_used: llmUsed }
    },
  })

  return createWorkflow({
    id: 'day-ahead-situation',
    inputSchema: SituationInput,
    outputSchema: SituationOutput,
  })
    .then(forecastStep)
    .then(analyzeStep)
    .commit()
}