add util to pick_cheapest_model_from_llm_provider

This commit is contained in:
Ishaan Jaff 2024-08-29 09:27:20 -07:00
parent 4b6a2fa4f3
commit 45774010f7

View file

@ -1,5 +1,7 @@
from typing import Dict, Optional
import litellm
def _ensure_extra_body_is_safe(extra_body: Optional[Dict]) -> Optional[Dict]:
"""
@ -26,3 +28,26 @@ def _ensure_extra_body_is_safe(extra_body: Optional[Dict]) -> Optional[Dict]:
extra_body["metadata"]["prompt"] = _prompt.__dict__
return extra_body
def pick_cheapest_model_from_llm_provider(custom_llm_provider: str):
"""
Pick a random model from the LLM provider.
"""
if custom_llm_provider not in litellm.models_by_provider:
raise ValueError(f"Unknown LLM provider: {custom_llm_provider}")
known_models = litellm.models_by_provider.get(custom_llm_provider, [])
min_cost = float("inf")
cheapest_model = None
for model in known_models:
model_info = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
_cost = model_info.get("input_cost_per_token", 0) + model_info.get(
"output_cost_per_token", 0
)
if _cost < min_cost:
min_cost = _cost
cheapest_model = model
return cheapest_model