Update agent.py

agent.py

@@ -1,4 +1,4 @@
-from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, LiteLLMModel # HfApiModel, OpenAIServerModel
+from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, LiteLLMModel, tool, load_tool # HfApiModel, OpenAIServerModel
 from tools.final_answer import FinalAnswerTool
 import asyncio
 import os
@@ -12,18 +12,14 @@ class MagAgent:
             model_id="gemini/gemini-2.0-flash",
             api_key= os.environ.get("GEMINI_KEY"),
             max_tokens=8192
-#        model = HfApiModel(
-#            max_tokens=2096,
-#            temperature=0.5,
-#            model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
-#            custom_role_conversions=None,
         )
+
         # Load prompt templates
         with open("prompts.yaml", 'r') as stream:
             prompt_templates = yaml.safe_load(stream)
         
         self.agent = CodeAgent(
-            model= model,  # OpenAIServerModel(model_id="gpt-4o"), 
+            model= model,
             tools=[
                 DuckDuckGoSearchTool(),
                 WikipediaSearchTool(),