SupplyAgent/api/services/find_agent_v2.py

"""Admin-facing read model and execution service for find_agent_v2."""

from __future__ import annotations

import json
from datetime import datetime, timedelta
from typing import Any
from zoneinfo import ZoneInfo

from sqlalchemy import func, or_, select, update

from find_agent_v2.demand_context import prepare_latest_v2_demand_run_by_name
from find_agent_v2.models import (
    FindAgentV2Candidate,
    FindAgentV2Evidence,
    FindAgentV2Round,
    FindAgentV2Run,
    FindAgentV2Search,
)
from find_agent_v2.runner import run_prepared_find_agent_v2
from find_agent_v2.service import (
    RUN_TIMEOUT_MINUTES,
    RUN_TIMEOUT_REASON,
    get_find_agent_v2_service,
)
from supply_infra.db.session import get_session


def _expire_overdue_runs(session, *, now: datetime | None = None) -> int:
    current = now or datetime.now(ZoneInfo("Asia/Shanghai")).replace(tzinfo=None)
    cutoff = current - timedelta(minutes=RUN_TIMEOUT_MINUTES)
    result = session.execute(
        update(FindAgentV2Run)
        .where(
            FindAgentV2Run.status == "running",
            FindAgentV2Run.create_time < cutoff,
        )
        .values(
            status="failed",
            outcome_status="failed",
            stop_reason=RUN_TIMEOUT_REASON,
            update_time=current,
        )
    )
    return int(result.rowcount or 0)


def _json(raw: str | None, default: Any) -> Any:
    try:
        return json.loads(raw) if raw else default
    except (TypeError, ValueError):
        return default


def _time(value: Any) -> str | None:
    return value.isoformat() if value is not None else None


def _number(value: Any) -> float | None:
    return float(value) if value is not None else None


def _run(row: FindAgentV2Run) -> dict[str, Any]:
    return {
        "id": int(row.id),
        "run_id": row.run_id,
        "demand_grade_id": row.demand_grade_id,
        "demand_word": row.demand_word,
        "status": row.status,
        "outcome_status": row.outcome_status,
        "current_round": row.current_round,
        "search_count": row.search_count,
        "candidate_count": row.candidate_count,
        "valid_primary_count": row.valid_primary_count,
        "input_tokens": row.input_tokens,
        "output_tokens": row.output_tokens,
        "total_tokens": row.total_tokens,
        "cost_usd": _number(row.cost_usd) or 0,
        "intent_summary": row.intent_summary,
        "stop_reason": row.stop_reason,
        "obagent_run_uid": row.obagent_run_uid,
        "triggered_by_user_id": row.triggered_by_user_id,
        "triggered_by_username": row.triggered_by_username,
        "trigger_source": row.trigger_source,
        "create_time": _time(row.create_time),
        "update_time": _time(row.update_time),
    }


def _evidence_item(row: FindAgentV2Evidence) -> dict[str, Any]:
    return {
        "id": int(row.id), "candidate_id": row.candidate_id,
        "evidence_type": row.evidence_type, "provider": row.provider,
        "status": row.status, "raw": _json(row.raw_json, {}),
        "normalized": _json(row.normalized_json, {}),
        "error_message": row.error_message, "create_time": _time(row.create_time),
    }


def _candidate_item(
    row: FindAgentV2Candidate, evidence: list[dict[str, Any]] | None = None,
) -> dict[str, Any]:
    return {
        "id": int(row.id), "first_search_id": row.first_search_id,
        "aweme_id": row.aweme_id, "title": row.title, "content_link": row.content_link,
        "author_name": row.author_name, "author_sec_uid": row.author_sec_uid,
        "source_keywords": _json(row.source_keywords_json, []), "tags": _json(row.tags_json, []),
        "publish_at": _time(row.publish_at), "duration_seconds": _number(row.duration_seconds),
        "play_count": row.play_count, "like_count": row.like_count,
        "comment_count": row.comment_count, "collect_count": row.collect_count,
        "share_count": row.share_count, "detail": _json(row.detail_json, {}),
        "portrait": _json(row.portrait_json, {}), "detail_status": row.detail_status,
        "portrait_status": row.portrait_status,
        "content_50_plus_ratio": _number(row.content_50_plus_ratio),
        "account_50_plus_ratio": _number(row.account_50_plus_ratio),
        "relevance_score": _number(row.relevance_score), "elder_score": _number(row.elder_score),
        "share_score": _number(row.share_score), "value_score": _number(row.value_score),
        "gate_status": row.gate_status, "gate_result": _json(row.gate_result_json, {}),
        "decision_bucket": row.decision_bucket, "decision_reason": row.decision_reason,
        "reject_reason_code": row.reject_reason_code, "evidence": evidence or [],
        "create_time": _time(row.create_time), "update_time": _time(row.update_time),
    }


