(testing - litellm.Router ) add unit test coverage for pattern matching / wildcard routing (#6250)

* add testing coverage for pattern match router

* fix add_pattern

* fix typo on router_cooldown_event_callback

* add testing for pattern match router

* fix add explanation for pattern match router

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")

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].append(llm_deployment)
-            self.patterns[regex].extend(llm_deployment)
-        else:
-            self.patterns[regex].append(llm_deployment)
 
     def _pattern_to_regex(self, pattern: str) -> str:
         """

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