mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
9b7ebb6a7d
* build(pyproject.toml): add new dev dependencies - for type checking * build: reformat files to fit black * ci: reformat to fit black * ci(test-litellm.yml): make tests run clear * build(pyproject.toml): add ruff * fix: fix ruff checks * build(mypy/): fix mypy linting errors * fix(hashicorp_secret_manager.py): fix passing cert for tls auth * build(mypy/): resolve all mypy errors * test: update test * fix: fix black formatting * build(pre-commit-config.yaml): use poetry run black * fix(proxy_server.py): fix linting error * fix: fix ruff safe representation error
264 lines
8.6 KiB
Python
264 lines
8.6 KiB
Python
"""
|
|
Class to handle llm wildcard routing and regex pattern matching
|
|
"""
|
|
|
|
import copy
|
|
import re
|
|
from re import Match
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
from litellm import get_llm_provider
|
|
from litellm._logging import verbose_router_logger
|
|
|
|
|
|
class PatternUtils:
|
|
@staticmethod
|
|
def calculate_pattern_specificity(pattern: str) -> Tuple[int, int]:
|
|
"""
|
|
Calculate pattern specificity based on length and complexity.
|
|
|
|
Args:
|
|
pattern: Regex pattern to analyze
|
|
|
|
Returns:
|
|
Tuple of (length, complexity) for sorting
|
|
"""
|
|
complexity_chars = ["*", "+", "?", "\\", "^", "$", "|", "(", ")"]
|
|
ret_val = (
|
|
len(pattern), # Longer patterns more specific
|
|
sum(
|
|
pattern.count(char) for char in complexity_chars
|
|
), # More regex complexity
|
|
)
|
|
return ret_val
|
|
|
|
@staticmethod
|
|
def sorted_patterns(
|
|
patterns: Dict[str, List[Dict]]
|
|
) -> List[Tuple[str, List[Dict]]]:
|
|
"""
|
|
Cached property for patterns sorted by specificity.
|
|
|
|
Returns:
|
|
Sorted list of pattern-deployment tuples
|
|
"""
|
|
return sorted(
|
|
patterns.items(),
|
|
key=lambda x: PatternUtils.calculate_pattern_specificity(x[0]),
|
|
reverse=True,
|
|
)
|
|
|
|
|
|
class PatternMatchRouter:
|
|
"""
|
|
Class to handle llm wildcard routing and regex pattern matching
|
|
|
|
doc: https://docs.litellm.ai/docs/proxy/configs#provider-specific-wildcard-routing
|
|
|
|
This class will store a mapping for regex pattern: List[Deployments]
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.patterns: Dict[str, List] = {}
|
|
|
|
def add_pattern(self, pattern: str, llm_deployment: Dict):
|
|
"""
|
|
Add a regex pattern and the corresponding llm deployments to the patterns
|
|
|
|
Args:
|
|
pattern: str
|
|
llm_deployment: str or List[str]
|
|
"""
|
|
# Convert the pattern to a regex
|
|
regex = self._pattern_to_regex(pattern)
|
|
if regex not in self.patterns:
|
|
self.patterns[regex] = []
|
|
self.patterns[regex].append(llm_deployment)
|
|
|
|
def _pattern_to_regex(self, pattern: str) -> str:
|
|
"""
|
|
Convert a wildcard pattern to a regex pattern
|
|
|
|
example:
|
|
pattern: openai/*
|
|
regex: openai/.*
|
|
|
|
pattern: openai/fo::*::static::*
|
|
regex: openai/fo::.*::static::.*
|
|
|
|
Args:
|
|
pattern: str
|
|
|
|
Returns:
|
|
str: regex pattern
|
|
"""
|
|
# # Replace '*' with '.*' for regex matching
|
|
# regex = pattern.replace("*", ".*")
|
|
# # Escape other special characters
|
|
# regex = re.escape(regex).replace(r"\.\*", ".*")
|
|
# return f"^{regex}$"
|
|
return re.escape(pattern).replace(r"\*", "(.*)")
|
|
|
|
def _return_pattern_matched_deployments(
|
|
self, matched_pattern: Match, deployments: List[Dict]
|
|
) -> List[Dict]:
|
|
new_deployments = []
|
|
for deployment in deployments:
|
|
new_deployment = copy.deepcopy(deployment)
|
|
new_deployment["litellm_params"][
|
|
"model"
|
|
] = PatternMatchRouter.set_deployment_model_name(
|
|
matched_pattern=matched_pattern,
|
|
litellm_deployment_litellm_model=deployment["litellm_params"]["model"],
|
|
)
|
|
new_deployments.append(new_deployment)
|
|
|
|
return new_deployments
|
|
|
|
def route(
|
|
self, request: Optional[str], filtered_model_names: Optional[List[str]] = None
|
|
) -> Optional[List[Dict]]:
|
|
"""
|
|
Route a requested model to the corresponding llm deployments based on the regex pattern
|
|
|
|
loop through all the patterns and find the matching pattern
|
|
if a pattern is found, return the corresponding llm deployments
|
|
if no pattern is found, return None
|
|
|
|
Args:
|
|
request: str - the received model name from the user (can be a wildcard route). If none, No deployments will be returned.
|
|
filtered_model_names: Optional[List[str]] - if provided, only return deployments that match the filtered_model_names
|
|
Returns:
|
|
Optional[List[Deployment]]: llm deployments
|
|
"""
|
|
try:
|
|
if request is None:
|
|
return None
|
|
|
|
sorted_patterns = PatternUtils.sorted_patterns(self.patterns)
|
|
regex_filtered_model_names = (
|
|
[self._pattern_to_regex(m) for m in filtered_model_names]
|
|
if filtered_model_names is not None
|
|
else []
|
|
)
|
|
for pattern, llm_deployments in sorted_patterns:
|
|
if (
|
|
filtered_model_names is not None
|
|
and pattern not in regex_filtered_model_names
|
|
):
|
|
continue
|
|
pattern_match = re.match(pattern, request)
|
|
if pattern_match:
|
|
return self._return_pattern_matched_deployments(
|
|
matched_pattern=pattern_match, deployments=llm_deployments
|
|
)
|
|
except Exception as e:
|
|
verbose_router_logger.debug(f"Error in PatternMatchRouter.route: {str(e)}")
|
|
|
|
return None # No matching pattern found
|
|
|
|
@staticmethod
|
|
def set_deployment_model_name(
|
|
matched_pattern: Match,
|
|
litellm_deployment_litellm_model: str,
|
|
) -> str:
|
|
"""
|
|
Set the model name for the matched pattern llm deployment
|
|
|
|
E.g.:
|
|
|
|
Case 1:
|
|
model_name: llmengine/* (can be any regex pattern or wildcard pattern)
|
|
litellm_params:
|
|
model: openai/*
|
|
|
|
if model_name = "llmengine/foo" -> model = "openai/foo"
|
|
|
|
Case 2:
|
|
model_name: llmengine/fo::*::static::*
|
|
litellm_params:
|
|
model: openai/fo::*::static::*
|
|
|
|
if model_name = "llmengine/foo::bar::static::baz" -> model = "openai/foo::bar::static::baz"
|
|
|
|
Case 3:
|
|
model_name: *meta.llama3*
|
|
litellm_params:
|
|
model: bedrock/meta.llama3*
|
|
|
|
if model_name = "hello-world-meta.llama3-70b" -> model = "bedrock/meta.llama3-70b"
|
|
"""
|
|
|
|
## BASE CASE: if the deployment model name does not contain a wildcard, return the deployment model name
|
|
if "*" not in litellm_deployment_litellm_model:
|
|
return litellm_deployment_litellm_model
|
|
|
|
wildcard_count = litellm_deployment_litellm_model.count("*")
|
|
|
|
# Extract all dynamic segments from the request
|
|
dynamic_segments = matched_pattern.groups()
|
|
|
|
if len(dynamic_segments) > wildcard_count:
|
|
return (
|
|
matched_pattern.string
|
|
) # default to the user input, if unable to map based on wildcards.
|
|
# Replace the corresponding wildcards in the litellm model pattern with extracted segments
|
|
for segment in dynamic_segments:
|
|
litellm_deployment_litellm_model = litellm_deployment_litellm_model.replace(
|
|
"*", segment, 1
|
|
)
|
|
|
|
return litellm_deployment_litellm_model
|
|
|
|
def get_pattern(
|
|
self, model: str, custom_llm_provider: Optional[str] = None
|
|
) -> Optional[List[Dict]]:
|
|
"""
|
|
Check if a pattern exists for the given model and custom llm provider
|
|
|
|
Args:
|
|
model: str
|
|
custom_llm_provider: Optional[str]
|
|
|
|
Returns:
|
|
bool: True if pattern exists, False otherwise
|
|
"""
|
|
if custom_llm_provider is None:
|
|
try:
|
|
(
|
|
_,
|
|
custom_llm_provider,
|
|
_,
|
|
_,
|
|
) = get_llm_provider(model=model)
|
|
except Exception:
|
|
# get_llm_provider raises exception when provider is unknown
|
|
pass
|
|
return self.route(model) or self.route(f"{custom_llm_provider}/{model}")
|
|
|
|
def get_deployments_by_pattern(
|
|
self, model: str, custom_llm_provider: Optional[str] = None
|
|
) -> List[Dict]:
|
|
"""
|
|
Get the deployments by pattern
|
|
|
|
Args:
|
|
model: str
|
|
custom_llm_provider: Optional[str]
|
|
|
|
Returns:
|
|
List[Dict]: llm deployments matching the pattern
|
|
"""
|
|
pattern_match = self.get_pattern(model, custom_llm_provider)
|
|
if pattern_match:
|
|
return pattern_match
|
|
return []
|
|
|
|
|
|
# Example usage:
|
|
# router = PatternRouter()
|
|
# router.add_pattern('openai/*', [Deployment(), Deployment()])
|
|
# router.add_pattern('openai/fo::*::static::*', Deployment())
|
|
# print(router.route('openai/gpt-4')) # Output: [Deployment(), Deployment()]
|
|
# print(router.route('openai/fo::hi::static::hi')) # Output: [Deployment()]
|
|
# print(router.route('something/else')) # Output: None
|