Update agent.py

agent.py

@@ -1,4 +1,5 @@
 import os
+from dotenv import load_dotenv
 from typing import TypedDict, List, Dict, Any, Optional
 from langchain.agents import create_tool_calling_agent, AgentExecutor, initialize_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
@@ -11,13 +12,16 @@ from langchain_community.tools import DuckDuckGoSearchResults
 from langchain_community.document_loaders import ImageCaptionLoader
 import requests, time
 import pandas as pd
-from pypdf import PdfReader
+from pathlib import Path
 from langchain_community.tools import WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
 from langchain_community.document_loaders import YoutubeLoader
 from langchain_community.document_loaders import UnstructuredExcelLoader
 from langchain_community.document_loaders import AssemblyAIAudioTranscriptLoader
 
+load_dotenv()
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+
 @tool
 def web_search(query: str) -> str:
     """Allows search through DuckDuckGo.
@@ -67,38 +71,79 @@ def youtube_transcript(video_url: str) -> str:
 
 # 4. File Reading
 @tool
-def read_file(dir: str) -> str:
-    """Read the content of the file
+def read_file(task_id: str) -> str:
+    """First download the file, then read its content
     Args:
-        dir: the filepath
+        dir: the task_id
     """
-    extension = dir.split['.'][-1]
-    if extension == 'pdf':
-        reader = PdfReader(dir)
-        contents = [p.extract_text() for p in reader.pages]
-        return "\n".join(contents)
-    else:
-        with open(dir) as f:
-            return f.read()
+    file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+    r = requests.get(file_url, timeout=15, allow_redirects=True)
+    with open('temp', "wb") as fp:
+        fp.write(r.content)
+    with open('temp') as f:
+        return f.read()
 
 @tool
-def excel_read(dir: str) -> str:
-    """Read the content of the excel file
+def excel_read(task_id: str) -> str:
+    """First download the excel file, then read its content
     Args:
-        dir: the filepath
+        dir: the task_id
     """
-    loader = UnstructuredExcelLoader(dir, mode="elements")
-    docs = loader.load()
-    contents = [doc.page_content for doc in docs]
-    return "\n".join(contents)
-
+    try:
+        file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+        r = requests.get(file_url, timeout=15, allow_redirects=True)
+        with open('temp.xlsx', "wb") as fp:
+            fp.write(r.content)
+        # Read the Excel file
+        df = pd.read_excel('temp.xlsx')
+        # Run various analyses based on the query
+        result = (
+            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+        )
+        result += f"Columns: {', '.join(df.columns)}\n\n"
+        # Add summary statistics
+        result += "Summary statistics:\n"
+        result += str(df.describe())
+        return result
+    except Exception as e:
+        return f"Error analyzing Excel file: {str(e)}"
+   
 @tool
-def mp3_listen(dir: str) -> str:
-    """Listen to the provided mp3 file
+def csv_read(task_id: str) -> str:
+    """First download the csv file, then read its content
     Args:
-        dir: the filepath
+        dir: the task_id
     """
-    loader = AssemblyAIAudioTranscriptLoader(file_path=dir)
+    try:
+        file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+        r = requests.get(file_url, timeout=15, allow_redirects=True)
+        with open('temp.csv', "wb") as fp:
+            fp.write(r.content)
+        # Read the CSV file
+        df = pd.read_csv(temp.csv)
+        # Run various analyses based on the query
+        result = (
+            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+        )
+        result += f"Columns: {', '.join(df.columns)}\n\n"
+        # Add summary statistics
+        result += "Summary statistics:\n"
+        result += str(df.describe())
+        return result
+    except Exception as e:
+        return f"Error analyzing CSV file: {str(e)}"
+        
+@tool
+def mp3_listen(task_id: str) -> str:
+    """First download the mp3 file, then listen to it
+    Args:
+        dir: the task_id
+    """
+    file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+    r = requests.get(file_url, timeout=15, allow_redirects=True)
+    with open('temp.mp3', "wb") as fp:
+        fp.write(r.content)
+    loader = AssemblyAIAudioTranscriptLoader(file_path="temp.mp3", api_key=os.getenv("AssemblyAI_API_KEY"))
     docs = loader.load()
     contents = [doc.page_content for doc in docs]
     return "\n".join(contents)
@@ -171,7 +216,7 @@ class BasicAgent:
             max_tokens=128,
             timeout=None,
             max_retries=2,
-            google_api_key="AIzaSyAxVUPaGJIgdxB46ZR0RWPKSjB9a63Z80o",
+            google_api_key=os.getenv("GEMINI_API_KEY"),
             # other params...
         )
         # System Prompt for few shot prompting
@@ -188,13 +233,16 @@ class BasicAgent:
                 - visit_webpage: visit the given webpage url by passing the url as input
                 - wiki_search: wiki search the content of the query by passing the query as input if the question asks for wiki search it
                 - youtube_transcript: fetch the transcript of the Youtube video by passing the video url as input if the question asks for watching a Youtube video
-                - read_file: read the content of the attached file by passing the file directory as input
-                - excel_read: read the content of the attached excel file by passing the file directory as input
-                - mp3_listen: listen to the content of the attached mp3 file by passing the file directory as input
-                - image_caption: understand the visual content of the attached image by passing the image directory as input
-                - python_tool: run the python code 
+                - read_file: read the content of the attached file by passing the TASK-ID as input
+                - excel_read: read the content of the attached excel file by passing the TASK-ID as input
+                - csv_read: read the content of the attached csv file by passing the TASK-ID as input
+                - mp3_listen: listen to the content of the attached mp3 file by passing the TASK-ID as input
+                - image_caption: understand the visual content of the attached image by passing the TASK-ID as input
+                - python_tool: run the python code
+
+                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]
         """
-        self.tools = [web_search, visit_webpage, wiki_search, youtube_transcript, read_file, excel_read, mp3_listen, image_caption, python_tool]
+        self.tools = [web_search, visit_webpage, wiki_search, youtube_transcript, read_file, excel_read, csv_read, mp3_listen, image_caption, python_tool]
         self.prompt = ChatPromptTemplate.from_messages([
             ("system", self.sys_prompt),
             ("human", "{input}")
@@ -209,12 +257,16 @@ class BasicAgent:
         )
         print("BasicAgent initialized.")
     
-    def __call__(self, question: str) -> str:
+    def __call__(self, task: dict) -> str:
+        task_id, question, Level, file_name = task["task_id"], task["question"], Level, task["file_name"]
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         # response = self.agent_exe.invoke({"input": f"Question: {question}"})
         # fixed_answer = response['message'][-1].content
         time.sleep(15)
-        fixed_answer = self.agent.run(question)
+        if file_name == "" or file_name is None:
+            fixed_answer = self.agent.run(question)
+        else:
+            fixed_answer = self.agent.run(f'{question} with TASK-ID: {task_id}')
         # fixed_answer = "This is a default answer."
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer