SupplyAgent/agents/demand_video_expand_agent/run.py

#!/usr/bin/env python3
"""单需求视频点位拓展判断 — 由任务层调用。"""
from __future__ import annotations

from dataclasses import dataclass

from supply_agent.types import AgentResult, Role

_POINT_TYPE_LABEL = {
    "inspiration": "灵感点",
    "purpose": "目的点",
    "key": "关键点",
}


@dataclass
class VideoPoint:
    video_id: str
    point_type: str
    point_data: str | None
    point_desc: str | None


@dataclass
class DemandExpandContext:
    biz_dt: str
    run_id: str
    demand_grade_id: int
    demand_name: str
    grade: str
    score: float | None
    video_ids: list[str]
    points: list[VideoPoint]


def build_expand_user_input(ctx: DemandExpandContext) -> str:
    """构建传给 Agent 的用户消息。"""
    score_text = f"{ctx.score:.2f}" if ctx.score is not None else "—"
    lines = [
        f"biz_dt={ctx.biz_dt}",
        f"run_id={ctx.run_id}",
        f"demand_grade_id={ctx.demand_grade_id}",
        f"demand_name={ctx.demand_name}",
        f"grade={ctx.grade}",
        f"score={score_text}",
        f"video_count={len(ctx.video_ids)}",
        "",
        "以下是从关联视频中提取的点位，请判断哪些可作为拓展需求：",
    ]
    for index, point in enumerate(ctx.points, start=1):
        type_label = _POINT_TYPE_LABEL.get(point.point_type, point.point_type)
        data_text = point.point_data or "—"
        desc_text = point.point_desc or "—"
        lines.append(
            f"[{index}|{type_label}|video={point.video_id}] "
            f"{data_text} | 描述: {desc_text}"
        )
    lines.extend(
        [
            "",
            "判断完成后调用 batch_save_demand_expansions 落库；",
            "将上述 biz_dt、demand_grade_id、demand_name、grade、run_id 原样传入工具。",
            "若无合适拓展，传 items=[]。",
        ]
    )
    return "\n".join(lines)


def judge_demand_expansion(ctx: DemandExpandContext) -> AgentResult:
    """对单个需求执行拓展判断。"""
    from agents.demand_video_expand_agent.agent import create_demand_video_expand_agent

    agent = create_demand_video_expand_agent()
    user_input = build_expand_user_input(ctx)
    return agent.run(user_input)


def extract_saved_count(result: AgentResult) -> int:
    """从 Agent 工具返回消息中解析保存条数。"""
    for msg in reversed(result.messages):
        if msg.role != Role.TOOL or not msg.content:
            continue
        text = str(msg.content)
        if "成功保存" not in text:
            continue
        try:
            part = text.split("成功保存", 1)[1].strip()
            return int(part.split("条", 1)[0].strip())
        except (IndexError, ValueError):
            continue
    return 0