Update agent.py

agent.py

@@ -1,69 +1,45 @@
-import os
 import json
-from smolagents import create_agent_executor, Agent, Task
-from langgraph.prebuilt import ToolExecutor
-from huggingface_hub import InferenceClient
-from tools import *
-
-# Load Hugging Face endpoint from environment
-HF_API_URL = os.getenv("HF_ENDPOINT_URL")  
-HF_API_TOKEN = os.getenv("HF_TOKEN") 
+from pathlib import Path
+from smolagents import Agent, Task
+from mistral_hf_wrapper import HuggingFaceModel
+from tools import (
+    chess_tools,
+    classifier_tool,
+    content_retriever_tool,
+    get_attachments_tool,
+    google_search_tools,
+    speech_recognition_tool,
+    youtube_video_tool,
+)
 
-if not HF_API_URL or not HF_API_TOKEN:
-    raise ValueError("Missing Hugging Face endpoint URL or token.")
+def load_tasks(path: str = "metadata.jsonl") -> list[Task]:
+    tasks = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            task_dict = json.loads(line.strip())
+            tasks.append(Task(**task_dict))
+    return tasks
 
-llm = InferenceClient(
-    model=HF_API_URL,
-    token=HF_API_TOKEN
+model = HuggingFaceModel(
+    endpoint_url="https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
 )
 
-def run_llm(prompt: str) -> str:
-    response = llm.text_generation(
-        prompt,
-        max_new_tokens=512,
-        do_sample=False,
-        temperature=0.0,
-        return_full_text=False,
-    )
-    return response.strip()
-
-tool_list = [
-    GetAttachmentTool(),
-    GoogleSearchTool(),
-    GoogleSiteSearchTool(),
-    ContentRetrieverTool(),
-    SpeechRecognitionTool(),
-    YouTubeVideoTool(),
-    ClassifierTool(),
-    ImageToChessBoardFENTool()
+toolset = [
+    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(),
 ]
 
-# Create tool executor for LangGraph
-tool_executor = ToolExecutor(tool_list)
-
-# Create agent instance
 agent = Agent(
-    llm=run_llm,
-    tools=tool_list,
+    model=model,
+    tools=toolset,
+    system_message="You are a helpful assistant solving GAIA benchmark tasks. Return only the final answer, nothing else."
 )
 
-agent_executor = create_agent_executor(
-    agent=agent,
-    tool_executor=tool_executor,
-    stream=False,
-)
-
-def load_tasks(metadata_path="metadata.jsonl") -> list[Task]:
-    with open(metadata_path, "r") as f:
-        tasks = []
-        for line in f:
-            data = json.loads(line)
-            tasks.append(Task(
-                task_id=data["question_id"],
-                input=data["answer"]  
-            ))
-        return tasks
-
 def solve_task(task: Task) -> str:
-    result = agent_executor.invoke(task.input)
-    return result.get("output", "")
+    return agent.run(task.question, task_id=task.task_id)