pick_cheapest_chat_models_from_llm_provider

This commit is contained in:
Ishaan Jaff 2025-03-26 10:59:02 -07:00
parent 8a72b67b18
commit fee20250a6

View file

@ -59,7 +59,9 @@ from litellm.litellm_core_utils.health_check_utils import (
_filter_model_params,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_request_utils import LitellmCoreRequestUtils
from litellm.litellm_core_utils.llm_request_utils import (
pick_cheapest_chat_models_from_llm_provider,
)
from litellm.litellm_core_utils.mock_functions import (
mock_embedding,
mock_image_generation,
@ -2348,6 +2350,8 @@ def completion( # type: ignore # noqa: PLR0915
or litellm.api_key
)
api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE")
new_params = deepcopy(optional_params)
response = vertex_chat_completion.completion( # type: ignore
model=model,
@ -2390,6 +2394,8 @@ def completion( # type: ignore # noqa: PLR0915
or get_secret("VERTEXAI_CREDENTIALS")
)
api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE")
new_params = deepcopy(optional_params)
if (
model.startswith("meta/")
@ -3655,6 +3661,8 @@ def embedding( # noqa: PLR0915
api_key or get_secret_str("GEMINI_API_KEY") or litellm.api_key
)
api_base = api_base or litellm.api_base or get_secret_str("GEMINI_API_BASE")
response = google_batch_embeddings.batch_embeddings( # type: ignore
model=model,
input=input,
@ -3669,6 +3677,8 @@ def embedding( # noqa: PLR0915
print_verbose=print_verbose,
custom_llm_provider="gemini",
api_key=gemini_api_key,
api_base=api_base,
client=client,
)
elif custom_llm_provider == "vertex_ai":
@ -3693,6 +3703,13 @@ def embedding( # noqa: PLR0915
or get_secret_str("VERTEX_CREDENTIALS")
)
api_base = (
api_base
or litellm.api_base
or get_secret_str("VERTEXAI_API_BASE")
or get_secret_str("VERTEX_API_BASE")
)
if (
"image" in optional_params
or "video" in optional_params
@ -3714,6 +3731,7 @@ def embedding( # noqa: PLR0915
print_verbose=print_verbose,
custom_llm_provider="vertex_ai",
client=client,
api_base=api_base,
)
else:
response = vertex_embedding.embedding(
@ -3731,6 +3749,8 @@ def embedding( # noqa: PLR0915
aembedding=aembedding,
print_verbose=print_verbose,
api_key=api_key,
api_base=api_base,
client=client,
)
elif custom_llm_provider == "oobabooga":
response = oobabooga.embedding(
@ -4693,6 +4713,14 @@ def image_generation( # noqa: PLR0915
or optional_params.pop("vertex_ai_credentials", None)
or get_secret_str("VERTEXAI_CREDENTIALS")
)
api_base = (
api_base
or litellm.api_base
or get_secret_str("VERTEXAI_API_BASE")
or get_secret_str("VERTEX_API_BASE")
)
model_response = vertex_image_generation.image_generation(
model=model,
prompt=prompt,
@ -4704,6 +4732,7 @@ def image_generation( # noqa: PLR0915
vertex_location=vertex_ai_location,
vertex_credentials=vertex_credentials,
aimg_generation=aimg_generation,
client=client,
)
elif (
custom_llm_provider in litellm._custom_providers
@ -5421,10 +5450,8 @@ async def ahealth_check_wildcard_models(
) -> dict:
# this is a wildcard model, we need to pick a random model from the provider
cheapest_models = (
LitellmCoreRequestUtils.pick_cheapest_chat_models_from_llm_provider(
custom_llm_provider=custom_llm_provider, n=3
)
cheapest_models = pick_cheapest_chat_models_from_llm_provider(
custom_llm_provider=custom_llm_provider, n=3
)
if len(cheapest_models) == 0:
raise Exception(