nits

src/ldp/alg/rollout.py

@@ -121,7 +121,9 @@ async def sample_trajectories(self, **kwargs):
                 kwargs["environment_factory"],
                 kwargs.get("batch_size", 1),
                 kwargs.get("max_steps"),
-                log_exceptions_immediately=kwargs.get("log_exceptions_immediately", True)
+                log_exceptions_immediately=kwargs.get(
+                    "log_exceptions_immediately", True
+                ),
             )
 
         if "environments" in kwargs:
@@ -131,7 +133,9 @@ async def sample_trajectories(self, **kwargs):
             return await self._sample_trajectories_from_envs(
                 kwargs["environments"],
                 kwargs.get("max_steps"),
-                log_exceptions_immediately=kwargs.get("log_exceptions_immediately", True),
+                log_exceptions_immediately=kwargs.get(
+                    "log_exceptions_immediately", True
+                ),
             )
 
         raise TypeError(

src/ldp/nn/agent/simple_local_agent.py

@@ -46,7 +46,7 @@ class AgentLMConfig(_LMConfig):
         ),
         validate_default=True,
     )
-    max_traj_token_count: int | None = Field(
+    max_messages_token_count: int | None = Field(
         default=None,
         description="If set, raise an error if the total tokens in the trajectory exceed this value.",
     )
@@ -128,7 +128,7 @@ async def get_asv(
 
     def _validate_token_count(self, messages: list[Message], tools: list[Tool]):
         """Asserts token count for the trajectory is within the limit."""
-        if self.llm_model.max_traj_token_count is None:
+        if self.llm_model.max_messages_token_count is None:
             return
         messages_for_tokenizer = self._llm_call_op.prep_messages_for_tokenizer(messages)
         tools_for_tokenizer = self._llm_call_op.prep_tools_for_tokenizer(tools)
@@ -138,12 +138,12 @@ def _validate_token_count(self, messages: list[Message], tools: list[Tool]):
             messages=messages_for_tokenizer,
             tools=tools_for_tokenizer,  # type: ignore[arg-type]
         )
-        if total_tokens > self.llm_model.max_traj_token_count:
+        if total_tokens > self.llm_model.max_messages_token_count:
             logger.error(
-                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
+                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_messages_token_count}"
             )
             raise ValueError(
-                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_traj_token_count}"
+                f"Token limit exceeded for trajectory: {total_tokens} > {self.llm_model.max_messages_token_count}"
             )
 
     # TODO: maybe remove these recomputation methods. I added them to debug some things. But idk,