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 typing import Any

from sqlalchemy import func, or_, select

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 get_find_agent_v2_service
from supply_infra.db.session import get_session


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 list_runs(
    *, status: str | None = None, keyword: str | None = None,
    limit: int = 30, offset: int = 0,
) -> dict[str, Any]:
    with get_session() as 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(limit).offset(offset)
        ))
        return {"items": [_run(row) for row in rows], "total": total, "limit": limit, "offset": offset}


def get_run_detail(run_id: str) -> dict[str, Any] | None:
    with get_session() as 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)
        ))
        candidates = list(session.scalars(
            select(FindAgentV2Candidate).where(FindAgentV2Candidate.run_id == run_id)
            .order_by(FindAgentV2Candidate.id)
        ))
        evidence = list(session.scalars(
            select(FindAgentV2Evidence).where(FindAgentV2Evidence.run_id == run_id)
            .order_by(FindAgentV2Evidence.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,
            "provider_state": _json(row.provider_state_json, {}),
            "result_count": row.result_count, "status": row.status,
            "error_message": row.error_message,
            "raw_response": _json(row.raw_response_json, {}),
            "create_time": _time(row.create_time),
        } for row in searches]
        evidence_by_candidate: dict[int, list[dict[str, Any]]] = {}
        evidence_items = []
        for row in evidence:
            item = {
                "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),
            }
            evidence_items.append(item)
            if row.candidate_id is not None:
                evidence_by_candidate.setdefault(int(row.candidate_id), []).append(item)
        candidate_items = [{
            "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_by_candidate.get(int(row.id), []),
            "create_time": _time(row.create_time), "update_time": _time(row.update_time),
        } for row in candidates]
        candidates_by_aweme = {item["aweme_id"]: item for item in candidate_items}
        for search in search_items:
            raw_results = search["raw_response"].get("search_results", [])
            result_ids = [
                str(item.get("aweme_id") or "")
                for item in raw_results
                if isinstance(item, dict) and item.get("aweme_id")
            ]
            search["result_aweme_ids"] = result_ids
            search["matched_candidates"] = [
                candidates_by_aweme[aweme_id]
                for aweme_id in result_ids
                if aweme_id in candidates_by_aweme
            ]
        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})
        for row in evidence_items:
            timeline.append({"type": "evidence", "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, "candidates": candidate_items,
            "evidence": evidence_items, "timeline": timeline,
        }


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