feat(router.py): support mock testing fallbacks flag on router

This commit is contained in:
Krrish Dholakia 2024-04-24 17:33:00 -07:00
parent f54510b6ee
commit 94cbe5516e
2 changed files with 97 additions and 1 deletions

View file

@ -1309,12 +1309,18 @@ class Router:
Try calling the function_with_retries
If it fails after num_retries, fall back to another model group
"""
mock_testing_fallbacks = kwargs.get("mock_testing_fallbacks", None)
model_group = kwargs.get("model")
fallbacks = kwargs.get("fallbacks", self.fallbacks)
context_window_fallbacks = kwargs.get(
"context_window_fallbacks", self.context_window_fallbacks
)
try:
if mock_testing_fallbacks is not None and mock_testing_fallbacks == True:
raise Exception(
f"This is a mock exception for model={model_group}, to trigger a fallback. Fallbacks={fallbacks}"
)
response = await self.async_function_with_retries(*args, **kwargs)
verbose_router_logger.debug(f"Async Response: {response}")
return response