demand-agent-new/examples/demand/local_output_sink.py

from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo

from examples.demand.data_query_tools import (
    build_dwd_multi_demand_pool_di_rows,
    build_feature_point_data_rows,
)


OUTPUT_FILENAMES = {
    "run_manifest": "run_manifest.json",
    "demand_task": "demand_task.json",
    "demand_items": "demand_items.json",
    "rejected_demand_items": "rejected_demand_items.json",
    "demand_content": "demand_content.json",
    "dwd_multi_demand_pool_di": "dwd_multi_demand_pool_di.json",
    "feature_point_data": "feature_point_data.json",
}

_CHINA_TZ = ZoneInfo("Asia/Shanghai")
_DT_FMT = "%Y%m%d"


def _today_dt() -> str:
    return datetime.now(_CHINA_TZ).strftime(_DT_FMT)


def _as_list(value: Any) -> list[Any]:
    if value is None:
        return []
    if isinstance(value, list):
        return value
    if isinstance(value, tuple):
        return list(value)
    return [value]


def _normalize_demand_content_row(row: Any) -> Any:
    if not isinstance(row, dict):
        return row

    normalized = dict(row)
    ext_data = normalized.get("ext_data")
    if isinstance(ext_data, str) and ext_data.strip():
        try:
            normalized["ext_data"] = json.loads(ext_data)
        except json.JSONDecodeError:
            pass
    return normalized


def _write_json(path: Path, payload: Any) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_name(f".{path.name}.tmp")
    tmp_path.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2, default=str) + "\n",
        encoding="utf-8",
    )
    tmp_path.replace(path)
    return path


class LocalOutputSink:
    """Writes DemandAgent local-json output mirrors for one run."""

    def __init__(self, output_dir: str | Path):
        self.output_dir = Path(output_dir)

    @classmethod
    def for_run(cls, base_dir: str | Path, run_id: str | int) -> "LocalOutputSink":
        return cls(Path(base_dir) / str(run_id))

    def path_for(self, name: str) -> Path:
        try:
            filename = OUTPUT_FILENAMES[name]
        except KeyError as exc:
            raise ValueError(f"unknown local output name: {name}") from exc
        return self.output_dir / filename

    def default_run_manifest(self) -> dict[str, Any]:
        return {
            "output_mode": "local_json",
            "output_dir": str(self.output_dir),
            "created_at": datetime.now(_CHINA_TZ).isoformat(),
        }

    def initialize(self, run_manifest: dict[str, Any] | None = None) -> dict[str, Path]:
        return self.write_all(run_manifest=run_manifest)

    def write_run_manifest(self, payload: dict[str, Any] | None = None) -> Path:
        if payload is None:
            payload = self.default_run_manifest()
        return _write_json(self.path_for("run_manifest"), payload)

    def write_demand_task(self, payload: Any = None) -> Path:
        if payload is None:
            payload = {}
        return _write_json(self.path_for("demand_task"), payload)

    def write_demand_items(self, items: Any = None) -> Path:
        return _write_json(self.path_for("demand_items"), _as_list(items))

    def write_rejected_demand_items(self, items: Any = None) -> Path:
        return _write_json(self.path_for("rejected_demand_items"), _as_list(items))

    def write_demand_content(self, rows: Any = None) -> Path:
        normalized_rows = [_normalize_demand_content_row(row) for row in _as_list(rows)]
        return _write_json(self.path_for("demand_content"), normalized_rows)

    def write_dwd_multi_demand_pool_di(self, rows: Any = None) -> Path:
        return _write_json(self.path_for("dwd_multi_demand_pool_di"), _as_list(rows))

    def write_dwd_multi_demand_pool_di_from_demand_content(
            self,
            rows: list[dict],
            partition_dt: str | None = None,
    ) -> Path:
        output_rows = build_dwd_multi_demand_pool_di_rows(rows=rows, partition_dt=partition_dt or _today_dt())
        return self.write_dwd_multi_demand_pool_di(output_rows)

    def write_feature_point_data(self, rows: Any = None) -> Path:
        return _write_json(self.path_for("feature_point_data"), _as_list(rows))

    def write_feature_point_data_from_names(self, names: list[str], dt: str | None = None) -> Path:
        output_rows = build_feature_point_data_rows(names=names, dt=dt or _today_dt())
        return self.write_feature_point_data(output_rows)

    def write_all(
            self,
            *,
            run_manifest: dict[str, Any] | None = None,
            demand_task: Any = None,
            demand_items: Any = None,
            rejected_demand_items: Any = None,
            demand_content_rows: Any = None,
            dwd_multi_demand_pool_di_rows: Any = None,
            feature_point_data_rows: Any = None,
    ) -> dict[str, Path]:
        paths = {
            "run_manifest": self.write_run_manifest(run_manifest),
            "demand_task": self.write_demand_task(demand_task),
            "demand_items": self.write_demand_items(demand_items),
            "rejected_demand_items": self.write_rejected_demand_items(rejected_demand_items),
            "demand_content": self.write_demand_content(demand_content_rows),
            "dwd_multi_demand_pool_di": self.write_dwd_multi_demand_pool_di(dwd_multi_demand_pool_di_rows),
            "feature_point_data": self.write_feature_point_data(feature_point_data_rows),
        }
        return paths


def write_local_outputs(
        output_dir: str | Path,
        *,
        run_manifest: dict[str, Any] | None = None,
        demand_task: Any = None,
        demand_items: Any = None,
        rejected_demand_items: Any = None,
        demand_content_rows: Any = None,
        dwd_multi_demand_pool_di_rows: Any = None,
        feature_point_data_rows: Any = None,
) -> dict[str, Path]:
    return LocalOutputSink(output_dir).write_all(
        run_manifest=run_manifest,
        demand_task=demand_task,
        demand_items=demand_items,
        rejected_demand_items=rejected_demand_items,
        demand_content_rows=demand_content_rows,
        dwd_multi_demand_pool_di_rows=dwd_multi_demand_pool_di_rows,
        feature_point_data_rows=feature_point_data_rows,
    )