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

"""Build local weight-score JSON files from PG Pattern V2 facts."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from examples.demand.pg_pattern_repository import PG_PATTERN_SOURCE_SYSTEM, query_weight_score_rows


DIMENSIONS = ("实质", "形式", "意图")
LEVELS = ("元素", "分类")


def _score(post_count: Any, occurrence_count: Any) -> float:
    """Stable MVP score from real PG coverage counts."""
    return float(post_count or 0) + float(occurrence_count or 0) / 1000000.0


def _normalize_element_row(row: dict[str, Any]) -> dict[str, Any]:
    occurrence_count = int(row.get("occurrence_count") or 0)
    post_count = int(row.get("post_count") or 0)
    return {
        "name": row.get("name"),
        "element_type": row.get("element_type"),
        "category_id": row.get("category_id"),
        "category_path": row.get("category_path"),
        "occurrence_count": occurrence_count,
        "post_count": post_count,
        "score": _score(post_count, occurrence_count),
        "source_system": PG_PATTERN_SOURCE_SYSTEM,
    }


def _normalize_category_row(row: dict[str, Any]) -> dict[str, Any]:
    occurrence_count = int(row.get("occurrence_count") or 0)
    post_count = int(row.get("post_count") or 0)
    return {
        "category": row.get("category"),
        "category_id": row.get("category_id"),
        "category_path": row.get("category_path"),
        "element_type": row.get("element_type"),
        "level": row.get("level"),
        "occurrence_count": occurrence_count,
        "post_count": post_count,
        "score": _score(post_count, occurrence_count),
        "source_system": PG_PATTERN_SOURCE_SYSTEM,
    }


def build_pg_weight_score_files(
        execution_id: int,
        *,
        base_dir: str | Path,
        limit_per_file: int = 5000,
) -> dict[str, int]:
    """Write old-shape weight JSON files under `{base_dir}/{execution_id}`."""
    output_dir = Path(base_dir) / str(execution_id)
    output_dir.mkdir(parents=True, exist_ok=True)

    counts: dict[str, int] = {}
    for dimension in DIMENSIONS:
        for level in LEVELS:
            rows = query_weight_score_rows(
                execution_id=execution_id,
                level=level,
                dimension=dimension,
                limit=limit_per_file,
            )
            if level == "元素":
                payload = [_normalize_element_row(dict(row)) for row in rows]
            else:
                payload = [_normalize_category_row(dict(row)) for row in rows]
            payload.sort(key=lambda item: float(item.get("score") or 0), reverse=True)
            path = output_dir / f"{dimension}_{level}.json"
            path.write_text(
                json.dumps(payload, ensure_ascii=False, indent=2, default=str) + "\n",
                encoding="utf-8",
            )
            counts[f"{dimension}_{level}"] = len(payload)
    return counts