MoviePilot/app/workflow/actions/fetch_medias.py

from typing import List, Optional

from pydantic import Field

from app.workflow.actions import BaseAction
from app.chain.recommend import RecommendChain
from app.application.configuration import get_chain_runtime_config_snapshot
from app.schemas.workflow import ActionParams
from app.schemas.workflow import ActionContext
from app.runtime.config import global_vars
from app.runtime.events import eventmanager
from app.runtime.log import logger
from app.schemas.event import RecommendSourceEventData
from app.schemas.workflow import MediaInfo
from app.schemas.types import ChainEventType
from app.adapters.network.http import RequestUtils


class FetchMediasParams(ActionParams):
    """
    获取媒体数据参数
    """
    source_type: Optional[str] = Field(default="ranking", description="来源")
    sources: Optional[List[str]] = Field(default=[], description="榜单")
    api_path: Optional[str] = Field(default=None, description="API路径")


class FetchMediasAction(BaseAction):
    """
    获取媒体数据
    """

    contract = {
        "outputs": [{"name": "medias", "label": "媒体", "kind": "list"}],
    }

    def __init__(self, action_id: str):
        super().__init__(action_id)

        self._medias = []
        self._has_error = False
        self.__inner_sources = [
            {
                "func": RecommendChain().tmdb_trending,
                "name": '流行趋势',
                "api_path": "recommend/tmdb_trending"
            },
            {
                "func": RecommendChain().douban_movie_showing,
                "name": '正在热映',
                "api_path": "recommend/douban_showing"
            },
            {
                "func": RecommendChain().bangumi_calendar,
                "name": 'Bangumi每日放送',
                "api_path": "recommend/bangumi_calendar"
            },
            {
                "func": RecommendChain().tmdb_movies,
                "name": 'TMDB热门电影',
                "api_path": "recommend/tmdb_movies"
            },
            {
                "func": RecommendChain().tmdb_tvs,
                "name": 'TMDB热门电视剧',
                "api_path": "recommend/tmdb_tvs?with_original_language=zh|en|ja|ko"
            },
            {
                "func": RecommendChain().douban_movie_hot,
                "name": '豆瓣热门电影',
                "api_path": "recommend/douban_movie_hot"
            },
            {
                "func": RecommendChain().douban_tv_hot,
                "name": '豆瓣热门电视剧',
                "api_path": "recommend/douban_tv_hot"
            },
            {
                "func": RecommendChain().douban_tv_animation,
                "name": '豆瓣热门动漫',
                "api_path": "recommend/douban_tv_animation"
            },
            {
                "func": RecommendChain().douban_movies,
                "name": '豆瓣最新电影',
                "api_path": "recommend/douban_movies"
            },
            {
                "func": RecommendChain().douban_tvs,
                "name": '豆瓣最新电视剧',
                "api_path": "recommend/douban_tvs"
            },
            {
                "func": RecommendChain().douban_movie_top250,
                "name": '豆瓣电影TOP250',
                "api_path": "recommend/douban_movie_top250"
            },
            {
                "func": RecommendChain().douban_tv_weekly_chinese,
                "name": '豆瓣国产剧集榜',
                "api_path": "recommend/douban_tv_weekly_chinese"
            },
            {
                "func": RecommendChain().douban_tv_weekly_global,
                "name": '豆瓣全球剧集榜',
                "api_path": "recommend/douban_tv_weekly_global"
            }
        ]

        # 广播事件，请示额外的推荐数据源支持
        event_data = RecommendSourceEventData()
        event = eventmanager.send_event(ChainEventType.RecommendSource, event_data)
        # 使用事件返回的上下文数据
        if event and event.event_data:
            event_data: RecommendSourceEventData = event.event_data
            if event_data.extra_sources:
                self.__inner_sources.extend([s.model_dump() for s in event_data.extra_sources])

    name = "获取媒体数据"
    description = "获取榜单等媒体数据列表"
    data = FetchMediasParams().model_dump()

    @property
    def success(self) -> bool:
        return not self._has_error

    def __get_source(self, source: str):
        """
        获取数据源
        """
        for s in self.__inner_sources:
            if s['api_path'] == source:
                return s
        return None

    def execute(self, workflow_id: int, params: dict, context: ActionContext) -> ActionContext:
        """
        获取媒体数据，填充到medias
        """
        params = FetchMediasParams(**params)
        try:
            if params.source_type == "ranking":
                for api_path in params.sources:
                    if global_vars.is_workflow_stopped(workflow_id):
                        break
                    source = self.__get_source(api_path)
                    if not source:
                        continue
                    logger.info(f"获取媒体数据 {source} ...")
                    name = source.get("name")
                    results = []
                    if source.get("func"):
                        results = source['func']()
                    else:
                        # 调用内部API获取数据
                        runtime_config = get_chain_runtime_config_snapshot()
                        api_url = (
                            f"http://127.0.0.1:{runtime_config.api_port}"
                            f"/api/v1/{source['api_path']}?token={runtime_config.api_token}"
                        )
                        res = RequestUtils(timeout=15).post_res(api_url)
                        if res:
                            results = res.json()
                    if results:
                        logger.info(f"{name} 获取到 {len(results)} 条数据")
                        self._medias.extend([MediaInfo(**r) for r in results])
                    else:
                        logger.error(f"{name} 获取数据失败")
            else:
                # 调用内部API获取数据
                runtime_config = get_chain_runtime_config_snapshot()
                api_url = (
                    f"http://127.0.0.1:{runtime_config.api_port}{params.api_path}"
                    f"?token={runtime_config.api_token}"
                )
                res = RequestUtils(timeout=15).post_res(api_url)
                if res:
                    results = res.json()
                    if results:
                        logger.info(f"{params.api_path} 获取到 {len(results)} 条数据")
                        self._medias.extend([MediaInfo(**r) for r in results])
        except Exception as e:
            logger.error(f"获取媒体数据失败: {e}")
            self._has_error = True

        if self._medias:
            context.medias.extend(self._medias)

        self.job_done(f"获取到 {len(self._medias)} 条媒数据")
        return context