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 pydantic import BaseModel, Field, model_validator

from find_agent_v2.qwen_video_understanding_30s import understand_candidate_video_30s_v2
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]


class CandidateEvaluation(BaseModel):
    """One candidate decision. Prefer candidate_id; aweme_id is a safe fallback."""

    candidate_id: int | None = Field(
        default=None,
        description="数据库候选编号，即输入中的 candidate_id；不要填写 aweme_id。",
    )
    aweme_id: str | None = Field(
        default=None,
        description="抖音视频号；仅在无法填写 candidate_id 时作为兼容定位字段。",
    )
    relevance_score: float = Field(ge=0, le=1, description="R 相关性评分，0~1。")
    elder_score: float = Field(ge=0, le=1, description="E 中老年适配评分，0~1。")
    share_score: float = Field(ge=0, le=1, description="S 传播力评分，0~1。")
    value_score: float = Field(ge=0, le=1, description="V 综合价值评分，0~1。")
    decision_bucket: str = Field(description="只能是 primary 或 rejected。")
    decision_reason: str = Field(min_length=1, description="基于真实证据的中文判断理由。")
    reject_reason_code: str | None = None

    @model_validator(mode="after")
    def validate_identity_and_bucket(self):
        if self.candidate_id is None and not str(self.aweme_id or "").strip():
            raise ValueError("candidate_id 和 aweme_id 至少提供一个")
        if self.decision_bucket not in {"primary", "rejected"}:
            raise ValueError("decision_bucket 只能是 primary/rejected")
        return self


class EvaluationBatch(BaseModel):
    """Structured evaluator output consumed and persisted by the host."""

    items: list[CandidateEvaluation] = Field(
        min_length=1,
        description="当前 Worker 分片内全部候选的评估结果，每个候选恰好一条。",
    )


def normalize_evaluation_items(
    items: list[dict[str, Any] | CandidateEvaluation],
    *,
    allowed_candidates: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """Validate a complete worker shard and resolve aweme_id without widening access."""
    allowed_ids = {int(item["candidate_id"]) for item in allowed_candidates}
    aweme_to_id = {
        str(item.get("aweme_id") or ""): int(item["candidate_id"])
        for item in allowed_candidates
    }
    normalized: list[dict[str, Any]] = []
    for raw in items:
        value = raw if isinstance(raw, CandidateEvaluation) else CandidateEvaluation.model_validate(raw)
        candidate_id = value.candidate_id
        if candidate_id is None:
            candidate_id = aweme_to_id.get(str(value.aweme_id or ""))
        if candidate_id not in allowed_ids:
            raise ValueError(
                "评估项不属于当前 Worker 分片；"
                f"candidate_id={candidate_id}, aweme_id={value.aweme_id}, "
                f"allowed_candidate_ids={sorted(allowed_ids)}"
            )
        payload = value.model_dump(exclude_none=True)
        payload["candidate_id"] = candidate_id
        payload.pop("aweme_id", None)
        normalized.append(payload)
    requested = [int(item["candidate_id"]) for item in normalized]
    if len(requested) != len(set(requested)):
        raise ValueError("同一 candidate_id 不能重复评估")
    missing = allowed_ids - set(requested)
    if missing:
        raise ValueError(f"必须一次评估完整分片，缺少 candidate_ids={sorted(missing)}")
    return normalized


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_video_understanding(
            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_video_understanding(
                run_id: str, candidate_id: int, prompt: str,
            ) -> str:
                if run_id != worker_run_id or int(candidate_id) not in worker_allowed:
                    return json.dumps({
                        "error": "candidate_id 或 run_id 超出当前评估 Worker 分片",
                        "allowed_candidate_ids": sorted(worker_allowed),
                    }, ensure_ascii=False)
                return await fn(run_id=run_id, candidate_id=int(candidate_id), prompt=prompt)

            return bound_video_understanding

        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[CandidateEvaluation]) -> str:
                if run_id != worker_run_id:
                    raise ValueError("run_id 不属于当前 Worker")
                state = get_find_agent_v2_service().get_full_state(
                    run_id, pending_only=True,
                )
                allowed_candidates = [
                    item for item in state.get("candidates", [])
                    if int(item["candidate_id"]) in worker_allowed
                ]
                normalized = normalize_evaluation_items(
                    items, allowed_candidates=allowed_candidates,
                )
                return fn(run_id=run_id, items=normalized)

            return bound_sync

        wrapper: ToolFn
        if name == "understand_candidate_video_30s_v2":
            wrapper = make_video_understanding(original, name, run_id, allowed)
        elif 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[CandidateEvaluation]) -> str:
    """按 candidate_id 写入完整分片的 R/E/S/V 与分池；不要用 aweme_id 代替。"""
    payload = [item.model_dump(exclude_none=True) if isinstance(item, CandidateEvaluation) else item for item in items]
    updated = get_find_agent_v2_service().evaluate(run_id, payload)
    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, ...] = (
    understand_candidate_video_30s_v2,
    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))