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",
}
_SEARCH_TOOLS = {
    "douyin_search",
    "douyin_search_tikhub",
    "douyin_user_videos",
}
_FAILURE_REPORT_PREFIX = "任务未完成（工具故障）"
_MIN_PRIMARY_TARGET = 5
_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,
        "主推荐",
        ("淘汰候选",),
    )
    rejected_section = _report_section(
        content,
        "淘汰候选",
        ("搜索树", "缺失数据", "总结"),
    )
    if primary_section is None or rejected_section is None:
        return "最终报告必须分别包含“主推荐”和“淘汰候选”段"

    primary_ids = set(_VIDEO_ID_PATTERN.findall(primary_section))
    rejected_ids = set(_VIDEO_ID_PATTERN.findall(rejected_section))
    wrong_primary = sorted(
        video_id
        for video_id in primary_ids
        if bucket_by_id.get(video_id) != "primary"
    )
    wrong_rejected = sorted(
        video_id
        for video_id in rejected_ids
        if bucket_by_id.get(video_id) != "rejected"
    )
    if wrong_primary:
        return (
            "主推荐段包含非 primary 候选: "
            + ", ".join(wrong_primary[:5])
            + "。必须按数据库 decision_bucket 输出"
        )
    if wrong_rejected:
        return (
            "淘汰候选段包含非 rejected 候选: "
            + ", ".join(wrong_rejected[:5])
            + "。必须按数据库 decision_bucket 输出"
        )

    primary_count = int(run_state.get("primary_count") or 0)
    if primary_count > 0 and not primary_ids:
        return "数据库存在 primary 候选，但主推荐段没有输出 aweme_id"
    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 _validate_report_primary_section(
    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,
        "主推荐",
        ("淘汰候选", "搜索树", "缺失数据", "总结"),
    )
    if primary_section is None:
        return "最终报告必须包含“主推荐”段"

    primary_ids = set(_VIDEO_ID_PATTERN.findall(primary_section))
    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 输出"
        )

    primary_count = int(run_state.get("primary_count") or 0)
    if primary_count > 0 and not primary_ids:
        return "数据库存在 primary 候选，但主推荐段没有输出 aweme_id"
    return None


def _guard_relaxed_primary_target(
    events: list[tuple[int, str, dict[str, object]]],
    messages: list[Message],
) -> str | None:
    """Allow finish when enough primary candidates are persisted."""
    evaluation_events = [
        (index, payload)
        for index, name, payload in events
        if name == "batch_save_video_candidate_evaluations"
    ]
    if not evaluation_events:
        return "尚未保存候选评估"
    _, 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 = state_events[-1]
    if _primary_count_from_state(state) < _MIN_PRIMARY_TARGET:
        return f"数据库 primary 数量不足 {_MIN_PRIMARY_TARGET} 条"

    last_assistant = _last_final_assistant_message(messages)
    if last_assistant is None or not (last_assistant.content or "").strip():
        return "最终回答为空"
    report_error = _validate_report_primary_section(last_assistant.content, state)
    if report_error:
        return report_error
    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 _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:
    """Enforce the final search → evidence → evaluation → audit → state → report order."""
    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 "尚未成功创建视频发现运行"

    state_events = [
        (index, payload)
        for index, name, payload in events
        if name == "query_video_discovery_state"
    ]
    if state_events:
        _, latest_state = state_events[-1]
        if _primary_count_from_state(latest_state) >= _MIN_PRIMARY_TARGET:
            return _guard_relaxed_primary_target(events, messages)

    search_indexes = [
        index
        for index, name, _ in events
        if name == "record_video_search_page"
    ]
    if not search_indexes:
        return "尚未持久化任何搜索页"
    last_search_index = max(search_indexes)
    raw_search_indexes = [
        index for index, name, _ in events if name in _SEARCH_TOOLS
    ]
    if raw_search_indexes and last_search_index < max(raw_search_indexes):
        return "最后一次搜索结果尚未通过 record_video_search_page 持久化"

    evidence_indexes = [
        index for index, name, _ in events if name in _EVIDENCE_TOOLS
    ]
    if not evidence_indexes:
        return "尚未获取并整理候选证据"
    last_evidence_index = max(evidence_indexes)
    if last_evidence_index < last_search_index:
        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 < last_evidence_index:
        return "最后一次证据获取发生在候选评估之后，证据尚未重新评估并保存"
    if evaluation.get("status") != "finished":
        return "最后一次候选保存尚未把运行状态设置为 finished"

    audit_events = [
        (index, payload)
        for index, name, payload in events
        if name == "audit_video_discovery_run"
    ]
    if not audit_events:
        return "尚未调用 audit_video_discovery_run 执行数据库完成审计"
    audit_index, audit = audit_events[-1]
    if audit_index < evaluation_index:
        return "候选评估晚于最后一次审计，请重新调用 audit_video_discovery_run"
    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]
    if state_index < audit_index:
        return "最终状态查询必须在审计通过之后执行"

    last_assistant = _last_final_assistant_message(messages)
    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=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