Create-knowledge-find-decode/acquisition/runner.py

"""Formal acquisition runner backed by the repository boundary."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable
from uuid import UUID

from acquisition.classification.coarse import ClassificationResult, coarse_classify_item
from acquisition.crawler import RateLimiter
from acquisition.domain import AcquisitionJob, Query
from acquisition.media.service import StabilizedMedia, stabilize_media_urls
from acquisition.platforms import PlatformAdapter, get_platform_adapter
from acquisition.repositories.base import AcquisitionRepository
from core.config import Settings

DEFAULT_PLATFORMS = ("xiaohongshu", "weixin", "douyin")

AdapterFactory = Callable[[str], PlatformAdapter]
Classifier = Callable[..., ClassificationResult]
MediaStabilizer = Callable[..., list[StabilizedMedia]]
RateLimiterFactory = Callable[[str], Any]


@dataclass(frozen=True)
class RunBatchResult:
    run_id: UUID
    total_jobs: int
    done: int
    partial: int
    failed: int
    skipped: int


class PlatformRateLimiter:
    """Share one rate bucket per platform across search and detail calls."""

    def __init__(self, platform: str, delegate: RateLimiter | None = None) -> None:
        self.platform = platform
        self.delegate = delegate or RateLimiter(
            min_interval_seconds=10.0,
            max_interval_seconds=12.0,
        )

    def wait(self, bucket: str) -> None:
        self.delegate.wait(self.platform)


def _query_id(query: Query) -> UUID:
    if query.id is None:
        raise RuntimeError("formal acquisition queries must have id before running")
    return query.id


def _job_id(job: AcquisitionJob) -> UUID:
    if job.id is None:
        raise RuntimeError("formal acquisition jobs must have id before running")
    return job.id


def _best_video_url(media: list[StabilizedMedia]) -> str:
    for row in media:
        if row.media_type == "video" and row.cdn_url:
            return row.cdn_url
    return ""


def _image_urls(media: list[StabilizedMedia]) -> list[str]:
    return [row.cdn_url for row in media if row.media_type == "image" and row.cdn_url]


def _source_payload(candidate: Any, detail: Any) -> dict[str, Any]:
    return {
        "candidate": candidate.model_dump() if hasattr(candidate, "model_dump") else {},
        "detail": detail.raw if isinstance(getattr(detail, "raw", None), dict) else {},
    }


def _record_candidate(
    repo: AcquisitionRepository,
    *,
    job: AcquisitionJob,
    query: Query,
    platform: str,
    candidate: Any,
    detail: Any,
    settings: Settings,
    media_stabilizer: MediaStabilizer,
    classifier: Classifier,
    classify: bool,
) -> bool:
    media_rows = media_stabilizer(
        image_urls=detail.image_urls,
        video_urls=detail.video_urls,
        settings=settings,
    )
    item = repo.upsert_candidate_item(
        platform=platform,
        job_id=_job_id(job),
        query_id=_query_id(query),
        platform_item_id=detail.source_id or candidate.source_id or None,
        canonical_url=detail.url or candidate.url or None,
        content_type=detail.content_type or None,
        title=detail.title or None,
        author_name=detail.author or None,
        raw_summary=(detail.body_text or "")[:1000],
        status="candidate",
        source_payload=_source_payload(candidate, detail),
        metadata={"candidate_rank": candidate.rank},
    )
    if item.id is None:
        raise RuntimeError("repository returned candidate without id")

    for row in media_rows:
        repo.add_media_asset(
            item_id=item.id,
            media_type=row.media_type,
            source_url=row.source_url,
            oss_url=row.cdn_url,
            cdn_url=row.cdn_url,
            position=row.position,
            status=row.status,
            source_payload={},
            metadata={},
        )

    if classify:
        result = classifier(
            platform=platform,
            title=detail.title,
            body_text=detail.body_text,
            image_urls=_image_urls(media_rows),
            video_url=_best_video_url(media_rows),
            settings=settings,
        )
        repo.add_item_classification(
            item_id=item.id,
            is_creation_knowledge=result.is_creation_knowledge,
            label=result.label,
            confidence=result.confidence,
            reason=result.reason,
            model_name=settings.video_model,
            prompt_version=result.prompt_version,
            result_payload=result.result_payload,
            status=result.status,
            error_message=result.error_message,
        )
    return True


def run_batch(
    repo: AcquisitionRepository,
    *,
    batch_id: UUID,
    settings: Settings,
    platforms: tuple[str, ...] | list[str] = DEFAULT_PLATFORMS,
    search_limit: int = 10,
    display_limit: int = 5,
    classify: bool = True,
    resume: bool = True,
    skip_done: bool = True,
    run_key: str | None = None,
    adapter_factory: AdapterFactory = get_platform_adapter,
    media_stabilizer: MediaStabilizer = stabilize_media_urls,
    classifier: Classifier = coarse_classify_item,
    rate_limiter_factory: RateLimiterFactory | None = None,
) -> RunBatchResult:
    """Run query x platform acquisition and write formal cloud-state rows."""
    queries = repo.list_queries_for_batch(batch_id, keep=True)
    run = repo.create_acquisition_run(
        batch_id=batch_id,
        run_key=run_key or f"acquisition:{batch_id}",
        status="running",
        metadata={
            "platforms": list(platforms),
            "search_limit": search_limit,
            "display_limit": display_limit,
            "classify": classify,
            "resume": resume,
        },
    )
    if run.id is None:
        raise RuntimeError("repository returned acquisition run without id")

    total_jobs = len(queries) * len(platforms)
    done = partial = failed = skipped = 0
    gates: dict[str, Any] = {}

    for query in queries:
        query_id = _query_id(query)
        for platform in platforms:
            job = repo.ensure_acquisition_job(
                run_id=run.id,
                query_id=query_id,
                platform=platform,
                search_limit=search_limit,
                display_limit=display_limit,
                status="pending",
                metadata={"query_text": query.query_text},
            )
            if skip_done and job.status == "done":
                skipped += 1
                continue

            attempts = job.attempt_count + 1
            job = repo.update_acquisition_job(
                _job_id(job),
                status="running",
                attempt_count=attempts,
                error_message=None,
            )
            errors: list[str] = []
            display_count = 0
            searched_count = 0
            try:
                adapter = adapter_factory(platform)
                gate = gates.get(platform)
                if gate is None:
                    gate = (
                        rate_limiter_factory(platform)
                        if rate_limiter_factory
                        else PlatformRateLimiter(platform)
                    )
                    gates[platform] = gate
                candidates = adapter.search(
                    query.query_text,
                    settings=settings,
                    limit=search_limit,
                    rate_limiter=gate,
                )
                searched_count = len(candidates)
                for candidate in candidates:
                    if display_count >= display_limit:
                        break
                    try:
                        detail = adapter.fetch_detail(
                            candidate,
                            settings=settings,
                            rate_limiter=gate,
                        )
                        if _record_candidate(
                            repo,
                            job=job,
                            query=query,
                            platform=platform,
                            candidate=candidate,
                            detail=detail,
                            settings=settings,
                            media_stabilizer=media_stabilizer,
                            classifier=classifier,
                            classify=classify,
                        ):
                            display_count += 1
                    except Exception as exc:
                        errors.append(str(exc)[:160])
                        continue

                status = "done" if display_count >= display_limit else (
                    "partial" if display_count else "failed"
                )
                if status == "done":
                    done += 1
                elif status == "partial":
                    partial += 1
                else:
                    failed += 1
                repo.update_acquisition_job(
                    _job_id(job),
                    status=status,
                    attempt_count=attempts,
                    error_message=None if status != "failed" else "; ".join(errors[-3:]),
                    metadata={
                        "query_text": query.query_text,
                        "searched_count": searched_count,
                        "display_count": display_count,
                        "errors": errors[-3:],
                    },
                )
            except Exception as exc:
                failed += 1
                repo.update_acquisition_job(
                    _job_id(job),
                    status="failed",
                    attempt_count=attempts,
                    error_message=str(exc)[:300],
                    metadata={
                        "query_text": query.query_text,
                        "searched_count": searched_count,
                        "display_count": display_count,
                        "errors": errors[-3:],
                    },
                )

    run_status = (
        "done"
        if failed == 0 and partial == 0
        else "partial"
        if done > 0 or partial > 0
        else "failed"
    )
    update_run = getattr(repo, "update_acquisition_run", None)
    if update_run:
        update_run(
            run.id,
            status=run_status,
            metadata={
                "total_jobs": total_jobs,
                "done": done,
                "partial": partial,
                "failed": failed,
                "skipped": skipped,
            },
        )

    return RunBatchResult(
        run_id=run.id,
        total_jobs=total_jobs,
        done=done,
        partial=partial,
        failed=failed,
        skipped=skipped,
    )