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
…n pre-built responses
2 participants
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