Update app.py

app.py

@@ -1,146 +1,321 @@
 import os
-from dotenv import load_dotenv
+from typing import Annotated, Optional, TypedDict
 import gradio as gr
+from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
+from langchain_openai import ChatOpenAI
+from langgraph.graph.message import add_messages
+from langgraph.graph import StateGraph, START
+from langgraph.prebuilt import tools_condition, ToolNode
 import requests
-
-from typing import List, Dict, Union, Optional
 import pandas as pd
-import wikipediaapi
-import requests
-#from bs4 import BeautifulSoup
-import random
-import re
-from typing import Optional
-from datetime import datetime
-import google.generativeai as genai
+from langchain.tools import Tool
+from dotenv import load_dotenv
+
+from arxiv_searcher import ArxivSearcher
+from chess_algebraic_notation_retriever import ChessAlgebraicNotationMoveRetriever
+from excel_file_reader import ExcelFileReader
+from image_question_answer_tool import ImageQuestionAnswerTool
+from python_code_question_answer_tool import PythonCodeQuestionAnswerTool
+from tavily_searcher import TavilySearcher
+from transcriber import Transcriber
+from wikipedia_searcher import WikipediaSearcher
+from youtube_video_question_answer_tool import YoutubeVideoQuestionAnswerTool
 
 load_dotenv()
 
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-
+ASSOCIATED_FILE_ENDPOINT = f"{DEFAULT_API_URL}/files/"
 
 # --- Basic Agent Definition ---
+# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+#search_tool = DuckDuckGoSearchRun()
 
-class BasicAgent:
-     def __init__(self, model_name: str = "gemini-pro"):
-        """
-        Multi-modal agent powered by Google Gemini with:
-        - Web search
-        - Wikipedia access
-        - Document processing
-        """
-        self.model = genai.GenerativeModel(model_name)
-        self.wiki = wikipediaapi.Wikipedia('en')
-        self.searx_url = "https://searx.space/search"  # Public Searx instance
-        
-        print("BasicAgent initialized.")
-
-     def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        fixed_answer = self.process_request(question)
-        print(f"Agent returning answer: {fixed_answer}")
-        return fixed_answer
-
+#search_tool = DuckDuckGoSearcherTool()
 
+def retrieve_task_file(task_id: str) -> Optional[bytes]:
+    """
+    Retrieve the task file for a given task ID.
+    """
+    try:
+        response = requests.get(ASSOCIATED_FILE_ENDPOINT + task_id, timeout=15)
+        response.raise_for_status()
+        if response.status_code != 200:
+            print(f"Error fetching file: {response.status_code}")
+            return None
+        #print(f"Fetched file: {response.content}")
+        return response.content
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching file: {e}")
+        return None
+    except Exception as e:
+        print(f"An unexpected error occurred fetching file: {e}")
+        return None
 
-     def generate_response(self, prompt: str) -> str:
-        """Get response from Gemini"""
-        try:
-            response = self.model.generate_content(prompt)
-            return response.text
-        except Exception as e:
-            return f"Error generating response: {str(e)}"
-
-     def web_search(self, query: str) -> List[Dict]:
-        """Use SearxNG meta-search engine"""
-        params = {
-            "q": query,
-            "format": "json",
-            "engines": "google,bing,duckduckgo"
-        }
-        try:
-            response = requests.get(self.searx_url, params=params)
-            response.raise_for_status()
-            return response.json().get("results", [])
-        except requests.RequestException:
-            return []
-
-     def wikipedia_search(self, query: str) -> str:
-        """Get Wikipedia summary"""
-        page = self.wiki.page(query)
-        return page.summary if page.exists() else "No Wikipedia page found"
-
-     def process_document(self, file_path: str) -> str:
-        """Handle PDF, Word, CSV, Excel files"""
-        if not os.path.exists(file_path):
-            return "File not found"
-        
-        ext = os.path.splitext(file_path)[1].lower()
-        
-        try:
-            if ext == '.pdf':
-                return self._process_pdf(file_path)
-            elif ext in ('.doc', '.docx'):
-                return self._process_word(file_path)
-            elif ext == '.csv':
-                return pd.read_csv(file_path).to_string()
-            elif ext in ('.xls', '.xlsx'):
-                return pd.read_excel(file_path).to_string()
-            else:
-                return "Unsupported file format"
-        except Exception as e:
-            return f"Error processing document: {str(e)}"
+def retrieve_next_chess_move_in_algebraic_notation(task_file_path: str, is_black_turn: bool) -> str:
+    """
+    Retrieve the next chess move in algebraic notation from an image path.
+    """
+    if task_file_path is None:
+        return "Error: Task file not found."
+    # Retrieve the next chess move in algebraic notation
+    next_chess_move = ChessAlgebraicNotationMoveRetriever().retrieve(task_file_path, is_black_turn)
+    return next_chess_move
+
+# Initialize the tool
+retrieve_next_chess_move_in_algebraic_notation_tool = Tool(
+    name="retrieve_next_chess_move_in_algebraic_notation",
+    func=retrieve_next_chess_move_in_algebraic_notation,
+    description="Retrieve the next chess move in algebraic notation from an image path."
+)
+
+def transcribe_audio(file_path: str) -> str:
+    if file_path is None:
+        return "Error: Audio path not found."
+    # Transcribe the audio
+    return Transcriber().transcribe(file_path)
+
+# Initialize the tool
+transcribe_audio_tool = Tool(
+    name="transcribe_audio",
+    func=transcribe_audio,
+    description="Transcribe the audio from an audio path."
+)
+
+# Initialize the tool
+answer_python_code_tool = PythonCodeQuestionAnswerTool()
+
+# Initialize the tool
+answer_image_question_tool = ImageQuestionAnswerTool()
+
+# Initialize the tool
+answer_youtube_video_question_tool = YoutubeVideoQuestionAnswerTool()
+
+'''def answer_youtube_video_question(youtube_video_url: str, question: str) -> str:
+    """
+    Answer the question based on the youtube video.
+    """
+    if youtube_video_url is None:
+        return "Error: Video not found."
+    # Download the video
+    video_path = YoutubeVideoDownloader().download_video(youtube_video_url)
+    # Answer the question
+    return VideoQuestionAnswer().answer(video_path, question)
+# Initialize the tool
+answer_youtube_video_question_tool = Tool(
+    name="answer_youtube_video_question",
+    func=answer_youtube_video_question,
+    description="Answer the question based on the youtube video."
+)'''
+
+def read_excel_file(file_path: str) -> str:
+    if file_path is None:
+        return "Error: File not found."
+    return ExcelFileReader().read_file(file_path)
+
+# Initialize the tool
+read_excel_file_tool = Tool(
+    name="read_excel_file",
+    func=read_excel_file,
+    description="Read the excel file."
+)
+
+# Initialize the tool
+wikipedia_search_tool = Tool(
+    name="wikipedia_search",
+    func=WikipediaSearcher().search,
+    description="Search Wikipedia for a given query."
+)
+
+# Initialize the tool
+arxiv_search_tool = Tool(
+    name="arxiv_search",
+    func=ArxivSearcher().search,
+    description="Search Arxiv for a given query."
+)
+
+tavily_search_tool = Tool(
+    name="tavily_search",
+    func=TavilySearcher().search,
+    description="Search the web for a given query."
+)
+
+def format_gaia_answer(answer: str) -> str:
+    llm = ChatOpenAI(model="o3-mini", openai_api_key=os.getenv("OPENAI_API_KEY"))
+    prompt = f"""
+    You are formatting answers for the GAIA benchmark, which requires responses to be concise and unambiguous.
+    Given the answer: {answer}
+    Return the answer in the correct GAIA format:
+    - If the answer is a single word or number, return it without any additional text or formatting.
+    - If the answer is a list, return a comma-separated list without any additional text or formatting.
+    - If the answer is a string, return it without any additional text or formatting.
+    Do not include any prefixes, dots, enumerations, explanations, or quotation marks.
+    Do not include any additional text or formatting.
+    """
+    response = llm.invoke(prompt)
+    # Delete double quotes
+    return response.content.strip().replace('"', '')
+
+class AgentState(TypedDict):
+    # The document provided
+    messages: Annotated[list[AnyMessage], add_messages]
+    file_path: Optional[str]
+    
+class BasicAgent:
+    def __init__(self):
+        tools = [
+            tavily_search_tool, 
+            arxiv_search_tool,
+            wikipedia_search_tool, 
+            transcribe_audio_tool, 
+            answer_python_code_tool, 
+            answer_image_question_tool, 
+            answer_youtube_video_question_tool,
+            read_excel_file_tool
+        ]
+        '''llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash", 
+            temperature=0.2,
+            api_key=os.getenv("GEMINI_API_KEY")
+        )'''
+        llm = ChatOpenAI(model="o3-mini", openai_api_key=os.getenv("OPENAI_API_KEY"))
+        self.llm_with_tools = llm.bind_tools(tools)
+        builder = StateGraph(AgentState)
+
+        # Define nodes: these do the work
+        builder.add_node("assistant", self.assistant)
+        builder.add_node("tools", ToolNode(tools))
+
+        # Define edges: these determine how the control flow moves
+        builder.add_edge(START, "assistant")
+        builder.add_conditional_edges(
+            "assistant",
+            # If the latest message requires a tool, route to tools
+            # Otherwise, provide a direct response
+            tools_condition,
+        )
+        builder.add_edge("tools", "assistant")
+        self.agent = builder.compile()
 
-     def _process_pdf(self, file_path: str) -> str:
-        """Process PDF using Gemini's vision capability"""
-        try:
-            # For Gemini 1.5 or later which supports file uploads
-            with open(file_path, "rb") as f:
-                file = genai.upload_file(f)
-                response = self.model.generate_content(
-                    ["Extract and summarize the key points from this document:", file]
-                )
-                return response.text
-        except:
-            # Fallback for older Gemini versions
-            try:
-                import PyPDF2
-                with open(file_path, 'rb') as f:
-                    reader = PyPDF2.PdfReader(f)
-                    return "\n".join([page.extract_text() for page in reader.pages])
-            except ImportError:
-                return "PDF processing requires PyPDF2 (pip install PyPDF2)"
-
-     def _process_word(self, file_path: str) -> str:
-        """Process Word documents"""
-        try:
-            from docx import Document
-            doc = Document(file_path)
-            return "\n".join([para.text for para in doc.paragraphs])
-        except ImportError:
-            return "Word processing requires python-docx (pip install python-docx)"
+        print("BasicAgent initialized.")
 
-     def process_request(self, request: Union[str, Dict]) -> str:
+    def assistant(self, state: AgentState):
+        # System message
+        textual_description_of_tools="""
+            tavily_search(query: str) -> str:
+                Search the web for a given query.
+                Args:
+                    query: Query to search the web for (string).
+                Returns:
+                    A single string containing the information found on the web.
+            arxiv_search(query: str) -> str:
+                Search Arxiv, that contains scientific papers, for a given query.
+                Args:
+                    query: Query to search Arxiv for (string).
+                Returns:
+                    A single string containing the answer to the question.
+            wikipedia_search(query: str) -> str:
+                Search Wikipedia for a given query.
+                Args:
+                    query: Query to search Wikipedia for (string).
+                Returns:
+                    A single string containing the answer to the question.
+            transcribe_audio(file_path: str) -> str:
+                Transcribe the audio from an audio path.
+                Args:
+                    file_path: File path of the audio file (string).
+                Returns:
+                    A single string containing the transcribed text from the audio.
+            
+            answer_python_code(file_path: str, question: str) -> str:
+                Answer the question based on the python code.
+                Args:
+                    file_path: File path of the python file (string).
+                    question: Question to answer (string).
+                Returns:
+                    A single string containing the answer to the question.
+            
+            answer_image_question(file_path: str, question: str) -> str:
+                Answer the question based on the image.
+                Args:
+                    file_path: File path of the image (string).
+                    question: Question to answer (string).
+                Returns:
+                    A single string containing the answer to the question.
+            
+            download_youtube_video(youtube_video_url: str) -> str:
+                Download the Youtube video into a local file based on the URL
+                Args:
+                    youtube_video_url: A youtube video url (string).
+                Returns:
+                    A single string containing the file path of the downloaded youtube video.
+            answer_youtube_video_question(file_path: str, question: str) -> str:
+                Answer the question based on file path of the downloaded youtube video
+                Args:
+                    file_path: File path of the downloaded youtube video (string).
+                    question: Question to answer (string).
+                Returns:
+                    A single string containing the answer to the question.
+            
+            read_excel_file(file_path: str) -> str:
+                Read the excel file.
+                Args:
+                    file_path: File path of the excel file (string).
+                Returns:
+                    A markdown formatted string containing the contents of the excel file.
         """
