Update app.py

app.py

@@ -6,7 +6,7 @@ import pandas as pd
 from typing import TypedDict, Annotated, Union, Dict, Any
 from smolagents import DuckDuckGoSearchTool
 from langchain_community.tools import DuckDuckGoSearchRun
-from tools import get_hub_stats, analyze_image, read_excel_file
+from tools import get_hub_stats, analyze_image, read_excel_file, download_file
 from langchain.tools import Tool
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from langchain_openai import ChatOpenAI
@@ -33,7 +33,7 @@ class BasicAgent:
 class NewAgent:
     def __init__(self):
         print("NewAgent initialized.")
-    def __call__(self, question: str, file_name: str = "") -> str:
+    def __call__(self, question: str, file_name: str = "", task_id: str = "") -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         
         # Initialize the web search tool
@@ -56,6 +56,12 @@ class NewAgent:
             func=read_excel_file,
             description="Reads an Excel file and returns structured information about its contents."
         )
+        # Initialize the Download File tool
+        download_file_tool = Tool(
+            name="download_file",
+            func=lambda task_id_filename: download_file(*task_id_filename.split(",", 1)),
+            description="Downloads a file associated with a task_id. Input should be 'task_id,filename'"
+        )
         # Generate the chat interface, including the tools
         tools = [
             search_tool,
@@ -63,6 +69,7 @@ class NewAgent:
             hub_stats_tool,
             image_analysis_tool,
             read_excel_tool,
+            download_file_tool,
         ]
         # llm = ChatOpenAI(model="gpt-4o")
         llm = ChatOpenAI(model="gpt-4.1")
@@ -76,12 +83,15 @@ class NewAgent:
                 content=f"""
                 You are a general AI assistant. I will ask you a question. 
                 If a file_name is provided, it indicates there's an associated file you may need to analyze.
+                You can download files using the download_file tool with the format 'task_id,filename'.
+                After downloading, you can analyze images with analyze_image or Excel files with read_excel_file.
                 If you cannot find an answer, you may report your thoughts.
                 If you find an answer, your response should only contain your final answer. Report nothing before or after this answer.
                 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.
+                Current task_id: {task_id}
                 Current file (if any): {file_name}"""
                                    )
             return {
@@ -173,7 +183,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
-            submitted_answer = agent(question_text, file_name)
+            submitted_answer = agent(question_text, file_name, task_id)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "File": file_name, "Submitted Answer": submitted_answer})
         except Exception as e: