search_context_cost_per_query

This commit is contained in:
Ishaan Jaff 2025-03-22 14:52:58 -07:00
parent 3a454d00df
commit cf22d31b2b
3 changed files with 35 additions and 10 deletions

View file

@ -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[

View file

@ -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),
)

View file

@ -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"]
)