-        Handle different request types:
-        - Direct text queries
-        - File processing requests
-        - Complex multi-step requests
+        file_path=state["file_path"]
+        prompt = f"""
+            You are a helpful assistant that can analyse images, videos, excel files and Python scripts and run computations with provided tools:
+            {textual_description_of_tools} 
+            You have access to the file path of the attached file in case it's informed. Currently the file path is: {file_path}
+            Be direct and specific. GAIA benchmark requires exact matching answers.
+            For example, if asked "What is the capital of France?", respond simply with "Paris".
+            Do not include any prefixes, dots, enumerations, explanations, or quotation marks.
+            Do not include any additional text or formatting.
+            If you are required a number, return a number, not the items.
         """
-        if isinstance(request, dict):
-            if 'steps' in request:
-                results = []
-                for step in request['steps']:
-                    if step['type'] == 'search':
-                        results.append(self.web_search(step['query']))
-                    elif step['type'] == 'process':
-                        results.append(self.process_document(step['file']))
-                return self.generate_response(f"Process these results: {results}")
-            return "Unsupported request format"
-        
-        return self.generate_response(request)
+        sys_msg = SystemMessage(content=prompt)
+
+        return {
+            "messages": [self.llm_with_tools.invoke([sys_msg] + state["messages"], config={"configurable": {"file_path": state["file_path"]}})],
+            "file_path": state["file_path"]
+        }
+        '''return {
+            "messages": [self.llm_with_tools.invoke(
+                state["messages"],
+                config={"configurable": {"file_path": state["file_path"]}}  # Aquí pasas el task_id
+            )],
+            "file_path": state["file_path"]
+        }'''
+    
+    def __call__(self, question: str, task_id: str, file_name: str) -> str:
+        print(f"######################### Agent received question (first 50 chars): {question[:50]}... with file_name: {file_name}")
+
+        # Get the file path
+        tmp_file_path = None
+        if file_name is not None and file_name != "":
+            file_content = retrieve_task_file(task_id)
+            if file_content is not None:
+                print(f"Saving file {file_name} to tmp folder")
+                tmp_file_path = f"tmp/{file_name}"
+                with open(tmp_file_path, "wb") as f:
+                    f.write(file_content)
+        # Show the file path
+        print(f"File path: {tmp_file_path}")
+
+        messages = self.agent.invoke({"messages": [HumanMessage(question)], "file_path": tmp_file_path})
+        # Show the messages
+        for m in messages['messages']:
+            m.pretty_print()
+        answer = messages["messages"][-1].content
+        answer = format_gaia_answer(answer)
+        print(f"######################### Agent returning answer: {answer}\n")
+        # Delete the file
+        if tmp_file_path is not None:
+            os.remove(tmp_file_path)
+        return answer