Adding new tools to the agent + downloading task files

app.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import suppress
 from pprint import pprint
 from typing import TypedDict, List, Dict, Any, Optional, Tuple
 from typing_extensions import Annotated
@@ -22,10 +23,6 @@ from langchain_core.tools import tool
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-# class State(TypedDict):
-#     question: str
-#     messages: Annotated[List[AnyMessage], add_messages]
-
 
 class State(MessagesState):
     question: str
@@ -33,7 +30,11 @@ class State(MessagesState):
 
 class BasicAgent:
     def __init__(self):
-        self.tools = [BasicAgent.search_tool]
+        self.tools = [
+            BasicAgent.search_tool,
+            BasicAgent.find_local_files_tool,
+            BasicAgent.read_text_file_tool
+        ]
 
         # Chat model with tool support
         self.model = ChatOpenAI(model="gpt-4o", temperature=0)
@@ -67,15 +68,15 @@ class BasicAgent:
         - Use tools when more information is needed.
         - Always ground your answer in tool results or reasoning.
         - If a tool fails or returns nothing, acknowledge it and try an alternative.
-        - If you can't find a confident answer, say so.
+        - If you can't find a confident answer, return your best guess.
         
         **Important:** 
+        - When giving the final answer, reply with only the answer — nothing else.
         - Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
         - If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
         - If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
         - If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
         - If there is no direct ask about the response format, think if it's possible to answer with a number, if yes - do it, if no - answer with as few words as possible, not the whole sentence.
-        - When giving the final answer, reply with only the answer — nothing else.
 
         **Examples:**
 
@@ -137,6 +138,25 @@ class BasicAgent:
         
         print("\nTool result:", result)
         return result[:max_length]
+
+    @staticmethod
+    @tool(
+      description="List task files.",
+    )
+    def find_local_files_tool() -> list[str]:
+        print(f"\nCalling find local files tool")
+        files = [f for f in os.listdir() if os.path.isfile(f) and f.startswith('task_file_')]
+        print(f"\nReturning", files)
+        return files
+
+    @staticmethod
+    @tool(
+      description="Read the text file and return it's content.",
+    )
+    def read_text_file_tool(file_name: str) -> str:
+        print(f"\nCalling read text file tool for", file_name)
+        with open(file_name, 'r') as f:
+            return f.read()
     
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
@@ -199,6 +219,16 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         if not task_id or question_text is None:
             print(f"Skipping item with missing task_id or question: {item}")
             continue
+
+        try:
+            file_url = f"{api_url}/files/{task_id}"
+            file_name = f"task_file_{task_id}"
+            with open(file_name, "wb") as file:
+                response = get(file_url)
+                file.write(response.content)
+        except:
+            print("Didn't manage to download a file, probably it's not expected for this task")
+        
         try:
             submitted_answer, logs = agent(question_text)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
@@ -206,6 +236,9 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         except Exception as e:
              print(f"Error running agent on task {task_id}: {e}")
              results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
+        finally:
+            with suppress(Exception):
+                os.remove(file_name)
 
     if not answers_payload:
         print("Agent did not produce any answers to submit.")