Add Bedrock/Anthropic prompt caching

Author: larryhudson

examples/pydantic_ai_examples/cache_processors.py

@@ -0,0 +1,181 @@
+"""Example message history processors for automatic cache point insertion.
+
+This module demonstrates how to use message history processors to automatically
+insert CachePoint objects for prompt caching optimization.
+"""
+
+from typing import Callable
+
+from pydantic_ai.messages import (
+    CachePoint,
+    ModelMessage,
+    ModelRequest,
+    SystemPromptPart,
+    UserPromptPart,
+)
+
+
+def cache_system_prompt_processor(messages: list[ModelMessage]) -> list[ModelMessage]:
+    """Add cache point after the last system prompt.
+
+    This processor finds the last system prompt in the message history and
+    adds a cache point to the beginning of the next user message, effectively
+    caching all system prompts.
+
+    Args:
+        messages: List of model messages to process
+
+    Returns:
+        Modified list of messages with cache points added
+    """
+    result = []
+    last_system_idx = -1
+
+    for i, message in enumerate(messages):
+        if isinstance(message, ModelRequest):
+            for part in message.parts:
+                if isinstance(part, SystemPromptPart):
+                    last_system_idx = i
+        result.append(message)
+
+    # Insert cache point after last system prompt
+    if last_system_idx >= 0 and last_system_idx < len(result) - 1:
+        next_message = result[last_system_idx + 1]
+        if isinstance(next_message, ModelRequest):
+            for part in next_message.parts:
+                if isinstance(part, UserPromptPart) and isinstance(part.content, list):
+                    part.content.insert(0, CachePoint())
+                    break
+                elif isinstance(part, UserPromptPart) and isinstance(part.content, str):
+                    # Convert string content to list and add cache point
+                    part.content = [CachePoint(), part.content]
+                    break
+
+    return result
+
+
+def cache_long_context_processor(
+    min_tokens: int = 1024,
+) -> Callable[[list[ModelMessage]], list[ModelMessage]]:
+    """Add cache points before content that likely exceeds token threshold.
+
+    This is a simplified example that estimates content length. In a real
+    implementation, you would want to use a proper tokenizer for accurate counts.
+
+    Args:
+        min_tokens: Minimum estimated tokens before adding a cache point
+
+    Returns:
+        A processor function that adds cache points for long content
+    """
+
+    def processor(messages: list[ModelMessage]) -> list[ModelMessage]:
+        result = []
+
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                for part in message.parts:
+                    if isinstance(part, UserPromptPart):
+                        if isinstance(part.content, str):
+                            # Simple estimation: ~4 characters per token
+                            if len(part.content) > min_tokens * 4:
+                                part.content = [CachePoint(), part.content]
+                        elif isinstance(part.content, list):
+                            # Look for large text blocks
+                            for i, item in enumerate(part.content):
+                                if isinstance(item, str) and len(item) > min_tokens * 4:
+                                    # Insert cache point before large text
+                                    part.content.insert(i, CachePoint())
+                                    break
+            result.append(message)
+
+        return result
+
+    return processor
+
+
+def cache_document_context_processor(
+    messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    """Add cache points after document content.
+
+    This processor adds cache points after any document or binary content
+    to cache large context documents.
+
+    Args:
+        messages: List of model messages to process
+
+    Returns:
+        Modified list of messages with cache points added
+    """
+    result = []
+
+    for message in messages:
+        if isinstance(message, ModelRequest):
+            for part in message.parts:
+                if isinstance(part, UserPromptPart) and isinstance(part.content, list):
+                    new_content = []
+                    for item in part.content:
+                        new_content.append(item)
+                        # Add cache point after document/binary content
+                        if hasattr(item, 'media_type') or hasattr(item, 'data'):
+                            new_content.append(CachePoint())
+                    part.content = new_content
+        result.append(message)
+
+    return result
+
+
+def cache_conversation_turns_processor(
+    messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    """Add cache points at regular conversation intervals.
+
+    This processor adds cache points every few conversation turns to cache
+    conversational context progressively.
+
+    Args:
+        messages: List of model messages to process
+
+    Returns:
+        Modified list of messages with cache points added
+    """
+    result = []
+    turn_count = 0
+
+    for message in messages:
+        if isinstance(message, ModelRequest):
+            for part in message.parts:
+                if isinstance(part, UserPromptPart):
+                    turn_count += 1
+                    # Add cache point every 3 turns
+                    if turn_count % 3 == 0:
+                        if isinstance(part.content, str):
+                            part.content = [CachePoint(), part.content]
+                        elif isinstance(part.content, list):
+                            part.content.insert(0, CachePoint())
+        result.append(message)
+
+    return result
+
+
+def multi_level_cache_processor(messages: list[ModelMessage]) -> list[ModelMessage]:
+    """Example of multiple cache points for hierarchical caching.
+
+    This processor demonstrates adding multiple cache points at different levels:
+    - After system prompts
+    - After large context
+    - At conversation intervals
+
+    Args:
+        messages: List of model messages to process
+
+    Returns:
+        Modified list of messages with cache points added
+    """
+    # Apply multiple processors in sequence
+    processed = cache_system_prompt_processor(messages)
+    processed = cache_long_context_processor(512)(processed)
+    processed = cache_conversation_turns_processor(processed)
+
+    return processed