(litellm sdk speedup) - use _model_contains_known_llm_provider in response_cost_calculator to check if the model contains a known litellm provider (#7721)

* define _cached_get_model_info_helper

* use _cached_get_model_info_helper

* speed up _select_model_name_for_cost_calc
This commit is contained in:
Ishaan Jaff 2025-01-12 15:40:05 -08:00 committed by GitHub
parent 6518bc70a0
commit 2c25ea5737
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,6 +1,7 @@
# What is this?
## File for 'response_cost' calculation in Logging
import time
from functools import lru_cache
from typing import Any, List, Literal, Optional, Tuple, Union
from pydantic import BaseModel
@ -51,7 +52,12 @@ from litellm.llms.vertex_ai.image_generation.cost_calculator import (
)
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.rerank import RerankResponse
from litellm.types.utils import CallTypesLiteral, PassthroughCallTypes, Usage
from litellm.types.utils import (
CallTypesLiteral,
LlmProvidersSet,
PassthroughCallTypes,
Usage,
)
from litellm.utils import (
CallTypes,
CostPerToken,
@ -392,23 +398,18 @@ def _select_model_name_for_cost_calc(
if base_model is not None:
return_model = base_model
completion_response_model: Optional[str] = None
if completion_response is not None and isinstance(completion_response, BaseModel):
completion_response_model = getattr(completion_response, "model", None)
hidden_params = getattr(completion_response, "_hidden_params", None)
if completion_response_model is None and hidden_params is not None:
if (
hidden_params.get("model", None) is not None
and len(hidden_params["model"]) > 0
):
return_model = hidden_params.get("model", model)
completion_response_model: Optional[str] = getattr(
completion_response, "model", None
)
hidden_params: Optional[dict] = getattr(completion_response, "_hidden_params", None)
if completion_response_model is None and hidden_params is not None:
if (
hidden_params is not None
and hidden_params.get("region_name", None) is not None
hidden_params.get("model", None) is not None
and len(hidden_params["model"]) > 0
):
region_name = hidden_params.get("region_name", None)
elif completion_response is not None and isinstance(completion_response, dict):
completion_response_model = completion_response.get("model", None)
return_model = hidden_params.get("model", model)
if hidden_params is not None and hidden_params.get("region_name", None) is not None:
region_name = hidden_params.get("region_name", None)
if return_model is None and completion_response_model is not None:
return_model = completion_response_model
@ -419,9 +420,7 @@ def _select_model_name_for_cost_calc(
if (
return_model is not None
and custom_llm_provider is not None
and not any(
return_model.startswith(provider) for provider in litellm.provider_list
)
and not _model_contains_known_llm_provider(return_model)
): # add provider prefix if not already present, to match model_cost
if region_name is not None:
return_model = f"{custom_llm_provider}/{region_name}/{return_model}"
@ -431,6 +430,15 @@ def _select_model_name_for_cost_calc(
return return_model
@lru_cache(maxsize=16)
def _model_contains_known_llm_provider(model: str) -> bool:
"""
Check if the model contains a known llm provider
"""
_provider_prefix = model.split("/")[0]
return _provider_prefix in LlmProvidersSet
def _get_usage_object(
completion_response: Any,
) -> Optional[Usage]: