diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index d1c2431fd..84a788bf2 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -25,7 +25,7 @@ async def router_cooldown_event_callback( """ Callback triggered when a deployment is put into cooldown by litellm - - Updates deploymen state on Prometheus + - Updates deployment state on Prometheus - Increments cooldown metric for deployment on Prometheus """ verbose_logger.debug("In router_cooldown_event_callback - updating prometheus") diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 6c4ec8e6b..814481f5e 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -32,10 +32,7 @@ class PatternMatchRouter: regex = self._pattern_to_regex(pattern) if regex not in self.patterns: self.patterns[regex] = [] - if isinstance(llm_deployment, list): - self.patterns[regex].extend(llm_deployment) - else: - self.patterns[regex].append(llm_deployment) + self.patterns[regex].append(llm_deployment) def _pattern_to_regex(self, pattern: str) -> str: """ diff --git a/tests/local_testing/test_router_pattern_matching.py b/tests/local_testing/test_router_pattern_matching.py new file mode 100644 index 000000000..d7e76b88b --- /dev/null +++ b/tests/local_testing/test_router_pattern_matching.py @@ -0,0 +1,157 @@ +""" +This tests the pattern matching router + +Pattern matching router is used to match patterns like openai/*, vertex_ai/*, anthropic/* etc. (wildcard matching) +""" + +import sys, os, time +import traceback, asyncio +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +from litellm import Router +from litellm.router import Deployment, LiteLLM_Params, ModelInfo +from concurrent.futures import ThreadPoolExecutor +from collections import defaultdict +from dotenv import load_dotenv + +load_dotenv() + +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + + +def test_pattern_match_router_initialization(): + router = PatternMatchRouter() + assert router.patterns == {} + + +def test_add_pattern(): + """ + Tests that openai/* is added to the patterns + + when we try to get the pattern, it should return the deployment + """ + router = PatternMatchRouter() + deployment = Deployment( + model_name="openai-1", + litellm_params=LiteLLM_Params(model="gpt-3.5-turbo"), + model_info=ModelInfo(), + ) + router.add_pattern("openai/*", deployment.to_json(exclude_none=True)) + assert len(router.patterns) == 1 + assert list(router.patterns.keys())[0] == "^openai/.*$" + + # try getting the pattern + assert router.route(request="openai/gpt-15") == [ + deployment.to_json(exclude_none=True) + ] + + +def test_add_pattern_vertex_ai(): + """ + Tests that vertex_ai/* is added to the patterns + + when we try to get the pattern, it should return the deployment + """ + router = PatternMatchRouter() + deployment = Deployment( + model_name="this-can-be-anything", + litellm_params=LiteLLM_Params(model="vertex_ai/gemini-1.5-flash-latest"), + model_info=ModelInfo(), + ) + router.add_pattern("vertex_ai/*", deployment.to_json(exclude_none=True)) + assert len(router.patterns) == 1 + assert list(router.patterns.keys())[0] == "^vertex_ai/.*$" + + # try getting the pattern + assert router.route(request="vertex_ai/gemini-1.5-flash-latest") == [ + deployment.to_json(exclude_none=True) + ] + + +def test_add_multiple_deployments(): + """ + Tests adding multiple deployments for the same pattern + + when we try to get the pattern, it should return the deployment + """ + router = PatternMatchRouter() + deployment1 = Deployment( + model_name="openai-1", + litellm_params=LiteLLM_Params(model="gpt-3.5-turbo"), + model_info=ModelInfo(), + ) + deployment2 = Deployment( + model_name="openai-2", + litellm_params=LiteLLM_Params(model="gpt-4"), + model_info=ModelInfo(), + ) + router.add_pattern("openai/*", deployment1.to_json(exclude_none=True)) + router.add_pattern("openai/*", deployment2.to_json(exclude_none=True)) + assert len(router.route("openai/gpt-4o")) == 2 + + +def test_pattern_to_regex(): + """ + Tests that the pattern is converted to a regex + """ + router = PatternMatchRouter() + assert router._pattern_to_regex("openai/*") == "^openai/.*$" + assert ( + router._pattern_to_regex("openai/fo::*::static::*") + == "^openai/fo::.*::static::.*$" + ) + + +def test_route_with_none(): + """ + Tests that the router returns None when the request is None + """ + router = PatternMatchRouter() + assert router.route(None) is None + + +def test_route_with_multiple_matching_patterns(): + """ + Tests that the router returns the first matching pattern when there are multiple matching patterns + """ + router = PatternMatchRouter() + deployment1 = Deployment( + model_name="openai-1", + litellm_params=LiteLLM_Params(model="gpt-3.5-turbo"), + model_info=ModelInfo(), + ) + deployment2 = Deployment( + model_name="openai-2", + litellm_params=LiteLLM_Params(model="gpt-4"), + model_info=ModelInfo(), + ) + router.add_pattern("openai/*", deployment1.to_json(exclude_none=True)) + router.add_pattern("openai/gpt-*", deployment2.to_json(exclude_none=True)) + assert router.route("openai/gpt-3.5-turbo") == [ + deployment1.to_json(exclude_none=True) + ] + + +# Add this test to check for exception handling +def test_route_with_exception(): + """ + Tests that the router returns None when there is an exception calling router.route() + """ + router = PatternMatchRouter() + deployment = Deployment( + model_name="openai-1", + litellm_params=LiteLLM_Params(model="gpt-3.5-turbo"), + model_info=ModelInfo(), + ) + router.add_pattern("openai/*", deployment.to_json(exclude_none=True)) + + router.patterns = ( + [] + ) # this will cause router.route to raise an exception, since router.patterns should be a dict + + result = router.route("openai/gpt-3.5-turbo") + assert result is None