mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
fix(openai): carry cache_control to prompt_cache_breakpoint on chat completions to gpt-5.6 (#46101)
* fix(openai): carry cache_control to prompt_cache_breakpoint on chat completions to gpt-5.6 A cache_control mark on a system, developer, user or tool message sent to /v1/chat/completions now reaches an OpenAI deployment that reads prompt_cache_breakpoint (gpt-5.6+ on api.openai.com, or one opted in through prompt_cache_options) as a prompt_cache_breakpoint on the same content part, instead of being stripped. Marks on assistant turns and tool definitions are still dropped, and deployments that do not read breakpoints keep today's behavior. * test(openai): rewrite chat prompt_cache_breakpoint wire test as sent-vs-received literals Each test now shows the request the caller sends and the exact body OpenAI receives as plain JSON, with one helper that returns the upstream body. Drops the per-test marker plumbing and the spend-log check, which did not test the cache mark translation. * test(openai): cover a marked chat message with null content on the breakpoint path A system, developer, user or tool message that carries cache_control but whose content is neither a string nor a list of parts is handed to the existing strip untouched instead of raising inside the transform. The new case pins that the request still goes out with the mark stripped and no breakpoint added. * refactor(openai): reuse the cache control hook's breakpoint placement on chat completions The chat transform had its own copy of how a message-level cache_control lands on a part for OpenAI's prompt_cache_breakpoint. It now calls one hook helper that the injection path also uses, so which roles and part types can carry a breakpoint is decided in one place. A message-level mark on list content now lands on the last part OpenAI accepts a breakpoint on, matching injection. * test(e2e): cover cache_control reaching gpt-5.6 explicit prompt caching on chat completions Four live cases through the proxy: a marked system part, a marked string system message, a marked user part and a marked tool result. Each must write the cache on the first call and read it on the second. * fix(openai): parse content parts instead of a banned TypeGuard when renaming cache marks The strict ruff gate bans typing.TypeGuard as unverified narrowing. The part rename now validates an accepted part with pydantic before handing it to the shared rename helper.
This commit is contained in:
ryan-crabbe-berri
GitHub
parent
b77108a409
commit
365c4885de
6 changed files with 651 additions and 30 deletions
|
|
@ -23,6 +23,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import PromptManagementClient
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
cache_control_as_prompt_cache_breakpoint,
|
||||
with_prompt_cache_breakpoint,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
|
|
@ -72,6 +73,7 @@ OPENAI_API_HOST: Final = "api.openai.com"
|
|||
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
|
||||
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
_CONTENT_PART_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
|
||||
|
||||
|
|
@ -90,6 +92,12 @@ def _validated_object_list(value: object) -> list[object] | None:
|
|||
return None
|
||||
|
||||
|
||||
def _part_with_cache_control_renamed(part: object) -> object:
|
||||
if not _accepts_prompt_cache_breakpoint(part):
|
||||
return part
|
||||
return cache_control_as_prompt_cache_breakpoint(_CONTENT_PART_ADAPTER.validate_python(part))
|
||||
|
||||
|
||||
def configured_injection_points(value: object) -> Sequence[CacheControlInjectionPoint]:
|
||||
if not isinstance(value, (list, tuple)):
|
||||
return ()
|
||||
|
|
@ -242,6 +250,10 @@ def _accepts_prompt_cache_breakpoint(block: object) -> bool:
|
|||
return isinstance(block, dict) and block.get("type") in OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES
|
||||
|
||||
|
||||
def _may_carry_prompt_cache_breakpoint(message: Mapping[str, object]) -> bool:
|
||||
return message.get("role") != "assistant"
|
||||
|
||||
|
||||
# Set by a caller whose message list is not the one that goes upstream -- today the
|
||||
# Responses API layer, whose `instructions` only becomes a system message further down.
|
||||
# Tells this hook to hand role-targeted points to the pass holding the final messages
|
||||
|
|
@ -620,9 +632,32 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
|
||||
return message
|
||||
|
||||
@staticmethod
|
||||
def cache_control_as_prompt_cache_breakpoints(message: AllMessageValues) -> AllMessageValues:
|
||||
"""Every Anthropic ``cache_control`` on the message or its parts moved to OpenAI's ``prompt_cache_breakpoint``,
|
||||
placed where the injection hook would put one. Marks with no OpenAI slot are left for the caller to strip."""
|
||||
fields: Final[Mapping[str, object]] = dict(message)
|
||||
if not _may_carry_prompt_cache_breakpoint(fields):
|
||||
return message
|
||||
parts: Final = _validated_object_list(fields.get("content"))
|
||||
renamed: Final = (
|
||||
message
|
||||
if parts is None
|
||||
else cast( # cast-ok: the same message with its part marks renamed
|
||||
AllMessageValues,
|
||||
{
|
||||
**fields,
|
||||
"content": [_part_with_cache_control_renamed(part) for part in parts],
|
||||
},
|
||||
)
|
||||
)
|
||||
if fields.get("cache_control") is None:
|
||||
return renamed
|
||||
return AnthropicCacheControlHook._insert_prompt_cache_breakpoint_in_message(renamed)
|
||||
|
||||
@staticmethod
|
||||
def _insert_prompt_cache_breakpoint_in_message(message: AllMessageValues) -> AllMessageValues:
|
||||
if message.get("role") == "assistant":
|
||||
if not _may_carry_prompt_cache_breakpoint(message):
|
||||
return message
|
||||
message_content: Final = message.get("content", None)
|
||||
if isinstance(message_content, str):
|
||||
|
|
|
|||
|
|
@ -4,12 +4,13 @@ Support for gpt model family
|
|||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS
|
||||
|
|
@ -72,6 +73,7 @@ else:
|
|||
|
||||
|
||||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_OBJECT_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -> str | None:
|
||||
|
|
@ -79,6 +81,14 @@ def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) -
|
|||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def _prompt_cache_options(optional_params: Mapping[str, object]) -> object | None:
|
||||
try:
|
||||
extra_body: Final = _OBJECT_MAPPING.validate_python(optional_params.get("extra_body"))
|
||||
except ValidationError:
|
||||
return None
|
||||
return extra_body.get("prompt_cache_options")
|
||||
|
||||
|
||||
class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
||||
"""
|
||||
Reference: https://platform.openai.com/docs/api-reference/chat/create
|
||||
|
|
@ -452,6 +462,38 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
custom_llm_provider, api_base
|
||||
)
|
||||
|
||||
def _messages_and_tools_in_cache_mark_dialect(
|
||||
self,
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
tools: Sequence["ChatCompletionToolParam"] | None,
|
||||
prompt_cache_options: object,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[Sequence[AllMessageValues], Sequence["ChatCompletionToolParam"] | None]:
|
||||
"""
|
||||
A deployment that reads OpenAI's ``prompt_cache_breakpoint`` (gpt-5.6+ on
|
||||
api.openai.com, or one opted in through ``prompt_cache_options`` or the cost
|
||||
map) gets each Anthropic ``cache_control`` mark rewritten in place. Any other
|
||||
endpoint keeps or strips the marks exactly as before.
|
||||
"""
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
|
||||
custom_llm_provider: Final = _litellm_params_str(litellm_params, "custom_llm_provider")
|
||||
api_base: Final = _litellm_params_str(litellm_params, "api_base")
|
||||
reads_breakpoints: Final = AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(
|
||||
model, custom_llm_provider, api_base, prompt_cache_options
|
||||
)
|
||||
if not reads_breakpoints and self._should_preserve_cache_control_for_endpoint(custom_llm_provider, api_base):
|
||||
return messages, tools
|
||||
renamed: Final = (
|
||||
[AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(message) for message in messages]
|
||||
if reads_breakpoints
|
||||
else list(messages)
|
||||
)
|
||||
return self.remove_cache_control_flag_from_messages_and_tools(
|
||||
model=model, messages=renamed, tools=list(tools) if tools is not None else None
|
||||
)
|
||||
|
||||
def _sanitized_tools_update_for_openai(
|
||||
self,
|
||||
optional_params: Mapping[str, object],
|
||||
|
|
@ -504,17 +546,17 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
Returns:
|
||||
dict: The transformed request. Sent as the body of the API call.
|
||||
"""
|
||||
messages = self._transform_messages(
|
||||
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model
|
||||
transformed_messages, tools = self._messages_and_tools_in_cache_mark_dialect(
|
||||
model=model,
|
||||
messages=self._transform_messages(
|
||||
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model
|
||||
),
|
||||
tools=optional_params.get("tools", []),
|
||||
prompt_cache_options=_prompt_cache_options(optional_params),
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
if not self._should_preserve_cache_control_for_endpoint(
|
||||
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
|
||||
):
|
||||
messages, tools = self.remove_cache_control_flag_from_messages_and_tools(
|
||||
model=model, messages=messages, tools=optional_params.get("tools", [])
|
||||
)
|
||||
if tools is not None and len(tools) > 0:
|
||||
optional_params["tools"] = tools
|
||||
if tools is not None and len(tools) > 0:
|
||||
optional_params["tools"] = tools
|
||||
|
||||
optional_params.pop("max_retries", None)
|
||||
if not optional_params.get("tools") and not optional_params.get("functions"):
|
||||
|
|
@ -522,7 +564,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
|
||||
return {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"messages": transformed_messages,
|
||||
**optional_params,
|
||||
**self._sanitized_tools_update_for_openai(optional_params, litellm_params),
|
||||
}
|
||||
|
|
@ -535,22 +577,17 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
transformed_messages = await self._transform_messages(
|
||||
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model, is_async=True
|
||||
transformed_messages, tools = self._messages_and_tools_in_cache_mark_dialect(
|
||||
model=model,
|
||||
messages=await self._transform_messages(
|
||||
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model, is_async=True
|
||||
),
|
||||
tools=optional_params.get("tools", []),
|
||||
prompt_cache_options=_prompt_cache_options(optional_params),
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
if not self._should_preserve_cache_control_for_endpoint(
|
||||
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
|
||||
):
|
||||
(
|
||||
transformed_messages,
|
||||
tools,
|
||||
) = self.remove_cache_control_flag_from_messages_and_tools(
|
||||
model=model,
|
||||
messages=transformed_messages,
|
||||
tools=optional_params.get("tools", []),
|
||||
)
|
||||
if tools is not None and len(tools) > 0:
|
||||
optional_params["tools"] = tools
|
||||
if tools is not None and len(tools) > 0:
|
||||
optional_params["tools"] = tools
|
||||
if self.__class__._is_base_class:
|
||||
if not optional_params.get("tools") and not optional_params.get("functions"):
|
||||
optional_params.pop("tool_choice", None)
|
||||
|
|
|
|||
|
|
@ -26,6 +26,11 @@ Each case asserts the feature actually happened, not just a 200. Coverage matrix
|
|||
cacheable prefix goes out as a plain system string with a ``prompt_cache_key``
|
||||
and the second call must report ``prompt_tokens_details.cached_tokens`` > 0.
|
||||
|
||||
- OpenAI (gpt-5.6) in explicit cache mode: OpenAI caches nothing unless a
|
||||
breakpoint reaches it, so a ``cache_control`` mark on a system part, a string
|
||||
system message, a user part or a tool result must write the cache on the
|
||||
first call and read it on the second.
|
||||
|
||||
service_tier lives in test_provider_features_e2e.py.
|
||||
|
||||
The provider-native cache_control request shape is not expressible with the
|
||||
|
|
@ -46,7 +51,24 @@ from e2e_config import unique_marker
|
|||
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
|
||||
from e2e_http import Result, UnknownApiError, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from models import CacheControl, ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, RichMessage, TextBlock, Usage
|
||||
from models import (
|
||||
CacheControl,
|
||||
ChatAssistantTurn,
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatTool,
|
||||
ChatToolFunction,
|
||||
ChatToolResultTurn,
|
||||
LiteLLMParamsBody,
|
||||
PromptCacheOptions,
|
||||
RichMessage,
|
||||
TextBlock,
|
||||
TextContentPart,
|
||||
ToolCall,
|
||||
ToolCallFunction,
|
||||
Usage,
|
||||
)
|
||||
from passthrough_client import PassthroughClient
|
||||
import os
|
||||
|
||||
|
|
@ -157,6 +179,30 @@ def _assert_cache_read_on_second_call(
|
|||
)
|
||||
|
||||
|
||||
def _cache_write_tokens(usage: Usage | None) -> int:
|
||||
if usage is None or usage.prompt_tokens_details is None:
|
||||
return 0
|
||||
return usage.prompt_tokens_details.cache_write_tokens or 0
|
||||
|
||||
|
||||
def _assert_openai_writes_then_reads_cache(client: PassthroughClient, key: str, body: ChatBody) -> None:
|
||||
first: Final = unwrap(client.proxy.chat(key, body))
|
||||
assert _cache_write_tokens(first.usage) > 0, (
|
||||
f"{body.model}: the first call on a never-seen prefix wrote no cache ({first.usage}), "
|
||||
"so the cache_control mark never reached OpenAI as a breakpoint"
|
||||
)
|
||||
|
||||
deadline: Final = time.monotonic() + 30.0
|
||||
while True:
|
||||
second = unwrap(client.proxy.chat(key, body))
|
||||
read_tokens = _cached_read_tokens(second.usage)
|
||||
if read_tokens > 0 or time.monotonic() >= deadline:
|
||||
break
|
||||
time.sleep(3.0)
|
||||
|
||||
assert read_tokens > 0, f"{body.model}: the identical second call read no cache ({second.usage})"
|
||||
|
||||
|
||||
def _cold_cache_call(send: Callable[[str], Result[ChatResponse]]) -> ChatResponse | None:
|
||||
result: Final = send(_cacheable_prefix())
|
||||
match result:
|
||||
|
|
@ -412,3 +458,96 @@ class TestCacheControl:
|
|||
_assert_cache_read_on_second_call(
|
||||
model, lambda prefix: _plain_cache_chat(client, key, model, prefix, cache_key)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.covers("llm.chat_completions.openai.prompt_cache_5m.nonstream.works")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.OPENAI,),
|
||||
models=(OPENAI_MODEL,),
|
||||
capabilities=(Capability.PROMPT_CACHING,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
class TestCacheControlOnOpenAIExplicitCaching:
|
||||
@pytest.fixture
|
||||
def model(self, client: PassthroughClient, resources: ResourceManager) -> str:
|
||||
name: Final = f"e2e-openai-explicit-cache-{unique_marker()}"
|
||||
model_id: Final = client.proxy.create_model(
|
||||
name,
|
||||
LiteLLMParamsBody(
|
||||
model=OPENAI_MODEL,
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
prompt_cache_options=PromptCacheOptions(mode="explicit"),
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
return name
|
||||
|
||||
def test_cache_control_on_a_system_text_part(self, client: PassthroughClient, scoped_key: str, model: str) -> None:
|
||||
body: Final = ChatBody(
|
||||
model=model,
|
||||
max_tokens=64,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="system",
|
||||
content=[TextContentPart(text=_cacheable_prefix(), cache_control=CacheControl())],
|
||||
),
|
||||
ChatMessage(role="user", content="Reply with one word."),
|
||||
],
|
||||
)
|
||||
_assert_openai_writes_then_reads_cache(client, scoped_key, body)
|
||||
|
||||
def test_cache_control_on_a_string_system_message(
|
||||
self, client: PassthroughClient, scoped_key: str, model: str
|
||||
) -> None:
|
||||
body: Final = ChatBody(
|
||||
model=model,
|
||||
max_tokens=64,
|
||||
messages=[
|
||||
ChatMessage(role="system", content=_cacheable_prefix(), cache_control=CacheControl()),
|
||||
ChatMessage(role="user", content="Reply with one word."),
|
||||
],
|
||||
)
|
||||
_assert_openai_writes_then_reads_cache(client, scoped_key, body)
|
||||
|
||||
def test_cache_control_on_a_user_text_part(self, client: PassthroughClient, scoped_key: str, model: str) -> None:
|
||||
body: Final = ChatBody(
|
||||
model=model,
|
||||
max_tokens=64,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
content=[
|
||||
TextContentPart(text=_cacheable_prefix(), cache_control=CacheControl()),
|
||||
TextContentPart(text="Reply with one word."),
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
_assert_openai_writes_then_reads_cache(client, scoped_key, body)
|
||||
|
||||
def test_cache_control_on_a_tool_result(self, client: PassthroughClient, scoped_key: str, model: str) -> None:
|
||||
body: Final = ChatBody(
|
||||
model=model,
|
||||
max_tokens=64,
|
||||
reasoning_effort="none",
|
||||
tools=[ChatTool(function=ChatToolFunction(name="fetch_report", parameters={"type": "object"}))],
|
||||
messages=[
|
||||
ChatMessage(role="user", content="Fetch the report and summarize it in one sentence."),
|
||||
ChatAssistantTurn(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id="call_1", type="function", function=ToolCallFunction(name="fetch_report", arguments="{}")
|
||||
)
|
||||
]
|
||||
),
|
||||
ChatToolResultTurn(
|
||||
tool_call_id="call_1",
|
||||
content=[TextContentPart(text=_cacheable_prefix(), cache_control=CacheControl())],
|
||||
),
|
||||
],
|
||||
)
|
||||
_assert_openai_writes_then_reads_cache(client, scoped_key, body)
|
||||
|
|
|
|||
|
|
@ -233,6 +233,7 @@ ContentPart = TextContentPart | ImageContentPart
|
|||
class ChatMessage(BaseModel):
|
||||
role: str
|
||||
content: str | list[ContentPart]
|
||||
cache_control: CacheControl | None = None
|
||||
|
||||
|
||||
class CacheControl(BaseModel):
|
||||
|
|
@ -337,7 +338,7 @@ class ChatAssistantTurn(BaseModel):
|
|||
class ChatToolResultTurn(BaseModel):
|
||||
role: Literal["tool"] = "tool"
|
||||
tool_call_id: str
|
||||
content: str
|
||||
content: str | list[ContentPart]
|
||||
|
||||
|
||||
type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn
|
||||
|
|
@ -465,6 +466,7 @@ class ChatChoice(BaseModel):
|
|||
|
||||
class PromptTokensDetails(BaseModel):
|
||||
cached_tokens: int | None = None
|
||||
cache_write_tokens: int | None = None
|
||||
|
||||
|
||||
class CompletionTokensDetails(BaseModel):
|
||||
|
|
@ -1344,6 +1346,10 @@ class CacheControlInjectionPoint(BaseModel):
|
|||
index: int | None = None
|
||||
|
||||
|
||||
class PromptCacheOptions(BaseModel):
|
||||
mode: Literal["implicit", "explicit"]
|
||||
|
||||
|
||||
class LiteLLMParamsBody(BaseModel):
|
||||
"""POST /model/new litellm_params: `model` is the only required field; `api_key`
|
||||
et al may be an `os.environ/FOO` reference the proxy resolves at call time.
|
||||
|
|
@ -1408,6 +1414,7 @@ class LiteLLMParamsBody(BaseModel):
|
|||
cooldown_time: float | None = None
|
||||
extra_body: DeploymentExtraBody | None = None
|
||||
cache_control_injection_points: list[CacheControlInjectionPoint] | None = None
|
||||
prompt_cache_options: PromptCacheOptions | None = None
|
||||
tpm: int | None = None
|
||||
weight: int | None = None
|
||||
order: int | None = None
|
||||
|
|
|
|||
|
|
@ -0,0 +1,204 @@
|
|||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support.client import Gateway, gateway_from_environment, string_value
|
||||
from integration._support.openai_wire import answering_model_discovery, chat_reply
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(120)
|
||||
|
||||
_SYSTEM_AND_USER_WITH_CACHE_CONTROL: Final[list[JsonValue]] = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "Answer briefly"},
|
||||
{"type": "text", "text": "Stable rules", "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this", "cache_control": {"type": "ephemeral", "ttl": "1h"}},
|
||||
{"type": "text", "text": "and nothing else"},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _reply_ok(request: Request) -> Reply:
|
||||
body: Final = rv.JSON_OBJECT.validate_json(request.body)
|
||||
return chat_reply("chatcmpl-1", string_value(body["model"]), "ok", stream=body.get("stream") is True)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Proxy:
|
||||
gateway: Gateway
|
||||
openai: Wire
|
||||
gpt_5_6_with_prompt_cache_options: str
|
||||
gpt_5_6_without_prompt_cache_options: str
|
||||
|
||||
def openai_receives(
|
||||
self, model: str, messages: Sequence[JsonValue], *, stream: bool = False, tools: Sequence[JsonValue] = ()
|
||||
) -> dict[str, JsonValue]:
|
||||
_ = self.openai.drain()
|
||||
request: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"messages": list(messages),
|
||||
"stream": stream,
|
||||
"cache": {"no-cache": True},
|
||||
}
|
||||
response: Final = self.gateway.request(
|
||||
"POST", "/v1/chat/completions", {**request, "tools": list(tools)} if tools else request
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return self.received_by_openai()
|
||||
|
||||
def received_by_openai(self) -> dict[str, JsonValue]:
|
||||
posts: Final = tuple(request for request in self.openai.drain() if request.method == "POST")
|
||||
assert [post.target for post in posts] == ["/v1/chat/completions"], posts
|
||||
return rv.JSON_OBJECT.validate_json(posts[0].body)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def proxy() -> Iterator[_Proxy]:
|
||||
with (
|
||||
wire_server(answering_model_discovery(_reply_ok)) as openai_upstream,
|
||||
gateway_from_environment() as gateway,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
api_base: Final = f"{openai_upstream.url}/v1"
|
||||
yield _Proxy(
|
||||
gateway,
|
||||
openai_upstream,
|
||||
gpt_5_6_with_prompt_cache_options=scenario.model(
|
||||
model="openai/gpt-5.6", api_base=api_base, prompt_cache_options={"mode": "explicit"}
|
||||
),
|
||||
gpt_5_6_without_prompt_cache_options=scenario.model(model="openai/gpt-5.6", api_base=api_base),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"])
|
||||
def test_cache_control_on_a_part_becomes_prompt_cache_breakpoint_on_that_part(proxy: _Proxy, stream: bool) -> None:
|
||||
received: Final = proxy.openai_receives(
|
||||
proxy.gpt_5_6_with_prompt_cache_options, _SYSTEM_AND_USER_WITH_CACHE_CONTROL, stream=stream
|
||||
)
|
||||
|
||||
assert received["messages"] == [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "Answer briefly"},
|
||||
{"type": "text", "text": "Stable rules", "prompt_cache_breakpoint": {"mode": "explicit"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}},
|
||||
{"type": "text", "text": "and nothing else"},
|
||||
],
|
||||
},
|
||||
]
|
||||
assert received["prompt_cache_options"] == {"mode": "explicit"}
|
||||
assert received.get("stream", False) is stream
|
||||
assert "cache_control" not in str(received)
|
||||
|
||||
|
||||
def test_cache_control_on_a_string_message_sent_with_the_openai_sdk(proxy: _Proxy) -> None:
|
||||
with openai.OpenAI(
|
||||
api_key=proxy.gateway.key, base_url=f"{str(proxy.gateway.client.base_url).rstrip('/')}/v1", max_retries=0
|
||||
) as sdk:
|
||||
completion: Final = sdk.chat.completions.create(
|
||||
model=proxy.gpt_5_6_with_prompt_cache_options,
|
||||
messages=[
|
||||
{"role": "system", "content": "Answer briefly", "cache_control": {"type": "ephemeral"}}, # pyright: ignore[reportArgumentType] # Anthropic mark the SDK forwards verbatim
|
||||
{"role": "user", "content": "Say hi"},
|
||||
],
|
||||
extra_body={"cache": {"no-cache": True}},
|
||||
)
|
||||
assert completion.choices[0].message.content == "ok"
|
||||
|
||||
assert proxy.received_by_openai()["messages"] == [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "Answer briefly", "prompt_cache_breakpoint": {"mode": "explicit"}}],
|
||||
},
|
||||
{"role": "user", "content": "Say hi"},
|
||||
]
|
||||
|
||||
|
||||
def test_cache_control_on_a_tool_result_is_kept_and_on_a_tool_definition_is_dropped(proxy: _Proxy) -> None:
|
||||
received: Final = proxy.openai_receives(
|
||||
proxy.gpt_5_6_with_prompt_cache_options,
|
||||
[
|
||||
{"role": "user", "content": "Look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": [{"type": "text", "text": "found it", "cache_control": {"type": "ephemeral"}}],
|
||||
},
|
||||
],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "parameters": {}},
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
assert received["messages"] == [
|
||||
{"role": "user", "content": "Look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": [{"type": "text", "text": "found it", "prompt_cache_breakpoint": {"mode": "explicit"}}],
|
||||
},
|
||||
], "a breakpoint on a tool message part writes the cache on live gpt-5.6 Chat Completions (2026-10-12)"
|
||||
assert received["tools"] == [{"type": "function", "function": {"name": "lookup", "parameters": {}}}]
|
||||
assert "cache_control" not in str(received)
|
||||
|
||||
|
||||
def test_cache_control_on_an_assistant_turn_is_dropped(proxy: _Proxy) -> None:
|
||||
received: Final = proxy.openai_receives(
|
||||
proxy.gpt_5_6_with_prompt_cache_options,
|
||||
[
|
||||
{"role": "user", "content": "First question"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "Earlier answer", "cache_control": {"type": "ephemeral"}}],
|
||||
},
|
||||
{"role": "user", "content": "Second question"},
|
||||
],
|
||||
)
|
||||
|
||||
assert received["messages"] == [
|
||||
{"role": "user", "content": "First question"},
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "Earlier answer"}]},
|
||||
{"role": "user", "content": "Second question"},
|
||||
]
|
||||
assert "cache_control" not in str(received)
|
||||
assert "prompt_cache_breakpoint" not in str(received)
|
||||
|
||||
|
||||
def test_custom_api_base_without_prompt_cache_options_forwards_cache_control_unchanged(proxy: _Proxy) -> None:
|
||||
received: Final = proxy.openai_receives(
|
||||
proxy.gpt_5_6_without_prompt_cache_options, _SYSTEM_AND_USER_WITH_CACHE_CONTROL
|
||||
)
|
||||
|
||||
assert received["messages"] == _SYSTEM_AND_USER_WITH_CACHE_CONTROL
|
||||
assert "prompt_cache_options" not in received
|
||||
assert "prompt_cache_breakpoint" not in str(received)
|
||||
|
|
@ -940,6 +940,205 @@ class TestCacheControlPreservationForCustomEndpoint:
|
|||
assert all("cache_control" not in m for m in body["messages"])
|
||||
|
||||
|
||||
_EPHEMERAL: Final[dict[str, str]] = {"type": "ephemeral"}
|
||||
_EXPLICIT: Final[dict[str, str]] = {"mode": "explicit"}
|
||||
_LOOKUP_TOOL: Final[dict[str, object]] = {"type": "function", "function": {"name": "lookup", "parameters": {}}}
|
||||
|
||||
|
||||
def _text(value: str) -> dict[str, object]:
|
||||
return {"type": "text", "text": value}
|
||||
|
||||
|
||||
def _marked(part: dict[str, object], mark: dict[str, str] = _EPHEMERAL) -> dict[str, object]:
|
||||
return {**part, "cache_control": mark}
|
||||
|
||||
|
||||
def _breakpoint(part: dict[str, object]) -> dict[str, object]:
|
||||
return {**part, "prompt_cache_breakpoint": _EXPLICIT}
|
||||
|
||||
|
||||
class TestCacheControlAsPromptCacheBreakpoint:
|
||||
"""
|
||||
A deployment that reads OpenAI's prompt_cache_breakpoint (gpt-5.6+ on api.openai.com,
|
||||
or one opted in through prompt_cache_options) gets each Anthropic cache_control mark
|
||||
rewritten on the same part instead of stripped, so the cache the caller asked for is
|
||||
still written. Every other endpoint keeps or strips the marks exactly as before.
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
self.config = OpenAIGPTConfig()
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_openai_base_env(self, monkeypatch):
|
||||
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
|
||||
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
|
||||
monkeypatch.setattr(litellm, "api_base", None, raising=False)
|
||||
|
||||
@staticmethod
|
||||
def _marked_messages() -> list[dict[str, object]]:
|
||||
return [
|
||||
{"role": "system", "content": [_text("Answer briefly"), _marked(_text("Stable rules"))]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
_marked(_text("Describe this"), {"type": "ephemeral", "ttl": "1h"}),
|
||||
_marked({"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}),
|
||||
_text("and nothing else"),
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [_marked(_text("earlier answer"))]},
|
||||
{"role": "user", "content": "and again"},
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _rewritten_messages() -> list[dict[str, object]]:
|
||||
return [
|
||||
{"role": "system", "content": [_text("Answer briefly"), _breakpoint(_text("Stable rules"))]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
_breakpoint(_text("Describe this")),
|
||||
_breakpoint({"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}),
|
||||
_text("and nothing else"),
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [_text("earlier answer")]},
|
||||
{"role": "user", "content": "and again"},
|
||||
]
|
||||
|
||||
def _transform(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, object]],
|
||||
*,
|
||||
custom_llm_provider: str = "openai",
|
||||
api_base: str | None = None,
|
||||
optional_params: dict[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
return self.config.transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params or {},
|
||||
litellm_params={"custom_llm_provider": custom_llm_provider, "api_base": api_base},
|
||||
headers={},
|
||||
)
|
||||
|
||||
def test_gpt_5_6_rewrites_system_and_user_marks_and_drops_the_assistant_mark(self):
|
||||
body = self._transform("gpt-5.6", self._marked_messages())
|
||||
assert body["messages"] == self._rewritten_messages()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_gpt_5_6_rewrites_the_same_way(self):
|
||||
body = await self.config.async_transform_request(
|
||||
model="gpt-5.6",
|
||||
messages=self._marked_messages(),
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "openai", "api_base": None},
|
||||
headers={},
|
||||
)
|
||||
assert body["messages"] == self._rewritten_messages()
|
||||
|
||||
def test_gpt_5_6_turns_a_marked_string_message_into_one_marked_text_part(self):
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "system", "content": "You are helpful.", "cache_control": _EPHEMERAL},
|
||||
{"role": "user", "content": "Hello", "cache_control": {"type": "ephemeral", "ttl": "1h"}},
|
||||
{"role": "user", "content": "unmarked"},
|
||||
]
|
||||
body = self._transform("gpt-5.6", messages)
|
||||
assert body["messages"] == [
|
||||
{"role": "system", "content": [_breakpoint(_text("You are helpful."))]},
|
||||
{"role": "user", "content": [_breakpoint(_text("Hello"))]},
|
||||
{"role": "user", "content": "unmarked"},
|
||||
]
|
||||
|
||||
def test_gpt_5_6_message_level_mark_on_list_content_marks_the_last_part(self):
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "user", "content": [_text("first"), _text("last")], "cache_control": _EPHEMERAL},
|
||||
]
|
||||
body = self._transform("gpt-5.6", messages)
|
||||
assert body["messages"] == [{"role": "user", "content": [_text("first"), _breakpoint(_text("last"))]}]
|
||||
|
||||
def test_gpt_5_6_marked_message_with_null_content_still_goes_out_with_the_mark_stripped(self):
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "user", "content": None, "cache_control": _EPHEMERAL},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
body = self._transform("gpt-5.6", messages)
|
||||
assert body["messages"] == [{"role": "user", "content": None}, {"role": "user", "content": "Hello"}]
|
||||
|
||||
def test_gpt_5_6_keeps_the_mark_on_a_tool_message_and_drops_the_tool_definition_mark(self):
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "user", "content": "look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [_marked(_text("found it"))]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "found it again", "cache_control": _EPHEMERAL},
|
||||
]
|
||||
body = self._transform("gpt-5.6", messages, optional_params={"tools": [_marked(dict(_LOOKUP_TOOL))]})
|
||||
assert body["messages"][2:] == [
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [_breakpoint(_text("found it"))]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [_breakpoint(_text("found it again"))]},
|
||||
], "a breakpoint on a tool message part writes the cache on live gpt-5.6 Chat Completions (2026-10-12)"
|
||||
assert body["tools"] == [_LOOKUP_TOOL]
|
||||
|
||||
def test_gpt_4o_on_openai_still_strips_every_mark(self):
|
||||
body = self._transform(
|
||||
"gpt-4o", self._marked_messages(), optional_params={"tools": [_marked(dict(_LOOKUP_TOOL))]}
|
||||
)
|
||||
assert body["messages"] == [
|
||||
{"role": "system", "content": [_text("Answer briefly"), _text("Stable rules")]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
_text("Describe this"),
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},
|
||||
_text("and nothing else"),
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [_text("earlier answer")]},
|
||||
{"role": "user", "content": "and again"},
|
||||
]
|
||||
assert body["tools"] == [_LOOKUP_TOOL]
|
||||
|
||||
def test_custom_api_base_without_opt_in_keeps_cache_control_untouched(self):
|
||||
body = self._transform("gpt-5.6", self._marked_messages(), api_base="http://localhost:4000/v1")
|
||||
assert body["messages"] == self._marked_messages()
|
||||
|
||||
def test_custom_api_base_opted_in_through_prompt_cache_options_rewrites(self):
|
||||
body = self._transform(
|
||||
"gpt-5.6",
|
||||
self._marked_messages(),
|
||||
api_base="http://localhost:4000/v1",
|
||||
optional_params={"extra_body": {"prompt_cache_options": {"mode": "explicit"}}},
|
||||
)
|
||||
assert body["messages"] == self._rewritten_messages()
|
||||
assert body["extra_body"] == {"prompt_cache_options": {"mode": "explicit"}}
|
||||
|
||||
def test_non_openai_provider_still_strips(self):
|
||||
body = self._transform(
|
||||
"gpt-5.6",
|
||||
self._marked_messages(),
|
||||
custom_llm_provider="fireworks_ai",
|
||||
api_base="https://api.fireworks.ai/inference/v1",
|
||||
)
|
||||
assert body["messages"] == [
|
||||
{"role": "system", "content": [_text("Answer briefly"), _text("Stable rules")]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
_text("Describe this"),
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},
|
||||
_text("and nothing else"),
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [_text("earlier answer")]},
|
||||
{"role": "user", "content": "and again"},
|
||||
]
|
||||
|
||||
|
||||
class TestToolChoiceWithoutToolsDropped:
|
||||
def setup_method(self):
|
||||
self.config = OpenAIGPTConfig()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue