From cf22d31b2b01f220756d909e38dc6556c8496c2b Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 22 Mar 2025 14:52:58 -0700
Subject: [PATCH] search_context_cost_per_query

---
 litellm/types/utils.py                        |  9 +++++
 litellm/utils.py                              |  3 ++
 .../test_token_counting.py                    | 33 +++++++++++++------
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 2af0c791d88..a7894e424c8 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -100,6 +100,12 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
     supports_web_search: Optional[bool]
 
 
+class SearchContextCostPerQuery(TypedDict, total=False):
+    search_context_size_low: float
+    search_context_size_medium: float
+    search_context_size_high: float
+
+
 class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     key: Required[str]  # the key in litellm.model_cost which is returned
 
@@ -136,6 +142,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     output_cost_per_video_per_second: Optional[float]  # only for vertex ai models
     output_cost_per_audio_per_second: Optional[float]  # only for vertex ai models
     output_cost_per_second: Optional[float]  # for OpenAI Speech models
+    search_context_cost_per_query: Optional[
+        SearchContextCostPerQuery
+    ]  # Cost for using web search tool
 
     litellm_provider: Required[str]
     mode: Required[
diff --git a/litellm/utils.py b/litellm/utils.py
index 317599f864b..dc97c4d898f 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -4545,6 +4545,9 @@ def _get_model_info_helper(  # noqa: PLR0915
                     "supports_native_streaming", None
                 ),
                 supports_web_search=_model_info.get("supports_web_search", False),
+                search_context_cost_per_query=_model_info.get(
+                    "search_context_cost_per_query", None
+                ),
                 tpm=_model_info.get("tpm", None),
                 rpm=_model_info.get("rpm", None),
             )
diff --git a/tests/logging_callback_tests/test_token_counting.py b/tests/logging_callback_tests/test_token_counting.py
index d4da465041a..6821c7d345c 100644
--- a/tests/logging_callback_tests/test_token_counting.py
+++ b/tests/logging_callback_tests/test_token_counting.py
@@ -21,16 +21,18 @@ sys.path.insert(
 import litellm
 import asyncio
 from typing import Optional
-from litellm.types.utils import StandardLoggingPayload, Usage
+from litellm.types.utils import StandardLoggingPayload, Usage, ModelInfoBase
 from litellm.integrations.custom_logger import CustomLogger
 
 
 class TestCustomLogger(CustomLogger):
     def __init__(self):
         self.recorded_usage: Optional[Usage] = None
+        self.standard_logging_payload: Optional[StandardLoggingPayload] = None
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         standard_logging_payload = kwargs.get("standard_logging_object")
+        self.standard_logging_payload = standard_logging_payload
         print(
             "standard_logging_payload",
             json.dumps(standard_logging_payload, indent=4, default=str),
@@ -246,15 +248,6 @@ async def test_stream_token_counting_anthropic_with_include_usage():
     )
 
 
-class TestCustomLogger(CustomLogger):
-    def __init__(self):
-        self.standard_logging_payload: Optional[StandardLoggingPayload] = None
-
-    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
-        print("kwargs: ", kwargs)
-        self.standard_logging_payload = kwargs.get("standard_logging_object", None)
-
-
 @pytest.mark.asyncio
 async def test_openai_web_search_logging_cost_tracking():
     """Makes a simple web search request and validates the response contains web search annotations and all expected fields are present"""
@@ -278,3 +271,23 @@ async def test_openai_web_search_logging_cost_tracking():
         "logged standard logging payload: ",
         json.dumps(test_custom_logger.standard_logging_payload, indent=4),
     )
+    standard_logging_payload = test_custom_logger.standard_logging_payload
+    response_cost = standard_logging_payload.get("response_cost")
+    assert response_cost is not None
+
+    # Assert the cost = Token Usage + Web Search Cost
+    model_map_information = standard_logging_payload["model_map_information"]
+    model_map_value: ModelInfoBase = model_map_information["model_map_value"]
+    total_token_cost = (
+        standard_logging_payload["prompt_tokens"]
+        * model_map_value["input_cost_per_token"]
+    ) + (
+        standard_logging_payload["completion_tokens"]
+        * model_map_value["output_cost_per_token"]
+    )
+    print("total token cost:", total_token_cost)
+    assert (
+        response_cost
+        == total_token_cost
+        + model_map_value["search_context_cost_per_query"]["search_context_size_low"]
+    )