SupplyAgent/scripts/verify_find_agent_live.py

"""Run a real find_agent acceptance case and clean its database records by default."""

from __future__ import annotations

import argparse
import asyncio
import json
from typing import Any

from sqlalchemy import delete, func, select

from agents.find_agent import create_find_agent
from agents.find_agent.agent import _validate_report_buckets
from supply_agent.config import Settings
from supply_agent.types import AgentEventType
from supply_infra.db.models.video_discovery import (
    VideoDiscoveryCandidate,
    VideoDiscoveryRun,
    VideoDiscoverySearch,
)
from supply_infra.db.session import get_session


def _parse_json(raw: str) -> dict[str, Any]:
    try:
        value = json.loads(raw)
    except (TypeError, json.JSONDecodeError):
        return {}
    return value if isinstance(value, dict) else {}


def _cleanup_run(run_id: str) -> int:
    with get_session() as session:
        session.execute(
            delete(VideoDiscoveryCandidate).where(
                VideoDiscoveryCandidate.run_id == run_id
            )
        )
        session.execute(
            delete(VideoDiscoverySearch).where(
                VideoDiscoverySearch.run_id == run_id
            )
        )
        session.execute(
            delete(VideoDiscoveryRun).where(VideoDiscoveryRun.run_id == run_id)
        )
    with get_session() as session:
        remaining = session.scalar(
            select(func.count(VideoDiscoveryRun.id)).where(
                VideoDiscoveryRun.run_id == run_id
            )
        )
    return int(remaining or 0)


