scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

fix(auto_router): bill the routing embedding to the caller's key and team (#39532)

Browse source 
* fix(auto_router): bill the routing embedding to the caller's key and team

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(auto_router): validate the forwarded caller metadata with a pydantic model

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-03 13:53:30 -07:00 • committed by GitHub
parent aff6b7e212
commit eb6c24a2a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 146 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/litellm_core_utils/internal_call_metadata.py
View file

@ -18,9 +18,11 @@ caller's identity metadata, minus two things that must never be forwarded as-is:
from __future__ import annotations
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES
from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
@ -142,6 +144,19 @@ def forwarded_internal_call_metadata(
}
def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]:
kwargs: Final = request_kwargs or MappingProxyType({})
return MappingProxyType(
{k: v for k in ("litellm_session_id", "litellm_trace_id") if isinstance(v := kwargs.get(k), str)}
)
def effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None:
return initialize_standard_callback_dynamic_params(dict(request_kwargs) if request_kwargs else None).get(
"turn_off_message_logging"
)
def sanitized_forwardable_call_metadata(
parent_metadata: Mapping[str, object],
call_origin: InternalCallOrigin,

59
litellm/router_strategy/auto_router/auto_router.py
View file

@ -2,23 +2,41 @@
Auto-Routing Strategy that works with a Semantic Router Config
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Optional
from pydantic import BaseModel, ConfigDict
from litellm._logging import verbose_router_logger
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.internal_call_metadata import (
effective_turn_off_message_logging,
forwarded_internal_call_metadata,
parent_session_kwargs,
)
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN
if TYPE_CHECKING:
from semantic_router.routers import SemanticRouter
from semantic_router.routers.base import Route
from litellm.router import Router
from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder
from litellm.types.router import PreRoutingHookResponse
else:
Router = Any
PreRoutingHookResponse = Any
Route = Any
SemanticRouter = Any
LiteLLMRouterEncoder = Any
class _CallerMetadata(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
metadata: Mapping[str, object] | None = None
litellm_metadata: Mapping[str, object] | None = None
class AutoRouter(CustomLogger):
@ -50,6 +68,8 @@ class AutoRouter(CustomLogger):
"""
from semantic_router.routers import SemanticRouter
from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder
self.auto_router_config_path: str | None = auto_router_config_path
self.auto_router_config: str | None = auto_router_config
self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE
@ -59,6 +79,11 @@ class AutoRouter(CustomLogger):
self.embedding_model: str = embedding_model
self.max_input_chars: int = max_input_chars
self.litellm_router_instance: Router = litellm_router_instance
self.encoder: LiteLLMRouterEncoder = LiteLLMRouterEncoder(
litellm_router_instance=litellm_router_instance,
model_name=embedding_model,
max_input_chars=max_input_chars,
)
def _load_semantic_routing_routes(self) -> list[Route]:
from semantic_router.routers import SemanticRouter
@ -129,9 +154,6 @@ class AutoRouter(CustomLogger):
from semantic_router.routers import SemanticRouter
from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
from litellm.router_strategy.auto_router.litellm_encoder import (
LiteLLMRouterEncoder,
)
from litellm.types.router import PreRoutingHookResponse
resolved_messages: Final = (
@ -149,34 +171,47 @@ class AutoRouter(CustomLogger):
#######################
routelayer = SemanticRouter(
routes=self.loaded_routes,
encoder=LiteLLMRouterEncoder(
litellm_router_instance=self.litellm_router_instance,
model_name=self.embedding_model,
max_input_chars=self.max_input_chars,
),
encoder=self.encoder,
auto_sync=self.auto_sync_value,
)
self.routelayer = routelayer
message_content: Final = self._extract_text_from_messages(resolved_messages)
route_name: Final = self._matched_route_name(routelayer, message_content)
route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs)
return PreRoutingHookResponse(
model=route_name or self.default_model,
messages=messages,
)
def _matched_route_name(self, routelayer: "SemanticRouter", text: str) -> str | None:
async def _matched_route_name(
self, routelayer: "SemanticRouter", text: str, request_kwargs: Mapping[str, object]
) -> str | None:
"""Name of the route `text` matches, or None when nothing matched or the match failed.
The route layer embeds `text` to compare it against the routes, and that embedding call can
`text` is embedded here rather than by `routelayer(text=...)` so the caller's metadata reaches
`aembedding()` and the embedding's spend lands on the key/team that sent the request;
SemanticRouter has no way to pass kwargs through to its encoder. That embedding call can
fail (context limit, timeout, provider error). Choosing a model is a routing decision, so a
failure here falls back to the default model rather than failing the user's request.
"""
from semantic_router.schema import RouteChoice
try:
route_choice: Final = routelayer(text=text)
caller: Final = _CallerMetadata.model_validate(request_kwargs)
query_vector: Final = (
await self.encoder.aencode_queries(
[text],
metadata=forwarded_internal_call_metadata(caller.metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN),
litellm_metadata=forwarded_internal_call_metadata(
caller.litellm_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN
),
proxy_server_request={"body": {"model": self.embedding_model, "input": [text]}},
turn_off_message_logging=effective_turn_off_message_logging(request_kwargs),
**parent_session_kwargs(request_kwargs),
)
)[0]
route_choice: Final = await routelayer.acall(vector=query_vector)
except Exception as e: # noqa: BLE001 -- the embedding call behind the route layer can fail many ways (context limit, timeout, provider/network error); none of them may fail the request
verbose_router_logger.warning(
"AutoRouter: semantic routing failed (%s), falling back to default model %s", e, self.default_model

120
tests/test_litellm/router_strategy/test_auto_router.py
View file

@ -330,36 +330,46 @@ ROUTER_CONFIG: Final = json.dumps(
)
class FailingRouteLayer:
"""Route layer whose embedding call fails, as it does when the prompt exceeds the encoder's window."""
def __call__(self, text: str) -> Any:
raise ValueError(
"Internal_litellm_router API call failed. Error: litellm.InternalServerError: "
"input is too large to process. increase the physical batch size"
)
class FixedRouteLayer:
"""Route layer that returns whatever the test tells it to, recording the text it was asked about."""
"""Route layer that returns whatever the test tells it to for the query vector it is handed."""
def __init__(self, route_choice: Any) -> None:
self.route_choice = route_choice
self.seen_text: str | None = None
def __call__(self, text: str) -> Any:
self.seen_text = text
async def acall(self, vector: Any) -> Any:
return self.route_choice
def _embedding_response(input: List[str]) -> Any:
import litellm
return litellm.EmbeddingResponse(
data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))]
)
class StubEmbeddingRouter:
"""Stands in for the LiteLLM Router when the route index has to be built for real."""
"""Stands in for the LiteLLM Router, recording the text and kwargs each query embedding was made with."""
def __init__(self) -> None:
self.seen_text: str | None = None
self.aembedding_kwargs: Dict[str, Any] | None = None
def embedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
import litellm
return _embedding_response(input)
return litellm.EmbeddingResponse(
data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))]
async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
self.seen_text = input[0]
self.aembedding_kwargs = kwargs
return _embedding_response(input)
class FailingEmbeddingRouter(StubEmbeddingRouter):
"""Router whose query embedding fails, as it does when the prompt exceeds the encoder's window."""
async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
raise ValueError(
"litellm.InternalServerError: input is too large to process. increase the physical batch size"
)
@ -369,7 +379,7 @@ def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs:
auto_router_config=ROUTER_CONFIG,
default_model="fallback-model",
embedding_model="text-embedding-3-small",
litellm_router_instance=litellm_router_instance or MagicMock(),
litellm_router_instance=litellm_router_instance or StubEmbeddingRouter(),
**kwargs,
)
auto_router.routelayer = routelayer
@ -381,7 +391,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
@pytest.mark.asyncio
async def test_should_fall_back_to_default_model_when_the_embedding_call_fails(self):
auto_router: Final = _auto_router(FailingRouteLayer())
auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=FailingEmbeddingRouter())
result: Final = await auto_router.async_pre_routing_hook(
model="my-auto-router",
@ -440,8 +450,8 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
async def test_should_still_route_to_the_matched_route_when_one_matches(self):
from semantic_router.schema import RouteChoice
layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
auto_router: Final = _auto_router(layer)
router: Final = StubEmbeddingRouter()
auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
result: Final = await auto_router.async_pre_routing_hook(
model="my-auto-router",
@ -451,7 +461,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
assert result is not None
assert result.model == "code-model"
assert layer.seen_text == "fix this stack trace"
assert router.seen_text == "fix this stack trace"
class TestAutoRouterEmbeddingInputCap:
@ -483,8 +493,8 @@ class TestAutoRouterRoutesResponsesApiInput:
async def test_should_route_a_string_input_when_messages_is_none(self):
from semantic_router.schema import RouteChoice
layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
auto_router: Final = _auto_router(layer)
router: Final = StubEmbeddingRouter()
auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
result: Final = await auto_router.async_pre_routing_hook(
model="my-auto-router",
@ -498,14 +508,14 @@ class TestAutoRouterRoutesResponsesApiInput:
assert result is not None
assert result.model == "code-model"
assert result.messages is None
assert layer.seen_text == "fix this stack trace"
assert router.seen_text == "fix this stack trace"
@pytest.mark.asyncio
async def test_should_route_a_list_input_with_instructions_when_messages_is_none(self):
from semantic_router.schema import RouteChoice
layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
auto_router: Final = _auto_router(layer)
router: Final = StubEmbeddingRouter()
auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
result: Final = await auto_router.async_pre_routing_hook(
model="my-auto-router",
@ -525,13 +535,13 @@ class TestAutoRouterRoutesResponsesApiInput:
assert result is not None
assert result.model == "code-model"
assert layer.seen_text is not None
assert "fix this stack trace" in layer.seen_text
assert router.seen_text is not None
assert "fix this stack trace" in router.seen_text
@pytest.mark.asyncio
async def test_should_skip_routing_when_neither_messages_nor_input_is_present(self):
layer: Final = FixedRouteLayer(None)
auto_router: Final = _auto_router(layer)
router: Final = StubEmbeddingRouter()
auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router)
result: Final = await auto_router.async_pre_routing_hook(
model="my-auto-router",
@ -540,12 +550,12 @@ class TestAutoRouterRoutesResponsesApiInput:
)
assert result is None
assert layer.seen_text is None
assert router.seen_text is None
@pytest.mark.asyncio
async def test_should_keep_routing_an_empty_messages_list_to_the_default_model(self):
layer: Final = FixedRouteLayer(None)
auto_router: Final = _auto_router(layer)
router: Final = StubEmbeddingRouter()
auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router)
result: Final = await auto_router.async_pre_routing_hook(
model="my-auto-router",
@ -555,4 +565,42 @@ class TestAutoRouterRoutesResponsesApiInput:
assert result is not None
assert result.model == "fallback-model"
assert layer.seen_text == ""
assert router.seen_text == ""
class TestAutoRouterAttributesItsEmbeddingSpend:
"""The query embedding is billed to the key that sent the request, like any other call it made."""
@pytest.mark.asyncio
async def test_should_forward_the_callers_identity_to_the_query_embedding_minus_its_budget_reservation(self):
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
router: Final = StubEmbeddingRouter()
auto_router: Final = _auto_router(None, litellm_router_instance=router)
request_kwargs: Final = {
"metadata": {
"user_api_key": "hashed-key",
"user_api_key_team_id": "team-1",
"user_api_key_budget_reservation": {"reservation_id": "r-1"},
},
"litellm_session_id": "session-1",
}
result: Final = await auto_router.async_pre_routing_hook(
model="my-auto-router",
request_kwargs=request_kwargs,
messages=[{"role": "user", "content": "fix this stack trace"}],
)
assert result is not None
assert router.seen_text == "fix this stack trace"
assert router.aembedding_kwargs is not None
forwarded: Final = router.aembedding_kwargs["metadata"]
assert forwarded["user_api_key"] == "hashed-key"
assert forwarded["user_api_key_team_id"] == "team-1"
assert forwarded[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "autorouter_classifier"
assert "user_api_key_budget_reservation" not in forwarded
assert router.aembedding_kwargs["litellm_session_id"] == "session-1"
assert router.aembedding_kwargs["proxy_server_request"] == {
"body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]}
}