需求分配完整流程修改

supply_infra/scheduler/grade_assignment.py

@@ -0,0 +1,143 @@
+"""调度执行层：把已落库计划任务转换为最小化的需求分类分配。"""
+from __future__ import annotations
+
+from agents.demand_grade_agent.tools.demand_priority import build_demand_priority_index
+from supply_infra.db.repositories.demand_belong_category_repo import DemandBelongCategoryRepository
+from supply_infra.db.repositories.demand_belong_pool_rel_repo import DemandBelongPoolRelRepository
+from supply_infra.db.repositories.multi_demand_pool_di_repo import MultiDemandPoolDiRepository
+from supply_infra.db.session import get_session
+
+
+def _category_ids_by_name(
+    pool_rows: list,
+    belong_ids_by_pool: dict[int, list[int]],
+    category_by_belong: dict[int, int],
+) -> dict[str, set[int]]:
+    """把需求池行解析为需求词到全部有效分类 ID 的映射。"""
+    result: dict[str, set[int]] = {}
+    for row in pool_rows:
+        if not row.demand_name:
+            continue
+        category_ids = result.setdefault(str(row.demand_name), set())
+        for belong_id in belong_ids_by_pool.get(int(row.id), []):
+            category_id = category_by_belong.get(int(belong_id))
+            if category_id is not None:
+                category_ids.add(category_id)
+    return result
+
+
+def _load_category_ids_by_name(pool_rows: list) -> dict[str, set[int]]:
+    with get_session() as session:
+        belong_ids_by_pool = DemandBelongPoolRelRepository(session).get_belong_ids_by_pool_ids(
+            [int(row.id) for row in pool_rows]
+        )
+        all_belong_ids = {
+            belong_id
+            for values in belong_ids_by_pool.values()
+            for belong_id in values
+        }
+        category_by_belong = {
+            int(row.id): int(row.category_id)
+            for row in DemandBelongCategoryRepository(session).get_by_ids(all_belong_ids)
+            if row.category_id is not None
+        }
+    return _category_ids_by_name(pool_rows, belong_ids_by_pool, category_by_belong)
+
+
+def build_supplement_grade_assignment(biz_dt: str, demand_names: list[str]) -> dict:
+    """为仍未分级的需求生成需求词与所属分类的最小交接件。"""
+    requested_names = list(
+        dict.fromkeys(
+            str(name).strip()
+            for name in demand_names
+            if name is not None and str(name).strip()
+        )
+    )
+    if not requested_names:
+        return {"biz_dt": biz_dt, "items": []}
+
+    requested_set = set(requested_names)
+    with get_session() as session:
+        pool_repo = MultiDemandPoolDiRepository(session)
+        all_pool_rows = pool_repo.list_by_biz_dt(biz_dt)
+        matched_pool_rows = [
+            row for row in all_pool_rows if row.demand_name in requested_set
+        ]
+        priority_index = build_demand_priority_index(all_pool_rows)
+
+    category_ids_by_name = _load_category_ids_by_name(matched_pool_rows)
+    ordered_names = sorted(
+        requested_names,
+        key=lambda name: (
+            priority_index.get(name, {}).get("source_rank_score") is None,
+            -float(priority_index.get(name, {}).get("source_rank_score") or 0),
+            float(priority_index.get(name, {}).get("global_demand_rank") or float("inf")),
+            name,
+        ),
+    )
+    return {
+        "biz_dt": biz_dt,
+        "items": [
+            {
+                "demand_name": name,
+                "category_ids": sorted(category_ids_by_name.get(name, set())),
+            }
+            for name in ordered_names
+        ],
+    }
+
+
+def build_grade_batch_assignment(
+    biz_dt: str,
+    category_ids: list[int],
+    max_demands: int = 20,
+    excluded_demand_names: list[str] | None = None,
+) -> dict:
+    """按计划节点取一小批需求，只返回需求词与其全部所属分类。"""
+    selected_ids = list(dict.fromkeys(int(value) for value in category_ids))
+    excluded = set(excluded_demand_names or [])
+    with get_session() as session:
+        belongs = DemandBelongCategoryRepository(session).list_by_category_ids(selected_ids)
+        pool_ids_by_belong = DemandBelongPoolRelRepository(session).get_pool_ids_by_belong_ids(
+            [int(row.id) for row in belongs]
+        )
+        pool_ids = sorted({pool_id for values in pool_ids_by_belong.values() for pool_id in values})
+        pool_repo = MultiDemandPoolDiRepository(session)
+        pool_rows = pool_repo.get_by_ids(pool_ids)
+        all_day_pool_rows = pool_repo.list_by_biz_dt(biz_dt)
+        priority_index = build_demand_priority_index(all_day_pool_rows)
+        candidate_names = {
+            str(row.demand_name)
+            for row in pool_rows
+            if row.biz_dt == biz_dt and row.demand_name and row.demand_name not in excluded
+        }
+
+    names = sorted(
+        candidate_names,
+        key=lambda name: (
+            priority_index.get(name, {}).get("source_rank_score") is None,
+            -float(priority_index.get(name, {}).get("source_rank_score") or 0),
+            float(priority_index.get(name, {}).get("global_demand_rank") or float("inf")),
+            name,
+        ),
+    )[: max(1, min(int(max_demands), 100))]
+    if not names:
+        return {"biz_dt": biz_dt, "items": []}
+
+    selected_name_set = set(names)
+    selected_pool_rows = [
+        row
+        for row in all_day_pool_rows
+        if row.demand_name in selected_name_set
+    ]
+    category_ids_by_name = _load_category_ids_by_name(selected_pool_rows)
+    return {
+        "biz_dt": biz_dt,
+        "items": [
+            {
+                "demand_name": name,
+                "category_ids": sorted(category_ids_by_name.get(name, set())),
+            }
+            for name in names
+        ],
+    }