akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

style(proxy): apply ruff format to the route-advice files

Browse source 
CI runs `ruff format --check` as a separate step after `ruff check`, and
only the latter was clean. Nothing here is a behaviour change: renaming
the dispatch variable to request_backend shortened two call sites enough
that their wrapping was no longer needed, and route_advice.py and its
tests were written to a hand-rolled width.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra 2026-08-05 16:31:09 -07:00
parent 47c9b6e8cc
commit bc48d2daf0
4 changed files with 77 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

8
headroom/proxy/handlers/anthropic.py
View file

@ -2657,9 +2657,7 @@ class AnthropicHandlerMixin:
)
else:
async with stage_timer.measure("upstream_connect"):
backend_response = await request_backend.send_message(
body, headers
)
backend_response = await request_backend.send_message(body, headers)
self.pipeline_extensions.emit(
PipelineStage.POST_SEND,
operation="proxy.request",
@ -2711,9 +2709,7 @@ class AnthropicHandlerMixin:
usage = backend_response.body.get("usage", {})
output_tokens = usage.get("output_tokens", 0)
_backend_name = (
request_backend.name if request_backend else "anthropic"
)
_backend_name = request_backend.name if request_backend else "anthropic"
# Eligible-only denominator for the active
# compression ratio: tokens in the live zone we
# actually attempted to compress. Frozen prefix

4
headroom/proxy/handlers/openai.py
View file

@ -3809,9 +3809,7 @@ class OpenAIHandlerMixin:
)
else:
# Non-streaming: use send_openai_message() → JSON
backend_response = await request_backend.send_openai_message(
body, headers
)
backend_response = await request_backend.send_openai_message(body, headers)
self.pipeline_extensions.emit(
PipelineStage.POST_SEND,
operation="proxy.request",

30
headroom/proxy/route_advice.py
View file

@ -75,9 +75,11 @@ def advice_from(request: Any) -> RouteAdvice | None:
if not isinstance(model, str) or not model:
return None
provider = getattr(obj, "provider", "") or ""
return RouteAdvice(model=model,
provider=provider if isinstance(provider, str) else "",
reason=str(getattr(obj, "reason", "") or "")[:400])
return RouteAdvice(
model=model,
provider=provider if isinstance(provider, str) else "",
reason=str(getattr(obj, "reason", "") or "")[:400],
)
class BackendResolver:
@ -134,8 +136,12 @@ class BackendResolver:
# backend will serve the request, and would log the same decision
# a second time.
body["model"] = advice.model
log.info("route advice: %s via %s (%s)", advice.model, provider,
advice.reason or "no reason given")
log.info(
"route advice: %s via %s (%s)",
advice.model,
provider,
advice.reason or "no reason given",
)
return backend
def _build(self, provider: str) -> Any:
@ -144,16 +150,19 @@ class BackendResolver:
# config -- so a typo builds a backend that only fails later, at
# request time, with an error that points nowhere near the typo.
if not _known_provider(provider):
log.warning("route advice: %r is not a litellm provider; ignoring",
provider)
log.warning("route advice: %r is not a litellm provider; ignoring", provider)
return None
try:
from headroom.backends.litellm import LiteLLMBackend
return LiteLLMBackend(provider=provider)
except Exception as exc: # noqa: BLE001 - never fail a request for this
log.warning("route advice: cannot build a backend for %r (%s); "
"falling back to the configured backend", provider, exc)
log.warning(
"route advice: cannot build a backend for %r (%s); "
"falling back to the configured backend",
provider,
exc,
)
return None
@ -178,8 +187,7 @@ def _known_provider(provider: str) -> bool:
try:
import litellm
names = {getattr(p, "value", None) or str(p)
for p in getattr(litellm, "provider_list", [])}
names = {getattr(p, "value", None) or str(p) for p in getattr(litellm, "provider_list", [])}
return provider in names
except Exception: # noqa: BLE001
return False

77
tests/test_route_advice.py
View file

@ -23,20 +23,26 @@ class _Req:
self.state = SimpleNamespace(**state)
DEFAULT = object() # stands in for the configured backend
DEFAULT = object() # stands in for the configured backend
# --- reading what an extension published ------------------------------------
def test_no_extension_means_no_advice():
assert advice_from(_Req()) is None
assert advice_from(object()) is None # no .state at all
assert advice_from(object()) is None # no .state at all
assert advice_from(None) is None
def test_advice_is_duck_typed_so_extensions_need_not_import_us():
a = advice_from(_Req(headroom_route=SimpleNamespace(
model="moonshot/kimi-k2", provider="moonshot", reason="cheaper")))
a = advice_from(
_Req(
headroom_route=SimpleNamespace(
model="moonshot/kimi-k2", provider="moonshot", reason="cheaper"
)
)
)
assert a == RouteAdvice("moonshot/kimi-k2", "moonshot", "cheaper")
@ -46,8 +52,12 @@ def test_an_extension_may_omit_everything_but_the_model():
def test_malformed_advice_is_ignored_rather_than_raised():
for bad in (SimpleNamespace(), SimpleNamespace(model=""),
SimpleNamespace(model=123), "not an object"):
for bad in (
SimpleNamespace(),
SimpleNamespace(model=""),
SimpleNamespace(model=123),
"not an object",
):
assert advice_from(_Req(headroom_route=bad)) is None
@ -58,6 +68,7 @@ def test_advice_needs_a_model():
# --- absent means unchanged, which is the whole safety argument -------------
def test_no_advice_returns_the_configured_backend():
r = BackendResolver(DEFAULT)
assert r.for_request(_Req()) is DEFAULT
@ -72,8 +83,7 @@ def test_a_native_provider_does_not_switch_backends():
"""Anthropic is the shape the proxy already holds, so a model rewrite is
enough and the extension has already done it."""
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="claude-haiku-4-5",
provider="anthropic"))
req = _Req(headroom_route=SimpleNamespace(model="claude-haiku-4-5", provider="anthropic"))
assert r.for_request(req) is DEFAULT
@ -85,6 +95,7 @@ def test_a_bare_anthropic_model_resolves_its_provider_and_stays_put():
# --- switching, and refusing to switch --------------------------------------
def test_a_foreign_provider_gets_its_own_backend(monkeypatch):
built = []
@ -93,12 +104,10 @@ def test_a_foreign_provider_gets_its_own_backend(monkeypatch):
built.append(provider)
self.provider = provider
monkeypatch.setattr(BackendResolver, "_build",
lambda self, p: FakeBackend(p))
monkeypatch.setattr(BackendResolver, "_build", lambda self, p: FakeBackend(p))
r = BackendResolver(DEFAULT)
body = {"model": "claude-opus-5"}
req = _Req(headroom_route=SimpleNamespace(model="moonshot/kimi-k2",
provider="moonshot"))
req = _Req(headroom_route=SimpleNamespace(model="moonshot/kimi-k2", provider="moonshot"))
got = r.for_request(req, body=body)
assert isinstance(got, FakeBackend) and got.provider == "moonshot"
# The extension could not safely write a foreign model id; we do it.
@ -107,8 +116,7 @@ def test_a_foreign_provider_gets_its_own_backend(monkeypatch):
def test_backends_are_built_once_per_provider(monkeypatch):
built = []
monkeypatch.setattr(BackendResolver, "_build",
lambda self, p: built.append(p) or object())
monkeypatch.setattr(BackendResolver, "_build", lambda self, p: built.append(p) or object())
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="x", provider="moonshot"))
for _ in range(5):
@ -118,8 +126,7 @@ def test_backends_are_built_once_per_provider(monkeypatch):
def test_a_backend_that_will_not_build_falls_back_and_stops_retrying(monkeypatch):
calls = []
monkeypatch.setattr(BackendResolver, "_build",
lambda self, p: calls.append(p) or None)
monkeypatch.setattr(BackendResolver, "_build", lambda self, p: calls.append(p) or None)
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="x", provider="nope"))
for _ in range(5):
@ -130,8 +137,10 @@ def test_a_backend_that_will_not_build_falls_back_and_stops_retrying(monkeypatch
def test_a_routing_preference_can_never_take_traffic_down(monkeypatch):
"""Missing credentials, missing optional dependency -- whatever the reason,
the request still has to be served."""
def boom(self, provider):
raise RuntimeError("no credentials")
monkeypatch.setattr(BackendResolver, "_build", boom)
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="x", provider="moonshot"))
@ -145,8 +154,7 @@ def test_an_unknown_provider_is_rejected_at_resolve_time():
Validate the name up front instead."""
r = BackendResolver(DEFAULT)
assert r._build("definitely-not-a-provider-name") is None
req = _Req(headroom_route=SimpleNamespace(
model="x", provider="definitely-not-a-provider-name"))
req = _Req(headroom_route=SimpleNamespace(model="x", provider="definitely-not-a-provider-name"))
assert r.for_request(req) is DEFAULT
@ -156,6 +164,7 @@ def test_a_real_provider_name_is_accepted():
# --- streaming, which is the path agents actually take ----------------------
class _Backend:
"""Records that it, and not some other backend, served the request."""
@ -178,8 +187,19 @@ async def _drive(handler, **kw):
from headroom.proxy.handlers.streaming import StreamingMixin
resp = await StreamingMixin._stream_response_bedrock(
handler, {"messages": []}, {}, "anthropic", "m", "rid",
0, 0, 0, [], {}, 0.0, **kw,
handler,
{"messages": []},
{},
"anthropic",
"m",
"rid",
0,
0,
0,
[],
{},
0.0,
**kw,
)
async for _ in resp.body_iterator:
pass
@ -224,8 +244,19 @@ async def _drive_openai(handler, **kw):
from headroom.proxy.handlers.streaming import StreamingMixin
resp = await StreamingMixin._stream_openai_via_backend(
handler, {"messages": []}, {}, "m", "rid", 0.0,
0, 0, 0, [], {}, 0.0, **kw,
handler,
{"messages": []},
{},
"m",
"rid",
0.0,
0,
0,
0,
[],
{},
0.0,
**kw,
)
async for _ in resp.body_iterator:
pass
@ -248,8 +279,10 @@ def test_openai_streaming_without_a_route_uses_the_configured_backend():
def test_the_resolver_follows_a_reassigned_default():
class H:
anthropic_backend = None
h = H()
from headroom.proxy.route_advice import BackendResolver as BR
first = BR(h.anthropic_backend)
assert first.default is None
h.anthropic_backend = DEFAULT