feat router allow * models

This commit is contained in:
Ishaan Jaff 2024-04-01 19:00:24 -07:00
parent d5d800e141
commit aabd7eff1f

View file

@ -191,6 +191,8 @@ class Router:
redis_cache=redis_cache, in_memory_cache=InMemoryCache()
) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc.
self.default_deployment = None # use this to track the users default deployment, when they want to use model = *
if model_list:
model_list = copy.deepcopy(model_list)
self.set_model_list(model_list)
@ -252,7 +254,6 @@ class Router:
}
}
"""
### ROUTING SETUP ###
if routing_strategy == "least-busy":
self.leastbusy_logger = LeastBusyLoggingHandler(
@ -2078,6 +2079,11 @@ class Router:
),
)
# Check if user is trying to use model_name == "*"
# this is a catch all model for their specific api key
if model["model_name"] == "*":
self.default_deployment = model
# Azure GPT-Vision Enhancements, users can pass os.environ/
data_sources = model.get("litellm_params", {}).get("dataSources", [])
@ -2248,6 +2254,13 @@ class Router:
)
model = self.model_group_alias[model]
if model not in self.model_names and self.default_deployment is not None:
updated_deployment = copy.deepcopy(
self.default_deployment
) # self.default_deployment
updated_deployment["litellm_params"]["model"] = model
return updated_deployment
## get healthy deployments
### get all deployments
healthy_deployments = [m for m in self.model_list if m["model_name"] == model]