diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 0bfa87f2..fda318f6 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -1419,6 +1419,14 @@ def _convert_message_to_dict(message: BaseMessage) -> dict: # if (name := message.name or message.additional_kwargs.get("name")) is not None: # message_dict["name"] = name + # Forward an Anthropic prompt-caching breakpoint placed at the message level. + # LangChain stores a message-level `cache_control` in `additional_kwargs`. Without + # forwarding it here, the breakpoint reaches the endpoint only when it is nested + # inside a typed content block, so prompt caching silently never activates for the + # other natural placements. + if (cache_control := message.additional_kwargs.get("cache_control")) is not None: + message_dict["cache_control"] = cache_control + if isinstance(message, ChatMessage): return {"role": message.role, **message_dict} elif isinstance(message, HumanMessage): diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index d5f6f614..5bca8b9b 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -713,6 +713,23 @@ def test_convert_message_not_propagate_id() -> None: assert "id" not in result +def test_convert_message_forwards_message_level_cache_control() -> None: + # A message-level cache_control breakpoint (stored by LangChain in + # additional_kwargs) must be forwarded to the request payload, so Anthropic + # prompt caching activates without nesting the flag inside a content block. + cache_control = {"type": "ephemeral"} + message = SystemMessage(content="foo", additional_kwargs={"cache_control": cache_control}) + result = _convert_message_to_dict(message) + assert result == {"role": "system", "content": "foo", "cache_control": cache_control} + + +def test_convert_message_omits_cache_control_when_absent() -> None: + # No cache_control key should appear when the message does not carry one. + message = SystemMessage(content="foo") + result = _convert_message_to_dict(message) + assert "cache_control" not in result + + def test_convert_message_with_tool_calls() -> None: ID = "call_fb5f5e1a-bac0-4422-95e9-d06e6022ad12" tool_calls = [