SupplyAgent/agents/find_agent/output_sync.py

"""Persist the recommendation buckets stated by the model's final response."""
from __future__ import annotations

import json
import re
from typing import Any

from supply_agent.types import AgentResult, Role
from supply_infra.db.repositories.video_discovery_repo import (
    VideoDiscoveryRepository,
)
from supply_infra.db.session import get_session

_VIDEO_ID_PATTERN = re.compile(r"(?<!\d)\d{15,22}(?!\d)")
_PRIMARY_LABELS = ("主推荐", "正式推荐")
_BACKUP_LABELS = ("补充推荐", "人工备选")
_END_LABELS = (
    "淘汰",
    "最有竞争力",
    "搜索轨迹",
    "搜索树",
    "缺失数据",
    "局限",
    "总结",
)


def _normalized_heading(line: str) -> str:
    text = line.strip()
    text = re.sub(r"^#{1,6}\s*", "", text)
    text = re.sub(r"^\d+\s*[.、]\s*", "", text)
    return text.replace("*", "").strip()


def extract_model_recommendation_ids(
    content: str,
) -> tuple[list[str], list[str]]:
    """Extract video ids from the model's primary and backup sections."""
    primary: list[str] = []
    backup: list[str] = []
    section: str | None = None

    for line in (content or "").splitlines():
        heading = _normalized_heading(line)
        if any(heading.startswith(label) for label in _PRIMARY_LABELS):
            section = "primary"
        elif any(heading.startswith(label) for label in _BACKUP_LABELS):
            section = "backup"
        elif any(heading.startswith(label) for label in _END_LABELS):
            section = None

        if section is None:
            continue
        target = primary if section == "primary" else backup
        for aweme_id in _VIDEO_ID_PATTERN.findall(line):
            if aweme_id not in target:
                target.append(aweme_id)

    primary_set = set(primary)
    return primary, [aweme_id for aweme_id in backup if aweme_id not in primary_set]


def _latest_run_id(result: AgentResult) -> str | None:
    for message in reversed(result.messages):
        if message.role != Role.TOOL or not message.content:
            continue
        try:
            payload = json.loads(message.content)
        except (TypeError, json.JSONDecodeError):
            continue
        if not isinstance(payload, dict):
            continue
        run_id = payload.get("run_id")
        if run_id:
            return str(run_id)
    return None


def persist_model_recommendations(result: AgentResult) -> dict[str, Any]:
    """
    Make the model's final recommendation sections authoritative in the database.

    The final response itself is never validated or rewritten. Video ids listed
    under primary/backup are persisted to those buckets exactly as stated.
    """
    run_id = _latest_run_id(result)
    primary_ids, backup_ids = extract_model_recommendation_ids(result.content)
    if not run_id or not (primary_ids or backup_ids):
        return {
            "run_id": run_id,
            "primary_ids": primary_ids,
            "backup_ids": backup_ids,
            "saved_count": 0,
        }

    rows = [
        {"aweme_id": aweme_id, "decision_bucket": "primary"}
        for aweme_id in primary_ids
    ]
    rows.extend(
        {"aweme_id": aweme_id, "decision_bucket": "backup"}
        for aweme_id in backup_ids
    )
    with get_session() as session:
        repo = VideoDiscoveryRepository(session)
        if repo.get_run(run_id) is None:
            return {
                "run_id": run_id,
                "primary_ids": primary_ids,
                "backup_ids": backup_ids,
                "saved_count": 0,
            }
        saved_count, _ = repo.save_candidate_evaluations(run_id, rows)
    return {
        "run_id": run_id,
        "primary_ids": primary_ids,
        "backup_ids": backup_ids,
        "saved_count": saved_count,
    }