chore(models): simplify conditions and fix return types (SWE-agent#216)

* chore(models): simplify conditions and fix return types

* undo formatting

---------

Co-authored-by: pmprones <[email protected]>

Author: mspronesti

.gitignore

@@ -189,3 +189,6 @@ website/frontend/build
 trajectories/*
 
 .vscode/**
+
+# PyCharm
+.idea/

sweagent/agent/models.py

@@ -16,7 +16,7 @@
     wait_random_exponential,
     retry_if_not_exception_type,
 )
-from typing import Optional
+from typing import Optional, Union
 
 logger = logging.getLogger("api_models")
 
@@ -49,7 +49,6 @@ def __add__(self, other):
             field.name: getattr(self, field.name) + getattr(other, field.name)
             for field in fields(self)
         })
-    
     def replace(self, other):
         if not isinstance(other, APIStats):
             raise TypeError("Can only replace APIStats with APIStats")
@@ -148,19 +147,13 @@ def update_stats(self, input_tokens: int, output_tokens: int) -> float:
         )
 
         # Check whether total cost or instance cost limits have been exceeded
-        if (
-            self.args.total_cost_limit > 0
-            and self.stats.total_cost >= self.args.total_cost_limit
-        ):
+        if 0 < self.args.total_cost_limit <= self.stats.total_cost:
             logger.warning(
                 f"Cost {self.stats.total_cost:.2f} exceeds limit {self.args.total_cost_limit:.2f}"
             )
             raise CostLimitExceededError("Total cost limit exceeded")
 
-        if (
-            self.args.per_instance_cost_limit > 0
-            and self.stats.instance_cost >= self.args.per_instance_cost_limit
-        ):
+        if 0 < self.args.per_instance_cost_limit <= self.stats.instance_cost:
             logger.warning(
                 f"Cost {self.stats.instance_cost:.2f} exceeds limit {self.args.per_instance_cost_limit:.2f}"
             )
@@ -233,7 +226,7 @@ def __init__(self, args: ModelArguments, commands: list[Command]):
 
     def history_to_messages(
         self, history: list[dict[str, str]], is_demonstration: bool = False
-    ) -> list[dict[str, str]]:
+    ) -> Union[str, list[dict[str, str]]]:
         """
         Create `messages` by filtering out all keys except for role/content per `history` turn
         """
@@ -273,6 +266,7 @@ def query(self, history: list[dict[str, str]]) -> str:
         self.update_stats(input_tokens, output_tokens)
         return response.choices[0].message.content
 
+
 class AnthropicModel(BaseModel):
     MODELS = {
         "claude-instant": {
@@ -326,7 +320,7 @@ def __init__(self, args: ModelArguments, commands: list[Command]):
 
     def history_to_messages(
         self, history: list[dict[str, str]], is_demonstration: bool = False
-    ) -> list[dict[str, str]]:
+    ) -> Union[str, list[dict[str, str]]]:
         """
         Create `prompt` by filtering out all keys except for role/content per `history` turn
         Reference: https://docs.anthropic.com/claude/reference/complete_post
@@ -440,7 +434,7 @@ def __init__(self, args: ModelArguments, commands: list[Command]):
 
     def history_to_messages(
         self, history: list[dict[str, str]], is_demonstration: bool = False
-    ) -> list[dict[str, str]]:
+    ) -> Union[str, list[dict[str, str]]]:
         """
         Create `messages` by filtering out all keys except for role/content per `history` turn
         """
@@ -516,7 +510,7 @@ class TogetherModel(BaseModel):
             "max_context": 32768,
             "cost_per_input_token": 6e-07,
             "cost_per_output_token": 6e-07,
-        },           
+        },
     }
 
     SHORTCUTS = {
@@ -593,7 +587,7 @@ def __init__(self, args: ModelArguments, commands: list[Command]):
 
     def history_to_messages(
         self, history: list[dict[str, str]], is_demonstration: bool = False
-    ) -> list[dict[str, str]]:
+    ) -> Union[str, list[dict[str, str]]]:
         """
         Create `messages` by filtering out all keys except for role/content per `history` turn
         """
@@ -652,7 +646,7 @@ def query(self, history: list[dict[str, str]]) -> str:
                 break
             thought_all += thought
             thought = input("... ")
-        
+
         action = super().query(history, action_prompt="Action: ")
 
         return f"{thought_all}\n```\n{action}\n```"