From edfd2b954fe95f6602f567b5ac0738c0ff9d2379 Mon Sep 17 00:00:00 2001 From: hobostay Date: Wed, 11 Mar 2026 16:12:14 +0800 Subject: [PATCH 1/3] Fix TypeError when total_tokens is None in PortkeyModel ## Summary Fix a bug where `total_tokens` being None would cause a TypeError in cost calculation. This can happen with some Portkey API responses. ## Changes - **portkey_model.py**: Add check for None total_tokens and calculate from prompt_tokens + completion_tokens - **test_portkey_model.py**: Add test case for total_tokens=None scenario ## Testing Added test `test_portkey_model_total_tokens_none` that simulates the bug scenario and verifies the fix works correctly. Co-Authored-By: Claude Opus 4.6 --- src/minisweagent/models/portkey_model.py | 5 +++ tests/models/test_portkey_model.py | 49 ++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/src/minisweagent/models/portkey_model.py b/src/minisweagent/models/portkey_model.py index 6d23aef00..fba52491d 100644 --- a/src/minisweagent/models/portkey_model.py +++ b/src/minisweagent/models/portkey_model.py @@ -167,6 +167,11 @@ def _calculate_cost(self, response) -> dict[str, float]: f"Completion tokens are None for model {self.config.model_name}. Setting to 0. Full response: {response_for_cost_calc.model_dump()}" ) completion_tokens = 0 + if total_tokens is None: + logger.warning( + f"Total tokens are None for model {self.config.model_name}. Setting to sum of prompt and completion tokens. Full response: {response_for_cost_calc.model_dump()}" + ) + total_tokens = prompt_tokens + completion_tokens if total_tokens - prompt_tokens - completion_tokens != 0: # This is most likely related to how portkey treats cached tokens: It doesn't count them towards the prompt tokens (?) logger.warning( diff --git a/tests/models/test_portkey_model.py b/tests/models/test_portkey_model.py index f8df5dd78..0b0d13d17 100644 --- a/tests/models/test_portkey_model.py +++ b/tests/models/test_portkey_model.py @@ -191,3 +191,52 @@ def test_portkey_model_cost_validation_error(): assert "Error calculating cost" in str(exc_info.value) assert "MSWEA_COST_TRACKING='ignore_errors'" in str(exc_info.value) + + +def test_portkey_model_total_tokens_none(): + """Test that None total_tokens is handled gracefully in cost calculation.""" + mock_portkey_class = MagicMock() + mock_client = MagicMock() + mock_response = MagicMock() + mock_choice = MagicMock() + mock_message = MagicMock() + mock_usage = MagicMock() + mock_tool_call = MagicMock() + + mock_tool_call.id = "call_999" + mock_tool_call.function.name = "bash" + mock_tool_call.function.arguments = json.dumps({"command": "echo test"}) + mock_message.tool_calls = [mock_tool_call] + mock_message.content = None + mock_message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_999", "function": {"name": "bash", "arguments": '{"command": "echo test"}'}}], + } + mock_choice.message = mock_message + mock_response.choices = [mock_choice] + mock_response.model_dump.return_value = {"test": "response"} + mock_response.model_copy.return_value = mock_response + mock_response.usage = mock_usage + mock_usage.prompt_tokens = 10 + mock_usage.completion_tokens = 20 + # Simulate the bug: total_tokens is None + mock_usage.total_tokens = None + + mock_client.chat.completions.create.return_value = mock_response + mock_portkey_class.return_value = mock_client + + with patch("minisweagent.models.portkey_model.Portkey", mock_portkey_class): + with patch.dict(os.environ, {"PORTKEY_API_KEY": "test-key"}): + with patch("minisweagent.models.portkey_model.litellm.cost_calculator.completion_cost") as mock_cost: + mock_cost.return_value = 0.01 + + model = PortkeyModel(model_name="gpt-4o") + messages = [{"role": "user", "content": "test"}] + + # This should not raise a TypeError + result = model.query(messages) + + # Verify the result is still correct + assert result["extra"]["actions"] == [{"command": "echo test", "tool_call_id": "call_999"}] + assert result["extra"]["cost"] == 0.01 From 81a29370d3f4ce2fdd55e1134a7ef6a235a5f38c Mon Sep 17 00:00:00 2001 From: hobostay Date: Thu, 12 Mar 2026 12:06:06 +0800 Subject: [PATCH 2/3] refactor: add defensive check for None tokens in total_tokens calculation Add a helper function to handle potential None values when calculating total_tokens from prompt_tokens + completion_tokens, addressing the concern raised by @klieret. --- src/minisweagent/models/portkey_model.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/minisweagent/models/portkey_model.py b/src/minisweagent/models/portkey_model.py index fba52491d..19aac6aa3 100644 --- a/src/minisweagent/models/portkey_model.py +++ b/src/minisweagent/models/portkey_model.py @@ -171,7 +171,13 @@ def _calculate_cost(self, response) -> dict[str, float]: logger.warning( f"Total tokens are None for model {self.config.model_name}. Setting to sum of prompt and completion tokens. Full response: {response_for_cost_calc.model_dump()}" ) - total_tokens = prompt_tokens + completion_tokens + # Defense: prompt_tokens and completion_tokens are guaranteed to be non-None here + # because they were set to 0 above if they were None. But we use get_token() to + # be extra safe in case of unexpected API responses. + def get_token(val): + return val if val is not None else 0 + + total_tokens = get_token(prompt_tokens) + get_token(completion_tokens) if total_tokens - prompt_tokens - completion_tokens != 0: # This is most likely related to how portkey treats cached tokens: It doesn't count them towards the prompt tokens (?) logger.warning( From dc6e7ea74577cd4f4cd09f1cc577883f617a2d0f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Mar 2026 04:06:16 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/minisweagent/models/portkey_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/minisweagent/models/portkey_model.py b/src/minisweagent/models/portkey_model.py index 19aac6aa3..9950a2eb8 100644 --- a/src/minisweagent/models/portkey_model.py +++ b/src/minisweagent/models/portkey_model.py @@ -171,6 +171,7 @@ def _calculate_cost(self, response) -> dict[str, float]: logger.warning( f"Total tokens are None for model {self.config.model_name}. Setting to sum of prompt and completion tokens. Full response: {response_for_cost_calc.model_dump()}" ) + # Defense: prompt_tokens and completion_tokens are guaranteed to be non-None here # because they were set to 0 above if they were None. But we use get_token() to # be extra safe in case of unexpected API responses.