Update agent.py

agent.py

@@ -10,12 +10,12 @@ from langgraph.prebuilt import tools_condition
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from langchain_core.tools import tool
 from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import YoutubeLoader
+from youtube_transcript_api import YouTubeTranscriptApi
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_tavily import TavilySearch
 import tempfile
 import pandas as pd
-
+import numpy as np
 load_dotenv()
 
 # ReAct System Prompt
@@ -108,25 +108,14 @@ def load_youtube_transcript(url: str, add_video_info: bool = True, language: Lis
     
     Args:
         url: YouTube video URL
-        add_video_info: Whether to include video metadata
-        language: List of language codes in descending priority
-        translation: Language to translate transcript to
+       
     """
     try:
-        loader = YoutubeLoader.from_youtube_url(
-            url, 
-            add_video_info=add_video_info,
-            language=language,
-            translation=translation
-        )
-        docs = loader.load()
-        
-        formatted_transcript = "\n\n---\n\n".join([
-            f'<Document source="{doc.metadata.get("source", "")}" title="{doc.metadata.get("title", "")}" author="{doc.metadata.get("author", "")}" length="{doc.metadata.get("length", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in docs
-        ])
+        video_id = url.split("v=")[1]
+        ytt_api = YouTubeTranscriptApi()
+        docs = ytt_api.fetch(video_id)
         
-        return {"youtube_transcript": formatted_transcript}
+        return {"youtube_transcript": docs}
     except Exception as e:
         return f"Error loading YouTube transcript: {str(e)}"
 
@@ -645,7 +634,7 @@ class LangGraphAgent:
 
 if __name__ == "__main__":
     agent = LangGraphAgent()
-    question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
+    question = "The attached Excel file contains the sales of menu items for a local fast-food chain. What were the total sales that the chain made from food (not including drinks)? Express your answer in USD with two decimal places. task_id: 1234567890"
     answer = agent(question)