feat(latency): pass route-selection spend metadata to LiteLLM#54
Conversation
|
WalkthroughAdds correlated LiteLLM spend-log metadata for each latency-service attempt, records successful route selections in request context, and exposes sanitized ChangesRoute Selection Attribution
Estimated code review effort: 4 (Complex) | ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/test_route_selection_headers.py (1)
26-26: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
pytest-mock'smockerfixture instead ofunittest.mock.patch.
HealthPoller.start/stopare patched directly viaunittest.mock.patch.object. As per coding guidelines, tests should "use ... pytest-mock for general mocking." Swapping to themockerfixture would align with the repo convention (no functional difference here since patches are unconditional for the fixture's lifetime).♻️ Suggested refactor using pytest-mock
-from unittest.mock import patch +- with patch.object(HealthPoller, "start"), patch.object(HealthPoller, "stop"): + mocker.patch.object(HealthPoller, "start") + mocker.patch.object(HealthPoller, "stop") + if True:(requires adding a
mocker: MockerFixtureparameter to thelatency_appfixture)As per coding guidelines,
tests/**/*.pyshould "use respx for HTTP mocking and pytest-mock for general mocking." The custom_OpenAICompatStubloopback server (line 226) deviates fromrespxtoo, but that's explicitly justified in the module docstring for wire-level header assertions, so it's not flagged as a separate issue.Also applies to: 239-239
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_route_selection_headers.py` at line 26, Replace unittest.mock.patch usage in the tests, including the HealthPoller.start/stop patches, with pytest-mock's mocker fixture. Add the appropriate mocker fixture parameter to latency_app and use it for the unconditional patches, removing the patch import while preserving the existing test behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@switchyard/lib/backends/latency_service_llm_backend.py`:
- Around line 669-703: Update _call_endpoint’s extra_headers merge to remove or
overwrite every case variant of the protected spend-logs metadata header before
applying extra_headers, ensuring client-provided casing cannot spoof or
duplicate it. Preserve the protected header value from extra_headers and add a
test covering a differently-cased spoof key.
In `@switchyard/lib/endpoints/dispatch.py`:
- Around line 111-113: Update the response handling around
serialize_chain_result so route_selection_headers(ctx) is applied before
returning an existing Response. Merge the generated x-switchyard-* attribution
headers onto that response, preserving its existing headers and body, while
retaining the current serialization path for non-Response results.
---
Nitpick comments:
In `@tests/test_route_selection_headers.py`:
- Line 26: Replace unittest.mock.patch usage in the tests, including the
HealthPoller.start/stop patches, with pytest-mock's mocker fixture. Add the
appropriate mocker fixture parameter to latency_app and use it for the
unconditional patches, removing the patch import while preserving the existing
test behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 95878eb8-f38a-45b0-96eb-5e01868fcada
📒 Files selected for processing (13)
.agents/skills/switchyard-lib-core/SKILL.mddocs/internal/latency_service_routing.mdswitchyard/lib/backends/latency_service_llm_backend.pyswitchyard/lib/endpoints/anthropic_messages_endpoint.pyswitchyard/lib/endpoints/dispatch.pyswitchyard/lib/endpoints/openai_chat_endpoint.pyswitchyard/lib/endpoints/responses_endpoint.pyswitchyard/lib/endpoints/route_selection.pyswitchyard/lib/endpoints/upstream_error.pyswitchyard/lib/proxy_context.pytests/_chain_test_helpers.pytests/test_latency_service_llm_backend.pytests/test_route_selection_headers.py
| async def _call_endpoint( | ||
| self, | ||
| model_id: str, | ||
| target_request_type: ChatRequestType, | ||
| *, | ||
| api_key: str | None, | ||
| body: dict[str, object], | ||
| extra_headers: dict[str, str], | ||
| ) -> object: | ||
| # ``extra_headers`` rides the client wrapper's ``**kwargs`` into the | ||
| # OpenAI SDK's per-request header merge (over ``default_headers``). | ||
| # A passthrough body may itself carry an SDK-style ``extra_headers`` | ||
| # field (shape-preserving translation keeps unknown keys); merge it | ||
| # under ours rather than letting it collide with the keyword — it used | ||
| # to ride ``**body`` into the same SDK parameter, and the spend-logs | ||
| # header must win a name conflict so callers can't spoof it. A | ||
| # non-mapping value is dropped: it could only ever have broken the | ||
| # SDK call. | ||
| client_extra = body.pop("extra_headers", None) | ||
| headers: dict[str, object] = ( | ||
| {**client_extra, **extra_headers} | ||
| if isinstance(client_extra, Mapping) | ||
| else dict(extra_headers) | ||
| ) | ||
| if target_request_type == ChatRequestType.OPENAI_RESPONSES: | ||
| return await self._clients[model_id].aresponses( | ||
| api_key=api_key, | ||
| extra_headers=headers, | ||
| **body, | ||
| ) | ||
| return await self._clients[model_id].acompletion( | ||
| api_key=api_key, | ||
| extra_headers=headers, | ||
| **body, | ||
| ) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how the OpenAI SDK / httpx merge extra_headers with default headers,
# and whether case-differing duplicate keys collapse or survive as two headers.
python - <<'EOF'
import inspect
try:
import httpx
print(inspect.getsource(httpx.Headers.__init__))
except Exception as e:
print("httpx introspection failed:", e)
try:
import openai
from openai._base_client import BaseClient
print(inspect.getsource(BaseClient._build_headers))
except Exception as e:
print("openai SDK introspection failed:", e)
EOFRepository: NVIDIA-NeMo/Switchyard
Length of output: 1118
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and inspect the surrounding implementation.
ast-grep outline switchyard/lib/backends/latency_service_llm_backend.py --view expanded || true
sed -n '640,730p' switchyard/lib/backends/latency_service_llm_backend.py
# Find tests and call sites related to extra_headers / spend-logs metadata.
rg -n "extra_headers|spend-logs|SPEND_LOGS_METADATA_HEADER|test_client_body_extra_headers_forwarded_not_collided" switchyard -g '!**/__pycache__/**'Repository: NVIDIA-NeMo/Switchyard
Length of output: 10424
🌐 Web query:
OpenAI Python SDK BaseClient _build_headers extra_headers case-insensitive httpx.Headers merge behavior
💡 Result:
In the OpenAI Python SDK, the _build_headers method is responsible for constructing the final request headers by merging default client headers with any user-provided extra_headers [1][2][3]. The merge behavior operates as follows: 1. Initial Merging: The SDK first performs a dictionary-based merge of the client's default headers and the provided custom headers using a helper function (_merge_mappings) [1][2][3]. 2. Case-Insensitivity: The result of this dictionary merge is then passed to the httpx.Headers class constructor [1][2][3]. Because httpx.Headers is a case-insensitive multi-dict [4], it normalizes the header keys. 3. Precedence: The extra_headers provided by the user take precedence over client-level default headers when keys clash [5]. This is consistent with the standard behavior of merging dictionaries where the later dictionary overwrites the former, and further enforced by how the httpx.Headers object subsequently handles the provided key-value pairs [4][5]. Essentially, the process relies on the fact that while the preliminary dictionary merge is case-sensitive, the conversion into the httpx.Headers object creates a case-insensitive store that correctly manages the final header set for the request [1][2][3].
Citations:
- 1: https://github.com/openai/openai-python/blob/main/src/openai/_base_client.py
- 2: https://github.com/openai/openai-python/blob/v1.98.0/src/openai/_base_client.py
- 3: https://github.com/openai/openai-python/blob/722d3fff/src/openai/_base_client.py
- 4: https://www.python-httpx.org/api/
- 5: https://github.com/openai/openai-python/blob/722d3fff/tests/test_client.py
🏁 Script executed:
python3 - <<'PY'
import inspect
import httpx
print("httpx version:", httpx.__version__)
for name in ["__contains__", "__getitem__", "get_list", "items", "multi_items", "keys", "values", "raw"]:
try:
attr = getattr(httpx.Headers, name)
print(f"\n### Headers.{name}")
print(inspect.getsource(attr))
except Exception as e:
print(f"\n### Headers.{name} introspection failed: {e}")
h = httpx.Headers([
("x-litellm-spend-logs-metadata", "safe"),
("X-LiteLLM-Spend-Logs-Metadata", "spoof"),
])
print("\nconstructed headers repr:", h)
print("contains lowercase:", "x-litellm-spend-logs-metadata" in h)
print("contains spoof-cased:", "X-LiteLLM-Spend-Logs-Metadata" in h)
print("getitem lowercase:", h["x-litellm-spend-logs-metadata"])
print("get_list lowercase:", h.get_list("x-litellm-spend-logs-metadata"))
print("multi_items:", list(h.multi_items()))
print("raw:", h.raw)
PYRepository: NVIDIA-NeMo/Switchyard
Length of output: 3990
Normalize the protected header case-insensitively before merging. httpx.Headers treats header names case-insensitively and keeps duplicate raw entries, so a client can still send X-LiteLLM-Spend-Logs-Metadata alongside the protected lowercase key. Drop or rewrite any case variant of the spend-logs header, and add a test for a differently-cased spoof key.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@switchyard/lib/backends/latency_service_llm_backend.py` around lines 669 -
703, Update _call_endpoint’s extra_headers merge to remove or overwrite every
case variant of the protected spend-logs metadata header before applying
extra_headers, ensuring client-provided casing cannot spoof or duplicate it.
Preserve the protected header value from extra_headers and add a test covering a
differently-cased spoof key.
| if isinstance(result, Response): | ||
| return result | ||
| headers = route_selection_headers(ctx) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== dispatch.py outline ==\n'
ast-grep outline switchyard/lib/endpoints/dispatch.py --view expanded || true
printf '\n== relevant slice of dispatch.py ==\n'
nl -ba switchyard/lib/endpoints/dispatch.py | sed -n '1,220p'
printf '\n== search for serialize_chain_result usages/tests ==\n'
rg -n "serialize_chain_result|route_selection_headers|CTX_ROUTE_SELECTION|x-switchyard" switchyard -g '!**/__pycache__/**'
printf '\n== tests around route selection headers ==\n'
fd -a "test_route_selection_headers.py|route_selection" switchyard || trueRepository: NVIDIA-NeMo/Switchyard
Length of output: 820
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the endpoint code paths that return Response objects directly.
printf '\n== anthropic_messages_endpoint.py relevant lines ==\n'
nl -ba switchyard/lib/endpoints/anthropic_messages_endpoint.py | sed -n '1,260p'
printf '\n== responses_endpoint.py relevant lines ==\n'
nl -ba switchyard/lib/endpoints/responses_endpoint.py | sed -n '1,260p'
printf '\n== route_selection_headers implementation and tests ==\n'
rg -n "def route_selection_headers|class ProxyContext|CTX_ROUTE_SELECTION|route_selection_headers\\(" switchyard/lib switchyard/tests -g '!**/__pycache__/**'Repository: NVIDIA-NeMo/Switchyard
Length of output: 253
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== dispatch.py relevant slice ==\n'
sed -n '1,220p' switchyard/lib/endpoints/dispatch.py
printf '\n== anthropic_messages_endpoint.py relevant slice ==\n'
sed -n '1,260p' switchyard/lib/endpoints/anthropic_messages_endpoint.py
printf '\n== responses_endpoint.py relevant slice ==\n'
sed -n '1,260p' switchyard/lib/endpoints/responses_endpoint.py
printf '\n== search for route selection header tests ==\n'
rg -n "route_selection_headers|CTX_ROUTE_SELECTION|x-switchyard|serialize_chain_result" switchyard -g '!**/__pycache__/**'Repository: NVIDIA-NeMo/Switchyard
Length of output: 18660
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== route_selection.py ==\n'
sed -n '1,220p' switchyard/lib/endpoints/route_selection.py
printf '\n== search for chain implementations returning Response ==\n'
rg -n "return .*Response|StreamingResponse\\(|JSONResponse\\(" switchyard/lib -g '!**/__pycache__/**'
printf '\n== tests mentioning serialize_chain_result or route selection headers ==\n'
rg -n "serialize_chain_result|route_selection_headers|CTX_ROUTE_SELECTION|x-switchyard-router-model|x-switchyard-selected-model|x-switchyard-selected-provider" switchyard/tests switchyard -g '!*__pycache__*'Repository: NVIDIA-NeMo/Switchyard
Length of output: 8990
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== openai_chat_endpoint.py relevant slice ==\n'
sed -n '1,220p' switchyard/lib/endpoints/openai_chat_endpoint.py
printf '\n== upstream_error.py header stamping slice ==\n'
sed -n '180,260p' switchyard/lib/endpoints/upstream_error.pyRepository: NVIDIA-NeMo/Switchyard
Length of output: 6655
Stamp route-selection headers before returning an existing Response. serialize_chain_result returns Response unchanged before route_selection_headers(ctx) runs, so any pre-built response drops the x-switchyard-* attribution headers even when ctx already has a recorded selection. Merge the headers onto the existing response instead of bypassing them.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@switchyard/lib/endpoints/dispatch.py` around lines 111 - 113, Update the
response handling around serialize_chain_result so route_selection_headers(ctx)
is applied before returning an existing Response. Merge the generated
x-switchyard-* attribution headers onto that response, preserving its existing
headers and body, while retaining the current serialization path for
non-Response results.
1 participant
Summary
For tokenomics/spend reporting (SWITCH-914), a LiteLLM front proxy needs to tie its provider spend-log rows back to the Switchyard logical route that selected them. This adds both halves of that contract to the latency-service backend:
x-litellm-spend-logs-metadatarequest header whose JSON payload records the route selection (router_model,router_strategy: "latency",router_selected_endpoint,router_selected_model,router_selected_provider) plus a per-request uuid4router_correlation_id. One correlation id is shared across failover attempts; the selection fields are re-stamped per attempt so each provider spend-log row records the endpoint it actually hit.x-switchyard-router-model,x-switchyard-selected-model,x-switchyard-selected-provider, andx-switchyard-router-correlation-id(streaming included). Error responses also carry them when the failure happened after a billed upstream success (e.g. a response-translation error following an upstream 200), so the provider row's correlation id stays joinable.Design notes
X-Switchyard-Versionoutbound telemetry header.ctx.metadata[CTX_ROUTE_SELECTION]; the endpoint layer maps it to response headers via the newswitchyard/lib/endpoints/route_selection.py, shared by the success serializer (serialize_chain_result) and the error path (handle_chain_exception).router_modelis skipped when not latin-1-encodable or containing control characters, so a hostile model string can neither fail a billed response nor inject response headers. The outbound JSON header is immune (json.dumpsescapes it) and still records the raw value.extra_headersfield keeps working: it is merged under the spend-logs header, which wins name conflicts so it cannot be spoofed.serialize_chain_resultnow requiresctx, so a future endpoint cannot silently opt out of spend attribution.docs/internal/latency_service_routing.md;switchyard-lib-coreskill Quick Reference row added.Testing
tests/test_route_selection_headers.py): the realOpenAILLMClient/OpenAI SDK against a loopback HTTP stub, asserting the outbound header on the wire and the response headers back — chat (non-streaming + streaming), Anthropic, and Responses ingress; hostile-model and body-extra_headersregressions; error-path semantics in both directions.route_modelfallback) intests/test_latency_service_llm_backend.py.uv run ruff check .clean,uv run mypy switchyard(strict) clean, full hermetic suite: 2000 passed.switchyard --routing-profiles … -- serve(real uvicorn, loopback stub upstream) verified both header surfaces, streaming included.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
x-switchyard-*response headers identifying the selected model, provider, route, and correlation ID.Documentation