fix tests

All-Hands-AI · Feb 22, 2025 · 801b134 · 801b134
1 parent 5b063cc
commit 801b134
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tests/unit/test_llm.py b/tests/unit/test_llm.py
@@ -2,6 +2,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from litellm import PromptTokensDetails
 from litellm.exceptions import (
     RateLimitError,
 )
@@ -441,7 +442,7 @@ def test_llm_token_usage(mock_litellm_completion, default_config):
         'usage': {
             'prompt_tokens': 12,
             'completion_tokens': 3,
-            'prompt_tokens_details': {'cached_tokens': 2},
+            'prompt_tokens_details': PromptTokensDetails(cached_tokens=2),
             'model_extra': {'cache_creation_input_tokens': 5},
         },
     }
@@ -453,7 +454,7 @@ def test_llm_token_usage(mock_litellm_completion, default_config):
         'usage': {
             'prompt_tokens': 7,
             'completion_tokens': 2,
-            'prompt_tokens_details': {'cached_tokens': 1},
+            'prompt_tokens_details': PromptTokensDetails(cached_tokens=1),
             'model_extra': {'cache_creation_input_tokens': 3},
         },
     }

diff --git a/tests/unit/test_message_utils.py b/tests/unit/test_message_utils.py
@@ -315,7 +315,7 @@ def test_get_single_tokens_usage_for_event():
     assert found.response_id == 'test-response-id'
 
     # If we change the event's response ID, we won't find anything
-    mock_tool_call_metadata.model_response['id'] = 'some-other-id'
+    mock_tool_call_metadata.model_response.id = 'some-other-id'
     found2 = get_single_tokens_usage_for_event(event, metrics)
     assert found2 is None