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]


def bound_candidate_tools(
    functions: tuple[ToolFn, ...], *, run_id: str, candidate_ids: list[int],
) -> tuple[ToolFn, ...]:
    """Physically constrain worker tools to one run and one candidate shard."""
    allowed = {int(value) for value in candidate_ids}
    output: list[ToolFn] = []
    for original in functions:
        name = str(getattr(original, "_tool_name", original.__name__))
        if name == "query_pending_candidates_v2":
            def make_query(fn_name: str, worker_run_id: str, worker_allowed: set[int]):
                @tool(name=fn_name, description="仅查询当前 Worker 分片中的待评估候选。")
                def bound_query(run_id: str, limit: int = 100) -> str:
                    if run_id != worker_run_id:
                        return json.dumps({"error": "run_id 不属于当前 Worker"}, ensure_ascii=False)
                    state = get_find_agent_v2_service().get_full_state(
                        run_id, limit=limit, pending_only=True,
                    )
                    state["candidates"] = [
                        item for item in state["candidates"]
                        if int(item["candidate_id"]) in worker_allowed
                    ]
                    return json.dumps(state, ensure_ascii=False, default=str)

                return bound_query

            output.append(make_query(name, run_id, allowed))
            continue

        def make_async(fn: ToolFn, fn_name: str, worker_run_id: str, worker_allowed: set[int]):
            @tool(name=fn_name, description=getattr(fn, "_tool_description", ""))
            async def bound_async(run_id: str, candidate_ids: list[int]) -> str:
                if run_id != worker_run_id:
                    return json.dumps({"error": "run_id 不属于当前 Worker"}, ensure_ascii=False)
                requested = [int(value) for value in candidate_ids]
                if not requested or not set(requested) <= worker_allowed:
                    return json.dumps({
                        "error": "candidate_ids 超出当前 Worker 分片",
                        "allowed_candidate_ids": sorted(worker_allowed),
                    }, ensure_ascii=False)
                return await fn(run_id=run_id, candidate_ids=requested)

            return bound_async

        def make_sync(fn: ToolFn, fn_name: str, worker_run_id: str, worker_allowed: set[int]):
            @tool(name=fn_name, description=getattr(fn, "_tool_description", ""))
            def bound_sync(run_id: str, items: list[dict[str, Any]]) -> str:
                if run_id != worker_run_id:
                    return json.dumps({"error": "run_id 不属于当前 Worker"}, ensure_ascii=False)
                requested = {int(item.get("candidate_id") or 0) for item in items}
                if not requested or not requested <= worker_allowed:
                    return json.dumps({
                        "error": "items 超出当前 Worker 分片",
                        "allowed_candidate_ids": sorted(worker_allowed),
                    }, ensure_ascii=False)
                return fn(run_id=run_id, items=items)

            return bound_sync

        wrapper: ToolFn
        if name in {"fetch_candidate_details_v2", "fetch_candidate_portraits_v2"}:
            wrapper = make_async(original, name, run_id, allowed)
        elif name == "evaluate_candidates_v2":
            wrapper = make_sync(original, name, run_id, allowed)
        else:
            wrapper = original
        output.append(wrapper)
    return tuple(output)


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