SupplyAgent/find_agent_v2/service.py

"""Transactional service for the isolated ``find_agent_v2_*`` tables."""

from __future__ import annotations

import json
import uuid
from dataclasses import asdict
from decimal import Decimal
from typing import Any

from sqlalchemy import func, select

from find_agent_v2.models import (
    FindAgentV2Candidate,
    FindAgentV2Evidence,
    FindAgentV2Round,
    FindAgentV2Run,
    FindAgentV2Search,
)
from find_agent_v2.state import DiscoverySnapshot
from supply_infra.db.session import get_session
from find_agent_v2.gates import (
    build_rule_snapshot,
    evaluate_candidate_gate,
    parse_datetime_value,
)


class FindAgentV2RunNotFound(LookupError):
    pass


def _json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, default=str)


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


def _ratio(value: Any) -> Decimal | None:
    if value in (None, "") or isinstance(value, bool):
        return None
    number = Decimal(str(value))
    if number > 1 and number <= 100:
        number /= 100
    if number < 0 or number > 1:
        raise ValueError("ratio 必须在 0~1")
    return number.quantize(Decimal("0.000001"))


def _candidate_dict(row: FindAgentV2Candidate) -> dict[str, Any]:
    return {
        "candidate_id": int(row.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": _loads(row.source_keywords_json, []),
        "tags": _loads(row.tags_json, []),
        "publish_at": row.publish_at.isoformat() if row.publish_at else None,
        "duration_seconds": float(row.duration_seconds) if row.duration_seconds is not None else None,
        "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_status": row.detail_status,
        "portrait_status": row.portrait_status,
        "content_50_plus_ratio": float(row.content_50_plus_ratio) if row.content_50_plus_ratio is not None else None,
        "account_50_plus_ratio": float(row.account_50_plus_ratio) if row.account_50_plus_ratio is not None else None,
        "relevance_score": float(row.relevance_score) if row.relevance_score is not None else None,
        "elder_score": float(row.elder_score) if row.elder_score is not None else None,
        "share_score": float(row.share_score) if row.share_score is not None else None,
        "value_score": float(row.value_score) if row.value_score is not None else None,
        "gate_status": row.gate_status,
        "gate_result": _loads(row.gate_result_json, {}),
        "decision_bucket": row.decision_bucket,
        "decision_reason": row.decision_reason,
        "reject_reason_code": row.reject_reason_code,
    }


class FindAgentV2Service:
    def create_run(
        self,
        *,
        user_input: str,
        demand_word: str,
        demand_grade_id: int | None = None,
        run_id: str | None = None,
        rule_config: dict[str, Any] | None = None,
    ) -> str:
        run_key = str(run_id or uuid.uuid4().hex)[:64]
        with get_session() as session:
            exists = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_key))
            if exists is not None:
                raise ValueError(f"find_agent_v2 run_id 已存在: {run_key}")
            rules = rule_config or build_rule_snapshot()
            session.add(FindAgentV2Run(
                run_id=run_key,
                demand_grade_id=demand_grade_id,
                demand_word=str(demand_word)[:256],
                input_json=_json({"user_input": user_input}),
                rule_config_json=_json(rules),
                status="running",
            ))
        return run_key

    def lookup_run(self, run_id: str) -> dict[str, Any] | None:
        with get_session() as session:
            row = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if row is None:
                return None
            return {
                "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": float(row.cost_usd or 0),
                "intent_summary": row.intent_summary,
                "stop_reason": row.stop_reason,
                "obagent_run_uid": row.obagent_run_uid,
                "rule_config": _loads(row.rule_config_json, {}),
            }

    def set_obagent_run_uid(self, run_id: str, run_uid: str | None) -> None:
        if not run_uid:
            return
        with get_session() as session:
            row = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if row is None:
                raise FindAgentV2RunNotFound(run_id)
            row.obagent_run_uid = str(run_uid)[:64]

    def add_usage(self, run_id: str, usage: dict[str, Any]) -> None:
        """Accumulate one execution attempt; resume never erases earlier usage."""
        with get_session() as session:
            row = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if row is None:
                raise FindAgentV2RunNotFound(run_id)
            row.input_tokens += int(usage.get("input_tokens") or 0)
            row.output_tokens += int(usage.get("output_tokens") or 0)
            row.total_tokens += int(usage.get("total_tokens") or 0)
            row.cost_usd = (
                Decimal(str(row.cost_usd or 0)) + Decimal(str(usage.get("cost") or 0))
            ).quantize(Decimal("0.00000001"))

    def prepare_resume(self, run_id: str) -> dict[str, Any]:
        """Reset a terminal run for a recovery round without deleting audit rows."""
        with get_session() as session:
            row = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if row is None:
                raise FindAgentV2RunNotFound(run_id)
            if row.status == "running":
                raise ValueError(f"run_id={run_id} 已处于 running，不能重复 resume")
            row.status = "running"
            row.outcome_status = None
            row.stop_reason = None
        return self.require_run(run_id)

    def require_run(self, run_id: str) -> dict[str, Any]:
        run = self.lookup_run(run_id)
        if run is None:
            raise FindAgentV2RunNotFound(f"find_agent_v2 run_id 不存在: {run_id}")
        return run

    def get_run_user_input(self, run_id: str) -> str:
        """Return the immutable task input stored when a v2 run was prepared."""
        with get_session() as session:
            row = session.scalar(select(FindAgentV2Run).where(
                FindAgentV2Run.run_id == run_id
            ))
            if row is None:
                raise FindAgentV2RunNotFound(run_id)
            user_input = _loads(row.input_json, {}).get("user_input")
            if not isinstance(user_input, str) or not user_input.strip():
                raise ValueError(f"run_id={run_id} 缺少有效 user_input")
            return user_input

    def begin_round(self, run_id: str, round_index: int, snapshot: DiscoverySnapshot) -> None:
        with get_session() as session:
            run = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if run is None:
                raise FindAgentV2RunNotFound(run_id)
            run.current_round = int(round_index)
            session.add(FindAgentV2Round(
                run_id=run_id,
                round_index=int(round_index),
                phase="planning",
                status="open",
                start_snapshot_json=_json(asdict(snapshot)),
            ))

    def update_round(
        self,
        run_id: str,
        round_index: int,
        *,
        phase: str | None = None,
        plan: str | None = None,
        status: str | None = None,
        snapshot: DiscoverySnapshot | None = None,
        error: str | None = None,
    ) -> None:
        with get_session() as session:
            row = session.scalar(select(FindAgentV2Round).where(
                FindAgentV2Round.run_id == run_id,
                FindAgentV2Round.round_index == int(round_index),
            ))
            if row is None:
                raise FindAgentV2RunNotFound(f"round 不存在: {run_id}/{round_index}")
            if phase is not None:
                row.phase = phase
            if plan is not None:
                row.plan_json = plan
            if status is not None:
                row.status = status
            if snapshot is not None:
                row.end_snapshot_json = _json(asdict(snapshot))
            if error is not None:
                row.error_message = error[:2000]

    def save_search(
        self,
        *,
        run_id: str,
        round_index: int,
        keyword: str,
        query_reason: str,
        source_type: str,
        provider: str,
        cursor: str,
        page_no: int,
        payload: dict[str, Any],
    ) -> dict[str, Any]:
        results = list(payload.get("search_results") or [])
        with get_session() as session:
            run = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if run is None:
                raise FindAgentV2RunNotFound(run_id)
            search = FindAgentV2Search(
                run_id=run_id,
                round_index=round_index,
                keyword=keyword[:256],
                query_reason=query_reason,
                source_type=source_type[:32],
                provider=provider[:32],
                cursor=str(cursor)[:128],
                page_no=page_no,
                has_more=int(bool(payload.get("has_more"))),
                next_cursor=str(payload.get("next_cursor") or "")[:128] or None,
                provider_state_json=_json({
                    "search_id": payload.get("search_id"),
                    "backtrace": payload.get("backtrace"),
                }),
                result_count=len(results),
                status="failed" if payload.get("error") else "success",
                error_message=str(payload.get("error") or "") or None,
                raw_response_json=_json(payload),
            )
            session.add(search)
            session.flush()
            new_count = 0
            for item in results:
                aweme_id = str(item.get("aweme_id") or "").strip()
                if not aweme_id:
                    continue
                candidate = session.scalar(select(FindAgentV2Candidate).where(
                    FindAgentV2Candidate.run_id == run_id,
                    FindAgentV2Candidate.aweme_id == aweme_id,
                ))
                author = item.get("author") if isinstance(item.get("author"), dict) else {}
                stats = item.get("statistics") if isinstance(item.get("statistics"), dict) else {}
                if candidate is None:
                    candidate = FindAgentV2Candidate(
                        run_id=run_id,
                        first_search_id=search.id,
                        aweme_id=aweme_id,
                        decision_bucket="pending_evaluation",
                    )
                    session.add(candidate)
                    new_count += 1
                keywords = _loads(candidate.source_keywords_json, [])
                if keyword not in keywords:
                    keywords.append(keyword)
                candidate.source_keywords_json = _json(keywords)
                candidate.title = str(item.get("desc") or item.get("title") or candidate.title or "")[:512] or None
                candidate.content_link = str(item.get("url") or candidate.content_link or "")[:1024] or None
                candidate.author_name = str(author.get("nickname") or candidate.author_name or "")[:256] or None
                candidate.author_sec_uid = str(author.get("sec_uid") or candidate.author_sec_uid or "")[:256] or None
                candidate.tags_json = _json(item.get("topics") or item.get("tags") or [])
                duration_ms = item.get("duration_ms")
                if duration_ms:
                    candidate.duration_seconds = Decimal(str(duration_ms)) / 1000
                candidate.play_count = stats.get("play_count") or candidate.play_count
                candidate.like_count = stats.get("digg_count") or stats.get("like_count") or candidate.like_count
                candidate.comment_count = stats.get("comment_count") or candidate.comment_count
                candidate.collect_count = stats.get("collect_count") or candidate.collect_count
                candidate.share_count = stats.get("share_count") or candidate.share_count
            run.search_count = int(session.scalar(select(func.count()).select_from(FindAgentV2Search).where(FindAgentV2Search.run_id == run_id)) or 0)
            session.flush()
            run.candidate_count = int(session.scalar(select(func.count()).select_from(FindAgentV2Candidate).where(FindAgentV2Candidate.run_id == run_id)) or 0)
            return {"search_id": int(search.id), "new_candidate_count": new_count, "result_count": len(results)}

    def candidate_inputs(self, run_id: str, candidate_ids: list[int]) -> list[dict[str, Any]]:
        with get_session() as session:
            rows = list(session.scalars(select(FindAgentV2Candidate).where(
                FindAgentV2Candidate.run_id == run_id,
                FindAgentV2Candidate.id.in_([int(v) for v in candidate_ids]),
            )))
            return [_candidate_dict(row) for row in rows]

    def save_details(self, run_id: str, details: list[dict[str, Any]], errors: list[dict[str, Any]]) -> None:
        by_id = {str(item.get("content_id") or ""): item for item in details}
        error_by_id = {str(item.get("content_id") or ""): item for item in errors}
        with get_session() as session:
            rows = list(session.scalars(select(FindAgentV2Candidate).where(
                FindAgentV2Candidate.run_id == run_id,
                FindAgentV2Candidate.aweme_id.in_(list(by_id) + list(error_by_id)),
            )))
            for row in rows:
                detail = by_id.get(row.aweme_id)
                error = error_by_id.get(row.aweme_id)
                if detail:
                    row.detail_status = "success"
                    row.detail_json = _json(detail)
                    row.title = str(detail.get("title") or detail.get("body_text") or row.title or "")[:512] or None
                    row.content_link = str(detail.get("content_link") or row.content_link or "")[:1024] or None
                    row.author_name = str(detail.get("channel_account_name") or row.author_name or "")[:256] or None
                    row.author_sec_uid = str(detail.get("channel_account_id") or row.author_sec_uid or "")[:256] or None
                    row.tags_json = _json(detail.get("topic_list") or [])
                    parsed = parse_datetime_value(detail.get("publish_at"))
                    row.publish_at = parsed.replace(tzinfo=None) if parsed else None
                    row.duration_seconds = detail.get("duration_seconds") or None
                    for key in ("play_count", "like_count", "comment_count", "collect_count", "share_count"):
                        value = detail.get(key)
                        if value is not None:
                            setattr(row, key, value)
                    status, raw = "success", detail
                else:
                    row.detail_status = "failed"
                    status, raw = "failed", error or {}
                session.add(FindAgentV2Evidence(
                    run_id=run_id,
                    candidate_id=row.id,
                    evidence_type="detail",
                    provider="crawler",
                    status=status,
                    raw_json=_json(raw),
                    error_message=str((error or {}).get("error") or "") or None,
                ))

    def save_portraits(self, run_id: str, results: list[dict[str, Any]]) -> None:
        with get_session() as session:
            for item in results:
                aweme_id = str(item.get("aweme_id") or "")
                row = session.scalar(select(FindAgentV2Candidate).where(
                    FindAgentV2Candidate.run_id == run_id,
                    FindAgentV2Candidate.aweme_id == aweme_id,
                ))
                if row is None:
                    continue
                normalization = item.get("age_normalization") or {}
                content = normalization.get("content") or {}
                account = normalization.get("account") or {}
                row.portrait_json = _json(item)
                row.portrait_status = "failed" if item.get("error") else "success"
                row.content_50_plus_ratio = _ratio(
                    content.get("older_ratio") if content.get("has_age_portrait") else None
                )
                row.account_50_plus_ratio = _ratio(
                    account.get("older_ratio") if account.get("has_age_portrait") else None
                )
                session.add(FindAgentV2Evidence(
                    run_id=run_id,
                    candidate_id=row.id,
                    evidence_type="portrait",
                    provider="douhot",
                    status=row.portrait_status,
                    raw_json=_json(item),
                    normalized_json=_json(normalization),
                    error_message=str(item.get("error") or "") or None,
                ))

    def evaluate(self, run_id: str, items: list[dict[str, Any]]) -> list[dict[str, Any]]:
        run_data = self.require_run(run_id)
        output: list[dict[str, Any]] = []
        with get_session() as session:
            run = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            assert run is not None
            for item in items:
                row = session.scalar(select(FindAgentV2Candidate).where(
                    FindAgentV2Candidate.run_id == run_id,
                    FindAgentV2Candidate.id == int(item.get("candidate_id") or 0),
                ))
                if row is None:
                    raise ValueError(f"candidate_id 不属于 run: {item.get('candidate_id')}")
                for key in ("relevance_score", "elder_score", "share_score", "value_score"):
                    setattr(row, key, _ratio(item.get(key)))
                requested = str(item.get("decision_bucket") or "rejected")
                if requested not in {"primary", "rejected"}:
                    raise ValueError("decision_bucket 只能是 primary/rejected")
                gate_input = _candidate_dict(row)
                gate = evaluate_candidate_gate(gate_input, run_data["rule_config"])
                row.gate_status = gate["status"]
                row.gate_result_json = _json(gate)
                row.decision_bucket = "primary" if requested == "primary" and gate["status"] == "pass" else "rejected"
                row.decision_reason = str(item.get("decision_reason") or "")
                failed = list(gate.get("failed_reason_codes") or [])
                row.reject_reason_code = (str(item.get("reject_reason_code") or "") or (failed[0] if failed else None))
                output.append({"candidate_id": int(row.id), "decision_bucket": row.decision_bucket, "gate_status": row.gate_status})
            session.flush()
            primaries = list(session.scalars(select(FindAgentV2Candidate).where(
                FindAgentV2Candidate.run_id == run_id,
                FindAgentV2Candidate.decision_bucket == "primary",
            )))
            run.valid_primary_count = len({row.aweme_id for row in primaries if row.gate_status == "pass"})
        return output

    def recount_valid_primary(self, run_id: str) -> int:
        """Recompute the denormalized primary counter after parallel worker writes."""
        with get_session() as session:
            run = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if run is None:
                raise FindAgentV2RunNotFound(run_id)
            count = int(session.scalar(select(func.count(func.distinct(
                FindAgentV2Candidate.aweme_id,
            ))).where(
                FindAgentV2Candidate.run_id == run_id,
                FindAgentV2Candidate.decision_bucket == "primary",
                FindAgentV2Candidate.gate_status == "pass",
            )) or 0)
            run.valid_primary_count = count
            return count

    def get_full_state(
        self, run_id: str, *, limit: int = 100, pending_only: bool = False,
    ) -> dict[str, Any]:
        run = self.require_run(run_id)
        with get_session() as session:
            searches = list(session.scalars(select(FindAgentV2Search).where(
                FindAgentV2Search.run_id == run_id,
            ).order_by(FindAgentV2Search.id)))
            candidate_query = select(FindAgentV2Candidate).where(
                FindAgentV2Candidate.run_id == run_id,
            )
            if pending_only:
                candidate_query = candidate_query.where(
                    FindAgentV2Candidate.decision_bucket == "pending_evaluation",
                )
            candidates = list(session.scalars(candidate_query.order_by(
                FindAgentV2Candidate.value_score.desc(), FindAgentV2Candidate.id,
            ).limit(max(1, min(limit, 500)))))
            return {
                "run": run,
                "searches": [{
                    "search_id": int(row.id), "round_index": row.round_index,
                    "keyword": row.keyword, "query_reason": row.query_reason,
                    "provider": row.provider, "page_no": row.page_no,
                    "has_more": bool(row.has_more), "next_cursor": row.next_cursor,
                    "status": row.status, "result_count": row.result_count,
                } for row in searches],
                "candidates": [_candidate_dict(row) for row in candidates],
            }

    def snapshot(self, run_id: str) -> DiscoverySnapshot:
        state = self.get_full_state(run_id)
        candidates = state["candidates"]
        buckets = [item["decision_bucket"] for item in candidates]
        run = state["run"]
        return DiscoverySnapshot(
            status=run["status"],
            search_count=run["search_count"],
            candidate_count=run["candidate_count"],
            pending_count=sum(value == "pending_evaluation" for value in buckets),
            primary_count=sum(value == "primary" for value in buckets),
            valid_primary_count=run["valid_primary_count"],
            rejected_count=sum(value == "rejected" for value in buckets),
            outcome_status=run.get("outcome_status") or "",
        )

    def finalize(
        self,
        run_id: str,
        *,
        failed: bool = False,
        reason: str = "",
        target_primary_count: int = 5,
    ) -> dict[str, Any]:
        snapshot = self.snapshot(run_id)
        if failed:
            outcome, status = "failed", "failed"
        elif snapshot.valid_primary_count >= max(1, int(target_primary_count)):
            outcome, status = "goal_met", "finished"
        elif snapshot.valid_primary_count > 0:
            outcome, status = "partial", "finished"
        else:
            outcome, status = "no_match", "finished"
        with get_session() as session:
            run = session.scalar(select(FindAgentV2Run).where(FindAgentV2Run.run_id == run_id))
            if run is None:
                raise FindAgentV2RunNotFound(run_id)
            run.status = status
            run.outcome_status = outcome
            run.stop_reason = reason[:2000] or None
        return self.require_run(run_id)


_SERVICE = FindAgentV2Service()


def get_find_agent_v2_service() -> FindAgentV2Service:
    return _SERVICE