Anthropic prompt caching (inc. Anthropic on Bedrock)

[heyjohnlim]

Description

Hello, i am testing a port of my agentic app to Pydantic AI, and i would like to request easy support prompt caching ["cache-control"] = {"type": "ephemeral"} for Anthropic models.

I find that my costs are reduced by half when i enable this feature in my current app, and i think lots of people would find it beneficial. At the moment i am caching tool calls and system prompt.

Thanks in advance.

References

https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching

[pydanticai-bot]

PydanticAI Github Bot Found 1 issues similar to this one:

1. "Prompt caching support #138" (90% similar)

[voberoi]

+1. I originated #138 but my most pressing need is to handle prompt caching with Anthropic models.

[ninapool]

@Kludex any idea if this is something that you'd be able to work on?

[Kludex]

At some point, yes.

[eluo28]

is there a reason passing the explicit prompt caching as part of the message is not supported? looking to use pydantic ai vs instructor for my production usecase with bedrock as provider but really need prompt caching for costs and time

[Kludex]

I'll check this soon.

[jnishiyama]

For anyone who needs this ASAP, you can just override the _map_message method in the AnthropicModel. This is what I did to implement prompt caching.

[Kludex]

For anyone who needs this ASAP, you can just override the _map_message method in the AnthropicModel. This is what I did to implement prompt caching.

Can you share how u did it?

[jnishiyama]

For anyone who needs this ASAP, you can just override the _map_message method in the AnthropicModel. This is what I did to implement prompt caching.

Can you share how u did it?

Yes, I am sure there is a better approach, but I needed a quick solution.

Notes:

• The reason for the isintance(..., str) checks is because I simply don't understand in what scenarios this would be the case.
• I also changed the return type, which is why I'm ignoring the type error. It ultimately doesn't matter as messages.create takes a str or Iterable[TextBlockParam] for the system argument anyways.

class AnthropicWithCache(AnthropicModel):
    async def _map_message(  # type: ignore
        self, messages: list[ModelMessage]
    ) -> tuple[list[TextBlockParam], list[MessageParam]]:
        system_prompt, anthropic_messages = await super()._map_message(messages)
        cached_system_prompt = [
            TextBlockParam(
                type="text", text=system_prompt, cache_control={"type": "ephemeral"}
            )
        ]
        for anthropic_message in reversed(anthropic_messages):
            if isinstance(anthropic_message, str):
                continue

            last_message_content = anthropic_message["content"]
            if isinstance(last_message_content, str):
                continue

            *last_message_content_items, last_content_item = last_message_content
            if isinstance(last_content_item, str):
                continue

            elif hasattr(last_content_item, "cache_control"):
                last_content_item = cast(
                    TextBlockParam
                    | ImageBlockParam
                    | ToolUseBlockParam
                    | ToolResultBlockParam
                    | DocumentBlockParam,
                    last_content_item,
                )
                cached_last_content_item = last_content_item
                cached_last_content_item["cache_control"] = CacheControlEphemeralParam(
                    type="ephemeral"
                )
                last_message_content_items.append(cached_last_content_item)
                anthropic_message["content"] = last_message_content_items
                break
        return cached_system_prompt, anthropic_messages

FYI: If you want the usage to include cached tokens you need to override those as well

[mwildehahn]

from anthropic.types import (
    MessageParam,
    TextBlockParam,
)
from pydantic_ai.messages import ModelMessage, ModelRequest, SystemPromptPart
from pydantic_ai.models.anthropic import AnthropicModel


class AnthropicModelWithCache(AnthropicModel):
    async def _map_message(  # type: ignore
        self, messages: list[ModelMessage]
    ) -> tuple[list[TextBlockParam], list[MessageParam]]:
        _, anthropic_messages = await super()._map_message(messages)
        system_prompt: list[TextBlockParam] = []
        is_cached = False
        for message in reversed(messages):
            if isinstance(message, ModelRequest):
                for part in reversed(message.parts):
                    if isinstance(part, SystemPromptPart):
                        if not part.dynamic_ref and not is_cached:
                            block = TextBlockParam(
                                type="text",
                                text=part.content,
                                cache_control={"type": "ephemeral"},
                            )
                            is_cached = True
                        else:
                            block = TextBlockParam(
                                type="text",
                                text=part.content,
                            )

                        system_prompt.append(block)

        system_prompt.reverse()
        return system_prompt, anthropic_messages

I tweaked the above slightly so it avoids caching dynamic system prompts.

[jerry-reevo]

Anyone happen to have a workaround for caching Anthropic tools?

[jnishiyama]

Anyone happen to have a workaround for caching Anthropic tools?

@jerry-reevo You can add something like this method to the above I believe

def _get_tools(
        self, model_request_parameters: ModelRequestParameters
    ) -> list[ToolParam]:
        tools = [
            self._map_tool_definition(r)
            for r in model_request_parameters.function_tools
        ]
        if model_request_parameters.output_tools:
            tools += [
                self._map_tool_definition(r)
                for r in model_request_parameters.output_tools
            ]

        tools[-1]["cache_control"] = CacheControlEphemeralParam(type="ephemeral")
        return tools

[DouweM]

Let's use #138 as the canonical issue for this feature.