mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(router): keep the request's metadata on rust core fallback hop kwargs (#50263)
The fallback loggers got the walk's own metadata in place of the request's, so prometheus fallback counters lost the caller's key and team labels. Lay the walk's stamps over a copy of the request's bucket, as run_async_fallback did
This commit is contained in:
parent
3b20b48473
commit
41cbbfeb8d
2 changed files with 59 additions and 2 deletions
|
|
@ -264,10 +264,18 @@ class WalkMetadata:
|
|||
|
||||
def hop_kwargs(self, event: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a fresh dict for the loggers
|
||||
"""The kwargs the level held after a hop (``kwargs["model"] = target``), as the
|
||||
fallback loggers see them."""
|
||||
fallback loggers see them: the request's own metadata (the proxy's ``user_api_key_*``
|
||||
labels) with the level's walk stamps laid over it, as ``run_async_fallback``'s
|
||||
``_hop_metadata`` copies it."""
|
||||
scope: Final = self._scope(event.get("scope"))
|
||||
model: Final = scope.metadata.get("model_group", event.get("to"))
|
||||
return {**self._request, **scope.kwargs, "model": model, self._field: dict(scope.metadata)}
|
||||
request_bucket: Final = as_str_mapping(self._request.get(self._field)) or {}
|
||||
return {
|
||||
**self._request,
|
||||
**scope.kwargs,
|
||||
"model": model,
|
||||
self._field: {**request_bucket, **scope.metadata},
|
||||
}
|
||||
|
||||
def log_retry(self, event: Mapping[str, object], error: BaseException | None) -> None:
|
||||
"""A ``retry_logged`` event: the failure joins the scope's ``previous_models``."""
|
||||
|
|
|
|||
49
tests/unit/router_rust/test_walk.py
Normal file
49
tests/unit/router_rust/test_walk.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
import pytest
|
||||
|
||||
from litellm.router_rust.invoke.walk import WalkMetadata
|
||||
|
||||
|
||||
def _no_base_model(_model_id: str) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"])
|
||||
def test_hop_kwargs_keep_the_requests_key_labels_under_the_walk_stamps(field: str) -> None:
|
||||
"""The fallback loggers (prometheus' fallback counters) read the caller's key and team off
|
||||
the hop's metadata, so a hop keeps the request's bucket and lays its own stamps over it."""
|
||||
request = {
|
||||
"model": "primary",
|
||||
field: {
|
||||
"user_api_key_hash": "hashed-key",
|
||||
"user_api_key_alias": "key-alias",
|
||||
"user_api_key_team_id": "team-id",
|
||||
"user_api_key_team_alias": "team-alias",
|
||||
"attempted_fallbacks": "stale",
|
||||
},
|
||||
}
|
||||
walk = WalkMetadata(request, field, _no_base_model)
|
||||
walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5})
|
||||
|
||||
kwargs = walk.hop_kwargs({"scope": 1, "to": "fallback"})
|
||||
|
||||
assert kwargs["model"] == "fallback"
|
||||
metadata = kwargs[field]
|
||||
assert isinstance(metadata, dict)
|
||||
assert metadata["user_api_key_hash"] == "hashed-key"
|
||||
assert metadata["user_api_key_alias"] == "key-alias"
|
||||
assert metadata["user_api_key_team_id"] == "team-id"
|
||||
assert metadata["user_api_key_team_alias"] == "team-alias"
|
||||
assert metadata["model_group"] == "fallback"
|
||||
assert metadata["original_model_group"] == "primary"
|
||||
assert metadata["attempted_fallbacks"] == 1
|
||||
assert request[field]["attempted_fallbacks"] == "stale"
|
||||
assert "model_group" not in request[field]
|
||||
|
||||
|
||||
def test_hop_kwargs_without_request_metadata_hold_the_walk_stamps() -> None:
|
||||
walk = WalkMetadata({"model": "primary"}, "metadata", _no_base_model)
|
||||
walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5})
|
||||
|
||||
metadata = walk.hop_kwargs({"scope": 1, "to": "fallback"})["metadata"]
|
||||
|
||||
assert metadata == {"original_model_group": "primary", "model_group": "fallback", "attempted_fallbacks": 1}
|
||||
Loading…
Add table
Reference in a new issue