add math agent.

basic_agent.py

@@ -18,47 +18,67 @@ Take the time to plan the steps to reach the solution. Show the steps and then e
 
 class BasicAgent:
 
-  def __init__(self, model_id=None):
+  def __init__(self):
     print("BasicAgent initializing.")
     # Logs appear to be swallowed.
     LOG.warning("logging BasicAgent initialized.")
 
-    # Assume we will use the default model creation
-    self.model = None
-
-    if model_id:
-      self.model_id = model_id
-
-      # Handle the special cases
-      if model_id.lower() == "google":
-        self.model_id = "google"
-
-        # Use Google gemini free tier
-        GEM_KEY=os.environ["GOOGLE_API_KEY"]
-        self.model = smolagents.OpenAIServerModel(
-            model_id="gemini-2.0-flash",
-            api_base="https://generativelanguage.googleapis.com/v1beta/openai/",
-            api_key=GEM_KEY,
-            temperature=0.3)
-      elif model_id.lower() == "local":
-        self.model_id = "Qwen/Qwen3-4B-FP8"
-        # Run locally.
-        self.model = smolagents.TransformersModel(
-          model_id=self.model_id,
-          max_new_tokens=32000,
-          temperature=0.3
-          )
-    else:
-      self.model_id = "Qwen/Qwen3-32B"
-
-    if not self.model:
-      self.model = smolagents.HfApiModel(
+    self.manager_model_id = "Qwen/Qwen3-32B"
+    self.manager_model = None
+
+    self.code_model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
+    self.code_model = None
+
+    """
+    # Handle the special cases
+    if self.model_id.lower() == "google":
+      self.model_id = "google"
+
+      # Use Google gemini free tier
+      GEM_KEY=os.environ["GOOGLE_API_KEY"]
+      self.model = smolagents.OpenAIServerModel(
+          model_id="gemini-2.0-flash",
+          api_base="https://generativelanguage.googleapis.com/v1beta/openai/",
+          api_key=GEM_KEY,
+          temperature=0.3)
+    elif self.model_id.lower() == "local":
+      self.model_id = "Qwen/Qwen3-4B-FP8"
+      # Run locally.
+      self.model = smolagents.TransformersModel(
+        model_id=self.model_id,
+        max_new_tokens=32000,
+        temperature=0.3
+        )
+     """
+
+    if not self.manager_model:
+      self.manager_model = smolagents.HfApiModel(
           max_tokens=32000,
           temperature=0.3,
-          model_id=self.model_id,
+          model_id=self.manager_model_id,
           custom_role_conversions=None,
       )
-    print(f"NEW2: BasicAgent {self.model_id=} {self.model=}")
+
+    if not self.code_model:
+      self.code_model = smolagents.HfApiModel(
+          max_tokens=32000,
+          temperature=0.3,
+          model_id=self.code_model_id,
+          custom_role_conversions=None,
+      )
+
+    print(f"NEW3: BasicAgent {self.manager_model_id=} {self.code_model_id=}")
+
+    self.math_agent = smolagents.CodeAgent(
+        name="date_time_math_agent",
+        description="Handle date, time and math caclulations",
+        model=self.code_model,
+        tools=[],
+        max_steps=6,
+        verbosity_level=2,
+        planning_interval=None,
+        additional_authorized_imports=["datetime", "time", "calendar"],
+    )
 
     web_search_tools = [
         smolagents.DuckDuckGoSearchTool(),
@@ -70,7 +90,7 @@ class BasicAgent:
     self.web_search_agent = smolagents.CodeAgent(
         name="web_search_agent",
         description="Search the web",
-        model=self.model,
+        model=self.code_model,
         tools=web_search_tools,
         max_steps=6,
         verbosity_level=2,
@@ -78,17 +98,17 @@ class BasicAgent:
         additional_authorized_imports=["duckduckgo_search"],
     )
 
-    wiki_search_tools = [
+    wikipedia_search_tools = [
         smolagents.WikipediaSearchTool(),
         smolagents.VisitWebpageTool(),
         smolagents.FinalAnswerTool()
     ]
 
-    self.wiki_search_agent = smolagents.CodeAgent(
+    self.wikipedia_search_agent = smolagents.CodeAgent(
         name="wikipedia_search_agent",
-        description="Search wikipedia",
-        model=self.model,
-        tools=wiki_search_tools,
+        description="Best search of wikipedia pages",
+        model=self.code_model,
+        tools=wikipedia_search_tools,
         max_steps=6,
         verbosity_level=2,
         planning_interval=None,
@@ -103,18 +123,15 @@ class BasicAgent:
         name="manager_agent",
         description="Manger of other agents",
         tools=[smolagents.FinalAnswerTool()],
-        model=self.model,
+        model=self.manager_model,
+        add_base_tools=True,
         max_steps=10,
         verbosity_level=2,
         planning_interval=None,
         additional_authorized_imports=["duckduckgo_search", "wikipedia-api"],
-        managed_agents=[self.web_search_agent, self.wiki_search_agent])
+        managed_agents=[self.web_search_agent, self.wikipedia_search_agent, self.math_agent])
 
   def __call__(self, question: str) -> str:
-    # Avoid rate limiting issues
-    if self.model_id == "google":
-      time.sleep(1)
-
     print(f"NEW Agent received question (first 50 chars): {question[:50]}...")
     prompt = f"{SYSTEM_PROMPT}\n\n{question}"
     answer = self.manager_agent.run(prompt)