fix - router add model logic

This commit is contained in:
Ishaan Jaff 2024-05-10 12:32:16 -07:00
parent 8a35354dd6
commit 9d3f01c6ae

View file

@ -2558,20 +2558,20 @@ class Router:
self.set_client(model=deployment.to_json(exclude_none=True)) self.set_client(model=deployment.to_json(exclude_none=True))
# set region (if azure model) # set region (if azure model)
try: # try:
if "azure" in deployment.litellm_params.model: # if "azure" in deployment.litellm_params.model:
region = litellm.utils.get_model_region( # region = litellm.utils.get_model_region(
litellm_params=deployment.litellm_params, mode=None # litellm_params=deployment.litellm_params, mode=None
) # )
deployment.litellm_params.region_name = region # deployment.litellm_params.region_name = region
except Exception as e: # except Exception as e:
verbose_router_logger.error( # verbose_router_logger.error(
"Unable to get the region for azure model - {}, {}".format( # "Unable to get the region for azure model - {}, {}".format(
deployment.litellm_params.model, str(e) # deployment.litellm_params.model, str(e)
) # )
) # )
pass # [NON-BLOCKING] # pass # [NON-BLOCKING]
return deployment return deployment
@ -2610,7 +2610,6 @@ class Router:
- The added/updated deployment - The added/updated deployment
""" """
# check if deployment already exists # check if deployment already exists
if deployment.model_info.id in self.get_model_ids(): if deployment.model_info.id in self.get_model_ids():
# remove the previous deployment # remove the previous deployment
removal_idx: Optional[int] = None removal_idx: Optional[int] = None
@ -2620,16 +2619,9 @@ class Router:
if removal_idx is not None: if removal_idx is not None:
self.model_list.pop(removal_idx) self.model_list.pop(removal_idx)
else:
# add to model list # if the model_id is not in router
_deployment = deployment.to_json(exclude_none=True) self.add_deployment(deployment=deployment)
self.model_list.append(_deployment)
# initialize client
self._add_deployment(deployment=deployment)
# add to model names
self.model_names.append(deployment.model_name)
return deployment return deployment
def delete_deployment(self, id: str) -> Optional[Deployment]: def delete_deployment(self, id: str) -> Optional[Deployment]: