Pacama95/chatbot_agent

1from langchain_community.document_loaders import YoutubeLoader2from langchain_community.document_loaders.youtube import TranscriptFormat3from langchain.docstore.document import Document4 5from pydantic import BaseModel, Field6from langchain.tools import BaseTool7 8from typing import Type, Any, List9    10 11class YoutubeVideoTranscript(BaseModel):12    video_url: str = Field(description='URL to the video to get the transcription for.')13    language: str = Field(description="The transcription language.")14 15class YoutubeVideoTranscriptTool(BaseTool):16 17    name: str = "youtubue_video_transcript"18    description: str = (19        "Given an URL to a YouTube video, returns the transcription for that video."20    )21    args_schema: Type[BaseModel] = YoutubeVideoTranscript22 23    def __init__(self, **kwargs: Any) -> None:24        super().__init__()25 26    def _run(self, video_url: str, language: str) -> str:27        """28            Given an URL to a YouTube video, returns the transcription for that video.29            Args:30                video_url (str): URL to the video to get the transcription for.31                language (str): The transcription language.32 33            Returns:34                List[Document]: The video transcriptions in chunks of 30 seconds.35        """36        loader = YoutubeLoader.from_youtube_url(video_url, add_video_info = False, language = [language], transcript_format=TranscriptFormat.CHUNKS, translation=language, chunk_size_seconds=30)37 38        transcriptions: List[Document] = loader.load()39 40        return transcriptions41