async def _run(args: argparse.Namespace) -> None:
    agent = create_find_agent(settings=Settings.from_env(LOG_ENABLED=False))
    relevant_points = [{"point": point} for point in args.relevant_point]
    prompt = (
        "这是一次真实验收测试，请完整执行并严格遵守完成条件。\n"
        f"demand_word：{args.demand_word}\n"
        f"seed_video_title：{args.seed_video_title}\n"
        f"relevant_points：{json.dumps(relevant_points, ensure_ascii=False)}\n"
        "请寻找老年人或临近退休人群可能喜欢观看和分享的视频。"
    )

    tool_calls: list[str] = []
    tool_payloads: list[dict[str, Any]] = []
    run_ids: list[str] = []
    domain_errors: list[dict[str, Any]] = []
    audit_results: list[dict[str, Any]] = []
    states: list[dict[str, Any]] = []
    final_content = ""
    done_data: dict[str, Any] = {}

    try:
        async for event in agent.astream(prompt):
            if event.type == AgentEventType.TOOL_CALL:
                name = str(event.data.get("name") or "")
                call_args = _parse_json(str(event.data.get("arguments") or "{}"))
                tool_calls.append(name)
                tool_payloads.append({"name": name, "args": call_args})
                print(f"TOOL_CALL {len(tool_calls)} {name}", flush=True)
            elif event.type == AgentEventType.TOOL_RESULT:
                name = str(event.data.get("name") or "")
                data = _parse_json(str(event.data.get("content") or ""))
                if name == "create_video_discovery_run" and data.get("run_id"):
                    run_ids.append(str(data["run_id"]))
                if name in {
                    "audit_video_discovery_process",
                    "audit_video_discovery_run",
                } and not data.get("error"):
                    audit_results.append(data)
                if name == "query_video_discovery_state" and not data.get("error"):
                    states.append(data)
                if data.get("error"):
                    domain_errors.append(
                        {
                            "tool": name,
                            "budget_exhausted": bool(data.get("budget_exhausted")),
                            "error": str(data["error"])[:300],
                        }
                    )
                status = "DOMAIN_ERROR" if data.get("error") else "OK"
                if event.data.get("is_error") and not data.get("error"):
                    status = "ERROR"
                extra = ""
                if data.get("error"):
                    extra = " error=" + json.dumps(
                        str(data["error"])[:240],
                        ensure_ascii=False,
                    )
                if name in {
                    "audit_video_discovery_process",
                    "audit_video_discovery_run",
                } and not data.get("error"):
                    audit_flag = data.get("can_finish")
                    extra = f" can_finish={audit_flag}"
                    if audit_flag is False:
                        extra += " violations=" + json.dumps(
                            data.get("critical_violations"),
                            ensure_ascii=False,
                        )
                print(f"TOOL_RESULT {name} {status}{extra}", flush=True)
            elif event.type == AgentEventType.MESSAGE:
                final_content = str(event.data.get("content") or "")
            elif event.type == AgentEventType.DONE:
                done_data = event.data
                if not final_content:
                    final_content = str(event.data.get("content") or "")

        unique_keywords: list[str] = []
        recorded_sources: list[Any] = []
        pagination_calls = 0
        for entry in tool_payloads:
            name = entry["name"]
            call_args = entry["args"]
            if name in {"douyin_search", "douyin_search_tikhub"}:
                keyword = str(call_args.get("keyword") or "").strip()
                if keyword and keyword not in unique_keywords:
                    unique_keywords.append(keyword)
                if (
                    str(call_args.get("cursor") or "0") not in {"", "0"}
                    or call_args.get("search_id")
                ):
                    pagination_calls += 1
            if name == "record_video_search_page":
                recorded_sources.append(call_args.get("source_type"))

        last_state = states[-1] if states else {}
        final_bucket_error = (
            _validate_report_buckets(final_content, last_state)
            if final_content and last_state
            else "缺少最终文本或最终状态"
        )
        report = {
            "model": agent.model,
            "iterations": done_data.get("iterations"),
            "tool_calls": len(tool_calls),
            "tool_names": tool_calls,
            "unique_search_keywords": unique_keywords,
            "pagination_calls": pagination_calls,
            "recorded_source_types": recorded_sources,
            "qwen_calls": tool_calls.count("qwen_video_analyze"),
            "audit_calls": len(audit_results),
            "last_audit_can_finish": (
                audit_results[-1].get("can_finish") if audit_results else None
            ),
            "last_audit_violations": (
                audit_results[-1].get("critical_violations")
                if audit_results
                else None
            ),
            "domain_errors": domain_errors,
            "state_search_count": len(last_state.get("searches", [])),
            "state_candidate_count": len(last_state.get("candidates", [])),
            "state_run": last_state.get("run"),
            "final_content_chars": len(final_content),
            "final_has_primary_section": "主推荐" in final_content,
            "final_has_backup_section": "补充推荐" in final_content,
            "final_bucket_consistent": final_bucket_error is None,
            "final_bucket_error": final_bucket_error,
            "max_iterations_failure": final_content.startswith(
                "Max iterations reached"
            ),
            "run_ids": run_ids,
        }
        print("AGENT_REPORT " + json.dumps(report, ensure_ascii=False), flush=True)
        print("FINAL_CONTENT_BEGIN", flush=True)
        print(final_content, flush=True)
        print("FINAL_CONTENT_END", flush=True)
    finally:
        if not args.keep_data:
            for run_id in set(run_ids):
                remaining = _cleanup_run(run_id)
                print(f"CLEANUP {run_id} remaining={remaining}", flush=True)


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--demand-word",
        default="个人养老金税收优惠",
    )
    parser.add_argument(
        "--seed-video-title",
        default="个人养老金制度全面实施，退休前这样缴纳可以享受税收优惠",
    )
    parser.add_argument(
        "--relevant-point",
        action="append",
        default=[
            "个人养老金缴费如何抵扣个税",
            "适合临近退休人群转发给家人了解",
        ],
    )
    parser.add_argument(
        "--keep-data",
        action="store_true",
        help="保留本次运行的数据库记录；默认清理。",
    )
    return parser.parse_args()


if __name__ == "__main__":
    asyncio.run(_run(_arguments()))