mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
search_context_cost_per_query
This commit is contained in:
parent
3a454d00df
commit
cf22d31b2b
3 changed files with 35 additions and 10 deletions
|
|
@ -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[
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue