next-tjwater-server/app/mcp/analysis_models.py

"""Bounded calculation inputs; identity, routing and output paths are host-owned."""

from datetime import datetime, timedelta
from typing import Annotated, Literal
from uuid import UUID

from pydantic import Field, model_validator

from app.mcp.models import Id, StrictInput


MAX_CALCULATION_TIMEOUT_SECONDS = 86400


class CalculationInput(StrictInput):
    timeout_seconds: int = Field(
        default=120, ge=1, le=MAX_CALCULATION_TIMEOUT_SECONDS, strict=True,
        description="本次计算等待预算，秒；由Agent按模型规模和任务时长选择，受工具目录中的服务端上限约束；与模拟duration不同，超时不保证已启动的计算停止或回滚",
    )


class BurstDetectionInput(CalculationInput):
    target_time: datetime = Field(description="带时区的侦测时刻；读取此前15天同周期压力数据，不是单小时定位")
    sensor_nodes: list[Id] = Field(min_length=5, max_length=20, description="已核实的压力监测节点ID，至少5个完整传感器")
    sampling_interval_minutes: int = Field(default=15, ge=1, le=60)

    @model_validator(mode="after")
    def valid_target(self):
        if self.target_time.tzinfo is None:
            raise ValueError("target_time requires explicit timezone")
        if 1440 % self.sampling_interval_minutes:
            raise ValueError("sampling interval must divide a day")
        if len(set(self.sensor_nodes)) != len(self.sensor_nodes):
            raise ValueError("sensor_nodes must be distinct")
        return self


class BurstLocationInput(CalculationInput):
    data_source: Literal["monitoring", "simulation"] = "monitoring"
    simulation_run_id: UUID | None = Field(default=None, description="已有分析运行UUID，不是方案名称；simulation模式只读取已有异常结果及模型基线")
    burst_leakage: float = Field(gt=0, allow_inf_nan=False, description="有依据的假定漏水流量，WNTR SI单位m³/s；不是已知爆管管段，不能随意猜测")
    pressure_scada_ids: list[Id] | None = Field(default=None, min_length=1, max_length=20, description="压力监测设备关联的模型节点ID，不是设备查询编号；省略时自动读取当前项目SCADA配置中的压力测点，不需要猜测节点ID")
    flow_scada_ids: list[Id] | None = Field(default=None, min_length=1, max_length=20, description="流量监测设备关联的模型要素ID，不是设备查询编号")
    scada_burst_start: datetime
    scada_burst_end: datetime
    scada_normal_start: datetime | None = Field(default=None, description="monitoring模式正常对照窗开始，须与结束时间成对提供；均省略时用爆管窗前一天相同时段，不代表该时段已确认正常")
    scada_normal_end: datetime | None = Field(default=None, description="monitoring模式正常对照窗结束，须与开始时间成对提供；simulation模式不填写，使用同时间窗模型基线")

    @model_validator(mode="after")
    def valid_windows(self):
        pairs = [(self.scada_burst_start, self.scada_burst_end)]
        if self.data_source == "monitoring":
            if self.simulation_run_id is not None:
                raise ValueError("monitoring cannot use simulation_run_id")
            if self.scada_normal_start is not None or self.scada_normal_end is not None:
                pairs.append((self.scada_normal_start, self.scada_normal_end))
        elif self.simulation_run_id is None or self.scada_normal_start is not None or self.scada_normal_end is not None:
            raise ValueError("simulation requires run_id and uses the same-window model baseline")
        for start, end in pairs:
            if start is None or end is None or start.tzinfo is None or end.tzinfo is None:
                raise ValueError("explicit timezone-aware burst and normal windows are required")
            if not timedelta(0) < end - start <= timedelta(days=1):
                raise ValueError("each window must be positive and at most 24 hours")
        for ids in (self.pressure_scada_ids, self.flow_scada_ids):
            if ids is not None and len(set(ids)) != len(ids):
                raise ValueError("sensor IDs must be distinct")
        return self


class ValveIsolationInput(CalculationInput):
    accident_element: list[Id] = Field(min_length=1, max_length=20, description="经核实或明确作为候选的事故管段/节点ID")
    disabled_valves: list[Id] = Field(default_factory=list, max_length=200, description="已知无法关闭的阀门；未提供不代表现场确认全部可用")


class LeakageIdentificationInput(CalculationInput):
    scada_start: datetime
    scada_end: datetime
    sensor_nodes: list[Id] = Field(min_length=1, max_length=20)
    q_sum: float = Field(gt=0, allow_inf_nan=False, description="有依据的总漏损流量约束，m³/s，不得使用无依据默认值")
    start_time: float = Field(ge=0, lt=24, allow_inf_nan=False, description="模型起算小时；需与SCADA时间窗及模型时钟对齐")
    timestep: int = Field(default=15, ge=5, le=60, description="计算步长，分钟；需与观测采样对应")
    dma_count: int = Field(ge=1, le=20)

    @model_validator(mode="after")
    def valid_window(self):
        if self.scada_start.tzinfo is None or self.scada_end.tzinfo is None:
            raise ValueError("timestamps require explicit timezone")
        duration = self.scada_end - self.scada_start
        if not timedelta(0) < duration <= timedelta(days=1):
            raise ValueError("window must be positive and at most 24 hours")
        if duration.total_seconds() % (self.timestep * 60):
            raise ValueError("window must contain whole calculation steps")
        if len(set(self.sensor_nodes)) != len(self.sensor_nodes):
            raise ValueError("sensor_nodes must be distinct")
        return self


class ScenarioInput(CalculationInput):
    start_time: datetime
    duration: int = Field(default=900, ge=60, le=86400, description="模拟时长，秒；首次优先短时间窗")
    output_node_ids: list[Id] = Field(min_length=1, max_length=20, description="需要返回结果的节点；不把局部结果解释为全网无影响")
    output_link_ids: list[Id] = Field(default_factory=list, max_length=20)

    @model_validator(mode="after")
    def timezone_required(self):
        if self.start_time.tzinfo is None:
            raise ValueError("start_time requires explicit timezone")
        return self


class ContaminantSimulationInput(ScenarioInput):
    source: Id = Field(description="假设污染源节点；这是正向传播模拟，不是自动反演污染源")
    concentration: float = Field(gt=0, allow_inf_nan=False, description="假设源浓度，mg/L")
    pattern: Id | None = Field(default=None, description="已存在的污染源时间模式ID；不填为持续恒定源")


class ValveClosureInput(ScenarioInput):
    valves: list[Id] = Field(min_length=1, max_length=20, description="在临时模型中关闭的阀门，不下发现场控制")


class FlushingInput(ScenarioInput):
    drainage_node_id: Id
    flushing_flow_m3h: float = Field(gt=0, allow_inf_nan=False, description="明确的冲洗流量，m³/h；不使用无依据的自动排放量")
    closed_valves: list[Id] = Field(default_factory=list, max_length=20)


NonnegativeValue = Annotated[float, Field(ge=0, allow_inf_nan=False)]
PumpSeries = Annotated[list[NonnegativeValue], Field(min_length=1, max_length=96)]


class PressureRegulationInput(ScenarioInput):
    pump_control: dict[Id, PumpSeries] = Field(min_length=1, max_length=20, description="已登记泵ID及运行序列；工频值为0/1，变频值为频率Hz，每项对应模型水力步长，不自动优化")
    tank_init_level: dict[Id, Annotated[float, Field(gt=0, allow_inf_nan=False)]] = Field(default_factory=dict, max_length=20, description="水箱初始水位，须在模型允许范围内")


class SensorPlacementInput(CalculationInput):
    method: Literal["sensitivity", "kmeans"]
    sensor_count: int = Field(ge=1, le=20)
    min_diameter: int = Field(default=0, ge=0, le=5000, description="最小管径，mm")