Create-knowledge-find-decode/acquisition/platforms/douyin.py

"""Douyin formal acquisition adapter."""
from __future__ import annotations

from typing import Any, Iterable

from acquisition.content_mode import infer_content_mode
from acquisition.crawler import fetch_post_detail
from acquisition.platforms.base import PlatformCandidate, PlatformItem, PlatformSearchPage
from acquisition.search import SearchHit, search_douyin_hits, search_douyin_page
from core.config import Settings


class DouyinAdapter:
    platform = "douyin"

    def search(
        self,
        query: str,
        *,
        settings: Settings,
        limit: int,
        rate_limiter: Any,
    ) -> list[PlatformCandidate]:
        hits = search_douyin_hits(
            query,
            content_type="视频",
            limit=limit,
            settings=settings,
            rate_limiter=rate_limiter,
        )
        return [
            PlatformCandidate(
                rank=i,
                platform=self.platform,
                provider=hit.provider,
                source_id=hit.content_id,
                raw={
                    "id": hit.content_id,
                    "search_provider": hit.provider,
                    "original": hit.raw,
                    "request_payload": hit.request_payload,
                },
            )
            for i, hit in enumerate(hits, start=1)
        ]

    def search_pages(
        self,
        query: str,
        *,
        settings: Settings,
        limit: int,
        rate_limiter: Any,
        max_pages: int = 2,
    ) -> Iterable[PlatformSearchPage]:
        provider = "piaoquantv"
        cursors: dict[str, str] = {}
        rank = 1
        seen: set[str] = set()
        for page_index in range(1, max_pages + 1):
            page = search_douyin_page(
                query,
                provider=provider,
                cursors=cursors,
                content_type="视频",
                limit=limit,
                settings=settings,
                rate_limiter=rate_limiter,
            )
            candidates: list[PlatformCandidate] = []
            for page_rank, hit in enumerate(page.rows, start=1):
                if not isinstance(hit, SearchHit):
                    continue
                if not hit.content_id or hit.content_id in seen:
                    continue
                seen.add(hit.content_id)
                candidates.append(
                    PlatformCandidate(
                        rank=rank,
                        platform=self.platform,
                        provider=hit.provider,
                        source_id=hit.content_id,
                        raw={
                            "id": hit.content_id,
                            "search_provider": hit.provider,
                            "original": hit.raw,
                            "request_payload": hit.request_payload,
                            "page_index": page_index,
                            "page_rank": page_rank,
                            "source_cursor": page.cursor,
                        },
                    )
                )
                rank += 1
            yield PlatformSearchPage(
                page_index=page_index,
                candidates=candidates,
                raw_count=page.raw_count,
                cursor=page.cursor,
                next_cursor=page.next_cursor,
                has_more=page.has_more,
                provider=page.provider,
                request_payload=page.request_payload,
                raw_metadata=page.raw_metadata,
            )
            provider = page.provider
            cursors = dict(page.raw_metadata.get("cursors") or cursors)
            current_cursor = page.cursor
            if not page.has_more or not page.next_cursor or page.next_cursor == current_cursor:
                break

    def fetch_detail(
        self,
        candidate: PlatformCandidate,
        *,
        settings: Settings,
        rate_limiter: Any,
    ) -> PlatformItem:
        if not candidate.provider:
            raise RuntimeError("douyin candidate missing search provider")
        post = fetch_post_detail(
            candidate.source_id,
            platform=self.platform,
            provider=candidate.provider,
            settings=settings,
            rate_limiter=rate_limiter,
        )
        content_mode = infer_content_mode(
            platform=self.platform,
            content_type=post.content_type,
            body_text=post.body_text,
            image_urls=post.image_urls,
            video_urls=post.video_urls,
            raw=post.raw if isinstance(post.raw, dict) else {},
        )
        return PlatformItem(
            platform=self.platform,
            provider=post.provider or candidate.provider,
            source_id=post.content_id or candidate.source_id,
            url=post.url or candidate.url,
            content_type=post.content_type or "",
            content_mode=content_mode,
            title=post.title or candidate.title,
            author=post.author_name or candidate.author,
            body_text=post.body_text or "",
            image_urls=post.image_urls,
            video_urls=post.video_urls,
            raw=post.raw if isinstance(post.raw, dict) else {},
        )