run lint

ServiceNow · Jan 22, 2025 · 56a3b1d · 56a3b1d
1 parent 644cf78
commit 56a3b1d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/tapeagents/llms.py b/tapeagents/llms.py
@@ -208,7 +208,9 @@ def get_token_costs(self) -> dict:
         """
         return {"input": 0, "output": 0}
 
-    def log_output(self, prompt: Prompt, message: LLMOutput, cached: bool = False, count_tokens: bool = True) -> None | LLMCall:
+    def log_output(
+        self, prompt: Prompt, message: LLMOutput, cached: bool = False, count_tokens: bool = True
+    ) -> None | LLMCall:
         """
         Logs the output of an LLM (Language Model) call along with its metadata.
 
@@ -233,7 +235,7 @@ def log_output(self, prompt: Prompt, message: LLMOutput, cached: bool = False, c
         else:
             # -1 is the default value of prompt and output length tokens when token counting is disabled
             prompt_length_tokens = -1
-            output_length_tokens = -1 
+            output_length_tokens = -1
 
         llm_call = LLMCall(
             timestamp=datetime.datetime.now().isoformat(),
@@ -735,7 +737,9 @@ def batch_generate(self, prompts: list[Prompt]) -> list[LLMCall]:
                 llm_call.output_length_tokens = len(chat_completion_logprobs)
                 self._stats["prompt_length_tokens"].append(llm_call.prompt_length_tokens)
                 self._stats["output_length_tokens"].append(llm_call.output_length_tokens)
-                assert llm_call.output_length_tokens <= self.parameters["max_tokens"], f"output_length_tokens: {llm_call.output_length_tokens}, max_tokens: {self.parameters['max_tokens']}"
+                assert (
+                    llm_call.output_length_tokens <= self.parameters["max_tokens"]
+                ), f"output_length_tokens: {llm_call.output_length_tokens}, max_tokens: {self.parameters['max_tokens']}"
             else:
                 llm_call = self.log_output(prompts[i], output, count_tokens=True)
             llm_call.logprobs = logprobs

diff --git a/tapeagents/nodes.py b/tapeagents/nodes.py
@@ -90,7 +90,7 @@ def make_prompt(self, agent: Any, tape: Tape) -> Prompt:
         cleaned_tape = self.prepare_tape(tape)
         steps_description = self.get_steps_description(tape, agent)
         messages = self.tape_to_messages(cleaned_tape, steps_description)
-        #TODO: benchmark counting the token for every single call
+        # TODO: benchmark counting the token for every single call
         if self.trim_tape_when_too_long and agent.llm.count_tokens(messages) > (agent.llm.context_size - 500):
             cleaned_tape = self.trim_tape(cleaned_tape)
         messages = self.tape_to_messages(cleaned_tape, steps_description)