Source code for langchain_community.tools.youtube.search

"""
自https://github.com/venuv/langchain_yt_tools 改编

CustomYTSearchTool搜索与某人相关的YouTube视频，并返回指定数量的视频URL。
此工具的输入应为逗号分隔的列表，
 - 第一部分包含人名
 - 第二部分（可选）是要返回的视频结果的最大数量"""
import json
from typing import Optional

from langchain_core.callbacks import CallbackManagerForToolRun

from langchain_community.tools import BaseTool


class YouTubeSearchTool(BaseTool):
    """查询YouTube的工具。"""

    name: str = "youtube_search"
    description: str = (
        "search for youtube videos associated with a person. "
        "the input to this tool should be a comma separated list, "
        "the first part contains a person name and the second a "
        "number that is the maximum number of video results "
        "to return aka num_results. the second part is optional"
    )

    def _search(self, person: str, num_results: int) -> str:
        from youtube_search import YoutubeSearch

        results = YoutubeSearch(person, num_results).to_json()
        data = json.loads(results)
        url_suffix_list = [
            "https://www.youtube.com" + video["url_suffix"] for video in data["videos"]
        ]
        return str(url_suffix_list)

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        """使用这个工具。"""
        values = query.split(",")
        person = values[0]
        if len(values) > 1:
            num_results = int(values[1])
        else:
            num_results = 2
        return self._search(person, num_results)