SupplyAgent/find_agent_v2/tools.py

"""Stage tools that persist exclusively to ``find_agent_v2_*`` tables."""

from __future__ import annotations

import json
from collections.abc import Callable, Iterable
from typing import Any

from find_agent_v2.providers import (
    fetch_details,
    fetch_portraits,
    search_internal,
    search_tikhub,
)
from find_agent_v2.service import get_find_agent_v2_service
from supply_agent.tools import tool
from supply_agent.tools.registry import ToolRegistry

ToolFn = Callable[..., Any]


@tool
async def search_videos_v2(run_id: str, round_index: int, searches: list[dict[str, Any]]) -> str:
    """批量搜索并仅写入 find_agent_v2_search/candidate；searches 最多 6 项。"""
    if not searches or len(searches) > 6:
        return json.dumps({"error": "searches 必须为 1~6 项"}, ensure_ascii=False)
    service = get_find_agent_v2_service()
    outputs: list[dict[str, Any]] = []
    for raw_task in searches:
        keyword = str(raw_task.get("keyword") or "").strip()
        reason = str(raw_task.get("query_reason") or "").strip()
        provider = str(raw_task.get("provider") or "internal_keyword")
        if not keyword or not reason:
            outputs.append({"error": "keyword/query_reason 不能为空"})
            continue
        max_pages = max(1, min(int(raw_task.get("max_pages") or 1), 2))
        cursor: str | int = raw_task.get("cursor") or 0
        provider_search_id = str(raw_task.get("search_id") or "")
        backtrace = str(raw_task.get("backtrace") or "")
        for page_no in range(1, max_pages + 1):
            common = {
                "keyword": keyword,
                "content_type": str(raw_task.get("content_type") or "视频"),
                "sort_type": str(raw_task.get("sort_type") or "综合排序"),
                "publish_time": str(raw_task.get("publish_time") or "不限"),
                "min_duration_seconds": int(raw_task.get("min_duration_seconds") or 30),
            }
            if provider == "tikhub":
                payload = await search_tikhub(
                    **common,
                    cursor=int(cursor or 0),
                    filter_duration=str(raw_task.get("filter_duration") or "不限"),
                    search_id=provider_search_id,
                    backtrace=backtrace,
                )
            else:
                provider = "internal_keyword"
                payload = await search_internal(**common, cursor=str(cursor or "0"))
            saved = service.save_search(
                run_id=run_id,
                round_index=int(round_index),
                keyword=keyword,
                query_reason=reason,
                source_type=str(raw_task.get("source_type") or "mixed"),
                provider=provider,
                cursor=str(cursor),
                page_no=page_no,
                payload=payload,
            )
            outputs.append({
                "keyword": keyword,
                "provider": provider,
                "page_no": page_no,
                "error": payload.get("error"),
                "has_more": bool(payload.get("has_more")),
                "next_cursor": payload.get("next_cursor"),
                **saved,
            })
            if payload.get("error") or not payload.get("has_more"):
                break
            cursor = payload.get("next_cursor") or cursor
            provider_search_id = str(payload.get("search_id") or provider_search_id)
            backtrace = str(payload.get("backtrace") or backtrace)
    return json.dumps({"run_id": run_id, "searches": outputs}, ensure_ascii=False)


@tool
async def fetch_candidate_details_v2(run_id: str, candidate_ids: list[int]) -> str:
    """批量获取候选详情并仅写入 find_agent_v2_candidate/evidence；最多 8 项。"""
    service = get_find_agent_v2_service()
    candidates = service.candidate_inputs(run_id, candidate_ids[:8])
    payload = await fetch_details([item["aweme_id"] for item in candidates])
    service.save_details(run_id, list(payload.get("details") or []), list(payload.get("errors") or []))
    return json.dumps({
        "run_id": run_id,
        "success_count": int(payload.get("success_count") or len(payload.get("details") or [])),
        "failed_count": int(payload.get("failed_count") or len(payload.get("errors") or [])),
        "errors": payload.get("errors") or [],
    }, ensure_ascii=False)


@tool
async def fetch_candidate_portraits_v2(run_id: str, candidate_ids: list[int]) -> str:
    """批量获取候选双侧年龄画像并仅写入 find_agent_v2_candidate/evidence。"""
    service = get_find_agent_v2_service()
    candidates = service.candidate_inputs(run_id, candidate_ids[:8])
    payload = await fetch_portraits([{
            "aweme_id": item["aweme_id"],
            "author_sec_uid": item.get("author_sec_uid"),
        } for item in candidates])
    results = list(payload.get("results") or [])
    service.save_portraits(run_id, results)
    return json.dumps({"run_id": run_id, "count": len(results), "results": results}, ensure_ascii=False)


@tool
def evaluate_candidates_v2(run_id: str, items: list[dict[str, Any]]) -> str:
    """写入 R/E/S/V 与 primary/rejected；程序会对 primary 强制执行 P0 门禁。"""
    updated = get_find_agent_v2_service().evaluate(run_id, items)
    return json.dumps({"run_id": run_id, "updated": updated}, ensure_ascii=False)


@tool
def query_find_agent_v2_state(run_id: str, limit: int = 100) -> str:
    """查询完全隔离的 find_agent_v2 运行、搜索与候选状态。"""
    state = get_find_agent_v2_service().get_full_state(run_id, limit=limit)
    return json.dumps(state, ensure_ascii=False, default=str)


@tool
def query_pending_candidates_v2(run_id: str, limit: int = 100) -> str:
    """仅查询当前 run 尚未分池的 pending_evaluation 候选。"""
    state = get_find_agent_v2_service().get_full_state(
        run_id, limit=limit, pending_only=True,
    )
    return json.dumps(state, ensure_ascii=False, default=str)


SEARCH_TOOLS: tuple[ToolFn, ...] = (search_videos_v2, query_find_agent_v2_state)
EVIDENCE_TOOLS: tuple[ToolFn, ...] = (
    fetch_candidate_details_v2,
    fetch_candidate_portraits_v2,
    query_pending_candidates_v2,
)
EVALUATION_TOOLS: tuple[ToolFn, ...] = (
    evaluate_candidates_v2,
    query_pending_candidates_v2,
)
REPORT_TOOLS: tuple[ToolFn, ...] = (query_find_agent_v2_state,)


def build_tool_registry(functions: Iterable[ToolFn]) -> ToolRegistry:
    return ToolRegistry().from_decorated(*tuple(functions))