From 9aa2dc2e8e6e03e058df94def80188ad9ac02c9a Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Mon, 6 Oct 2025 18:30:01 -0700 Subject: [PATCH] [Refactor] Utils: extract inner function from client (#15234) * fix: remove func definition from inside client It makes the function bigger and harder to understand, I left just the wrappers. * fix: test_arouter_test_team_model failure - Added fallback to the model_name to index functionality. --- litellm/router.py | 14 ++ litellm/utils.py | 316 ++++++++++++++++++++++++---------------------- 2 files changed, 180 insertions(+), 150 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 260e7c3f8c0..ef7bcba635f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6249,6 +6249,20 @@ class Router: returned_models.append(alias_model) else: returned_models.append(model) + elif team_id is not None: + # Fallback: if team_id is provided and model_name not in index, + # check if model_name matches any team_public_model_name + # O(n) scan but only when team_id lookup fails + for idx, model in enumerate(self.model_list): + if self.should_include_deployment( + model_name=model_name, model=model, team_id=team_id + ): + if model_alias is not None: + alias_model = copy.deepcopy(model) + alias_model["model_name"] = model_alias + returned_models.append(alias_model) + else: + returned_models.append(model) return returned_models diff --git a/litellm/utils.py b/litellm/utils.py index 3c6c3ac86e4..bb6638f781c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -910,159 +910,169 @@ def _get_wrapper_timeout( return timeout +def check_coroutine(value) -> bool: + return get_coroutine_checker().is_async_callable(value) + + +async def async_pre_call_deployment_hook(kwargs: Dict[str, Any], call_type: str): + """ + Allow modifying the request just before it's sent to the deployment. + + Use this instead of 'async_pre_call_hook' when you need to modify the request AFTER a deployment is selected, but BEFORE the request is sent. + """ + try: + typed_call_type = CallTypes(call_type) + except ValueError: + typed_call_type = None # unknown call type + + modified_kwargs = kwargs.copy() + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + result = await callback.async_pre_call_deployment_hook( + modified_kwargs, typed_call_type + ) + if result is not None: + modified_kwargs = result + + return modified_kwargs + + +async def async_post_call_success_deployment_hook( + request_data: dict, response: Any, call_type: Optional[CallTypes] +) -> Optional[Any]: + """ + Allow modifying / reviewing the response just after it's received from the deployment. + """ + try: + typed_call_type = CallTypes(call_type) + except ValueError: + typed_call_type = None # unknown call type + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + result = await callback.async_post_call_success_deployment_hook( + request_data, cast(LLMResponseTypes, response), typed_call_type + ) + if result is not None: + return result + + return response + + +def post_call_processing( + original_response, + model, + optional_params: Optional[dict], + original_function, + rules_obj, +): + try: + if original_response is None: + pass + else: + call_type = original_function.__name__ + if ( + call_type == CallTypes.completion.value + or call_type == CallTypes.acompletion.value + ): + is_coroutine = check_coroutine(original_response) + if is_coroutine is True: + pass + else: + if ( + isinstance(original_response, ModelResponse) + and len(original_response.choices) > 0 + ): + model_response: Optional[str] = original_response.choices[ + 0 + ].message.content # type: ignore + if model_response is not None: + ### POST-CALL RULES ### + rules_obj.post_call_rules( + input=model_response, model=model + ) + ### JSON SCHEMA VALIDATION ### + if litellm.enable_json_schema_validation is True: + try: + if ( + optional_params is not None + and "response_format" in optional_params + and optional_params["response_format"] + is not None + ): + json_response_format: Optional[dict] = None + if ( + isinstance( + optional_params["response_format"], + dict, + ) + and optional_params[ + "response_format" + ].get("json_schema") + is not None + ): + json_response_format = optional_params[ + "response_format" + ] + elif _parsing._completions.is_basemodel_type( + optional_params["response_format"] # type: ignore + ): + json_response_format = ( + type_to_response_format_param( + response_format=optional_params[ + "response_format" + ] + ) + ) + if json_response_format is not None: + litellm.litellm_core_utils.json_validation_rule.validate_schema( + schema=json_response_format[ + "json_schema" + ]["schema"], + response=model_response, + ) + except TypeError: + pass + if ( + optional_params is not None + and "response_format" in optional_params + and isinstance( + optional_params["response_format"], dict + ) + and "type" in optional_params["response_format"] + and optional_params["response_format"]["type"] + == "json_object" + and "response_schema" + in optional_params["response_format"] + and isinstance( + optional_params["response_format"][ + "response_schema" + ], + dict, + ) + and "enforce_validation" + in optional_params["response_format"] + and optional_params["response_format"][ + "enforce_validation" + ] + is True + ): + # schema given, json response expected, and validation enforced + litellm.litellm_core_utils.json_validation_rule.validate_schema( + schema=optional_params["response_format"][ + "response_schema" + ], + response=model_response, + ) + + except Exception as e: + raise e + def client(original_function): # noqa: PLR0915 rules_obj = Rules() - def check_coroutine(value) -> bool: - return get_coroutine_checker().is_async_callable(value) - - async def async_pre_call_deployment_hook(kwargs: Dict[str, Any], call_type: str): - """ - Allow modifying the request just before it's sent to the deployment. - - Use this instead of 'async_pre_call_hook' when you need to modify the request AFTER a deployment is selected, but BEFORE the request is sent. - """ - try: - typed_call_type = CallTypes(call_type) - except ValueError: - typed_call_type = None # unknown call type - - modified_kwargs = kwargs.copy() - - for callback in litellm.callbacks: - if isinstance(callback, CustomLogger): - result = await callback.async_pre_call_deployment_hook( - modified_kwargs, typed_call_type - ) - if result is not None: - modified_kwargs = result - - return modified_kwargs - - async def async_post_call_success_deployment_hook( - request_data: dict, response: Any, call_type: Optional[CallTypes] - ) -> Optional[Any]: - """ - Allow modifying / reviewing the response just after it's received from the deployment. - """ - try: - typed_call_type = CallTypes(call_type) - except ValueError: - typed_call_type = None # unknown call type - - for callback in litellm.callbacks: - if isinstance(callback, CustomLogger): - result = await callback.async_post_call_success_deployment_hook( - request_data, cast(LLMResponseTypes, response), typed_call_type - ) - if result is not None: - return result - - return response - - def post_call_processing(original_response, model, optional_params: Optional[dict]): - try: - if original_response is None: - pass - else: - call_type = original_function.__name__ - if ( - call_type == CallTypes.completion.value - or call_type == CallTypes.acompletion.value - ): - is_coroutine = check_coroutine(original_response) - if is_coroutine is True: - pass - else: - if ( - isinstance(original_response, ModelResponse) - and len(original_response.choices) > 0 - ): - model_response: Optional[str] = original_response.choices[ - 0 - ].message.content # type: ignore - if model_response is not None: - ### POST-CALL RULES ### - rules_obj.post_call_rules( - input=model_response, model=model - ) - ### JSON SCHEMA VALIDATION ### - if litellm.enable_json_schema_validation is True: - try: - if ( - optional_params is not None - and "response_format" in optional_params - and optional_params["response_format"] - is not None - ): - json_response_format: Optional[dict] = None - if ( - isinstance( - optional_params["response_format"], - dict, - ) - and optional_params[ - "response_format" - ].get("json_schema") - is not None - ): - json_response_format = optional_params[ - "response_format" - ] - elif _parsing._completions.is_basemodel_type( - optional_params["response_format"] # type: ignore - ): - json_response_format = ( - type_to_response_format_param( - response_format=optional_params[ - "response_format" - ] - ) - ) - if json_response_format is not None: - litellm.litellm_core_utils.json_validation_rule.validate_schema( - schema=json_response_format[ - "json_schema" - ]["schema"], - response=model_response, - ) - except TypeError: - pass - if ( - optional_params is not None - and "response_format" in optional_params - and isinstance( - optional_params["response_format"], dict - ) - and "type" in optional_params["response_format"] - and optional_params["response_format"]["type"] - == "json_object" - and "response_schema" - in optional_params["response_format"] - and isinstance( - optional_params["response_format"][ - "response_schema" - ], - dict, - ) - and "enforce_validation" - in optional_params["response_format"] - and optional_params["response_format"][ - "enforce_validation" - ] - is True - ): - # schema given, json response expected, and validation enforced - litellm.litellm_core_utils.json_validation_rule.validate_schema( - schema=optional_params["response_format"][ - "response_schema" - ], - response=model_response, - ) - - except Exception as e: - raise e @wraps(original_function) def wrapper(*args, **kwargs): # noqa: PLR0915 @@ -1273,6 +1283,8 @@ def client(original_function): # noqa: PLR0915 original_response=result, model=model or None, optional_params=kwargs, + original_function=original_function, + rules_obj=rules_obj, ) # [OPTIONAL] ADD TO CACHE @@ -1489,7 +1501,11 @@ def client(original_function): # noqa: PLR0915 return result ### POST-CALL RULES ### post_call_processing( - original_response=result, model=model, optional_params=kwargs + original_response=result, + model=model, + optional_params=kwargs, + original_function=original_function, + rules_obj=rules_obj, ) # Only run if call_type is a valid value in CallTypes if call_type in [ct.value for ct in CallTypes]: