From 69a6eb6256460bae99fdcb35e2279ed391730a61 Mon Sep 17 00:00:00 2001
From: Praveen11558 <44603409+Praveen11558@users.noreply.github.com>
Date: Tue, 10 Mar 2026 17:13:46 +0530
Subject: [PATCH] (fix): token_count calculation during inline data sent during
 multimodal inputs

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 .../llm_cost_calc/test_llm_cost_calc_utils.py                   | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index 600cfb4ac08..c30aa89f113 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -44,10 +44,12 @@ from litellm.types.utils import CacheCreationTokenDetails, Usage
 @pytest.fixture(autouse=True)
 def setup_local_model_cost():
     """Use the local model cost map for all tests."""
+    original_model_cost = litellm.model_cost
     os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
     litellm.model_cost = litellm.get_model_cost_map(url="")
     yield
     os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None)
+    litellm.model_cost = original_model_cost
  
 def test_reasoning_tokens_no_price_set():
     model = "o1-mini"