From 1d00b1dff9a18288dca8152613d944401b3eda11 Mon Sep 17 00:00:00 2001 From: Laura Barcziova Date: Mon, 26 Oct 2026 08:46:13 +0100 Subject: [PATCH] Re-enable Vertex Anthropic prompt caching BeeAI 0.1.84+ defaults allow_prompt_caching to False for every VertexAIChatModel, which silently stripped the cache_control_injection_points the reasoning agent sets. Every request then re-sent the full context as fresh input: cache read/write dropped to zero and Vertex input cost spiked. Re-enable it for Vertex Anthropic by passing allow_prompt_caching=True at construction (the adapter honors an explicit value), keyed off the public CHAT_MODEL string rather than a private attribute. Add a unit test that the flag is requested only for Vertex Anthropic, plus an end-to-end test that the transformed request retains its cache injection points. Assisted-by: Claude Opus 4.8 --- ymir/agents/tests/unit/test_utils.py | 82 +++++++++++++++++++--------- ymir/agents/utils.py | 16 +++--- 2 files changed, 62 insertions(+), 36 deletions(-) diff --git a/ymir/agents/tests/unit/test_utils.py b/ymir/agents/tests/unit/test_utils.py index c51de80d1..997d132d6 100644 --- a/ymir/agents/tests/unit/test_utils.py +++ b/ymir/agents/tests/unit/test_utils.py @@ -6,44 +6,72 @@ @pytest.mark.parametrize( - ( - "chat_model", - "litellm_provider_id", - "beeai_default_cache_setting", - "ymir_expected_cache_setting", - ), + ("chat_model", "expected_caching_kwarg"), [ - ("vertexai:claude-opus-4-6", "vertexai", False, True), - ("vertexai:claude-sonnet-4-6", "vertexai", False, True), - ("vertexai:gemini-2.5-pro", "vertexai", False, False), - ("anthropic:claude-opus-4-6", "anthropic", True, True), + # Vertex Anthropic: we must explicitly request caching at construction, + # because BeeAI 0.1.84+ defaults allow_prompt_caching=False for Vertex. + ("vertexai:claude-opus-4-6", True), + ("vertexai:claude-sonnet-4-6", True), + # Vertex Gemini: not requested (caching is also forced off below). + ("vertexai:gemini-2.5-pro", None), + # Direct Anthropic: not requested — BeeAI enables caching here by default. + ("anthropic:claude-opus-4-6", None), ], ) -def test_get_chat_model_only_enables_prompt_caching_for_vertex_anthropic( +def test_get_chat_model_requests_caching_only_for_vertex_anthropic( monkeypatch: pytest.MonkeyPatch, chat_model: str, - litellm_provider_id: str, - beeai_default_cache_setting: bool, - ymir_expected_cache_setting: bool, + expected_caching_kwarg: bool | None, ) -> None: - """Preserve caching for Vertex Anthropic after BeeAI 0.1.84 disables it - by default for every VertexAIChatModel. - - The mock model simulates the BeeAI default without Vertex credentials. - """ + """get_chat_model must pass allow_prompt_caching=True to ChatModel.from_name + for Vertex Anthropic models and leave every other provider at its BeeAI + default (no explicit override).""" captured: dict = {} - mock_model = SimpleNamespace( - _litellm_provider_id=litellm_provider_id, - allow_prompt_caching=beeai_default_cache_setting, - ) def from_name(*args, **kwargs): captured.update(kwargs) - return mock_model + return SimpleNamespace(allow_prompt_caching=None, tool_choice_support=None) monkeypatch.setenv("CHAT_MODEL", chat_model) monkeypatch.setattr(utils.ChatModel, "from_name", from_name) - assert utils.get_chat_model() is mock_model - assert "allow_prompt_caching" not in captured - assert mock_model.allow_prompt_caching is ymir_expected_cache_setting + utils.get_chat_model() + + assert captured.get("allow_prompt_caching") is expected_caching_kwarg + + +def test_vertex_anthropic_request_retains_cache_injection_points( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """End-to-end guard against the Oct-2026 caching regression: a real Vertex + Claude model built by get_chat_model must forward cache_control_injection_points + through request transformation, and drop them only when caching is disabled. + + This exercises the real VertexAIChatModel adapter — whose 0.1.84+ default of + allow_prompt_caching=False silently strips the points — so it catches the bug + regardless of how the provider id is spelled or where the flag is read. + """ + from beeai_framework.backend.message import UserMessage + from beeai_framework.backend.types import ChatModelInput + + # VertexAIChatModel construction requires project/location; credentials may + # be empty (no network call is made during construction or transformation). + monkeypatch.setenv("CHAT_MODEL", "vertexai:claude-opus-4-6") + monkeypatch.setenv("GOOGLE_VERTEX_PROJECT", "test-project") + monkeypatch.setenv("GOOGLE_VERTEX_LOCATION", "us-east5") + monkeypatch.setenv("GOOGLE_VERTEX_CREDENTIALS", "") + + model = utils.get_chat_model() + assert model.allow_prompt_caching is True + + points = [{"location": "message", "index": 0}] + model._settings["cache_control_injection_points"] = points + + def transform() -> dict: + return model._transform_input(ChatModelInput(messages=[UserMessage("hi")])) + + assert transform().get("cache_control_injection_points") == points + + # Confirm the flag is what controls it: with caching off BeeAI strips them. + model.allow_prompt_caching = False + assert "cache_control_injection_points" not in transform() diff --git a/ymir/agents/utils.py b/ymir/agents/utils.py index d887c8801..71083c75e 100644 --- a/ymir/agents/utils.py +++ b/ymir/agents/utils.py @@ -58,6 +58,12 @@ def get_chat_model() -> ChatModel: settings["drop_params"] = True settings["additional_drop_params"] = ["temperature"] + # BeeAI 0.1.84+ defaults prompt caching off for Vertex models; re-enable it + # for Vertex Anthropic (set at construction so the adapter keeps the value). + model_kwargs: dict[str, Any] = {} + if chat_model.lower().startswith("vertexai") and "claude" in chat_model.lower(): + model_kwargs["allow_prompt_caching"] = True + model = ChatModel.from_name( chat_model, # this the preferred way to set parameters, don't do options=... @@ -73,16 +79,8 @@ def get_chat_model() -> ChatModel: # for transient 429 / rate-limit errors from the provider. settings=settings, allow_parallel_tool_calls=bool(reasoning_effort), + **model_kwargs, ) - if getattr(model, "_litellm_provider_id", None) == "vertexai": - # BeeAI uses "vertexai" as the litellm provider ID, but litellm only - # recognizes "vertex_ai" for adaptive thinking detection (Claude 4.6+). - # Without this, litellm sends thinking.type=enabled which opus-4-8 rejects. - model._litellm_provider_id = "vertex_ai" - if "claude" in chat_model.lower(): - # Preserve caching for Anthropic models accessed through Vertex AI. - # BeeAI 0.1.84 disables it by default for all Vertex models. - model.allow_prompt_caching = True if "gemini" in chat_model: # disable `required` for Gemini models model.tool_choice_support = {"single", "none", "auto"}