SupplyAgent/agents/find_agent/output_sync.py

"""Parse recommendation buckets from the model's final response."""
from __future__ import annotations

import re

_VIDEO_ID_PATTERN = re.compile(r"(?<!\d)\d{15,22}(?!\d)")
_PRIMARY_LABELS = ("主推荐", "正式推荐")
_REJECTED_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 rejected sections."""
    primary: list[str] = []
    rejected: 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 _REJECTED_LABELS):
            section = "rejected"
        elif any(heading.startswith(label) for label in _END_LABELS):
            section = None

        if section is None:
            continue
        target = primary if section == "primary" else rejected
        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 rejected if aweme_id not in primary_set
    ]