From 801b134c7f73eefed6266879a79c8133ed673c5f Mon Sep 17 00:00:00 2001 From: Engel Nyst Date: Sat, 22 Feb 2025 20:36:59 +0100 Subject: [PATCH] fix tests --- tests/unit/test_llm.py | 5 +++-- tests/unit/test_message_utils.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_llm.py b/tests/unit/test_llm.py index 86a9a1b618de..0f9a8c7362a7 100644 --- 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 index 34c52d861a23..e50ed08f9513 100644 --- 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