Support for prompt caching

[ashwin153]

• I want to send cache_control={"type": "ephemeral"} when using Anthropic models and magentic.Chat.
• I think a potential workaround is to create a CacheControlMessage that subclasses magentic.SystemMessage.
• What do you think of this workaround?
• Is this something you think should be better supported by magentic?

[ashwin153]

Maybe if model_config = pydantic.ConfigDict(extra='allow') was set on Message, I could do magentic.UserMessage("...", cache_control={"type": "ephemeral"}) and then add logic to my ChatModel.complete and ChatModel.acomplete to serialize this extra field in my Anthropic requests?

[ashwin153]

@jackmpcollins My bad I accidentally closed this can you repopen it?

[ashwin153]

import magentic.chat_model.message

# All messages before a cache breakpoint in a chat are eligible for prompt caching by Anthropic.
CACHE_BREAKPOINT = magentic.chat_model.message._RawMessage(
    {
        "cache_control": {"type": "ephemeral"},
        "content": "Remember everything I've told you up to this point.",
        "role": "system",
    },
)

here's my workaround

[jackmpcollins]

Hi @ashwin153 Your approach sounds right to me! I think magentic could have custom Message objects for different providers to support the params that they differ on. So in this case you could make a new message class

from magentic import AssistantMessage
from magentic.chat_model.message import ContentT

class AnthropicAssistantMessage(AssistantMessage[ContentT], Generic[ContentT]):
    cache_control: Literal["ephemeral", ...] | None = None

and you can register how to convert this magentic message into an anthropic message param

from anthropic.types import MessageParam
from magentic.chat_model.anthropic_chat_model import message_to_anthropic_message

@message_to_anthropic_message.register(AnthropicAssistantMessage)
def _(message: AnthropicAssistantMessage[Any]) -> MessageParam:
    # Need to call the implementation of this function for AssistantMessage to reuse it
    # I'm not sure how to do that off the top of my head, so inventing "impl_for"
    # Another option is to copy all the code from the existing implementation again here
    message_param = message_to_anthropic_message.impl_for(AssistantMessage)(message)
    # Add the extra field
    if message.cache_control is not None:
        message_param["cache_control"] = message.cache_control
    return message_param

see:

magentic/src/magentic/chat_model/anthropic_chat_model.py

Lines 163 to 210 in 50c7aee

	@message_to_anthropic_message.register(AssistantMessage)
	def _(message: AssistantMessage[Any]) -> MessageParam:
	if isinstance(message.content, str):
	return {
	"role": AnthropicMessageRole.ASSISTANT.value,
	"content": message.content,
	}
	if isinstance(message.content, FunctionCall):
	return {
	"role": AnthropicMessageRole.ASSISTANT.value,
	"content": [_function_call_to_tool_call_block(message.content)],
	}
	if isinstance(message.content, ParallelFunctionCall):
	return {
	"role": AnthropicMessageRole.ASSISTANT.value,
	"content": [
	_function_call_to_tool_call_block(function_call)
	for function_call in message.content
	],
	}
	if isinstance(message.content, StreamedResponse):
	content_blocks: list[TextBlockParam | ToolUseBlockParam] = []
	for item in message.content:
	if isinstance(item, StreamedStr):
	content_blocks.append({"type": "text", "text": item.to_string()})
	elif isinstance(item, FunctionCall):
	content_blocks.append(_function_call_to_tool_call_block(item))
	return {
	"role": AnthropicMessageRole.ASSISTANT.value,
	"content": content_blocks,
	}
	function_schema = function_schema_for_type(type(message.content))
	return {
	"role": AnthropicMessageRole.ASSISTANT.value,
	"content": [
	{
	"type": "tool_use",
	# Can be random because no result will be inserted back into the chat
	"id": _create_unique_id(),
	"name": function_schema.name,
	"input": json.loads(function_schema.serialize_args(message.content)),
	}
	],
	}

[jackmpcollins]

Actually seeing now that this is in the system section of the Anthropic request, so it needs an update in magentic itself. https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching

You could create an AnthropicSystemMessage, no need to register it since it will just be treated the same as SystemMessage by default, but in AnthropicChatModel the _extract_system_message function would need to be updated to extract multiple messages and format these appropriately.

magentic/src/magentic/chat_model/anthropic_chat_model.py

Line 445 in 50c7aee

system, messages = _extract_system_message(messages)