-
Notifications
You must be signed in to change notification settings - Fork 20
feat(latency): pass route-selection spend metadata to LiteLLM #54
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,6 +9,7 @@ | |
| from fastapi.responses import JSONResponse, Response, StreamingResponse | ||
|
|
||
| from switchyard.lib.endpoints.error_envelope import error_response | ||
| from switchyard.lib.endpoints.route_selection import route_selection_headers | ||
| from switchyard.lib.endpoints.upstream_error import record_upstream_attempt_success | ||
| from switchyard.lib.proxy_context import ProxyContext | ||
| from switchyard.lib.roles import TranslatedResponse | ||
|
|
@@ -96,16 +97,24 @@ def serialize_chain_result( | |
| *, | ||
| stream: bool, | ||
| sse_iter: Callable[[Any], AsyncIterator[str]], | ||
| ctx: ProxyContext, | ||
| ) -> Response: | ||
| """Serialize a chain result to the appropriate HTTP response. | ||
|
|
||
| Returns the result as-is if it is already a ``Response``, wraps it in a | ||
| ``StreamingResponse`` when streaming is requested, or JSON-serializes it. | ||
| Any route selection recorded on *ctx* is stamped as ``x-switchyard-*`` | ||
| response headers (streaming included — the backend call completed before | ||
| the response object is built, so the selection is final). ``ctx`` is | ||
| required so a new endpoint cannot silently opt out of spend attribution. | ||
| """ | ||
| if isinstance(result, Response): | ||
| return result | ||
| headers = route_selection_headers(ctx) | ||
|
Comment on lines
111
to
+113
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🗄️ 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 🤖 Prompt for AI Agents |
||
| if stream and hasattr(result, "__aiter__"): | ||
| return StreamingResponse(sse_iter(result), media_type="text/event-stream") | ||
| return StreamingResponse( | ||
| sse_iter(result), media_type="text/event-stream", headers=headers | ||
| ) | ||
| if hasattr(result, "model_dump"): | ||
| return JSONResponse(content=result.model_dump()) | ||
| return JSONResponse(content=result) | ||
| return JSONResponse(content=result.model_dump(), headers=headers) | ||
| return JSONResponse(content=result, headers=headers) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| """Route-selection response headers for spend/tokenomics attribution. | ||
|
|
||
| Maps the :data:`CTX_ROUTE_SELECTION` record a routing backend stored on | ||
| ``ctx`` (see :mod:`switchyard.lib.proxy_context`) to the ``x-switchyard-*`` | ||
| response headers a front proxy such as LiteLLM copies into its parent | ||
| spend-log row. Shared by the success serializer (``dispatch``) and the | ||
| error path (``upstream_error``) — a failure that happens *after* a billed | ||
| upstream success must still expose the selection, or the provider spend-log | ||
| row's correlation id becomes unjoinable. | ||
| """ | ||
|
|
||
| from collections.abc import Mapping | ||
|
|
||
| from switchyard.lib.proxy_context import CTX_ROUTE_SELECTION, ProxyContext | ||
|
|
||
| #: Response headers exposing the route selection behind an upstream call, | ||
| #: carrying the same correlation id the provider row received via the | ||
| #: outbound ``x-litellm-spend-logs-metadata`` header. | ||
| ROUTER_MODEL_HEADER = "x-switchyard-router-model" | ||
| SELECTED_MODEL_HEADER = "x-switchyard-selected-model" | ||
| SELECTED_PROVIDER_HEADER = "x-switchyard-selected-provider" | ||
| ROUTER_CORRELATION_ID_HEADER = "x-switchyard-router-correlation-id" | ||
|
|
||
| _ROUTE_SELECTION_RESPONSE_HEADERS = ( | ||
| (ROUTER_MODEL_HEADER, "router_model"), | ||
| (SELECTED_MODEL_HEADER, "router_selected_model"), | ||
| (SELECTED_PROVIDER_HEADER, "router_selected_provider"), | ||
| (ROUTER_CORRELATION_ID_HEADER, "router_correlation_id"), | ||
| ) | ||
|
|
||
|
|
||
| def _is_header_value_safe(value: str) -> bool: | ||
| """Whether *value* can be emitted as an HTTP/1.1 response-header value. | ||
|
|
||
| ``router_model`` echoes the client-supplied model string, so it must be | ||
| re-validated as header material: Starlette encodes response-header values | ||
| as latin-1 (a non-encodable value would fail response construction after | ||
| the upstream call already succeeded and was billed), and CTL characters — | ||
| CR/LF above all — would be a response-splitting vector on permissive | ||
| ASGI stacks. | ||
| """ | ||
| try: | ||
| value.encode("latin-1") | ||
| except UnicodeEncodeError: | ||
| return False | ||
| return not any(ord(ch) < 0x20 or ord(ch) == 0x7F for ch in value) | ||
|
|
||
|
|
||
| def route_selection_headers(ctx: ProxyContext) -> dict[str, str]: | ||
| """Response headers for the route selection recorded on *ctx*, if any. | ||
|
|
||
| Empty when no routing backend recorded a selection (passthrough chains, | ||
| failures before any upstream success). A recorded field that is absent or | ||
| not emittable as a header value is skipped — headers never carry | ||
| placeholder or unsafe values. | ||
| """ | ||
| selection = ctx.metadata.get(CTX_ROUTE_SELECTION) | ||
| if not isinstance(selection, Mapping): | ||
| return {} | ||
| headers: dict[str, str] = {} | ||
| for header_name, selection_key in _ROUTE_SELECTION_RESPONSE_HEADERS: | ||
| value = selection.get(selection_key) | ||
| if isinstance(value, str) and value and _is_header_value_safe(value): | ||
| headers[header_name] = value | ||
| return headers |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
Repository: NVIDIA-NeMo/Switchyard
Length of output: 1118
🏁 Script executed:
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:
🏁 Script executed:
Repository: NVIDIA-NeMo/Switchyard
Length of output: 3990
Normalize the protected header case-insensitively before merging.
httpx.Headerstreats header names case-insensitively and keeps duplicate raw entries, so a client can still sendX-LiteLLM-Spend-Logs-Metadataalongside 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