def _candidate_rows_with_evidence(session, query) -> list[dict[str, Any]]:
    rows = list(session.scalars(query))
    ids = [int(row.id) for row in rows]
    evidence_rows = list(session.scalars(
        select(FindAgentV2Evidence).where(FindAgentV2Evidence.candidate_id.in_(ids))
        .order_by(FindAgentV2Evidence.id)
    )) if ids else []
    evidence_by_candidate: dict[int, list[dict[str, Any]]] = {}
    for row in evidence_rows:
        if row.candidate_id is not None:
            evidence_by_candidate.setdefault(int(row.candidate_id), []).append(
                _evidence_item(row)
            )
    return [
        _candidate_item(row, evidence_by_candidate.get(int(row.id), []))
        for row in rows
    ]


def list_runs(
    *, status: str | None = None, keyword: str | None = None,
    page: int = 1, page_size: int = 30,
) -> dict[str, Any]:
    normalized_page = max(1, int(page))
    normalized_page_size = min(100, max(1, int(page_size)))
    offset = (normalized_page - 1) * normalized_page_size
    with get_session() as session:
        _expire_overdue_runs(session)
        conditions = []
        if status:
            conditions.append(FindAgentV2Run.status == status)
        normalized = str(keyword or "").strip().lower()
        if normalized:
            pattern = f"%{normalized}%"
            conditions.append(or_(
                func.lower(FindAgentV2Run.demand_word).like(pattern),
                func.lower(FindAgentV2Run.run_id).like(pattern),
                func.lower(func.coalesce(FindAgentV2Run.triggered_by_username, "")).like(pattern),
            ))
        total = int(session.scalar(
            select(func.count()).select_from(FindAgentV2Run).where(*conditions)
        ) or 0)
        rows = list(session.scalars(
            select(FindAgentV2Run).where(*conditions)
            .order_by(FindAgentV2Run.create_time.desc(), FindAgentV2Run.id.desc())
            .limit(normalized_page_size).offset(offset)
        ))
        total_pages = max(1, (total + normalized_page_size - 1) // normalized_page_size)
        return {
            "items": [_run(row) for row in rows],
            "total": total,
            "page": normalized_page,
            "page_size": normalized_page_size,
            "total_pages": total_pages,
            "has_previous": normalized_page > 1,
            "has_next": normalized_page < total_pages,
        }


def get_run_detail(run_id: str) -> dict[str, Any] | None:
    with get_session() as session:
        _expire_overdue_runs(session)
        run = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
        if run is None:
            return None
        rounds = list(session.scalars(
            select(FindAgentV2Round).where(FindAgentV2Round.run_id == run_id)
            .order_by(FindAgentV2Round.round_index, FindAgentV2Round.id)
        ))
        searches = list(session.scalars(
            select(FindAgentV2Search).where(FindAgentV2Search.run_id == run_id)
            .order_by(FindAgentV2Search.round_index, FindAgentV2Search.id)
        ))
        candidate_summaries = list(session.execute(select(
            FindAgentV2Candidate.aweme_id, FindAgentV2Candidate.decision_bucket,
        ).where(FindAgentV2Candidate.run_id == run_id)))
        round_items = [{
            "id": int(row.id), "round_index": row.round_index, "phase": row.phase,
            "status": row.status, "plan": _json(row.plan_json, {}),
            "start_snapshot": _json(row.start_snapshot_json, {}),
            "end_snapshot": _json(row.end_snapshot_json, {}),
            "error_message": row.error_message, "create_time": _time(row.create_time),
            "update_time": _time(row.update_time),
        } for row in rounds]
        search_items = [{
            "id": int(row.id), "round_index": row.round_index, "keyword": row.keyword,
            "query_reason": row.query_reason, "source_type": row.source_type,
            "provider": row.provider, "cursor": row.cursor, "page_no": row.page_no,
            "has_more": bool(row.has_more), "next_cursor": row.next_cursor,
            "result_count": row.result_count, "status": row.status,
            "error_message": row.error_message,
            "create_time": _time(row.create_time),
        } for row in searches]
        bucket_by_aweme = {
            str(aweme_id): str(bucket) for aweme_id, bucket in candidate_summaries
        }
        search_rows_by_id = {int(row.id): row for row in searches}
        for search in search_items:
            row = search_rows_by_id[search["id"]]
            raw_results = _json(row.raw_response_json, {}).get("search_results", [])
            result_ids = list(dict.fromkeys(
                str(item.get("aweme_id"))
                for item in raw_results
                if isinstance(item, dict) and item.get("aweme_id")
            ))
            buckets = [bucket_by_aweme.get(item, "pending_evaluation") for item in result_ids]
            search["video_count"] = len(result_ids)
            search["primary_count"] = sum(item == "primary" for item in buckets)
            search["rejected_count"] = sum(item == "rejected" for item in buckets)
            search["pending_count"] = (
                search["video_count"] - search["primary_count"] - search["rejected_count"]
            )
        timeline = []
        for row in round_items:
            timeline.append({"type": "round", "time": row["create_time"], "data": row})
        for row in search_items:
            timeline.append({"type": "search", "time": row["create_time"], "data": row})
        timeline.sort(key=lambda item: str(item.get("time") or ""))
        return {
            "run": {**_run(run), "input": _json(run.input_json, {}),
                    "rule_config": _json(run.rule_config_json, {})},
            "rounds": round_items, "searches": search_items, "timeline": timeline,
        }


def get_run_candidates(run_id: str) -> list[dict[str, Any]] | None:
    with get_session() as session:
        if session.scalar(select(FindAgentV2Run.id).where(FindAgentV2Run.run_id == run_id)) is None:
            return None
        return _candidate_rows_with_evidence(session, select(FindAgentV2Candidate).where(
            FindAgentV2Candidate.run_id == run_id,
        ).order_by(FindAgentV2Candidate.id))


def get_search_candidates(run_id: str, search_id: int) -> list[dict[str, Any]] | None:
    with get_session() as session:
        search = session.scalar(select(FindAgentV2Search).where(
            FindAgentV2Search.run_id == run_id, FindAgentV2Search.id == int(search_id),
        ))
        if search is None:
            return None
        raw_results = _json(search.raw_response_json, {}).get("search_results", [])
        aweme_ids = list(dict.fromkeys(
            str(item.get("aweme_id")) for item in raw_results
            if isinstance(item, dict) and item.get("aweme_id")
        ))
        if not aweme_ids:
            return []
        items = _candidate_rows_with_evidence(session, select(FindAgentV2Candidate).where(
            FindAgentV2Candidate.run_id == run_id,
            FindAgentV2Candidate.aweme_id.in_(aweme_ids),
        ))
        by_aweme = {item["aweme_id"]: item for item in items}
        return [by_aweme[item] for item in aweme_ids if item in by_aweme]


def prepare_test_run(demand_word: str, current_user: dict[str, Any]) -> dict[str, Any]:
    prepared = prepare_latest_v2_demand_run_by_name(
        demand_word,
        triggered_by_user_id=int(current_user["id"]),
        triggered_by_username=str(current_user["username"]),
        trigger_source="web_admin",
    )
    return prepared.summary()


def execute_test_run(run_id: str) -> None:
    """Background task entry point. Always persist bootstrap failures."""
    try:
        run_prepared_find_agent_v2(run_id)
    except Exception as exc:
        get_find_agent_v2_service().fail_run(run_id, f"{type(exc).__name__}: {exc}")