mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(chatgpt): derive the upstream session id from prompt_cache_key (#42014)
* fix(chatgpt): derive Responses session id from prompt_cache_key ChatGPT partitions the Codex prompt cache by the Responses session-id header (the Codex CLI sends a stable per-conversation key derived from prompt_cache_key). The provider only consulted explicit session ids and otherwise stamped a per-request uuid4, so no turn ever reused the prompt cache -- 0% cache hits on repeat turns, an order-of-magnitude quota/cost multiplier on subscription accounts. Resolution order after this change: explicit session config (session_id / litellm_session_id / metadata.session_id) > prompt_cache_key (sanitized with _safe_header_value) > litellm-internal ids > uuid4. Two guards, mirroring get_fireworks_session_id: - proxy-generated session ids (missing_session_id: "generate") are per-request and are skipped, checked in both "metadata" and "litellm_metadata" (LITELLM_METADATA_ROUTES, responses included) - internal per-request ids (litellm_trace_id / litellm_call_id) lose to a stable conversation key Supersedes the body-key forwarding of #37630: a random session header overrides the body key on the provider side, so the header must carry the caller's key. Related: #29978, #29993. Verified against a real ChatGPT Codex account: stable session id -> 97-100% cached on repeat turns; random per-request -> 0%. * style: ruff format * fix(chatgpt): address review - hash unsafe cache keys, move tests, drop comments - hash prompt_cache_keys containing header-unsafe chars instead of mapping them through _safe_header_value, whose replacement char could collapse distinct keys onto one cache shard and serialize unrelated requests (greptile review) - strip the explanatory comment blocks per AGENTS.md comment policy - move the tests into the existing mapped responses transformation test file per AGENTS.md test-placement rule * fix(chatgpt): keep session typing and test diff scoped * fix(chatgpt): hash surrogate cache keys, type session params, flatten tests * test(chatgpt): record upstream requests with a mock instead of appending to a list * fix(chatgpt): keep proxy-mapped session headers ahead of body session_id and ignore non-string cache keys * fix(chatgpt): hash cache keys with leading or trailing whitespace so the session header stays valid * test(chatgpt): type the session id test parameters * fix(chatgpt): keep trace ids copied into litellm_session_id from outranking prompt_cache_key * fix(chatgpt): mark session ids copied from trace_id so explicit ones keep precedence * fix(chatgpt): keep the trace-id session marker out of provider params The marker now uses the _litellm_ internal prefix, so bridges that pass litellm_params back into completion() drop it instead of sending it as a provider param --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Onur C. Cakmak
GitHub
parent
74392efbd4
commit
d9aae590e3
7 changed files with 202 additions and 35 deletions
|
|
@ -1737,6 +1737,7 @@ PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-"
|
|||
TYPESAFE_DANGER_ZONE_HEADER_PREFIX: Final = "typesafe-danger-zone-"
|
||||
INTERNAL_KWARG_PREFIX: Final = "_litellm_"
|
||||
CONTROL_OPTIONS_KEY: Final = f"{INTERNAL_KWARG_PREFIX}control"
|
||||
SESSION_ID_FROM_TRACE_ID_PARAM_KEY: Final = f"{INTERNAL_KWARG_PREFIX}session_id_from_trace_id"
|
||||
|
||||
AZURE_SPEECH_CUSTOM_LLM_PROVIDER: Final = "azure_speech"
|
||||
AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech"
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import Final, cast
|
|||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.constants import CONTROL_OPTIONS_KEY
|
||||
from litellm.constants import CONTROL_OPTIONS_KEY, SESSION_ID_FROM_TRACE_ID_PARAM_KEY
|
||||
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
|
||||
from litellm.llms.openai.data_residency import infer_openai_data_residency
|
||||
from litellm.types.litellm_params import MAX_CONTROL_INT_DIGITS, ControlOptions
|
||||
|
|
@ -212,6 +212,9 @@ def get_litellm_params(
|
|||
|
||||
# Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining)
|
||||
_meta: Final = resolved_metadata or {}
|
||||
session_id_from_trace_id: Final = (
|
||||
litellm_session_id is None and not _meta.get("session_id") and bool(_meta.get("trace_id"))
|
||||
)
|
||||
if litellm_session_id is None:
|
||||
litellm_session_id = _meta.get("session_id") or _meta.get("trace_id")
|
||||
if litellm_trace_id is None:
|
||||
|
|
@ -254,6 +257,7 @@ def get_litellm_params(
|
|||
or (get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None),
|
||||
"litellm_trace_id": litellm_trace_id,
|
||||
"litellm_session_id": litellm_session_id,
|
||||
**({SESSION_ID_FROM_TRACE_ID_PARAM_KEY: True} if session_id_from_trace_id else {}),
|
||||
"hf_model_name": hf_model_name,
|
||||
"custom_prompt_dict": custom_prompt_dict,
|
||||
"litellm_metadata": litellm_metadata,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import re
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
|
||||
from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY
|
||||
|
||||
_HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$")
|
||||
_FORBIDDEN_AFFINITY_HEADERS: Final = frozenset(
|
||||
|
|
@ -53,36 +53,51 @@ def _get_provider_affinity_header_name(litellm_params: object | None, default_he
|
|||
return validate_provider_affinity_header_name(header)
|
||||
|
||||
|
||||
def get_stable_session_id(litellm_params: object | None) -> str | None:
|
||||
if litellm_params is None:
|
||||
return None
|
||||
|
||||
direct_session_id: Final = _get_value(litellm_params, "session_id")
|
||||
if direct_session_id:
|
||||
return str(direct_session_id)
|
||||
|
||||
metadata_values: Final[tuple[object, ...]] = tuple(
|
||||
def _metadata_values(litellm_params: object) -> tuple[object, ...]:
|
||||
return tuple(
|
||||
value for key in ("metadata", "litellm_metadata") if (value := _get_value(litellm_params, key)) is not None
|
||||
)
|
||||
|
||||
|
||||
def _non_generated_litellm_session_id(litellm_params: object) -> str | None:
|
||||
has_generated_session_id: Final = any(
|
||||
isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY)
|
||||
for metadata in metadata_values
|
||||
_get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) for metadata in _metadata_values(litellm_params)
|
||||
)
|
||||
|
||||
litellm_session_id: Final = _get_value(litellm_params, "litellm_session_id")
|
||||
if litellm_session_id and not has_generated_session_id:
|
||||
return str(litellm_session_id)
|
||||
return str(litellm_session_id) if litellm_session_id and not has_generated_session_id else None
|
||||
|
||||
for metadata in metadata_values:
|
||||
if (
|
||||
isinstance(metadata, Mapping)
|
||||
and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY)
|
||||
and (value := metadata.get("session_id"))
|
||||
|
||||
def _metadata_session_id(litellm_params: object) -> str | None:
|
||||
for metadata in _metadata_values(litellm_params):
|
||||
if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) and (
|
||||
value := _get_value(metadata, "session_id")
|
||||
):
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
|
||||
def get_caller_litellm_session_id(litellm_params: object | None) -> str | None:
|
||||
if litellm_params is None or _get_value(litellm_params, SESSION_ID_FROM_TRACE_ID_PARAM_KEY):
|
||||
return None
|
||||
return _non_generated_litellm_session_id(litellm_params)
|
||||
|
||||
|
||||
def get_requested_session_id(litellm_params: object | None) -> str | None:
|
||||
if litellm_params is None:
|
||||
return None
|
||||
direct_session_id: Final = _get_value(litellm_params, "session_id")
|
||||
return str(direct_session_id) if direct_session_id else _metadata_session_id(litellm_params)
|
||||
|
||||
|
||||
def get_stable_session_id(litellm_params: object | None) -> str | None:
|
||||
if litellm_params is None:
|
||||
return None
|
||||
direct_session_id: Final = _get_value(litellm_params, "session_id")
|
||||
if direct_session_id:
|
||||
return str(direct_session_id)
|
||||
return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params)
|
||||
|
||||
|
||||
def add_provider_affinity_header(
|
||||
headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None
|
||||
) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Constants and helpers for ChatGPT subscription OAuth.
|
|||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
|
|
@ -17,6 +18,7 @@ import httpx
|
|||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
|
||||
from litellm.litellm_core_utils.provider_affinity import get_caller_litellm_session_id, get_requested_session_id
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.utils import ProviderSpecificHeader
|
||||
|
||||
|
|
@ -395,22 +397,31 @@ def _normalize_litellm_params(litellm_params: object) -> Mapping[str, object]:
|
|||
return MappingProxyType({})
|
||||
|
||||
|
||||
def _prompt_cache_session_id(prompt_cache_key: object) -> str | None:
|
||||
if not isinstance(prompt_cache_key, str) or not prompt_cache_key:
|
||||
return None
|
||||
return (
|
||||
prompt_cache_key
|
||||
if _safe_header_value(prompt_cache_key.strip()) == prompt_cache_key
|
||||
else hashlib.sha256(prompt_cache_key.encode("utf-8", "surrogatepass")).hexdigest()
|
||||
)
|
||||
|
||||
|
||||
def get_chatgpt_session_id(litellm_params: object) -> str | None:
|
||||
params: Final = _normalize_litellm_params(litellm_params)
|
||||
for key in ("litellm_session_id", "session_id"):
|
||||
value = params.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
metadata: Final = _normalize_litellm_params(params.get("metadata"))
|
||||
if metadata:
|
||||
value = metadata.get("session_id")
|
||||
if value:
|
||||
return str(value)
|
||||
for key in ("litellm_trace_id", "litellm_call_id"):
|
||||
value = params.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
return None
|
||||
return (
|
||||
get_caller_litellm_session_id(params)
|
||||
or get_requested_session_id(params)
|
||||
or _prompt_cache_session_id(params.get("prompt_cache_key"))
|
||||
or next(
|
||||
(
|
||||
str(value)
|
||||
for key in ("litellm_session_id", "litellm_trace_id", "litellm_call_id")
|
||||
if (value := params.get(key))
|
||||
),
|
||||
None,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def ensure_chatgpt_session_id(litellm_params: object) -> str:
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.litellm_core_utils.get_litellm_params import (
|
|||
stored_control_options,
|
||||
)
|
||||
from litellm.types.litellm_params import ControlOptions
|
||||
from litellm.utils import get_non_default_completion_params
|
||||
|
||||
|
||||
def _funnel_kwargs_completion_forwards(monkeypatch, wif_kwargs: dict[str, object]) -> dict[str, object]:
|
||||
|
|
@ -411,6 +412,14 @@ class TestMetadataFallsBackToLitellmMetadata:
|
|||
result["metadata"].pop("trace_id")
|
||||
assert litellm_metadata == {"trace_id": "trace-1"}
|
||||
|
||||
def test_session_id_derived_from_trace_id_adds_no_provider_params(self):
|
||||
from_trace_id: Final = get_litellm_params(metadata={"trace_id": "trace-1"})
|
||||
from_session_id: Final = get_litellm_params(metadata={"session_id": "session-1"})
|
||||
added_by_trace_id: Final = from_trace_id.keys() - from_session_id.keys()
|
||||
|
||||
assert added_by_trace_id
|
||||
assert not added_by_trace_id & get_non_default_completion_params(from_trace_id).keys()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, expected",
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.llms.chatgpt.authenticator import Authenticator
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
|
@ -584,3 +585,54 @@ class TestChatGPTResponsesMultiAccount:
|
|||
assert headers_b["ChatGPT-Account-Id"] == "acct-b"
|
||||
assert headers_default["Authorization"] == "Bearer token-default"
|
||||
assert headers_default["ChatGPT-Account-Id"] == "acct-default"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call", ["responses", "completion"])
|
||||
def test_prompt_cache_key_keeps_one_session_header_across_requests(
|
||||
call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
(tmp_path / "auth.json").write_text(
|
||||
json.dumps({"access_token": "token", "account_id": "acct", "expires_at": 4102444800})
|
||||
)
|
||||
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
|
||||
completed: Final = {
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-5.5",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "OK", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
}
|
||||
upstream: Final = MagicMock(
|
||||
side_effect=lambda _request: httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
content=f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode(),
|
||||
)
|
||||
)
|
||||
client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(upstream)))
|
||||
for _ in range(2):
|
||||
if call == "responses":
|
||||
litellm.responses(model="chatgpt/gpt-5.5", input="hi", prompt_cache_key="conv-123", client=client)
|
||||
else:
|
||||
litellm.completion(
|
||||
model="chatgpt/gpt-5.5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
prompt_cache_key="conv-123",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert [call.args[0].headers["session_id"] for call in upstream.call_args_list] == ["conv-123", "conv-123"]
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import h11
|
||||
import pytest
|
||||
|
||||
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.llms.chatgpt.common_utils import (
|
||||
extract_chatgpt_account_id,
|
||||
get_chatgpt_client_credential,
|
||||
get_chatgpt_session_id,
|
||||
is_chatgpt_oauth_key,
|
||||
select_chatgpt_client_credential_headers,
|
||||
)
|
||||
|
|
@ -88,3 +94,72 @@ def test_bearer_scheme_is_case_insensitive(scheme):
|
|||
|
||||
def test_other_authentication_schemes_are_not_chatgpt_logins():
|
||||
assert not is_chatgpt_oauth_key(f"Basic {CHATGPT_JWT}")
|
||||
|
||||
|
||||
class TestGetChatGPTSessionId:
|
||||
@pytest.mark.parametrize(
|
||||
"explicit",
|
||||
[
|
||||
{"litellm_session_id": "chosen"},
|
||||
{"session_id": "chosen"},
|
||||
{"metadata": {"session_id": "chosen"}},
|
||||
get_litellm_params(litellm_session_id="chosen", metadata={"trace_id": "chosen"}),
|
||||
],
|
||||
)
|
||||
def test_explicit_session_id_wins_over_prompt_cache_key(self, explicit: Mapping[str, object]) -> None:
|
||||
assert get_chatgpt_session_id({**explicit, "prompt_cache_key": "conv-123"}) == "chosen"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_scoped",
|
||||
[
|
||||
{"litellm_trace_id": "trace-1"},
|
||||
{"litellm_call_id": "call-1"},
|
||||
get_litellm_params(metadata={"trace_id": "trace-1"}),
|
||||
get_litellm_params(litellm_metadata={"trace_id": "trace-1"}),
|
||||
{"litellm_session_id": "gen-1", "metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True}},
|
||||
{
|
||||
"litellm_session_id": "gen-1",
|
||||
"litellm_metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True},
|
||||
},
|
||||
],
|
||||
)
|
||||
def test_prompt_cache_key_wins_over_request_scoped_ids(self, request_scoped: Mapping[str, object]) -> None:
|
||||
assert get_chatgpt_session_id({**request_scoped, "prompt_cache_key": "conv-123"}) == "conv-123"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mapped_header",
|
||||
[
|
||||
{
|
||||
"litellm_session_id": "header-session",
|
||||
"metadata": {"trace_id": "header-session", "session_id": "header-session"},
|
||||
},
|
||||
get_litellm_params(metadata={"trace_id": "header-session", "session_id": "header-session"}),
|
||||
],
|
||||
)
|
||||
def test_proxy_mapped_session_header_wins_over_body_session_id(self, mapped_header: Mapping[str, object]) -> None:
|
||||
assert (
|
||||
get_chatgpt_session_id({**mapped_header, "session_id": "body-session", "prompt_cache_key": "conv-123"})
|
||||
== "header-session"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("prompt_cache_key", [{"tenant": "a"}, 12345, True, ["x", "y"]])
|
||||
def test_non_string_prompt_cache_key_is_ignored(self, prompt_cache_key: object) -> None:
|
||||
assert get_chatgpt_session_id({"prompt_cache_key": prompt_cache_key, "litellm_call_id": "call-1"}) == "call-1"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"first, second",
|
||||
[
|
||||
("a\nb", "a\tb"),
|
||||
(json.loads('"conv-\\ud800"'), json.loads('"conv-\\udc00"')),
|
||||
("conv-123 ", " conv-123"),
|
||||
],
|
||||
)
|
||||
def test_header_unsafe_cache_keys_map_to_distinct_stable_header_values(self, first: str, second: str) -> None:
|
||||
first_id: Final = get_chatgpt_session_id({"prompt_cache_key": first})
|
||||
second_id: Final = get_chatgpt_session_id({"prompt_cache_key": second})
|
||||
|
||||
assert first_id is not None and second_id is not None
|
||||
assert first_id == get_chatgpt_session_id({"prompt_cache_key": first})
|
||||
assert first_id != second_id
|
||||
for value in (first_id, second_id):
|
||||
h11.Request(method="POST", target="/", headers=[("host", "chatgpt.com"), ("session_id", value)])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue