Update agent.py

agent.py

@@ -1,102 +1,112 @@
-# Libs
-import os
-import yaml
-import importlib
-
-from smolagents import (
-    CodeAgent,
-    DuckDuckGoSearchTool,
-    VisitWebpageTool,
-    WikipediaSearchTool,
-    OpenAIServerModel,
-    SpeechToTextTool,
-    FinalAnswerTool,
-)
-
-# Local
-from tools import GetTaskFileTool, LoadXlsxFileTool, LoadTextFileTool
-
-sys_instruction = (
-    "You are a general AI assistant. Answer each question by reporting your thoughts, "
-    "then submit ONLY a concise text using the 'final_answer' tool. "
-    "Final answer MUST be a number, a few words, or a comma-separated list of numbers and/or strings. "
-    "For numbers, avoid commas and units unless specified. For strings, avoid articles and abbreviations, "
-    "and write digits in full unless stated otherwise. Apply these rules for list elements as well."
-)
-
-prompts = yaml.safe_load(
-    importlib.resources.files("smolagents.prompts").joinpath("code_agent.yaml").read_text()
-)
-prompts["system_prompt"] = sys_instruction + prompts["system_prompt"]
-
-# req_instruction = (
-req_instruction = (
-    "You are an expert and helpful agent named {{name}}.\n"
-    "A valued client has assigned you the following task:\n"
-    "---\n"
-    "Task:\n"
-    "{{task}}\n"
-    "---\n"
-    "Use your tools as needed. Before completing the task, plan your actions carefully.\n"
-    "While completing the task, think step by step. And after completing the task, carefully double check your solution.\n\n"
-    "If you respond correctly, you will be rewarded with a very high bonus.\n\n"
-    "Your final_answer MUST be:\n"
-    "- a number,\n"
-    "- a short phrase,\n"
-    "- or a comma-separated list of numbers or strings (no articles or abbreviations).\n\n"
-    "Only the content passed to the final_answer tool will be preserved—everything else will be discarded."
-)
-prompts['managed_agent']['task'] = req_instruction
-prompts['managed_agent']['report'] = "{{final_answer}}"
-
-# print(prompts["system_prompt"])
-# print(prompts['planning'])
-# print(prompts['managed_agent'])
-# print(prompts['final_answer'])
-
-def get_model(
-        model_id: str = "gpt-4.1-mini",
-        model_temperature: float = 0.7,
-    ):
-    """
-    Create and return an OpenAIServerModel instance with the specified model ID and temperature.
-    """
-    # Load the model
-    model = OpenAIServerModel(
-            model_id=model_id,
-            api_key=os.getenv("OPENAI_API_KEY"),
-            temperature=model_temperature
-    )
-    return model
-
-def get_agent(
-        model_id: str = "gpt-4.1-mini",
-        model_temperature: float = 0.7,
-        agent_max_steps: int = 15,
-    ):
-    """
-    Create and return a CodeAgent instance with the specified model and tools.
-    """
-
-    # Defne the agent with the tools
-    agent = CodeAgent(
-        tools=[
-            DuckDuckGoSearchTool(),
-            VisitWebpageTool(),
-            WikipediaSearchTool(),
-            GetTaskFileTool(),
-            SpeechToTextTool(),
-            LoadXlsxFileTool(),
-            LoadTextFileTool(),
-            FinalAnswerTool(),
-        ],
-        model=get_model(
-            model_id=model_id,
-            model_temperature=model_temperature,
-        ),
-        prompt_templates=prompts,
-        max_steps=agent_max_steps,
-        additional_authorized_imports = ["pandas"],
-        name="GAIAAgent",
-    )
+# Libs
+import os
+import yaml
+import importlib
+
+from smolagents import (
+    CodeAgent,
+    DuckDuckGoSearchTool,
+    VisitWebpageTool,
+    WikipediaSearchTool,
+    OpenAIServerModel,
+    SpeechToTextTool,
+    FinalAnswerTool,
+)
+
+# Local
+from tools import GetTaskFileTool, LoadXlsxFileTool, LoadTextFileTool
+
+sys_instruction = (
+    "You are a general AI assistant. Answer each question by reporting your thoughts, "
+    "then submit ONLY a concise text using the 'final_answer' tool. "
+    "Final answer MUST be a number, a few words, or a comma-separated list of numbers and/or strings. "
+    "For numbers, avoid commas and units unless specified. For strings, avoid articles and abbreviations, "
+    "and write digits in full unless stated otherwise. Apply these rules for list elements as well."
+)
+
+prompts = yaml.safe_load(
+    importlib.resources.files("smolagents.prompts").joinpath("code_agent.yaml").read_text()
+)
+prompts["system_prompt"] = sys_instruction + prompts["system_prompt"]
+
+# req_instruction = (
+req_instruction = (
+    "You are an expert and helpful agent named {{name}}.\n"
+    "A valued client has assigned you the following task:\n"
+    "---\n"
+    "Task:\n"
+    "{{task}}\n"
+    "---\n"
+    "Use your tools as needed. Before completing the task, plan your actions carefully.\n"
+    "While completing the task, think step by step. And after completing the task, carefully double check your solution.\n\n"
+    "If you respond correctly, you will be rewarded with a very high bonus.\n\n"
+    "Your final_answer MUST be:\n"
+    "- a number,\n"
+    "- a short phrase,\n"
+    "- or a comma-separated list of numbers or strings (no articles or abbreviations).\n\n"
+    "Only the content passed to the final_answer tool will be preserved—everything else will be discarded."
+)
+prompts['managed_agent']['task'] = req_instruction
+prompts['managed_agent']['report'] = "{{final_answer}}"
+
+# print(prompts["system_prompt"])
+# print(prompts['planning'])
+# print(prompts['managed_agent'])
+# print(prompts['final_answer'])
+
+def get_model(
+        model_id: str = "gpt-4.1-mini",
+        model_temperature: float = 0.7,
+    ):
+    """
+    Create and return an OpenAIServerModel instance with the specified model ID and temperature.
+    """
+    # Load the model
+    if "gpt" in model_id:
+        model = OpenAIServerModel(
+                model_id=model_id,
+                api_key=os.getenv("OPENAI_API_KEY"),
+                temperature=model_temperature
+        )
+    elif "gemini" in model_id:
+        model = OpenAIServerModel(
+                model_id=model_id,
+                api_key=os.getenv("GOOGLEAI_API_KEY"),
+                api_base="https://generativelanguage.googleapis.com/v1beta/openai/",
+                temperature=model_temperature
+        )
+    else:
+        raise ValueError(f"Unknown model_id: {model_id}. Supported models are gpt and gemini.")
+    return model
+
+def get_agent(
+        model_id: str = "gpt-4.1-mini",
+        model_temperature: float = 0.7,
+        agent_max_steps: int = 15,
+    ):
+    """
+    Create and return a CodeAgent instance with the specified model and tools.
+    """
+
+    # Defne the agent with the tools
+    agent = CodeAgent(
+        tools=[
+            DuckDuckGoSearchTool(),
+            VisitWebpageTool(),
+            WikipediaSearchTool(),
+            GetTaskFileTool(),
+            SpeechToTextTool(),
+            LoadXlsxFileTool(),
+            LoadTextFileTool(),
+            FinalAnswerTool(),
+        ],
+        model=get_model(
+            model_id=model_id,
+            model_temperature=model_temperature,
+        ),
+        prompt_templates=prompts,
+        max_steps=agent_max_steps,
+        additional_authorized_imports = ["pandas"],
+        name="GAIAAgent",
+    )
     return agent