Update agent.py

agent.py

@@ -1,49 +1,29 @@
+import os
 import json
-from pathlib import Path
-from smolagents import Agent
-from smolagents.core.task import Task
-from mistral_hf_wrapper import HuggingFaceModel
 
-# Import tools
-from tools import (
-    chess_tools,
-    classifier_tool,
-    content_retriever_tool,
-    get_attachments_tool,
-    google_search_tools,
-    speech_recognition_tool,
-    youtube_video_tool,
-)
+from tools import TOOLS
+from metadata import load_metadata  
+from mistral_inference import query_mistral  
 
-# Load GAIA task list from metadata.jsonl
-def load_tasks(path: str = "metadata.jsonl") -> list[Task]:
-    with open(path, "r", encoding="utf-8") as f:
-        return [Task(**json.loads(line)) for line in f]
+API_URL = os.getenv("HF_MISTRAL_URL")
+API_TOKEN = os.getenv("HF_TOKEN")
 
-# Configure Mistral Inference Endpoint wrapper
-model = HuggingFaceModel(
-    endpoint_url="https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
-)
+HEADERS = {
+    "Authorization": f"Bearer {API_TOKEN}",
+    "Content-Type": "application/json"
+}
 
-# List of tool instances
-tools = [
-    chess_tools.ImageToChessBoardFENTool(),
-    classifier_tool.ClassifierTool(),
-    content_retriever_tool.ContentRetrieverTool(),
-    get_attachments_tool.GetAttachmentTool(),
-    google_search_tools.GoogleSearchTool(),
-    google_search_tools.GoogleSiteSearchTool(),
-    speech_recognition_tool.SpeechRecognitionTool(),
-    youtube_video_tool.YouTubeVideoTool(),
-]
+# Load all tasks from metadata.jsonl
+def load_tasks():
+    return load_metadata("metadata.jsonl")
 
-# Build the agent
-agent = Agent(
-    model=model,
-    tools=tools,
-    system_message="You are a helpful AI agent solving GAIA benchmark tasks. Only return the answer, no extra explanation."
-)
+# Solve a single task
+def solve_task(task, tools=TOOLS):
+    system_prompt = "You are a helpful agent. Use reasoning, tools if needed, and return the answer only."
+    user_prompt = task["question"]
 
-# Core solve function
-def solve_task(task: Task) -> str:
-    return agent.run(task.question, task_id=task.task_id)
+    response = query_mistral(API_URL, HEADERS, system_prompt, user_prompt)
+    return {
+        "task_id": task["question_id"],
+        "submitted_answer": response.strip()
+    }