« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Provider-specific cache annotation strategies for prompt assembly."""
3from __future__ import annotations
5from typing import Any, Protocol, runtime_checkable
7from lexigram.contracts.ai.llm import ChatMessage, TokenCounterProtocol
8from lexigram.logging import (
9 get_logger,
10)
12logger = get_logger(__name__)
14ANTHROPIC_MIN_CACHE_TOKENS = 1024
15ANTHROPIC_HAIKU_MIN_CACHE_TOKENS = 2048
16ANTHROPIC_MAX_BREAKPOINTS = 4
17DEEPSEEK_PADDING_BOUNDARY = 64
18GEMINI_CONTEXT_CACHE_MIN_TOKENS = 32_000
21@runtime_checkable
22class CacheStrategy(Protocol):
23 """Protocol for provider-specific cache annotation strategies."""
25 def annotate(
26 self,
27 messages: list[ChatMessage],
28 static_count: int,
29 token_counter: TokenCounterProtocol | None = None,
30 ) -> list[ChatMessage]:
31 """Annotate messages with provider-specific cache hints.
33 Args:
34 messages: Full assembled message list.
35 static_count: Number of messages that are static (from beginning).
36 token_counter: Optional token counter for size-aware annotations.
38 Returns:
39 Messages with provider-specific cache annotations applied.
40 """
41 ...
44class PassthroughCacheStrategy:
45 """No-op strategy — returns messages unchanged.
47 Used for unknown providers or when caching is not applicable.
48 """
50 def annotate(
51 self,
52 messages: list[ChatMessage],
53 static_count: int,
54 token_counter: TokenCounterProtocol | None = None,
55 ) -> list[ChatMessage]:
56 """Return messages unchanged."""
57 return messages
60class AnthropicCacheStrategy:
61 """Anthropic prompt caching strategy.
63 Inserts cache_control: {"type": "ephemeral"} on up to 4 static blocks.
64 Minimum 1,024 tokens per cached block (2,048 for Haiku models).
65 """
67 def __init__(self, min_tokens: int = ANTHROPIC_MIN_CACHE_TOKENS) -> None:
68 """Initialize with token minimum threshold.
70 Args:
71 min_tokens: Minimum tokens required to mark a block as cacheable.
72 """
73 self._min_tokens = min_tokens
75 def annotate(
76 self,
77 messages: list[ChatMessage],
78 static_count: int,
79 token_counter: TokenCounterProtocol | None = None,
80 ) -> list[ChatMessage]:
81 """Insert cache_control on eligible static blocks."""
82 result = list(messages)
83 breakpoints = 0
84 for i in range(min(static_count, len(result))):
85 if breakpoints >= ANTHROPIC_MAX_BREAKPOINTS:
86 break
87 msg = result[i]
88 token_count = (
89 token_counter.count(str(msg.content or ""))
90 if token_counter
91 else self._min_tokens
92 )
93 if token_count >= self._min_tokens:
94 # Add cache_control to the message's additional fields
95 updated = _set_cache_control(msg, {"type": "ephemeral"})
96 result[i] = updated
97 breakpoints += 1
98 logger.debug(
99 "anthropic_cache_breakpoint_added",
100 index=i,
101 tokens=token_count,
102 breakpoints=breakpoints,
103 )
104 return result
107class OpenAICacheStrategy:
108 """OpenAI prompt caching strategy.
110 No annotations needed — OpenAI caches automatically when first 1,024+
111 tokens are byte-identical. Warns if static prefix is too short.
112 """
114 def annotate(
115 self,
116 messages: list[ChatMessage],
117 static_count: int,
118 token_counter: TokenCounterProtocol | None = None,
119 ) -> list[ChatMessage]:
120 """Validate prefix length and return messages unchanged."""
121 if token_counter and static_count > 0:
122 static_messages = messages[:static_count]
123 static_tokens = token_counter.count_messages(static_messages) # type: ignore[arg-type]
124 if static_tokens < 1024:
125 logger.warning(
126 "openai_cache_prefix_too_short",
127 static_tokens=static_tokens,
128 minimum=1024,
129 )
130 return messages
133class DeepSeekCacheStrategy:
134 """DeepSeek cache strategy.
136 Pads static blocks to the nearest 64-token boundary using neutral
137 whitespace. Requires TokenCounterProtocol for accurate padding.
138 """
140 def annotate(
141 self,
142 messages: list[ChatMessage],
143 static_count: int,
144 token_counter: TokenCounterProtocol | None = None,
145 ) -> list[ChatMessage]:
146 """Pad static messages to 64-token boundaries."""
147 if token_counter is None:
148 logger.warning("deepseek_cache_no_token_counter", action="skipping_padding")
149 return messages
150 result = list(messages)
151 for i in range(min(static_count, len(result))):
152 msg = result[i]
153 content = str(msg.content or "")
154 current_tokens = token_counter.count(content)
155 remainder = current_tokens % DEEPSEEK_PADDING_BOUNDARY
156 if remainder != 0:
157 # Calculate target token count (next multiple of 64)
158 target_tokens = current_tokens + (DEEPSEEK_PADDING_BOUNDARY - remainder)
159 # Iteratively add spaces until we reach the target token count
160 padded_content = content
161 max_iterations = (
162 target_tokens * 10
163 ) # Upper bound to prevent infinite loop
164 for _ in range(max_iterations):
165 padded_content += " "
166 if token_counter.count(padded_content) >= target_tokens:
167 break
168 result[i] = _update_content(msg, padded_content)
169 logger.debug(
170 "deepseek_cache_padding_applied",
171 index=i,
172 original_tokens=current_tokens,
173 target_tokens=target_tokens,
174 padded_tokens=token_counter.count(padded_content),
175 )
176 return result
179class GeminiCacheStrategy:
180 """Gemini context caching strategy.
182 Flags static blocks >32k tokens for Context Caching API pre-creation.
183 Below 32k, passes through unchanged.
184 """
186 def annotate(
187 self,
188 messages: list[ChatMessage],
189 static_count: int,
190 token_counter: TokenCounterProtocol | None = None,
191 ) -> list[ChatMessage]:
192 """Flag large static blocks for context caching."""
193 if token_counter and static_count > 0:
194 static_messages = messages[:static_count]
195 total_tokens = token_counter.count_messages(static_messages) # type: ignore[arg-type]
196 if total_tokens > GEMINI_CONTEXT_CACHE_MIN_TOKENS:
197 logger.info(
198 "gemini_context_cache_recommended",
199 static_tokens=total_tokens,
200 threshold=GEMINI_CONTEXT_CACHE_MIN_TOKENS,
201 )
202 return messages
205class MistralCacheStrategy:
206 """Mistral prefix caching strategy.
208 Ensures stable [INST] prefix structure for prefix caching.
209 Passes messages through — structure is enforced at assembly time.
210 """
212 def annotate(
213 self,
214 messages: list[ChatMessage],
215 static_count: int,
216 token_counter: TokenCounterProtocol | None = None,
217 ) -> list[ChatMessage]:
218 """Return messages unchanged — structure enforced in assembler."""
219 return messages
222def _set_cache_control(msg: ChatMessage, cache_control: dict[str, Any]) -> ChatMessage:
223 """Return a new ChatMessage with cache_control added to metadata."""
224 from dataclasses import replace as _dc_replace
226 try:
227 metadata = dict(msg.metadata or {})
228 metadata["cache_control"] = cache_control
229 return _dc_replace(msg, metadata=metadata)
230 except (TypeError, AttributeError) as e:
231 logger.warning(
232 "cache_control_not_supported",
233 reason=f"Failed to set cache_control: {e}",
234 )
235 return msg
238def _update_content(msg: ChatMessage, content: str) -> ChatMessage:
239 """Return a new ChatMessage with updated content."""
240 from dataclasses import replace as _dc_replace
242 return _dc_replace(msg, content=content)
245class ProviderCacheStrategyRegistry:
246 """Registry mapping provider names to cache annotation strategies.
248 Uses registry-based dispatch — no if/elif chains.
250 Usage::
252 registry = ProviderCacheStrategyRegistry.with_defaults()
253 strategy = registry.for_provider("anthropic")
254 annotated = strategy.annotate(messages, static_count=3, token_counter=counter)
255 """
257 def __init__(self) -> None:
258 """Create an empty registry."""
259 self._strategies: dict[str, CacheStrategy] = {}
260 self._default: CacheStrategy = PassthroughCacheStrategy()
262 @classmethod
263 def with_defaults(cls) -> ProviderCacheStrategyRegistry:
264 """Create a registry pre-populated with all provider strategies."""
265 registry = cls()
266 registry.register("anthropic", AnthropicCacheStrategy())
267 registry.register("openai", OpenAICacheStrategy())
268 registry.register("azure", OpenAICacheStrategy())
269 registry.register("deepseek", DeepSeekCacheStrategy())
270 registry.register("gemini", GeminiCacheStrategy())
271 registry.register("mistral", MistralCacheStrategy())
272 registry.register("*", PassthroughCacheStrategy())
273 return registry
275 def register(self, provider: str, strategy: CacheStrategy) -> None:
276 """Register a strategy for a provider key.
278 Args:
279 provider: Provider name (e.g. "anthropic", "openai").
280 strategy: Cache annotation strategy.
281 """
282 self._strategies[provider] = strategy
284 def for_provider(self, provider: str) -> CacheStrategy:
285 """Get the strategy for the given provider.
287 Falls back to wildcard '*' strategy if provider not found, then to
288 PassthroughCacheStrategy default.
290 Args:
291 provider: Provider name.
293 Returns:
294 CacheStrategy for the provider (or wildcard/passthrough fallback).
295 """
296 return (
297 self._strategies.get(provider) or self._strategies.get("*") or self._default
298 )