[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.
This commit is contained in:
Alexsander Hamir 2025-10-06 18:30:01 -07:00 committed by GitHub
parent fbf7ca17ba
commit 9aa2dc2e8e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 180 additions and 150 deletions

View file

@ -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

View file

@ -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]: