mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
Merge 250b1a50d2 into 21881c5711
This commit is contained in:
devin-ai-integration[bot]
GitHub
commit
2b66eaf2d8
3 changed files with 1857 additions and 15 deletions
|
|
@ -151,6 +151,27 @@ def _session_id_from_baggage(baggage: str) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _caller_trace_field(data: Mapping[str, object], metadata_variable_name: str, field: str) -> str | None:
|
||||
"""The caller's value for a trace-control field, counted only when it is a
|
||||
usable id: a non-empty string. An explicitly empty/unusable value on the
|
||||
active metadata container still shadows the promoted requester value, but
|
||||
neither ever counts as "the caller supplied this field" on its own, so a
|
||||
numeric session id or an empty string cannot suppress the W3C header
|
||||
fallback or satisfy a missing-session-id policy."""
|
||||
active: Final = data.get(metadata_variable_name)
|
||||
if isinstance(active, Mapping) and field in active:
|
||||
active_map: Final = cast(Mapping[str, object], active) # cast-ok: isinstance above, free-form JSON values
|
||||
active_value: Final = active_map[field]
|
||||
return active_value if isinstance(active_value, str) and active_value else None
|
||||
promoted: Final = metadata_variable_name == "litellm_metadata" and field in LITELLM_TRACE_CONTROL_METADATA_FIELDS
|
||||
requester: Final = data.get("metadata")
|
||||
if not promoted or not isinstance(requester, Mapping):
|
||||
return None
|
||||
requester_map: Final = cast(Mapping[str, object], requester) # cast-ok: isinstance above, free-form JSON values
|
||||
requester_value: Final = requester_map.get(field)
|
||||
return requester_value if isinstance(requester_value, str) and requester_value else None
|
||||
|
||||
|
||||
def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None:
|
||||
"""Only proxy-validated keys are stamped, proven by the unforgeable
|
||||
via_virtual_key marker AND a known non-secret shape: the sha256 hex digest
|
||||
|
|
@ -833,11 +854,22 @@ def apply_missing_session_id_policy(
|
|||
):
|
||||
metadata["session_id"] = body_session_id
|
||||
return
|
||||
if data.get("litellm_session_id") or metadata.get("session_id"):
|
||||
caller_session_id: Final = _caller_trace_field(data, _metadata_variable_name, "session_id")
|
||||
if caller_session_id is not None:
|
||||
# Consumers that read the root field (router fallbacks, spend logs,
|
||||
# sandbox reuse) otherwise see no session and mint a uuid4 per request.
|
||||
if not data.get("litellm_session_id"):
|
||||
data["litellm_session_id"] = caller_session_id # rebind-ok: data is an out-param
|
||||
return
|
||||
if data.get("litellm_session_id"):
|
||||
return
|
||||
match policy:
|
||||
case "generate":
|
||||
session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4())
|
||||
session_id: Final = str(
|
||||
data.get("litellm_trace_id")
|
||||
or _caller_trace_field(data, _metadata_variable_name, "trace_id")
|
||||
or uuid.uuid4()
|
||||
)
|
||||
data["litellm_session_id"] = session_id # rebind-ok: data is an out-param
|
||||
data.setdefault("litellm_trace_id", session_id)
|
||||
metadata["session_id"] = session_id
|
||||
|
|
@ -1584,16 +1616,28 @@ class LiteLLMProxyRequestSetup:
|
|||
# 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
|
||||
# explicit litellm headers nor the Anthropic-metadata path found
|
||||
# anything - but lets a caller's existing traceparent/baggage headers
|
||||
# (from real OTel instrumentation) correlate with litellm's own logs
|
||||
# instead of generating an unrelated trace_id.
|
||||
# explicit litellm headers, the Anthropic-metadata path, nor the
|
||||
# caller's own request metadata set the field to a DIFFERENT usable id
|
||||
# - but lets a caller's existing traceparent/baggage headers (from
|
||||
# real OTel instrumentation) correlate with litellm's own logs instead
|
||||
# of generating an unrelated trace_id.
|
||||
normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)})
|
||||
if "litellm_trace_id" not in data:
|
||||
traceparent: Final = normalized_headers.get("traceparent")
|
||||
if isinstance(traceparent, str):
|
||||
trace_id_from_traceparent: Final = _trace_id_from_traceparent(traceparent)
|
||||
if trace_id_from_traceparent:
|
||||
# The caller's metadata wins over the header fallback unless
|
||||
# both carry the same id: stamping the root field then claims
|
||||
# nothing the caller did not already ask for, and keeps the
|
||||
# W3C-correlated root trace id instead of a generated uuid4.
|
||||
caller_trace_id: Final = _caller_trace_field(
|
||||
cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object
|
||||
_metadata_variable_name,
|
||||
"trace_id",
|
||||
)
|
||||
if trace_id_from_traceparent and (
|
||||
caller_trace_id is None or caller_trace_id == trace_id_from_traceparent
|
||||
):
|
||||
metadata_from_headers["trace_id"] = trace_id_from_traceparent
|
||||
data["litellm_trace_id"] = trace_id_from_traceparent # rebind-ok: data is an out-param
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -1603,7 +1647,14 @@ class LiteLLMProxyRequestSetup:
|
|||
baggage: Final = normalized_headers.get("baggage")
|
||||
if isinstance(baggage, str):
|
||||
session_id_from_baggage: Final = _session_id_from_baggage(baggage)
|
||||
if session_id_from_baggage:
|
||||
caller_session_id: Final = _caller_trace_field(
|
||||
cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object
|
||||
_metadata_variable_name,
|
||||
"session_id",
|
||||
)
|
||||
if session_id_from_baggage and (
|
||||
caller_session_id is None or caller_session_id == session_id_from_baggage
|
||||
):
|
||||
metadata_from_headers["session_id"] = session_id_from_baggage
|
||||
data["litellm_session_id"] = session_id_from_baggage # rebind-ok: data is an out-param
|
||||
verbose_proxy_logger.debug("Extracted session_id from W3C baggage header")
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -3674,6 +3674,106 @@ def test_add_litellm_metadata_from_request_headers_explicit_trace_id_beats_trace
|
|||
assert data["litellm_session_id"] == "explicit-trace-id-value"
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_body_trace_id_beats_traceparent():
|
||||
headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
|
||||
data = {"metadata": {"trace_id": "caller-chosen-trace-id"}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["metadata"]["trace_id"] == "caller-chosen-trace-id"
|
||||
assert "litellm_trace_id" not in data
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_body_session_id_beats_baggage():
|
||||
headers = {"baggage": "session.id=baggage-session-42"}
|
||||
data = {"metadata": {"session_id": "caller-chosen-session-id"}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["metadata"]["session_id"] == "caller-chosen-session-id"
|
||||
assert "litellm_session_id" not in data
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_body_steering_is_per_field():
|
||||
headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=baggage-session-42",
|
||||
}
|
||||
data = {"metadata": {"trace_id": "caller-chosen-trace-id"}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["metadata"]["trace_id"] == "caller-chosen-trace-id"
|
||||
assert data["litellm_session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_litellm_metadata_steering_honoured():
|
||||
headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
|
||||
data = {"litellm_metadata": {"trace_id": "caller-chosen-trace-id"}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
assert data["litellm_metadata"]["trace_id"] == "caller-chosen-trace-id"
|
||||
assert "litellm_trace_id" not in data
|
||||
|
||||
|
||||
@pytest.mark.parametrize("empty_session_id", ["", None])
|
||||
def test_add_litellm_metadata_from_request_headers_empty_body_session_id_falls_back_to_baggage(
|
||||
empty_session_id: str | None,
|
||||
):
|
||||
headers = {"baggage": "session.id=baggage-session-42"}
|
||||
data = {"metadata": {"session_id": empty_session_id}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["litellm_session_id"] == "baggage-session-42"
|
||||
assert data["metadata"]["session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["trace_id", "session_id"])
|
||||
def test_add_litellm_metadata_from_request_headers_promoted_metadata_beats_headers(field: str):
|
||||
headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=baggage-session-42",
|
||||
}
|
||||
data = {"metadata": {field: "caller-chosen"}, "litellm_metadata": {}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
assert field not in data["litellm_metadata"]
|
||||
assert f"litellm_{field}" not in data
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_equal_ids_still_stamp_root_fields():
|
||||
headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=matching-session-42",
|
||||
}
|
||||
data = {
|
||||
"metadata": {"trace_id": "4bf92f3577b34da6a3ce929d0e0e4736", "session_id": "matching-session-42"},
|
||||
}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
assert data["litellm_session_id"] == "matching-session-42"
|
||||
assert data["metadata"]["trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
assert data["metadata"]["session_id"] == "matching-session-42"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("non_string_session_id", [4815162342, True, {"session": "nested"}])
|
||||
def test_add_litellm_metadata_from_request_headers_non_string_body_session_id_falls_back_to_baggage(
|
||||
non_string_session_id: object,
|
||||
):
|
||||
headers = {"baggage": "session.id=header-session-42"}
|
||||
data = {"metadata": {"session_id": non_string_session_id}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["litellm_session_id"] == "header-session-42"
|
||||
assert data["metadata"]["session_id"] == "header-session-42"
|
||||
|
||||
|
||||
def _otel_span_with_trace_id(trace_id: int) -> NonRecordingSpan:
|
||||
return NonRecordingSpan(SpanContext(trace_id=trace_id, span_id=0x00F067AA0BA902B7, is_remote=False))
|
||||
|
||||
|
|
@ -8222,6 +8322,85 @@ async def test_missing_session_id_omit_keeps_client_supplied_session_id():
|
|||
assert _spend_log_session_id(updated) == "client-session-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("path", "client_body"),
|
||||
[
|
||||
("/v1/chat/completions", {"model": "gpt-4o", "messages": [], "metadata": {"session_id": ""}}),
|
||||
("/v1/responses", {"model": "gpt-4o", "input": "hi"}),
|
||||
],
|
||||
)
|
||||
async def test_missing_session_id_reject_accepts_baggage_session_id(path: str, client_body: dict[str, object]):
|
||||
request = _request_for(path)
|
||||
request.headers = {"baggage": "session.id=baggage-session-42"}
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data=client_body,
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "reject"},
|
||||
)
|
||||
|
||||
assert updated["litellm_session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
|
||||
@pytest.mark.parametrize("policy", [None, "reject", "generate"])
|
||||
async def test_promoted_caller_trace_ids_beat_traceparent_and_baggage(path: str, policy: str | None):
|
||||
request = _request_for(path)
|
||||
request.headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=baggage-session-42",
|
||||
}
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={"model": "gpt-4o", "metadata": {"trace_id": "caller-trace", "session_id": "caller-session"}},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": policy} if policy else {},
|
||||
)
|
||||
|
||||
assert updated["litellm_metadata"]["trace_id"] == "caller-trace"
|
||||
assert updated["litellm_metadata"]["session_id"] == "caller-session"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_session_id_reject_ignores_requester_session_id_shadowed_by_empty_litellm_metadata():
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-4o",
|
||||
"input": "hi",
|
||||
"metadata": {"session_id": "caller-session"},
|
||||
"litellm_metadata": {"session_id": ""},
|
||||
},
|
||||
request=_request_for("/v1/responses"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "reject"},
|
||||
)
|
||||
assert exc_info.value.code == "400"
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_empty_litellm_metadata_field_falls_back_to_headers():
|
||||
headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=baggage-session-42",
|
||||
}
|
||||
data = {
|
||||
"metadata": {"trace_id": "caller-trace", "session_id": "caller-session"},
|
||||
"litellm_metadata": {"trace_id": "", "session_id": ""},
|
||||
}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
assert data["litellm_session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"client_body",
|
||||
|
|
@ -8341,6 +8520,91 @@ async def test_missing_session_id_generate_reuses_traceparent_trace_id():
|
|||
assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[
|
||||
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
|
||||
{},
|
||||
],
|
||||
ids=["with_traceparent", "no_traceparent"],
|
||||
)
|
||||
async def test_missing_session_id_generate_reuses_promoted_caller_trace_id(path: str, headers: dict[str, str]):
|
||||
request = _request_for(path)
|
||||
request.headers = headers
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={"model": "gpt-4o", "input": "hi", "metadata": {"trace_id": "caller-trace"}},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "generate"},
|
||||
)
|
||||
|
||||
caller_trace_msg: Final = "generate must derive the session from the caller's un-promoted metadata.trace_id"
|
||||
assert updated["litellm_session_id"] == "caller-trace", caller_trace_msg
|
||||
assert updated["litellm_metadata"]["session_id"] == "caller-trace", caller_trace_msg
|
||||
assert updated["litellm_metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True, (
|
||||
"the derived session id must still be marked generated"
|
||||
)
|
||||
assert _spend_log_session_id(updated, "litellm_metadata") == "caller-trace", (
|
||||
"spend log and callback session ids must agree on the caller trace id"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("policy", ["generate", "reject"])
|
||||
async def test_missing_session_id_policy_promotes_caller_session_to_root_field(policy: str):
|
||||
"""A caller-supplied usable session id satisfies the missing_session_id policies on
|
||||
litellm_metadata routes (where it is not yet the managed metadata field) and must also
|
||||
land on the root ``litellm_session_id`` field: consumers that read the root field
|
||||
(router fallbacks, spend logs, sandbox reuse) otherwise mint a fresh uuid4 per request."""
|
||||
request = _request_for("/v1/responses")
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-4o",
|
||||
"input": "hi",
|
||||
"litellm_trace_id": "root-trace-42",
|
||||
"metadata": {"session_id": "caller-session-42"},
|
||||
},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": policy},
|
||||
)
|
||||
|
||||
assert updated["litellm_trace_id"] == "root-trace-42"
|
||||
assert updated["litellm_session_id"] == "caller-session-42"
|
||||
assert updated["litellm_metadata"]["session_id"] == "caller-session-42"
|
||||
assert SESSION_ID_GENERATED_METADATA_KEY not in updated["litellm_metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_session_id_generate_ignores_non_string_caller_session_id():
|
||||
"""A non-string session id is not a usable session: the generate policy must fall through
|
||||
to generation instead of letting an unusable value strand the root session field."""
|
||||
request = _request_for("/v1/responses")
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-4o",
|
||||
"input": "hi",
|
||||
"litellm_trace_id": "root-trace-42",
|
||||
"metadata": {"session_id": 4815162342},
|
||||
},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "generate"},
|
||||
)
|
||||
|
||||
assert updated["litellm_session_id"] == "root-trace-42"
|
||||
assert updated["litellm_metadata"]["session_id"] == "root-trace-42"
|
||||
assert updated["litellm_metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("policy", ["generate", "reject"])
|
||||
async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue