Update models and address auto import of agents (#35)

* Experimental changes

* updating models

* Dealing with auto import of agents/. Add opus 4.6 models

Author: MarcCote

agents/llm.py

@@ -25,6 +25,20 @@
     " When stuck, try using the `help` command to see what commands are available."
 )
 
+CLAUDE_MODELS = [
+    "claude-3.5-haiku",
+    "claude-3.5-sonnet",
+    "claude-3.5-sonnet-latest",
+    "claude-3.7-sonnet",
+    "claude-4-sonnet",
+    "claude-4-opus",
+    "claude-opus-4.5",
+    "claude-opus-4.6",
+    "claude-sonnet-4.5",
+    "claude-sonnet-4.6",
+    "claude-haiku-4.5",
+]
+
 
 class LLMAgent(tales.Agent):
 
@@ -98,12 +112,7 @@ def act(self, obs, reward, done, infos):
             "seed": self.seed,
             "stream": False,
         }
-        if self.llm in [
-            "claude-3.5-haiku",
-            "claude-3.5-sonnet",
-            "claude-3.5-sonnet-latest",
-            "claude-3.7-sonnet",
-        ]:
+        if self.llm in CLAUDE_MODELS:
             # For these models, we cannot set the seed.
             llm_kwargs.pop("seed")
 

agents/reasoning.py

@@ -28,6 +28,36 @@
 
 DEEPSEEK_CHAT_TEMPLATE_NO_THINK = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}{{'<｜Assistant｜>' + content + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜><think>\\n</think>\\n'}}{% endif %}"
 
+CLAUDE_MODELS = [
+    "claude-3.7-sonnet",
+    "claude-4-sonnet",
+    "claude-4-opus",
+    "claude-sonnet-4.5",
+    "claude-haiku-4.5",
+    "claude-opus-4.5",
+    "claude-opus-4.6",
+    "claude-sonnet-4.6",
+]
+
+OPENAI_MODELS = [
+    "o1",
+    "o1-mini",
+    "o1-preview",
+    "o3-mini",
+    "o4-mini",
+    "o3",
+    "gpt-5.1",
+    "gpt-5.2",
+    "gpt-5",
+    "gpt-5-mini",
+    "gpt-5-nano",
+]
+
+GEMINI_MODELS = [
+    "gemini-2.5-pro",
+    "gemini-3-pro-preview",
+]
+
 
 class ReasoningAgent(tales.Agent):
 
@@ -117,42 +147,25 @@ def act(self, obs, reward, done, infos):
             "stream": True,  # Should prevent openai.APITimeoutError
         }
         if isinstance(self.reasoning_effort, int):
-            if self.llm in ["claude-3.7-sonnet"]:
+            if self.llm in CLAUDE_MODELS:
                 llm_kwargs["thinking_budget"] = self.reasoning_effort
             else:
                 llm_kwargs["max_tokens"] = self.reasoning_effort
 
-        elif self.llm in [
-            "o1",
-            "o1-preview",
-            "o3-mini",
-            "o4-mini",
-            "o3",
-            "gpt-5",
-            "gpt-5-mini",
-            "gpt-5-nano",
-        ]:
+        elif self.llm in OPENAI_MODELS:
             llm_kwargs["reasoning_effort"] = self.reasoning_effort
 
-        if self.llm in [
-            "o1",
-            "o1-mini",
-            "o1-preview",
-            "o3-mini",
-            "o4-mini",
-            "o3",
-            "claude-3.7-sonnet",
-            "gpt-5",
-            "gpt-5-mini",
-            "gpt-5-nano",
-        ]:
+        elif self.llm in CLAUDE_MODELS:
+            llm_kwargs["thinking_effort"] = self.reasoning_effort
+
+        if self.llm in OPENAI_MODELS + CLAUDE_MODELS:
             # For these models, we cannot set the temperature.
             llm_kwargs.pop("temperature")
 
         if self.llm in ["o3-mini"]:
             llm_kwargs.pop("stream")
 
-        if self.llm in ["claude-3.7-sonnet"]:
+        if self.llm in CLAUDE_MODELS:
             llm_kwargs["thinking"] = 1
             llm_kwargs.pop("seed")
 
@@ -238,7 +251,7 @@ def act(self, obs, reward, done, infos):
             # Extract the action part from the response.
             action = action[reasoning_end:].strip()
 
-        elif self.llm in ["claude-3.7-sonnet"]:
+        elif self.llm in CLAUDE_MODELS:
             # Extract the thinking part from the response JSON.
             thinking = "".join(
                 [item.get("thinking", "") for item in response.json()["content"]]
@@ -253,24 +266,14 @@ def act(self, obs, reward, done, infos):
             "response": response_text,
         }
 
-        if self.llm in ["gemini-2.5-pro-preview-03-25", "gemini-2.5-pro-preview-05-06"]:
+        if self.llm in GEMINI_MODELS:
             stats["nb_tokens_prompt"] = response.usage().input
             stats["nb_tokens_thinking"] = response.usage().details.get(
                 "thoughtsTokenCount", 0
             )
             stats["nb_tokens_response"] = response.usage().output
 
-        elif self.llm in [
-            "o1",
-            "o1-mini",
-            "o1-preview",
-            "o3-mini",
-            "o4-mini",
-            "o3",
-            "gpt-5",
-            "gpt-5-mini",
-            "gpt-5-nano",
-        ]:
+        elif self.llm in OPENAI_MODELS:
             # stats["nb_tokens_prompt"] = self.token_counter(messages=messages),
             # stats["nb_tokens_response"] = self.token_counter(text=response_text)
             stats["nb_tokens_prompt"] = response.usage().input
@@ -281,8 +284,17 @@ def act(self, obs, reward, done, infos):
                 "completion_tokens_details"
             ]["reasoning_tokens"]
 
+        elif self.llm in CLAUDE_MODELS:
+            stats["nb_tokens_prompt"] = self.token_counter(messages=messages)
+            stats["nb_tokens_response"] = self.token_counter(text=response_text)
+            stats["nb_tokens_thinking"] = 0
+            if thinking:
+                stats["nb_tokens_thinking"] = (
+                    response.usage().output - self.token_counter(text=response_text)
+                )
+
         else:
-            stats["nb_tokens_prompt"] = (self.token_counter(messages=messages),)
+            stats["nb_tokens_prompt"] = self.token_counter(messages=messages)
             stats["nb_tokens_thinking"] = (
                 self.token_counter(text=thinking) if thinking else 0
             )

benchmark.py

@@ -405,6 +405,7 @@ def _maybe_load_agent_module():
     if args.agent:
         print(f"Importing agent(s) from {args.agent}.")
         for agent_file in glob.glob(args.agent):
+            print(f"Importing {agent_file}...")
             agent_dirname = os.path.dirname(agent_file)
             agent_filename, _ = os.path.splitext(os.path.basename(agent_file))
             if f"{agent_dirname}.{agent_filename}" in sys.modules:

tales/agent.py

@@ -33,7 +33,7 @@ def params(self):
 AGENTS = {}
 
 
-def register(name: str, desc: str, klass: callable, add_arguments: callable) -> None:
+def register(name: str, desc: str, klass: type, add_arguments: callable) -> None:
     """ Register a new type of Agent.
 
     Arguments:
@@ -60,7 +60,9 @@ def register(name: str, desc: str, klass: callable, add_arguments: callable) ->
         >>>          klass=RandomAgent,
         >>>          add_arguments=_add_arguments)
     """
-    if name in AGENTS:
-        raise ValueError(f"Agent '{name}' already registered.")
+    if name in AGENTS and str(klass) != str(AGENTS[name][1]):
+        raise ValueError(
+            f"Agent '{name}' already registered from {AGENTS.get(name)[1]}."
+        )
 
     AGENTS[name] = (desc, klass, add_arguments)

tales/token.py

@@ -46,7 +46,15 @@ def __init__(self, model: str):
         self.model = model
         if self.model in tiktoken.model.MODEL_TO_ENCODING:
             self.tokenize = tiktoken.encoding_for_model(self.model).encode
-        elif self.model in ("o4-mini", "o3", "gpt-5", "gpt-5-mini", "gpt-5-nano"):
+        elif self.model in (
+            "o4-mini",
+            "o3",
+            "gpt-5",
+            "gpt-5-mini",
+            "gpt-5-nano",
+            "gpt-5.1",
+            "gpt-5.2",
+        ):
             self.tokenize = tiktoken.encoding_for_model("o3-mini").encode
         elif self.model in ("gpt-4.1", "gpt-4.1-nano", "gpt-4.1-mini"):
             self.tokenize = tiktoken.encoding_for_model("gpt-4o").encode
@@ -99,19 +107,31 @@ def __call__(self, *, messages=None, text=None):
             system = messages[0]["content"]
             messages.pop(0)
 
-        return self.client.beta.messages.count_tokens(
+        nb_tokens = self.client.messages.count_tokens(
             model=self.model,
             messages=messages,
             system=system,
         ).input_tokens
 
+        if text is not None:
+            # Remove the boilerplate tokens needed for the count_tokens(...).
+            # i.e. [{"role": "assistant", "content": ""}]
+            nb_tokens -= (
+                16 - 3
+            )  # 16 tokens for the boilerplate, minus 3 for <|im_*|> tokens.
+
+        return nb_tokens
+
 
 class GeminiTokenCounter(TokenCounter):
 
     def __init__(self, model: Model):
         from google import genai
 
         self.model = model.model_id
+        # Strip 'gemini/' prefix if present
+        if self.model.startswith("gemini/"):
+            self.model = self.model[len("gemini/") :]
         self.client = genai.Client(api_key=model.get_key())
 
     def __call__(self, *, messages=None, text=None):