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")

_EVIDENCE_TOOLS = {
    "douyin_detail",
    "get_content_fans_portrait",
    "get_account_fans_portrait",
    "batch_fetch_portraits",
    "normalize_age_portraits",
}
_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_buckets(
    content: str,
    state: dict[str, object],
) -> str | None:
    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")
    }
    run = state.get("run")
    run_state = run if isinstance(run, dict) else {}

    primary_section = _report_section(
        content,
        "主推荐",
        ("补充推荐",),
    )
    backup_section = _report_section(
        content,
        "补充推荐",
        ("最有竞争力", "搜索树", "缺失数据", "总结"),
    )
    if primary_section is None or backup_section is None:
        return "最终报告必须分别包含“主推荐”和“补充推荐”段"

    primary_ids = set(_VIDEO_ID_PATTERN.findall(primary_section))
    backup_ids = set(_VIDEO_ID_PATTERN.findall(backup_section))
    wrong_primary = sorted(
        video_id
        for video_id in primary_ids
        if bucket_by_id.get(video_id) != "primary"
    )
    wrong_backup = sorted(
        video_id
        for video_id in backup_ids
        if bucket_by_id.get(video_id) != "backup"
    )
    if wrong_primary:
        return (
            "主推荐段包含非 primary 候选: "
            + ", ".join(wrong_primary[:5])
            + "。必须按数据库 decision_bucket 输出"
        )
    if wrong_backup:
        return (
            "补充推荐段包含非 backup 候选: "
            + ", ".join(wrong_backup[:5])
            + "。这些条目应移到淘汰候选，不得冒充 Agent 保留结果"
        )

    primary_count = int(run_state.get("primary_count") or 0)
    backup_count = int(run_state.get("backup_count") or 0)
    if primary_count > 0 and not primary_ids:
        return "数据库存在 primary 候选，但主推荐段没有输出 aweme_id"
    if primary_count == 0 and backup_count > 0 and not backup_ids:
        return "数据库只有 backup 保留候选，补充推荐段至少应输出一条"
    return None


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 find_agent_completion_guard(messages: list[Message]) -> str | None:
    """Reject final text until persisted state and the final audit prove completion."""
    events = _successful_tool_events(messages)
    if not any(name == "create_video_discovery_run" for _, name, _ in events):
        return "尚未成功创建视频发现运行"

    record_indexes = [
        index
        for index, name, _ in events
        if name == "record_video_search_page"
    ]
    if not record_indexes:
        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_index < max(record_indexes):
        return "最后一次搜索发生在候选评估之后，新增候选尚未重新评估"
    if evaluation.get("status") != "finished":
        return "最后一次候选保存尚未把运行状态设置为 finished"

    evidence_indexes = [
        index for index, name, _ in events if name in _EVIDENCE_TOOLS
    ]
    if evidence_indexes and evaluation_index < max(evidence_indexes):
        return "最后一次证据获取发生在候选评估之后，证据尚未重新保存"

    audit_events = [
        (index, payload)
        for index, name, payload in events
        if name in {
            "audit_video_discovery_process",
            "audit_video_discovery_run",
        }
    ]
    if not audit_events:
        return "尚未执行完成审计"
    audit_index, audit = audit_events[-1]
    if audit_index < evaluation_index:
        post_audit_evaluations = [
            payload
            for index, payload in evaluation_events
            if index > audit_index
        ]
        if not post_audit_evaluations or any(
            payload.get("audit_relevant_changed") is not False
            for payload in post_audit_evaluations
        ):
            return (
                "最后一次候选评估改变了审计相关状态，尚未重新审计；"
                "下一步只调用 audit_video_discovery_run，"
                "不要再次保存或查询"
            )
    if audit_index < max(record_indexes):
        return "最后一次搜索页保存尚未重新审计"
    if evidence_indexes and audit_index < max(evidence_indexes):
        return "最后一次证据获取尚未重新审计"
    if audit.get("can_finish") is not True:
        violations = audit.get("critical_violations")
        if isinstance(violations, list) and violations:
            summary = "；".join(str(item) for item in violations[:5])
            return f"最后一次审计未通过：{summary}"
        return "最后一次审计未通过"

    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]
    last_changed_evaluation_index = max(
        (
            index
            for index, payload in evaluation_events
            if payload.get("audit_relevant_changed") is not False
        ),
        default=evaluation_index,
    )
    if state_index < last_changed_evaluation_index:
        return "最终数据库状态早于最后一次有效候选变更，请重新查询"

    last_assistant = next(
        (
            message
            for message in reversed(messages)
            if message.role == Role.ASSISTANT and not message.tool_calls
        ),
        None,
    )
    if last_assistant is None or not (last_assistant.content or "").strip():
        return "最终回答为空"
    report_error = _validate_report_buckets(last_assistant.content, state)
    if report_error:
        return report_error
    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=None,
        tool_call_budgets={
            "search": (
                {
                    "douyin_search",
                    "douyin_search_tikhub",
                    "douyin_user_videos",
                },
                10,
            ),
            "detail": ({"douyin_detail"}, 2),
            "portrait": (
                {
                    "get_content_fans_portrait",
                    "get_account_fans_portrait",
                    "batch_fetch_portraits",
                },
                3,
            ),
            "evaluation_save": (
                {"batch_save_video_candidate_evaluations"},
                6,
            ),
            "state_query": ({"query_video_discovery_state"}, 8),
        },
        tool_repeat_requires_change={
            "query_video_discovery_state": {
                "record_video_search_page",
                "batch_save_video_candidate_evaluations",
            },
        },
    )

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

    return agent