SupplyAgent/agents/find_agent/agent.py

"""
find_agent 工厂 — 组装 Agent 实例。

每个业务 Agent 都应提供 create_xxx_agent() 工厂函数，
统一注册本 Agent 的工具 + 共享基础设施工具。
"""
from __future__ import annotations

import json
import re
from pathlib import Path

from supply_agent import Agent
from supply_agent.config import Settings
from supply_agent.types import Message, Role
from agents.find_agent.tools import register_all_tools

_PROMPT_PATH = Path(__file__).parent / "prompt" / "system_prompt.md"
FIND_AGENT_SYSTEM_PROMPT = _PROMPT_PATH.read_text(encoding="utf-8")

_FAILURE_REPORT_PREFIX = "任务未完成（工具故障）"
_VIDEO_ID_PATTERN = re.compile(r"(?<!\d)\d{15,22}(?!\d)")


def _report_section(
    content: str,
    start_marker: str,
    end_markers: tuple[str, ...],
) -> str | None:
    start = content.find(start_marker)
    if start < 0:
        return None
    ends = [
        position
        for marker in end_markers
        if (position := content.find(marker, start + len(start_marker))) >= 0
    ]
    end = min(ends) if ends else len(content)
    return content[start:end]


def _validate_report_primary_section(
    content: str,
    state: dict[str, object],
) -> str | None:
    """Only validate aweme_ids listed in the primary report section."""
    candidates = state.get("candidates")
    if not isinstance(candidates, list):
        return "最终状态缺少 candidates，无法校验报告分池"
    bucket_by_id = {
        str(item.get("aweme_id")): str(item.get("decision_bucket"))
        for item in candidates
        if isinstance(item, dict) and item.get("aweme_id")
    }

    primary_section = _report_section(
        content,
        "主推荐",
        ("淘汰候选", "搜索树", "缺失数据", "总结"),
    )
    if primary_section is None:
        return "最终报告必须包含“主推荐”段"

    primary_ids = set(_VIDEO_ID_PATTERN.findall(primary_section))
    if not primary_ids:
        return "主推荐段必须输出至少一个 aweme_id"

    wrong_primary = sorted(
        video_id
        for video_id in primary_ids
        if bucket_by_id.get(video_id) != "primary"
    )
    if wrong_primary:
        return (
            "主推荐段包含非 primary 候选: "
            + ", ".join(wrong_primary[:5])
            + "。必须按数据库 decision_bucket 输出"
        )
    return None


def _primary_count_from_state(state: dict[str, object]) -> int:
    run = state.get("run")
    if isinstance(run, dict):
        return int(run.get("primary_count") or 0)
    return 0


def _successful_tool_events(
    messages: list[Message],
) -> list[tuple[int, str, dict[str, object]]]:
    events: list[tuple[int, str, dict[str, object]]] = []
    for index, message in enumerate(messages):
        if message.role != Role.TOOL or not message.name or not message.content:
            continue
        try:
            payload = json.loads(message.content)
        except (TypeError, json.JSONDecodeError):
            continue
        if not isinstance(payload, dict) or payload.get("error"):
            continue
        events.append((index, message.name, payload))
    return events


def _last_final_assistant_message(messages: list[Message]) -> Message | None:
    return next(
        (
            message
            for message in reversed(messages)
            if message.role == Role.ASSISTANT and not message.tool_calls
        ),
        None,
    )


def find_agent_completion_guard(messages: list[Message]) -> str | None:
    """返回完成提示；默认仅提示、不阻断结束。"""
    last_assistant = _last_final_assistant_message(messages)
    content = (last_assistant.content or "").strip() if last_assistant else ""
    if content.startswith(_FAILURE_REPORT_PREFIX):
        return None

    events = _successful_tool_events(messages)
    if not any(name == "create_video_discovery_run" for _, name, _ in events):
        return "尚未成功创建视频发现运行"

    evaluation_events = [
        (index, payload)
        for index, name, payload in events
        if name == "batch_save_video_candidate_evaluations"
    ]
    if not evaluation_events:
        return "尚未保存候选评估"
    evaluation_index, evaluation = evaluation_events[-1]
    if evaluation.get("status") != "finished":
        return "最后一次候选保存尚未把运行状态设置为 finished"

    state_events = [
        (index, payload)
        for index, name, payload in events
        if name == "query_video_discovery_state"
    ]
    if not state_events:
        return "尚未查询最终数据库状态"
    state_index, state = state_events[-1]
    if state_index < evaluation_index:
        return "最终状态查询必须在候选保存为 finished 之后执行"

    run = state.get("run")
    run_state = run if isinstance(run, dict) else {}
    if run_state.get("status") != "finished":
        return "运行状态尚未持久化为 finished"

    if last_assistant is None or not content:
        return "最终回答为空"

    if _primary_count_from_state(state) > 0:
        return _validate_report_primary_section(content, state)
    return None


def create_find_agent(
    settings: Settings | None = None,
    *,
    model: str | None = None,
) -> Agent:
    """创建 find_agent 实例，注册所有相关工具。"""
    agent = Agent(
        settings=settings,
        name="find_agent",
        model="google/gemini-3-flash-preview",
        system_prompt=FIND_AGENT_SYSTEM_PROMPT,
        max_iterations=60,
        temperature=0.2,
        completion_guard=find_agent_completion_guard,
        tool_repeat_requires_change={
            "query_video_discovery_state": {
                "record_video_search_page",
                "batch_save_video_candidate_evaluations",
                "audit_video_discovery_run",
            },
        },
    )

    # 本 Agent 专属工具
    register_all_tools(agent.tools)

    return agent