scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-22 00:31:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 51

feat(router): add tag_regex support for header-based routing

Browse source 
Adds a new `tag_regex` field to litellm_params that lets operators route
requests based on regex patterns matched against request headers — primarily
User-Agent — without requiring per-developer tag configuration.

Use case: route all Claude Code traffic (User-Agent: claude-code/x.y.z) to
a dedicated deployment by setting:

  tag_regex:
    - "^User-Agent: claude-code\\/"

in the deployment's litellm_params. Works alongside existing `tags` routing;
exact tag match takes precedence over regex match. Unmatched requests fall
through to deployments tagged `default`.

The matched deployment, pattern, and user_agent are recorded in
`metadata["tag_routing"]` so they flow through to SpendLogs automatically.
This commit is contained in:
Ishaan Jaffer 2026-03-13 17:47:26 -07:00
parent 2b61f2a41b
commit a54b403097
4 changed files with 319 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/router.py
View file

@ -14,6 +14,7 @@ import hashlib
import inspect
import json
import logging
import re
import threading
import time
import traceback
@ -1551,7 +1552,9 @@ class Router:
# Drain any fire-and-forget tasks (e.g. alerting hooks)
# scheduled via asyncio.create_task during acompletion.
pending = asyncio.all_tasks()
pending.discard(asyncio.current_task())
current = asyncio.current_task()
if current is not None:
pending.discard(current)
if pending:
await asyncio.gather(*pending, return_exceptions=True)
@ -6526,6 +6529,18 @@ class Router:
deployment = self._add_deployment(deployment=deployment)
# Validate tag_regex patterns early so a bad regex fails at startup
# rather than silently misbehaving on the first matching request.
_tag_regex = deployment.litellm_params.get("tag_regex") or []
for pattern in _tag_regex:
try:
re.compile(pattern)
except re.error as exc:
raise ValueError(
f"Invalid regex in tag_regex for model '{deployment.model_name}': "
f"{pattern!r} — {exc}"
) from exc
model = deployment.to_json(exclude_none=True)
self._add_model_to_list_and_index_map(

90
litellm/router_strategy/tag_based_routing.py
View file

@ -6,6 +6,7 @@ Use this to route requests between Teams
- If no default_deployments are set, return all deployments
"""
import re
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
from litellm._logging import verbose_logger
@ -19,6 +20,29 @@ else:
LitellmRouter = Any
def _is_valid_deployment_tag_regex(
tag_regexes: List[str],
header_strings: List[str],
) -> Optional[str]:
"""
Test compiled regex patterns against "Header-Name: value" strings.
Returns the first matching pattern string, or None if nothing matches.
Uses re.compile() which has an internal LRU cache — no per-call overhead
after the first compile.
"""
for pattern in tag_regexes:
for header_str in header_strings:
try:
if re.search(pattern, header_str):
return pattern
except re.error:
verbose_logger.warning(
"tag_regex: invalid pattern %r — skipping", pattern
)
return None
def is_valid_deployment_tag(
deployment_tags: List[str], request_tags: List[str], match_any: bool = True
) -> bool:
@ -83,30 +107,74 @@ async def get_deployments_for_tag(
request_tags = metadata.get("tags")
match_any = llm_router_instance.tag_filtering_match_any
new_healthy_deployments = []
default_deployments = []
if request_tags:
# Build header strings for regex matching from what the proxy already stores.
# Currently we match against User-Agent; format matches "^User-Agent: claude-code/..."
user_agent = metadata.get("user_agent", "")
header_strings: List[str] = (
[f"User-Agent: {user_agent}"] if user_agent else []
)
new_healthy_deployments: List[Any] = []
default_deployments: List[Any] = []
has_tag_filter = bool(request_tags) or bool(header_strings)
if has_tag_filter:
verbose_logger.debug(
"get_deployments_for_tag routing: router_keys: %s", request_tags
"get_deployments_for_tag routing: request_tags=%s user_agent=%s",
request_tags,
user_agent,
)
# example this can be router_keys=["free", "custom"]
for deployment in healthy_deployments:
deployment_litellm_params = deployment.get("litellm_params")
deployment_tags = deployment_litellm_params.get("tags")
deployment_tag_regex = deployment_litellm_params.get("tag_regex")
verbose_logger.debug(
"deployment: %s, deployment_router_keys: %s",
deployment,
"deployment: %s tags: %s tag_regex: %s",
deployment.get("model_name"),
deployment_tags,
deployment_tag_regex,
)
if deployment_tags is None:
continue
matched_via: Optional[str] = None
matched_value: Optional[str] = None
if is_valid_deployment_tag(deployment_tags, request_tags, match_any):
# 1. Exact tag match (existing behaviour)
if deployment_tags and request_tags:
if is_valid_deployment_tag(deployment_tags, request_tags, match_any):
matched_via = "tags"
matched_value = next(
(t for t in deployment_tags if t in set(request_tags)),
deployment_tags[0],
)
# 2. Regex match against request headers (new)
if matched_via is None and deployment_tag_regex and header_strings:
regex_match = _is_valid_deployment_tag_regex(
deployment_tag_regex, header_strings
)
if regex_match is not None:
matched_via = "tag_regex"
matched_value = regex_match
if matched_via is not None:
verbose_logger.debug(
"tag routing match: deployment=%s matched_via=%s matched_value=%s",
deployment.get("model_name"),
matched_via,
matched_value,
)
# Record provenance in metadata so it flows to SpendLogs
metadata["tag_routing"] = {
"matched_deployment": deployment.get("model_name"),
"matched_via": matched_via,
"matched_value": matched_value,
"request_tags": request_tags or [],
"user_agent": user_agent,
}
new_healthy_deployments.append(deployment)
if "default" in deployment_tags:
if deployment_tags and "default" in deployment_tags:
default_deployments.append(deployment)
if len(new_healthy_deployments) == 0 and len(default_deployments) == 0:

7
litellm/types/router.py
View file

@ -198,6 +198,11 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
model_info: Optional[Dict] = None
mock_response: Optional[Union[str, ModelResponse, Exception, Any]] = None
# tag-based routing
tags: Optional[List[str]] = None
# regex patterns matched against request headers for tag routing
tag_regex: Optional[List[str]] = None
# auto-router params
auto_router_config_path: Optional[str] = None
auto_router_config: Optional[str] = None
@ -334,6 +339,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
# routing params
# use this for tag-based routing
tags: Optional[List[str]]
# regex patterns matched against request headers (e.g. "^User-Agent:\\s*claude-code\\/")
tag_regex: Optional[List[str]]
# deployment budgets
max_budget: Optional[float]

217
tests/test_litellm/router_strategy/test_router_tag_regex_routing.py Normal file
View file

@ -0,0 +1,217 @@
"""
Unit tests for tag_regex routing.
Tests _is_valid_deployment_tag_regex() and get_deployments_for_tag() with tag_regex
patterns, verifying that regex-based header matching works correctly alongside
existing tag-based routing.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../.."))
from unittest.mock import MagicMock
from litellm.router_strategy import tag_based_routing
from litellm.router_strategy.tag_based_routing import get_deployments_for_tag
_is_valid_deployment_tag_regex = tag_based_routing._is_valid_deployment_tag_regex
# ---------------------------------------------------------------------------
# _is_valid_deployment_tag_regex unit tests
# ---------------------------------------------------------------------------
def test_regex_matches_claude_code_user_agent():
"""^User-Agent: claude-code/ matches a claude-code UA string."""
result = _is_valid_deployment_tag_regex(
tag_regexes=[r"^User-Agent: claude-code\/"],
header_strings=["User-Agent: claude-code/1.2.3"],
)
assert result == r"^User-Agent: claude-code\/"
def test_regex_no_match_for_other_ua():
"""Pattern does not match a non-claude-code User-Agent."""
result = _is_valid_deployment_tag_regex(
tag_regexes=[r"^User-Agent: claude-code\/"],
header_strings=["User-Agent: Mozilla/5.0 (browser)"],
)
assert result is None
def test_regex_returns_first_matching_pattern():
"""When multiple patterns are provided, returns the first match."""
result = _is_valid_deployment_tag_regex(
tag_regexes=[r"^User-Agent: cursor\/", r"^User-Agent: claude-code\/"],
header_strings=["User-Agent: claude-code/2.0.0"],
)
assert result == r"^User-Agent: claude-code\/"
def test_regex_empty_inputs_return_none():
"""Empty lists return None without errors."""
assert _is_valid_deployment_tag_regex([], ["User-Agent: claude-code/1.0"]) is None
assert _is_valid_deployment_tag_regex([r"^User-Agent: claude-code\/"], []) is None
def test_invalid_regex_skipped_does_not_raise():
"""An invalid regex pattern is skipped (warning logged) — no exception raised."""
result = _is_valid_deployment_tag_regex(
tag_regexes=["[invalid(regex"],
header_strings=["User-Agent: claude-code/1.0"],
)
assert result is None
def test_regex_matches_version_range():
"""Semver-aware pattern matches multiple versions."""
pattern = r"^User-Agent: claude-code\/\d"
for ua in ["claude-code/1.0", "claude-code/2.0.0-beta.1", "claude-code/99.0"]:
result = _is_valid_deployment_tag_regex(
tag_regexes=[pattern],
header_strings=[f"User-Agent: {ua}"],
)
assert result == pattern, f"Expected match for UA: {ua}"
# ---------------------------------------------------------------------------
# get_deployments_for_tag integration tests
# ---------------------------------------------------------------------------
CLAUDE_CODE_DEPLOYMENT = {
"model_name": "claude-sonnet",
"litellm_params": {
"model": "openai/claude-code-deployment",
"api_key": "fake",
"mock_response": "cc",
"tag_regex": [r"^User-Agent: claude-code\/"],
},
"model_info": {"id": "claude-code-deployment"},
}
REGULAR_DEPLOYMENT = {
"model_name": "claude-sonnet",
"litellm_params": {
"model": "openai/regular-deployment",
"api_key": "fake",
"mock_response": "regular",
"tags": ["default"],
},
"model_info": {"id": "regular-deployment"},
}
ALL_DEPLOYMENTS = [CLAUDE_CODE_DEPLOYMENT, REGULAR_DEPLOYMENT]
def _make_router_mock(enable_tag_filtering=True, match_any=True):
mock = MagicMock()
mock.enable_tag_filtering = enable_tag_filtering
mock.tag_filtering_match_any = match_any
return mock
@pytest.mark.asyncio
async def test_claude_code_ua_routes_to_cc_deployment():
"""claude-code/x.y.z UA → claude-code-deployment via tag_regex."""
router = _make_router_mock()
result = await get_deployments_for_tag(
llm_router_instance=router,
model="claude-sonnet",
healthy_deployments=ALL_DEPLOYMENTS,
request_kwargs={"metadata": {"user_agent": "claude-code/1.2.3"}},
)
assert len(result) == 1
assert result[0]["model_info"]["id"] == "claude-code-deployment"
@pytest.mark.asyncio
async def test_regular_ua_routes_to_default_deployment():
"""Mozilla UA → regular-deployment via default tag fallback."""
router = _make_router_mock()
result = await get_deployments_for_tag(
llm_router_instance=router,
model="claude-sonnet",
healthy_deployments=ALL_DEPLOYMENTS,
request_kwargs={"metadata": {"user_agent": "Mozilla/5.0 (browser)"}},
)
assert len(result) == 1
assert result[0]["model_info"]["id"] == "regular-deployment"
@pytest.mark.asyncio
async def test_no_ua_routes_to_default_deployment():
"""No User-Agent → default deployment."""
router = _make_router_mock()
result = await get_deployments_for_tag(
llm_router_instance=router,
model="claude-sonnet",
healthy_deployments=ALL_DEPLOYMENTS,
request_kwargs={"metadata": {}},
)
assert len(result) == 1
assert result[0]["model_info"]["id"] == "regular-deployment"
@pytest.mark.asyncio
async def test_tag_routing_metadata_written_for_regex_match():
"""tag_routing metadata block is populated when regex matches."""
router = _make_router_mock()
metadata: dict = {"user_agent": "claude-code/2.0.0-beta.1"}
await get_deployments_for_tag(
llm_router_instance=router,
model="claude-sonnet",
healthy_deployments=ALL_DEPLOYMENTS,
request_kwargs={"metadata": metadata},
)
assert "tag_routing" in metadata
tr = metadata["tag_routing"]
assert tr["matched_via"] == "tag_regex"
assert tr["matched_value"] == r"^User-Agent: claude-code\/"
assert tr["user_agent"] == "claude-code/2.0.0-beta.1"
@pytest.mark.asyncio
async def test_tag_filtering_disabled_returns_all_deployments():
"""When enable_tag_filtering is False, all deployments returned regardless of UA."""
router = _make_router_mock(enable_tag_filtering=False)
result = await get_deployments_for_tag(
llm_router_instance=router,
model="claude-sonnet",
healthy_deployments=ALL_DEPLOYMENTS,
request_kwargs={"metadata": {"user_agent": "claude-code/1.0"}},
)
assert result == ALL_DEPLOYMENTS
@pytest.mark.asyncio
async def test_explicit_tag_match_takes_precedence_over_regex():
"""A deployment with both tags and tag_regex: exact tag match fires first."""
deployment_with_both = {
"model_name": "claude-sonnet",
"litellm_params": {
"model": "openai/both-deployment",
"api_key": "fake",
"tags": ["premium"],
"tag_regex": [r"^User-Agent: claude-code\/"],
},
"model_info": {"id": "both-deployment"},
}
router = _make_router_mock()
metadata: dict = {
"tags": ["premium"],
"user_agent": "claude-code/1.0",
}
result = await get_deployments_for_tag(
llm_router_instance=router,
model="claude-sonnet",
healthy_deployments=[deployment_with_both],
request_kwargs={"metadata": metadata},
)
assert len(result) == 1
tr = metadata.get("tag_routing", {})
assert tr.get("matched_via") == "tags"