phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat(router.py): support mock testing content policy + context window fallbacks

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-25 10:57:32 -07:00
parent 6f51da4e78
commit 0396d484fb
2 changed files with 76 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

26
litellm/router.py
View file

@ -2117,6 +2117,12 @@ class Router:
If it fails after num_retries, fall back to another model group
"""
mock_testing_fallbacks = kwargs.pop("mock_testing_fallbacks", None)
mock_testing_context_fallbacks = kwargs.pop(
"mock_testing_context_fallbacks", None
)
mock_testing_content_policy_fallbacks = kwargs.pop(
"mock_testing_content_policy_fallbacks", None
)
model_group = kwargs.get("model")
fallbacks = kwargs.get("fallbacks", self.fallbacks)
context_window_fallbacks = kwargs.get(
@ -2130,6 +2136,26 @@ class Router:
raise Exception(
f"This is a mock exception for model={model_group}, to trigger a fallback. Fallbacks={fallbacks}"
)
elif (
mock_testing_context_fallbacks is not None
and mock_testing_context_fallbacks is True
):
raise litellm.ContextWindowExceededError(
model=model_group,
llm_provider="",
message=f"This is a mock exception for model={model_group}, to trigger a fallback. \
Context_Window_Fallbacks={context_window_fallbacks}",
)
elif (
mock_testing_content_policy_fallbacks is not None
and mock_testing_content_policy_fallbacks is True
):
raise litellm.ContentPolicyViolationError(
model=model_group,
llm_provider="",
message=f"This is a mock exception for model={model_group}, to trigger a fallback. \
Context_Policy_Fallbacks={content_policy_fallbacks}",
)
response = await self.async_function_with_retries(*args, **kwargs)
verbose_router_logger.debug(f"Async Response: {response}")