"""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) service = get_find_agent_v2_service() state = { "run": service.require_run(run_id), "candidates": [ item for item in service.candidate_inputs( run_id, sorted(worker_allowed), ) if item.get("decision_bucket") == "pending_evaluation" ][:max(1, int(limit))], } 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) if set(requested) != worker_allowed: return json.dumps({ "error": "必须一次处理完整 Evidence Worker 分片", "requested_candidate_ids": requested, "required_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") service = get_find_agent_v2_service() allowed_candidates = [ item for item in service.candidate_inputs( run_id, sorted(worker_allowed), ) if item.get("decision_bucket") == "pending_evaluation" ] 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) def _search_has_results(payload: dict[str, Any]) -> bool: return bool(payload.get("search_results")) async def _search_provider_pages( *, service: Any, run_id: str, round_index: int, keyword: str, query_reason: str, source_type: str, provider: str, raw_task: dict[str, Any], common: dict[str, Any], max_pages: int, cursor: str | int, provider_search_id: str = "", backtrace: str = "", extra_output: dict[str, Any] | None = None, ) -> tuple[list[dict[str, Any]], bool]: outputs: list[dict[str, Any]] = [] has_results = False for page_no in range(1, max_pages + 1): 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: 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=query_reason, source_type=source_type, provider=provider, cursor=str(cursor), page_no=page_no, payload=payload, ) if _search_has_results(payload): has_results = True 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"), **(extra_output or {}), **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 outputs, has_results @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 项。 未指定 provider 时先打内部搜索，无候选再回退 TikHub。 显式指定 internal_keyword 只用内部搜索；显式指定 tikhub 只用 TikHub。 """ 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() requested_provider = str(raw_task.get("provider") or "").strip() allow_tikhub_fallback = not requested_provider provider = "tikhub" if requested_provider == "tikhub" else "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)) source_type = str(raw_task.get("source_type") or "mixed") 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), } pages, has_results = await _search_provider_pages( service=service, run_id=run_id, round_index=round_index, keyword=keyword, query_reason=reason, source_type=source_type, provider=provider, raw_task=raw_task, common=common, max_pages=max_pages, cursor=raw_task.get("cursor") or 0, provider_search_id=str(raw_task.get("search_id") or ""), backtrace=str(raw_task.get("backtrace") or ""), ) outputs.extend(pages) if allow_tikhub_fallback and not has_results: fallback, _ = await _search_provider_pages( service=service, run_id=run_id, round_index=round_index, keyword=keyword, query_reason=reason, source_type=source_type, provider="tikhub", raw_task=raw_task, common=common, max_pages=max_pages, cursor=0, extra_output={"fallback_from": "internal_keyword"}, ) outputs.extend(fallback) 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]) requested_aweme_ids = [str(item["aweme_id"]) for item in candidates] payload = await fetch_details(requested_aweme_ids) details = list(payload.get("details") or []) errors = list(payload.get("errors") or []) service.save_details( run_id, details, errors, requested_aweme_ids=requested_aweme_ids, ) successful_ids = { str(item.get("content_id") or "") for item in details } & set(requested_aweme_ids) error_by_id = {str(item.get("content_id") or ""): item for item in errors} normalized_errors = [ error_by_id.get(aweme_id) or { "content_id": aweme_id, "error": "上游详情响应未包含已请求候选", "error_code": "UPSTREAM_RESULT_MISSING", } for aweme_id in requested_aweme_ids if aweme_id not in successful_ids ] return json.dumps({ "run_id": run_id, "success_count": len(successful_ids), "failed_count": len(normalized_errors), "errors": normalized_errors, }, 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]) requested_aweme_ids = [str(item["aweme_id"]) for item in candidates] 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, requested_aweme_ids=requested_aweme_ids, ) result_by_id = {str(item.get("aweme_id") or ""): item for item in results} successful_ids = { aweme_id for aweme_id in requested_aweme_ids if aweme_id in result_by_id and not result_by_id[aweme_id].get("error") } errors = [ result_by_id.get(aweme_id) or { "aweme_id": aweme_id, "error": "上游画像响应未包含已请求候选", "error_code": "UPSTREAM_RESULT_MISSING", } for aweme_id in requested_aweme_ids if aweme_id not in successful_ids ] return json.dumps({ "run_id": run_id, "count": len(results), "success_count": len(successful_ids), "failed_count": len(errors), "errors": errors, "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))