Update agent.py

agent.py

@@ -1,11 +1,12 @@
 import os
 from dotenv import load_dotenv
 from typing import TypedDict, List, Dict, Any, Optional
+from urllib.parse import urlparse
 from langgraph.graph import StateGraph, START, END, MessagesState
 from langchain.agents import create_tool_calling_agent, ConversationalAgent, AgentExecutor, initialize_agent, create_react_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
-from langchain_core.tools import tool
+from langchain_core.tools import tool, Tool
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain.memory import ConversationBufferMemory
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
@@ -15,12 +16,12 @@ from langgraph.prebuilt import tools_condition
 # 1. Web Browsing
 from langchain_community.tools import DuckDuckGoSearchResults
 from langchain_community.document_loaders import ImageCaptionLoader
-import requests, time
+import requests, time, yt_dlp
 import pandas as pd
 from pathlib import Path
 from bs4 import BeautifulSoup
 from langchain_community.tools import WikipediaQueryRun
-from langchain_community.utilities import WikipediaAPIWrapper
+from langchain_community.utilities import WikipediaAPIWrapper, DuckDuckGoSearchAPIWrapper
 from langchain_community.document_loaders import YoutubeLoader
 from langchain_community.document_loaders import UnstructuredExcelLoader
 from langchain_community.document_loaders import AssemblyAIAudioTranscriptLoader
@@ -37,9 +38,14 @@ def duckduck_websearch(query: str) -> str:
         query: what you want to search
     """
     try:
-        search = DuckDuckGoSearchResults()
-        results = search.invoke(query)
-        return "\n".join(results)
+        # search = DuckDuckGoSearchResults()
+        # results = search.invoke(query)
+        search = search = DuckDuckGoSearchAPIWrapper(max_results=5)
+        results = search.run(query)
+        if not results or results.strip() == "":
+            return "No search results found."
+            
+        return results
     except Exception as e:
         print(str(e))
         print('Try to use request method for duckcudckgo Search')
@@ -301,13 +307,80 @@ class BasicAgent:
                 If Task ID is included in the question, remember to call the relevant read tools [ie. read_file, excel_read, csv_read, mp3_listen, image_caption]
                 Note: python_tool is called when the question mentions the term "Python" or any math calculation.
         """
-        self.tools = [duckduck_websearch, serper_websearch, visit_webpage, wiki_search, text_splitter, youtube_transcript, read_file, excel_read, csv_read, mp3_listen, image_caption, run_python]
+        # self.tools = [duckduck_websearch, serper_websearch, visit_webpage, wiki_search, text_splitter, self._analyze_video, youtube_transcript, read_file, excel_read, csv_read, mp3_listen, image_caption, run_python]
+        self.tools = [
+            Tool(
+                name="duckduck_websearch",
+                func=duckduck_websearch,
+                description="Search the web for information with DuckDuckGo"
+            ),
+            Tool(
+                name="serper_websearch",
+                func=serper_websearch,
+                description="Search the web for information with Serper"
+            ),
+            Tool(
+                name="visit_webpage",
+                func=visit_webpage,
+                description="Directly visit the webpage"
+            ),
+            Tool(
+                name="wiki_search",
+                func=wiki_search,
+                description="Search the information on Wikipedia"
+            ),
+            Tool(
+                name="text_splitter",
+                func=text_splitter,
+                description="Split text into chunks"
+            ),
+            Tool(
+                name="analyze_video",
+                func=self.analyze_video,
+                description+"Analyze YouTube video content directly"
+            ),
+            Tool(
+                name="youtube_transcript",
+                func=youtube_transcript,
+                description="Fetch the transcript of YouTube video"
+            ),
+            Tool(
+                name="read_file",
+                func=read_file,
+                description="Read the file content"
+            ),
+            Tool(
+                name="excel_read",
+                func=excel_read,
+                description="Read the content of Excel file"
+            ),
+            Tool(
+                name="csv_read",
+                func=csv_read,
+                description="Read the content of CSV file"
+            ),
+            Tool(
+                name='mp3_listen',
+                func=mp3_listen,
+                description="Listen to the MP3 file"
+            ),
+            Tool(
+                name="image_caption",
+                func=image_caption,
+                description="Understand the image content"
+            ),
+            Tool(
+                name="run_python",
+                func=run_python,
+                description="Run Python code"
+            )
+        ]
         # Setup memory
         self.memory = ConversationBufferMemory(
             memory_key="chat_history",
             return_messages=True
         )
-        self.agent = self.__setup__agent__()
+        self.agent = self.__setup_agent__()
         # self.prompt = ChatPromptTemplate.from_messages([
         #     ("system", self.sys_prompt),
         #     ("human", "{input}")
@@ -337,7 +410,7 @@ class BasicAgent:
         fixed_answer = "This is a default answer."
         
 
-        max_retries = 3
+        max_retries = 5
         base_sleep = 1
         for attempt in range(max_retries):
             try:
@@ -355,7 +428,67 @@ class BasicAgent:
                 return f"Error processing query after {max_retries} attempts: {str(e)}"
         return fixed_answer
 
-    def __setup__agent__(self) -> AgentExecutor:
+    @tool
+    def _analyze_video(self, url: str) -> str:
+        """Analyze video content using Gemini's video understanding capabilities."""
+        try:
+            # Validate URL
+            parsed_url = urlparse(url)
+            if not all([parsed_url.scheme, parsed_url.netloc]):
+                return "Please provide a valid video URL with http:// or https:// prefix."
+            
+            # Check if it's a YouTube URL
+            if 'youtube.com' not in url and 'youtu.be' not in url:
+                return "Only YouTube videos are supported at this time."
+
+            try:
+                # Configure yt-dlp with minimal extraction
+                ydl_opts = {
+                    'quiet': True,
+                    'no_warnings': True,
+                    'extract_flat': True,
+                    'no_playlist': True,
+                    'youtube_include_dash_manifest': False
+                }
+
+                with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+                    try:
+                        # Try basic info extraction
+                        info = ydl.extract_info(url, download=False, process=False)
+                        if not info:
+                            return "Could not extract video information."
+
+                        title = info.get('title', 'Unknown')
+                        description = info.get('description', '')
+                        
+                        # Create a detailed prompt with available metadata
+                        prompt = f"""Please analyze this YouTube video:
+Title: {title}
+URL: {url}
+Description: {description}
+Please provide a detailed analysis focusing on:
+1. Main topic and key points from the title and description
+2. Expected visual elements and scenes
+3. Overall message or purpose
+4. Target audience"""
+
+                        # Use the LLM with proper message format
+                        messages = [HumanMessage(content=prompt)]
+                        response = self.model.invoke(messages)
+                        return response.content if hasattr(response, 'content') else str(response)
+
+                    except Exception as e:
+                        if 'Sign in to confirm' in str(e):
+                            return "This video requires age verification or sign-in. Please provide a different video URL."
+                        return f"Error accessing video: {str(e)}"
+
+            except Exception as e:
+                return f"Error extracting video info: {str(e)}"
+
+        except Exception as e:
+            return f"Error analyzing video: {str(e)}"
+    
+    def __setup_agent__(self) -> AgentExecutor:
         PREFIX = """
                 You are a general AI assistant that can use various tools to answer question. I will ask you a question. Report your thoughts, and finish your answer with the following template: 
                 FINAL ANSWER: [YOUR FINAL ANSWER].