langchain-ai/deepagents
1
0
Fork 0
mirror of https://github.com/langchain-ai/deepagents.git synced 2026-07-21 17:25:26 -04:00
Code Issues 169 Packages Projects Releases 16 Wiki Activity

feat(code): set prompt_cache_key for OpenAI models (#4632)

Browse Source 
Set `prompt_cache_key` on OpenAI model calls (keyed to the thread ID) so
GPT-5.6+ models get reliable prompt-cache prefix matching.

---

OpenAI's GPT-5.6 and later model families require `prompt_cache_key` to
use the more reliable prompt-prefix matching for both implicit and
explicit caching. `ConfigurableModelMiddleware` now injects the active
thread ID as a top-level `prompt_cache_key` for OpenAI models
(`ls_provider == "openai"`), mirroring the existing Fireworks
session-affinity path. It's passed top-level (not via
`extra_body`/`extra_headers`), which is how
`langchain_openai.ChatOpenAI` forwards it for both the Chat Completions
and Responses APIs. Any user-supplied key is preserved.

Made by [Open
SWE](https://openswe.vercel.app/agents/d948dfb5-805e-b4cf-d194-9d579c86811e)

---------

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
Mason Daugherty
2026-07-13 02:14:30 -04:00
committed by GitHub
parent 70829c5846
commit 8cf57aca9f
2 changed files with 432 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/code/deepagents_code/configurable_model.py
+108 -9
View File
@@ -12,6 +12,7 @@ import logging
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from urllib.parse import urlsplit
from deepagents._models import ( # noqa: PLC2701
get_model_identifier,
@@ -102,6 +103,48 @@ def _is_fireworks_model(model: object) -> bool:
return _get_ls_provider(model) == "fireworks"
def _is_openai_model(model: object) -> bool:
"""Check whether a resolved model targets OpenAI's official API.
`ChatOpenAI` reports `'openai'` even when configured with a custom base URL,
so provider metadata alone cannot establish support for OpenAI-specific
request fields. Unknown endpoints are treated conservatively as
incompatible.
Returns:
`True` if the model reports `'openai'` and uses `api.openai.com` or an
official regional subdomain.
"""
if _get_ls_provider(model) != "openai":
return False
# Prefer the SDK client's resolved base URL: a default `ChatOpenAI()` leaves
# `openai_api_base` unset, but its `root_client.base_url` still resolves to
# the official `api.openai.com` default. Fall back to the constructor field
# only when the client is unavailable.
client = getattr(model, "root_client", None)
base_url = getattr(client, "base_url", None)
if base_url is None:
base_url = getattr(model, "openai_api_base", None)
if base_url is None:
# The provider is 'openai' yet no endpoint is discoverable. A genuine
# `ChatOpenAI` always exposes one, so this points to an unexpected model
# shape (e.g. an attribute renamed by an upstream upgrade). Skip the
# optimization instead of guessing, and leave a trace so the silent
# regression is diagnosable.
logger.debug("OpenAI model exposes no base URL; skipping prompt_cache_key")
return False
try:
hostname = urlsplit(str(base_url)).hostname
except ValueError:
logger.debug("OpenAI base URL is unparseable; skipping prompt_cache_key")
return False
return hostname == "api.openai.com" or (
hostname is not None and hostname.endswith(".api.openai.com")
)
_ANTHROPIC_ONLY_SETTINGS: set[str] = {"cache_control"}
"""Keys injected by Anthropic-specific middleware (e.g.
`AnthropicPromptCachingMiddleware`) that are not accepted by other providers and
@@ -162,6 +205,50 @@ def _with_fireworks_session_settings(
return {**model_settings, **updated}
def _with_openai_prompt_cache_key(
model: object, model_settings: dict[str, Any], thread_id: str
) -> dict[str, Any] | None:
"""Return model settings with an OpenAI `prompt_cache_key` added if needed.
Setting `prompt_cache_key` lets OpenAI route a conversation to a stable
prompt-cache prefix across turns, giving more reliable cache hits than the
automatic (keyless) matching; it is optional, and requests still cache
without it. It is a supported top-level `ChatOpenAI` invocation setting
forwarded to the OpenAI request payload, so passing it through
`model_settings` reaches the wire unchanged. `prompt_cache_key` is an
optional, additive request field: it sharpens prefix-cache routing on model
families that support it (GPT-5.6 and later) and is otherwise inert, so the
same key is sent to every OpenAI model without a version gate. This mirrors
the Fireworks path (`_with_fireworks_session_settings`), which sets
`prompt_cache_key` the same way and additionally sends an
`x-session-affinity` header.
A user-supplied `prompt_cache_key` is always preserved, whether it was
configured on the model (`model_kwargs`) or supplied for this invocation
(`model_settings`).
Returns:
A new `model_settings` dict with `prompt_cache_key` added, or `None` when
a key is already present on the model or in the settings (nothing to
add).
"""
model_kwargs = getattr(model, "model_kwargs", None)
if model_kwargs is not None and not isinstance(model_kwargs, Mapping):
# A non-mapping `model_kwargs` cannot carry a user-supplied key, so it is
# treated as "no key present" and injection proceeds. Trace the anomaly
# since a real `ChatOpenAI` always exposes a mapping here.
logger.debug(
"Ignoring non-mapping model_kwargs (%s) when checking for a "
"user-supplied prompt_cache_key",
type(model_kwargs).__name__,
)
if "prompt_cache_key" in model_settings or (
isinstance(model_kwargs, Mapping) and "prompt_cache_key" in model_kwargs
):
return None
return {**model_settings, "prompt_cache_key": thread_id}
def _get_context(request: ModelRequest) -> CLIContextSchema | None:
"""Return runtime context when it matches the CLI context shape."""
runtime = request.runtime
@@ -242,17 +329,29 @@ def _build_overrides(
if model_params:
overrides["model_settings"] = {**request.model_settings, **model_params}
# Inject the provider's prompt-cache routing hint from the active thread.
# Only one provider path applies per call; both share the fetch/guard/log
# tail below. `overrides.get` is side-effect-free, so resolving `settings`
# before the provider check is equivalent to doing it inside each branch.
effective_model = new_model if new_model is not None else request.model
if ctx.thread_id and _is_fireworks_model(effective_model):
if ctx.thread_id:
settings = overrides.get("model_settings", request.model_settings)
settings_with_session = _with_fireworks_session_settings(
settings, ctx.thread_id
)
if settings_with_session is not None:
overrides["model_settings"] = settings_with_session
# No thread ID in the message: it is treated as a sensitive session
# identifier. The line's presence alone confirms injection ran.
logger.debug("Injected Fireworks session settings")
if _is_fireworks_model(effective_model):
updated_settings = _with_fireworks_session_settings(settings, ctx.thread_id)
injected = "Fireworks session settings"
elif _is_openai_model(effective_model):
updated_settings = _with_openai_prompt_cache_key(
effective_model, settings, ctx.thread_id
)
injected = "OpenAI prompt_cache_key"
else:
updated_settings = None
injected = ""
if updated_settings is not None:
overrides["model_settings"] = updated_settings
# The thread ID is a sensitive session identifier, so it is kept out
# of the log line; the line firing at all confirms injection ran.
logger.debug("Injected %s", injected)
if not overrides:
return request
libs/code/tests/unit_tests/test_configurable_model.py
+324 -2
View File
@@ -23,6 +23,7 @@ from deepagents_code.configurable_model import (
_get_context,
_is_anthropic_model,
_is_fireworks_model,
_is_openai_model,
)
@@ -32,6 +33,7 @@ def _make_model(name: str) -> MagicMock:
model.model_name = name
model.model_dump.return_value = {"model_name": name}
model._get_ls_params.return_value = {"ls_provider": "openai"}
model.root_client = SimpleNamespace(base_url="https://api.openai.com/v1")
return model
@@ -576,9 +578,11 @@ class TestFireworksSessionSettings:
}
}
def test_non_fireworks_model_unchanged_with_thread_id(self) -> None:
def test_non_fireworks_non_openai_model_unchanged_with_thread_id(self) -> None:
model = _make_model("gemini-3.5-flash")
model._get_ls_params.return_value = {"ls_provider": "google_genai"}
request = _make_request(
_make_model("gpt-5.5"),
model,
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
@@ -719,6 +723,267 @@ class TestFireworksSessionSettings:
assert captured[0].model_settings["extra_headers"] is not original_headers
class TestOpenAIPromptCacheKey:
"""OpenAI model calls receive a `prompt_cache_key` from the thread ID."""
def test_openai_model_gets_prompt_cache_key(self) -> None:
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0].model is request.model
assert captured[0].model_settings == {"prompt_cache_key": "thread-123"}
def test_prompt_cache_key_merged_with_existing_settings(self) -> None:
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(thread_id="thread-123"),
model_settings={"temperature": 0.5},
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0].model_settings == {
"temperature": 0.5,
"prompt_cache_key": "thread-123",
}
def test_existing_prompt_cache_key_not_overwritten(self) -> None:
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(thread_id="thread-123"),
model_settings={"prompt_cache_key": "custom-cache"},
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0] is request
assert captured[0].model_settings == {"prompt_cache_key": "custom-cache"}
def test_model_prompt_cache_key_not_overwritten(self) -> None:
model = _make_model("gpt-5.6")
model.model_kwargs = {"prompt_cache_key": "model-cache"}
request = _make_request(
model,
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0] is request
assert captured[0].model_settings == {}
def test_non_mapping_model_kwargs_still_injects(self) -> None:
"""A non-mapping `model_kwargs` is treated as no key present."""
model = _make_model("gpt-5.6")
model.model_kwargs = ["not", "a", "mapping"]
request = _make_request(
model,
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0].model_settings == {"prompt_cache_key": "thread-123"}
def test_custom_openai_endpoint_skips_prompt_cache_key(self) -> None:
model = _make_model("gpt-5.6")
model.root_client = SimpleNamespace(base_url="https://proxy.example/v1")
request = _make_request(
model,
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0] is request
assert captured[0].model_settings == {}
def test_regional_openai_endpoint_gets_prompt_cache_key(self) -> None:
model = _make_model("gpt-5.6")
model.root_client = SimpleNamespace(base_url="https://eu.api.openai.com/v1")
request = _make_request(
model,
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0].model_settings == {"prompt_cache_key": "thread-123"}
@pytest.mark.parametrize(
"base_url",
[
"https://api.openai.com.example/v1",
"https://eu.api.openai.com.example/v1",
"https://fake-api.openai.com/v1",
],
)
def test_lookalike_openai_endpoint_skips_prompt_cache_key(
self, base_url: str
) -> None:
model = _make_model("gpt-5.6")
model.root_client = SimpleNamespace(base_url=base_url)
request = _make_request(
model,
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0] is request
assert captured[0].model_settings == {}
def test_empty_thread_id_skips_prompt_cache_key(self) -> None:
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(thread_id=""),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0] is request
def test_no_prompt_cache_key_without_thread_id(self) -> None:
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0] is request
def test_openai_swap_gets_prompt_cache_key(self) -> None:
base = _make_model("claude-sonnet-4-6")
base._get_ls_params.return_value = {"ls_provider": "anthropic"}
override = _make_model("gpt-5.6")
request = _make_request(
base,
context=CLIContext(model="openai:gpt-5.6", thread_id="thread-123"),
)
captured: list[ModelRequest] = []
with patch(_PATCH_CREATE, return_value=_make_model_result(override)):
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0].model is override
assert captured[0].model_settings == {"prompt_cache_key": "thread-123"}
def test_swap_to_openai_injects_key_and_strips_cache_control(self) -> None:
"""Anthropic→OpenAI swap injects the key and strips `cache_control`.
The real `/model` mid-thread scenario: a session running
`AnthropicPromptCachingMiddleware` (which sets `cache_control`) switches
to an OpenAI model. Injection and the Anthropic-only strip must both run
in the same pass, leaving only the cache key — otherwise `cache_control`
would reach the OpenAI SDK and raise `TypeError`.
"""
base = _make_model("claude-sonnet-4-6")
base._get_ls_params.return_value = {"ls_provider": "anthropic"}
override = _make_model("gpt-5.6")
request = _make_request(
base,
context=CLIContext(model="openai:gpt-5.6", thread_id="thread-123"),
model_settings={"cache_control": {"type": "ephemeral"}},
)
captured: list[ModelRequest] = []
with patch(_PATCH_CREATE, return_value=_make_model_result(override)):
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0].model is override
assert captured[0].model_settings == {"prompt_cache_key": "thread-123"}
def test_prompt_cache_key_layered_over_model_params(self) -> None:
"""The key is added on top of a `model_params` merge, not instead of it."""
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(
model_params={"temperature": 0.7}, thread_id="thread-123"
),
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert captured[0].model_settings == {
"temperature": 0.7,
"prompt_cache_key": "thread-123",
}
async def test_async_openai_model_gets_prompt_cache_key(self) -> None:
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(thread_id="thread-123"),
)
captured: list[ModelRequest] = []
async def handler(r: ModelRequest) -> ModelResponse[Any]: # noqa: RUF029
captured.append(r)
return _make_response()
await _mw.awrap_model_call(request, handler)
assert captured[0].model_settings == {"prompt_cache_key": "thread-123"}
def test_caller_model_settings_not_mutated(self) -> None:
"""Injection copies the caller's dict instead of mutating in place."""
model_settings = {"temperature": 0.5}
request = _make_request(
_make_model("gpt-5.6"),
context=CLIContext(thread_id="thread-123"),
model_settings=model_settings,
)
captured: list[ModelRequest] = []
_mw.wrap_model_call(
request, lambda r: (captured.append(r), _make_response())[1]
)
assert model_settings == {"temperature": 0.5}
assert captured[0].model_settings is not model_settings
class TestIsFireworksModel:
"""Direct tests for the `_is_fireworks_model` helper."""
@@ -744,6 +1009,63 @@ class TestIsFireworksModel:
assert _is_fireworks_model(model) is False
class TestIsOpenAIModel:
"""Direct tests for the `_is_openai_model` helper."""
def test_returns_true_for_openai(self) -> None:
assert _is_openai_model(_make_model("gpt-5.6")) is True
def test_returns_true_for_official_openai_endpoint(self) -> None:
model = _make_model("gpt-5.6")
model.root_client = SimpleNamespace(base_url="https://api.openai.com/v1")
assert _is_openai_model(model) is True
def test_returns_false_for_custom_openai_endpoint(self) -> None:
model = _make_model("gpt-5.6")
model.root_client = SimpleNamespace(base_url="https://proxy.example/v1")
assert _is_openai_model(model) is False
def test_returns_false_without_endpoint_metadata(self) -> None:
model = MagicMock(spec=BaseChatModel)
model._get_ls_params.return_value = {"ls_provider": "openai"}
assert _is_openai_model(model) is False
def test_falls_back_to_openai_api_base_for_official(self) -> None:
model = _make_model("gpt-5.6")
model.root_client = None
model.openai_api_base = "https://api.openai.com/v1"
assert _is_openai_model(model) is True
def test_falls_back_to_openai_api_base_for_custom(self) -> None:
model = _make_model("gpt-5.6")
model.root_client = None
model.openai_api_base = "https://proxy.example/v1"
assert _is_openai_model(model) is False
def test_returns_false_for_malformed_base_url(self) -> None:
model = _make_model("gpt-5.6")
model.root_client = SimpleNamespace(base_url="http://[::1")
assert _is_openai_model(model) is False
def test_returns_false_for_non_openai(self) -> None:
model = _make_model("accounts/fireworks/models/kimi-k2p7-code")
model._get_ls_params.return_value = {"ls_provider": "fireworks"}
assert _is_openai_model(model) is False
def test_returns_false_for_plain_object(self) -> None:
assert _is_openai_model(object()) is False
def test_returns_false_when_ls_params_returns_none(self) -> None:
model = MagicMock(spec=BaseChatModel)
model._get_ls_params.return_value = None
assert _is_openai_model(model) is False
def test_returns_false_when_ls_provider_not_str(self) -> None:
model = MagicMock(spec=BaseChatModel)
model._get_ls_params.return_value = {"ls_provider": object()}
assert _is_openai_model(model) is False
class TestIsAnthropicModel:
"""Direct tests for the `_is_anthropic_model` helper."""