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