ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix custom pricing - separate provider info from model info (#7990)

Browse Source 
* fix(utils.py): initial commit fixing custom cost tracking

refactors out provider specific model info from `get_model_info` - this was causing custom costs to be registered incorrectly

* fix(utils.py): cleanup `_supports_factory` to check provider info, if model info is None

some providers support features like vision across all models

* fix(utils.py): refactor to use _supports_factory

* test: update testing

* fix: fix linting errors

* test: fix testing
This commit is contained in:
Krish Dholakia 2025-01-25 21:49:28 -08:00 committed by GitHub
parent d9b8100cca
commit 03eef5a2a0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 118 additions and 91 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
litellm/cost_calculator.py
View File

@ -560,6 +560,10 @@ def completion_cost( # noqa: PLR0915
base_model=base_model,
)
verbose_logger.debug(
f"completion_response _select_model_name_for_cost_calc: {model}"
)
if completion_response is not None and (
isinstance(completion_response, BaseModel)
or isinstance(completion_response, dict)
@ -598,9 +602,6 @@ def completion_cost( # noqa: PLR0915
cache_read_input_tokens = prompt_tokens_details.get("cached_tokens", 0)
total_time = getattr(completion_response, "_response_ms", 0)
verbose_logger.debug(
f"completion_response response ms: {getattr(completion_response, '_response_ms', None)} "
)
hidden_params = getattr(completion_response, "_hidden_params", None)
if hidden_params is not None:

1
litellm/litellm_core_utils/litellm_logging.py
View File

@ -847,6 +847,7 @@ class Logging(LiteLLMLoggingBaseClass):
response_cost = litellm.response_cost_calculator(
**response_cost_calculator_kwargs
)
verbose_logger.debug(f"response_cost: {response_cost}")
return response_cost
except Exception as e: # error calculating cost
debug_info = StandardLoggingModelCostFailureDebugInformation(

10
litellm/llms/base_llm/base_utils.py
View File

@ -5,17 +5,15 @@ from typing import List, Optional, Type, Union
from openai.lib import _parsing, _pydantic
from pydantic import BaseModel
from litellm.types.utils import ModelInfoBase
from litellm.types.utils import ProviderSpecificModelInfo
class BaseLLMModelInfo(ABC):
@abstractmethod
def get_model_info(
def get_provider_info(
self,
model: str,
existing_model_info: Optional[ModelInfoBase] = None,
) -> Optional[ModelInfoBase]:
pass
) -> Optional[ProviderSpecificModelInfo]:
return None
@abstractmethod
def get_models(self) -> List[str]:

22
litellm/llms/fireworks_ai/chat/transformation.py
View File

@ -3,7 +3,7 @@ from typing import List, Literal, Optional, Tuple, Union, cast
import litellm
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject
from litellm.types.utils import ModelInfoBase, ProviderSpecificModelInfo
from litellm.types.utils import ProviderSpecificModelInfo
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
@ -159,30 +159,14 @@ class FireworksAIConfig(OpenAIGPTConfig):
)
return messages
def get_model_info(
self, model: str, existing_model_info: Optional[ModelInfoBase] = None
) -> ModelInfoBase:
def get_provider_info(self, model: str) -> ProviderSpecificModelInfo:
provider_specific_model_info = ProviderSpecificModelInfo(
supports_function_calling=True,
supports_prompt_caching=True, # https://docs.fireworks.ai/guides/prompt-caching
supports_pdf_input=True, # via document inlining
supports_vision=True, # via document inlining
)
if existing_model_info is not None:
return ModelInfoBase(
**{**existing_model_info, **provider_specific_model_info}
)
return ModelInfoBase(
key=model,
litellm_provider="fireworks_ai",
mode="chat",
input_cost_per_token=0.0,
output_cost_per_token=0.0,
max_tokens=None,
max_input_tokens=None,
max_output_tokens=None,
**provider_specific_model_info,
)
return provider_specific_model_info
def transform_request(
self,

19
litellm/llms/openai/chat/gpt_transformation.py
View File

@ -11,7 +11,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelInfoBase, ModelResponse
from litellm.types.utils import ModelResponse
from ..common_utils import OpenAIError
@ -255,23 +255,6 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
models = response.json()["data"]
return [model["id"] for model in models]
def get_model_info(
self, model: str, existing_model_info: Optional[ModelInfoBase] = None
) -> ModelInfoBase:
if existing_model_info is not None:
return existing_model_info
return ModelInfoBase(
key=model,
litellm_provider="openai",
mode="chat",
input_cost_per_token=0.0,
output_cost_per_token=0.0,
max_tokens=None,
max_input_tokens=None,
max_output_tokens=None,
)
@staticmethod
def get_api_key(api_key: Optional[str] = None) -> Optional[str]:
return (

6
litellm/llms/topaz/common_utils.py
View File

@ -1,7 +1,6 @@
from typing import List, Optional
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import ModelInfoBase
from ..base_llm.base_utils import BaseLLMModelInfo
from ..base_llm.chat.transformation import BaseLLMException
@ -12,11 +11,6 @@ class TopazException(BaseLLMException):
class TopazModelInfo(BaseLLMModelInfo):
def get_model_info(
self, model: str, existing_model_info: Optional[ModelInfoBase] = None
) -> Optional[ModelInfoBase]:
return existing_model_info
def get_models(self) -> List[str]:
return [
"topaz/Standard V2",

2
litellm/proxy/_new_secret_config.yaml
View File

@ -16,5 +16,3 @@ model_list:
litellm_settings:
callbacks: ["langsmith"]
default_internal_user_params:
available_teams: ["litellm_dashboard_54a81fa9-9c69-45e8-b256-0c36bf104e5f", "a29a2dc6-1347-4ebc-a428-e6b56bbba611", "test-group-12"]

84
litellm/utils.py
View File

@ -149,6 +149,7 @@ from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
ProviderField,
ProviderSpecificModelInfo,
StreamingChoices,
TextChoices,
TextCompletionResponse,
@ -1898,6 +1899,13 @@ def _supports_factory(model: str, custom_llm_provider: Optional[str], key: str)
verbose_logger.debug(
f"Model not found or error in checking {key} support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {str(e)}"
)
provider_info = get_provider_info(
model=model, custom_llm_provider=custom_llm_provider
)
if provider_info is not None and provider_info.get(key, False) is True:
return True
return False
@ -1958,23 +1966,11 @@ def supports_vision(model: str, custom_llm_provider: Optional[str] = None) -> bo
Returns:
bool: True if the model supports vision, False otherwise.
"""
try:
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=model, custom_llm_provider=custom_llm_provider
)
model_info = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
if model_info.get("supports_vision", False) is True:
return True
return False
except Exception as e:
verbose_logger.error(
f"Model not found or error in checking vision support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {str(e)}"
)
return False
return _supports_factory(
model=model,
custom_llm_provider=custom_llm_provider,
key="supports_vision",
)
def supports_embedding_image_input(
@ -2037,6 +2033,7 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915
},
}
"""
loaded_model_cost = {}
if isinstance(model_cost, dict):
loaded_model_cost = model_cost
@ -2054,6 +2051,9 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915
## override / add new keys to the existing model cost dictionary
updated_dictionary = _update_dictionary(existing_model, value)
litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary)
verbose_logger.debug(
f"added/updated model={model_cost_key} in litellm.model_cost: {model_cost_key}"
)
# add new model names to provider lists
if value.get("litellm_provider") == "openai":
if key not in litellm.open_ai_chat_completion_models:
@ -4048,6 +4048,26 @@ def _cached_get_model_info_helper(
return _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
def get_provider_info(
model: str, custom_llm_provider: Optional[str]
) -> Optional[ProviderSpecificModelInfo]:
## PROVIDER-SPECIFIC INFORMATION
# if custom_llm_provider == "predibase":
# _model_info["supports_response_schema"] = True
provider_config: Optional[BaseLLMModelInfo] = None
if custom_llm_provider and custom_llm_provider in LlmProvidersSet:
# Check if the provider string exists in LlmProviders enum
provider_config = ProviderConfigManager.get_provider_model_info(
model=model, provider=LlmProviders(custom_llm_provider)
)
model_info: Optional[ProviderSpecificModelInfo] = None
if provider_config:
model_info = provider_config.get_provider_info(model=model)
return model_info
def _get_model_info_helper( # noqa: PLR0915
model: str, custom_llm_provider: Optional[str] = None
) -> ModelInfoBase:
@ -4071,6 +4091,11 @@ def _get_model_info_helper( # noqa: PLR0915
potential_model_names = _get_potential_model_names(
model=model, custom_llm_provider=custom_llm_provider
)
verbose_logger.debug(
f"checking potential_model_names in litellm.model_cost: {potential_model_names}"
)
combined_model_name = potential_model_names["combined_model_name"]
stripped_model_name = potential_model_names["stripped_model_name"]
combined_stripped_model_name = potential_model_names[
@ -4111,7 +4136,6 @@ def _get_model_info_helper( # noqa: PLR0915
_model_info: Optional[Dict[str, Any]] = None
key: Optional[str] = None
provider_config: Optional[BaseLLMModelInfo] = None
if combined_model_name in litellm.model_cost:
key = combined_model_name
@ -4121,6 +4145,7 @@ def _get_model_info_helper( # noqa: PLR0915
):
_model_info = None
if _model_info is None and model in litellm.model_cost:
key = model
_model_info = _get_model_info_from_model_cost(key=key)
if not _check_provider_match(
@ -4131,6 +4156,7 @@ def _get_model_info_helper( # noqa: PLR0915
_model_info is None
and combined_stripped_model_name in litellm.model_cost
):
key = combined_stripped_model_name
_model_info = _get_model_info_from_model_cost(key=key)
if not _check_provider_match(
@ -4138,6 +4164,7 @@ def _get_model_info_helper( # noqa: PLR0915
):
_model_info = None
if _model_info is None and stripped_model_name in litellm.model_cost:
key = stripped_model_name
_model_info = _get_model_info_from_model_cost(key=key)
if not _check_provider_match(
@ -4145,6 +4172,7 @@ def _get_model_info_helper( # noqa: PLR0915
):
_model_info = None
if _model_info is None and split_model in litellm.model_cost:
key = split_model
_model_info = _get_model_info_from_model_cost(key=key)
if not _check_provider_match(
@ -4152,29 +4180,11 @@ def _get_model_info_helper( # noqa: PLR0915
):
_model_info = None
if custom_llm_provider and custom_llm_provider in LlmProvidersSet:
# Check if the provider string exists in LlmProviders enum
provider_config = ProviderConfigManager.get_provider_model_info(
model=model, provider=LlmProviders(custom_llm_provider)
)
if _model_info is None and provider_config is not None:
_model_info = cast(
Optional[Dict],
provider_config.get_model_info(
model=model, existing_model_info=_model_info
),
)
key = "provider_specific_model_info"
if _model_info is None or key is None:
raise ValueError(
"This model isn't mapped yet. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json"
)
## PROVIDER-SPECIFIC INFORMATION
if custom_llm_provider == "predibase":
_model_info["supports_response_schema"] = True
_input_cost_per_token: Optional[float] = _model_info.get(
"input_cost_per_token"
)
@ -4357,6 +4367,8 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod
custom_llm_provider=custom_llm_provider,
)
verbose_logger.debug(f"model_info: {_model_info}")
returned_model_info = ModelInfo(
**_model_info, supported_openai_params=supported_openai_params
)

48
tests/local_testing/test_completion_cost.py
View File

@ -2791,6 +2791,24 @@ def test_cost_calculator_with_base_model():
assert resp._hidden_params["response_cost"] > 0
@pytest.fixture
def model_item():
return {
"model_name": "random-model",
"litellm_params": {
"model": "openai/my-fake-model",
"api_key": "my-fake-key",
"api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
},
"model_info": {},
}
@pytest.mark.parametrize("base_model_arg", ["litellm_param", "model_info"])
def test_cost_calculator_with_base_model_with_router(base_model_arg, model_item):
from litellm import Router
@pytest.mark.parametrize("base_model_arg", ["litellm_param", "model_info"])
def test_cost_calculator_with_base_model_with_router(base_model_arg):
from litellm import Router
@ -2861,3 +2879,33 @@ def test_cost_calculator_with_custom_pricing():
)
assert resp.model == "random-model"
assert resp._hidden_params["response_cost"] > 0
@pytest.mark.parametrize(
"custom_pricing",
[
"litellm_params",
"model_info",
],
)
@pytest.mark.asyncio
async def test_cost_calculator_with_custom_pricing_router(model_item, custom_pricing):
from litellm import Router
litellm._turn_on_debug()
if custom_pricing == "litellm_params":
model_item["litellm_params"]["input_cost_per_token"] = 0.0000008
model_item["litellm_params"]["output_cost_per_token"] = 0.0000032
elif custom_pricing == "model_info":
model_item["model_info"]["input_cost_per_token"] = 0.0000008
model_item["model_info"]["output_cost_per_token"] = 0.0000032
router = Router(model_list=[model_item])
resp = await router.acompletion(
model="random-model",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="Hello, how are you?",
)
# assert resp.model == "random-model"
assert resp._hidden_params["response_cost"] > 0

10
tests/local_testing/test_get_model_info.py
View File

@ -32,12 +32,20 @@ def test_get_model_info_custom_llm_with_model_name():
litellm.get_model_info(model)
def test_get_model_info_custom_llm_with_same_name_vllm():
def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch):
"""
Tests if {custom_llm_provider}/{model_name} name given, and model exists in model info, the object is returned
"""
model = "command-r-plus"
provider = "openai" # vllm is openai-compatible
litellm.register_model(
{
"openai/command-r-plus": {
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
},
}
)
model_info = litellm.get_model_info(model, custom_llm_provider=provider)
print("model_info", model_info)
assert model_info["input_cost_per_token"] == 0.0