mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
feat(proxy)!: use Bedrock request metadata as the session id on pass-through routes (#49193)
This commit is contained in:
shrey
GitHub
parent
991ea7f7e8
commit
033c758299
2 changed files with 218 additions and 0 deletions
|
|
@ -112,6 +112,7 @@ _EXPLICIT_SESSION_HEADERS: Final = frozenset({"x-litellm-trace-id", "x-litellm-s
|
|||
# ``session-id``/``thread-id``; builds before the codex-api split sent
|
||||
# ``session_id``/``conversation_id``. Ordered session before thread.
|
||||
_CODEX_SESSION_ID_HEADERS: Final = ("session-id", "session_id", "thread-id", "conversation_id")
|
||||
_BEDROCK_SESSION_ID_METADATA_KEYS: Final = ("sessionId", "threadId", "conversationId")
|
||||
# Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores
|
||||
# (covers UUIDs and most common session-id formats).
|
||||
_SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
|
||||
|
|
@ -1055,6 +1056,22 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None:
|
|||
return session_id
|
||||
|
||||
|
||||
def _get_bedrock_session_id_from_request_metadata(
|
||||
normalized_headers: Mapping[str, str], request_body: Mapping[str, object]
|
||||
) -> str | None:
|
||||
"""Read the conversation id a Bedrock SDK client tags its calls with: Converse's ``requestMetadata``
|
||||
body field or InvokeModel's ``X-Amzn-Bedrock-Request-Metadata`` JSON header. Only the Bedrock
|
||||
pass-through route nests the client's body under ``data``."""
|
||||
if request_body.get("custom_llm_provider") != "bedrock" or not isinstance(request_body.get("data"), Mapping):
|
||||
return None
|
||||
header: Final = normalized_headers.get("x-amzn-bedrock-request-metadata")
|
||||
body_tags: Final = _request_fields(_request_fields(request_body.get("data")).get("requestMetadata"))
|
||||
header_tags: Final = _request_fields(safe_json_loads(header) if header else None)
|
||||
tag_sources: Final = itertools.product((body_tags, header_tags), _BEDROCK_SESSION_ID_METADATA_KEYS)
|
||||
values: Final = (tags.get(key) for tags, key in tag_sources)
|
||||
return next((value for value in values if isinstance(value, str) and value), None)
|
||||
|
||||
|
||||
def _is_llm_inference_route(request: Request) -> bool:
|
||||
route: Final = get_request_route(request)
|
||||
return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access(
|
||||
|
|
@ -1967,6 +1984,14 @@ class LiteLLMProxyRequestSetup:
|
|||
body_metadata["user_id"] = session_id
|
||||
verbose_proxy_logger.debug("Extracted session_id from Anthropic metadata.user_id")
|
||||
|
||||
if "litellm_session_id" not in data:
|
||||
bedrock_session_id: Final = _get_bedrock_session_id_from_request_metadata(normalized_headers, request_body)
|
||||
metadata_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id")
|
||||
if bedrock_session_id and metadata_session_id in (None, bedrock_session_id):
|
||||
metadata_from_headers["session_id"] = bedrock_session_id
|
||||
data["litellm_session_id"] = bedrock_session_id # rebind-ok: data is an out-param
|
||||
verbose_proxy_logger.debug("Extracted session_id from Bedrock request metadata")
|
||||
|
||||
# Last-resort fallback: the W3C standards for trace/session propagation
|
||||
# (https://www.w3.org/TR/trace-context/, https://www.w3.org/TR/baggage/).
|
||||
# Lower priority than everything above - only fires when neither the
|
||||
|
|
|
|||
|
|
@ -4323,6 +4323,186 @@ def test_add_litellm_metadata_from_request_headers_traceparent_and_baggage_toget
|
|||
assert data["litellm_session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
def _bedrock_passthrough_data(
|
||||
client_body: dict[str, object], custom_llm_provider: str = "bedrock"
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"method": "POST",
|
||||
"endpoint": "model/us.anthropic.claude-haiku-4-5-20251001-v1:0/converse",
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"data": client_body,
|
||||
"litellm_metadata": {},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("headers", "client_body"),
|
||||
[
|
||||
pytest.param({}, {"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}, id="converse_body"),
|
||||
pytest.param(
|
||||
{"X-Amzn-Bedrock-Request-Metadata": '{"threadId": "thread-abc-123", "team": "search"}'},
|
||||
{"messages": []},
|
||||
id="invoke_header",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_bedrock_request_metadata_thread_id_becomes_session_id(
|
||||
headers: dict[str, str], client_body: dict[str, object]
|
||||
):
|
||||
sent_body: Final = copy.deepcopy(client_body)
|
||||
data: Final = _bedrock_passthrough_data(client_body)
|
||||
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
|
||||
assert data["litellm_session_id"] == "thread-abc-123"
|
||||
assert data["litellm_metadata"] == {"session_id": "thread-abc-123"}
|
||||
assert "litellm_trace_id" not in data, "a conversation tag must not merge the turns into one trace"
|
||||
assert data["data"] == sent_body, "the tags must still reach Bedrock unchanged"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tags", "expected"),
|
||||
[
|
||||
({"threadId": "thread-abc-123", "sessionId": "session-xyz-789"}, "session-xyz-789"),
|
||||
({"conversationId": "conversation-42", "threadId": "thread-abc-123"}, "thread-abc-123"),
|
||||
({"conversationId": "conversation-42", "user": "shrey-1234"}, "conversation-42"),
|
||||
({"sessionId": "", "threadId": "thread-abc-123"}, "thread-abc-123"),
|
||||
({"sessionId": "conv-42"}, "conv-42"),
|
||||
({"sessionId": "user:123 / thread.7"}, "user:123 / thread.7"),
|
||||
],
|
||||
)
|
||||
def test_bedrock_request_metadata_session_key_priority(tags: dict[str, str], expected: str):
|
||||
data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": tags})
|
||||
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={}, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
|
||||
assert data["litellm_session_id"] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("headers", "data"),
|
||||
[
|
||||
pytest.param(
|
||||
{"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"'},
|
||||
_bedrock_passthrough_data({"messages": []}),
|
||||
id="malformed_header_json",
|
||||
),
|
||||
pytest.param(
|
||||
{"x-amzn-bedrock-request-metadata": '["thread-abc-123"]'},
|
||||
_bedrock_passthrough_data({"messages": []}),
|
||||
id="header_json_not_an_object",
|
||||
),
|
||||
pytest.param(
|
||||
{}, _bedrock_passthrough_data({"requestMetadata": {"threadId": 1234567890}}), id="non_string_value"
|
||||
),
|
||||
pytest.param({}, _bedrock_passthrough_data({"requestMetadata": {"threadId": ""}}), id="empty"),
|
||||
pytest.param({}, _bedrock_passthrough_data({"requestMetadata": {"user": "shrey-1234"}}), id="no_session_key"),
|
||||
pytest.param(
|
||||
{"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"}'},
|
||||
{"model": "claude", "custom_llm_provider": "bedrock", "litellm_metadata": {}},
|
||||
id="not_a_passthrough_body",
|
||||
),
|
||||
pytest.param(
|
||||
{},
|
||||
_bedrock_passthrough_data(
|
||||
{"requestMetadata": {"threadId": "thread-abc-123"}}, custom_llm_provider="vertex_ai"
|
||||
),
|
||||
id="not_bedrock",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_bedrock_request_metadata_without_usable_session_id_sets_no_session(
|
||||
headers: dict[str, str], data: dict[str, object]
|
||||
):
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
|
||||
assert "litellm_session_id" not in data
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body_tags", "header_tags", "expected"),
|
||||
[
|
||||
pytest.param({"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key"),
|
||||
pytest.param({"sessionId": "body-session"}, {"sessionId": "header-session"}, "body-session", id="body_wins"),
|
||||
],
|
||||
)
|
||||
def test_bedrock_request_metadata_reads_body_tags_then_header_tags(
|
||||
body_tags: dict[str, str], header_tags: dict[str, str], expected: str
|
||||
):
|
||||
data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": body_tags})
|
||||
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={"x-amzn-bedrock-request-metadata": json.dumps(header_tags)},
|
||||
data=data,
|
||||
_metadata_variable_name="litellm_metadata",
|
||||
)
|
||||
|
||||
assert data["litellm_session_id"] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_bucket", ["litellm_metadata", "metadata"])
|
||||
def test_bedrock_request_metadata_never_replaces_a_caller_metadata_session(metadata_bucket: str):
|
||||
data: Final = {
|
||||
**_bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-B"}}),
|
||||
metadata_bucket: {"session_id": "conversation-A"},
|
||||
}
|
||||
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={}, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
|
||||
assert "litellm_session_id" not in data
|
||||
assert data[metadata_bucket]["session_id"] == "conversation-A"
|
||||
|
||||
|
||||
def test_bedrock_request_metadata_stamps_a_tag_equal_to_the_caller_metadata_session():
|
||||
data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}})
|
||||
data["litellm_metadata"] = {"session_id": "thread-abc-123"}
|
||||
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={}, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
|
||||
assert data["litellm_session_id"] == "thread-abc-123"
|
||||
|
||||
|
||||
def test_bedrock_request_metadata_header_leaves_non_passthrough_caller_session_alone():
|
||||
data: Final = {"model": "claude", "custom_llm_provider": "bedrock", "metadata": {"session_id": "caller-session-1"}}
|
||||
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"}'},
|
||||
data=data,
|
||||
_metadata_variable_name="metadata",
|
||||
)
|
||||
|
||||
assert data["metadata"] == {"session_id": "caller-session-1"}
|
||||
assert "litellm_session_id" not in data
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("headers", "expected"),
|
||||
[
|
||||
pytest.param({"x-litellm-session-id": "explicit-session-1"}, "explicit-session-1", id="explicit_header_wins"),
|
||||
pytest.param({"baggage": "session.id=baggage-session-42"}, "thread-abc-123", id="beats_baggage"),
|
||||
],
|
||||
)
|
||||
def test_bedrock_request_metadata_session_precedence(headers: dict[str, str], expected: str):
|
||||
data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}})
|
||||
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
|
||||
assert data["litellm_session_id"] == expected
|
||||
assert data["litellm_metadata"]["session_id"] == expected
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_explicit_trace_id_beats_traceparent():
|
||||
"""x-litellm-trace-id must win over a traceparent header carrying a
|
||||
different trace-id - explicit litellm headers are always highest priority."""
|
||||
|
|
@ -9598,6 +9778,19 @@ async def test_missing_session_id_reject_accepts_body_metadata_session_id():
|
|||
assert updated["metadata"]["session_id"] == "body-session-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_session_id_reject_accepts_bedrock_request_metadata_session_id():
|
||||
updated: Final = await add_litellm_data_to_request(
|
||||
data=_bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}),
|
||||
request=_request_for("/bedrock/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/converse"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "reject"},
|
||||
)
|
||||
|
||||
assert _spend_log_session_id(updated, "litellm_metadata") == "thread-abc-123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_session_id_reject_returns_400_without_session_id():
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue