Caching Strategies for AI Systems#

from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Optional
import hashlib
import json
import redis
import numpy as np

# Educational example — not for clinical use

@dataclass
class CacheEntry:
    query: str
    query_embedding: list[float]
    response: str
    created_at: datetime
    ttl_seconds: int
    source_documents: list[str]     # For citation tracking


class SemanticCache:
    """
    Semantic cache for LLM query-response pairs.
    
    Uses Redis for storage with pgvector for similarity search,
    or in-process numpy for small deployments.
    
    Educational example — not for clinical use.
    NOT appropriate for patient-specific queries.
    """
    
    def __init__(
        self,
        similarity_threshold: float = 0.95,   # Very high threshold for clinical use
        default_ttl_seconds: int = 86400,     # 1 day
        embedding_client=None,
        redis_client=None,
    ):
        self.threshold = similarity_threshold
        self.default_ttl = default_ttl_seconds
        self.embedding_client = embedding_client
        self.redis = redis_client
    
    async def lookup(self, query: str) -> Optional[CacheEntry]:
        """
        Look up a query in the semantic cache.
        Returns a CacheEntry if a similar query was found above threshold,
        else returns None (cache miss).
        """
        query_embedding = await self.embedding_client.embed(query)
        
        # Retrieve candidate embeddings from store
        candidates = await self._get_candidates(query_embedding, top_k=10)
        
        for entry in candidates:
            similarity = self._cosine_similarity(
                np.array(query_embedding),
                np.array(entry.query_embedding)
            )
            if similarity >= self.threshold:
                # Check TTL
                if self._is_expired(entry):
                    await self._delete(entry.query)
                    continue
                return entry
        
        return None
    
    async def store(
        self,
        query: str,
        query_embedding: list[float],
        response: str,
        source_documents: list[str],
        ttl_seconds: Optional[int] = None,
    ) -> None:
        """Store a query-response pair in the semantic cache."""
        entry = CacheEntry(
            query=query,
            query_embedding=query_embedding,
            response=response,
            created_at=datetime.utcnow(),
            ttl_seconds=ttl_seconds or self.default_ttl,
            source_documents=source_documents,
        )
        
        key = self._cache_key(query)
        value = json.dumps({
            "query": entry.query,
            "query_embedding": entry.query_embedding,
            "response": entry.response,
            "created_at": entry.created_at.isoformat(),
            "ttl_seconds": entry.ttl_seconds,
            "source_documents": entry.source_documents,
        })
        
        await self.redis.setex(key, entry.ttl_seconds, value)
    
    async def invalidate_by_source(self, source_document_id: str) -> int:
        """
        Invalidate all cached responses that cite a specific source document.
        
        Call this when a clinical guideline is updated to prevent stale
        recommendations from being served from cache.
        """
        # This requires scanning all cached entries — expensive at scale.
        # Production implementation: maintain a reverse index (source_id → cache_keys).
        invalidated = 0
        
        async for key in self.redis.scan_iter("semantic_cache:*"):
            value = await self.redis.get(key)
            if value:
                entry_data = json.loads(value)
                if source_document_id in entry_data.get("source_documents", []):
                    await self.redis.delete(key)
                    invalidated += 1
        
        return invalidated
    
    def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
    
    def _is_expired(self, entry: CacheEntry) -> bool:
        age = (datetime.utcnow() - entry.created_at).total_seconds()
        return age > entry.ttl_seconds
    
    def _cache_key(self, query: str) -> str:
        content_hash = hashlib.sha256(query.encode()).hexdigest()[:16]
        return f"semantic_cache:{content_hash}"
    
    async def _get_candidates(self, embedding: list[float], top_k: int) -> list[CacheEntry]:
        # In production: use pgvector ANN search on the stored embeddings
        # This example shows simplified key-scan (suitable only for small caches)
        entries = []
        async for key in self.redis.scan_iter("semantic_cache:*"):
            value = await self.redis.get(key)
            if value:
                data = json.loads(value)
                entry = CacheEntry(
                    query=data["query"],
                    query_embedding=data["query_embedding"],
                    response=data["response"],
                    created_at=datetime.fromisoformat(data["created_at"]),
                    ttl_seconds=data["ttl_seconds"],
                    source_documents=data["source_documents"],
                )
                entries.append(entry)
        return entries[:top_k]
    
    async def _delete(self, query: str) -> None:
        await self.redis.delete(self._cache_key(query))

import anthropic

# Educational example — not for clinical use

client = anthropic.Anthropic()

def query_with_prompt_caching(
    system_prompt: str,
    retrieved_chunks: list[str],
    user_query: str,
) -> str:
    """
    Query the Anthropic API using prompt caching for system prompt
    and retrieved context chunks.
    
    Prompt caching reduces cost and time-to-first-token for requests
    sharing a common prefix. The system prompt and retrieved chunks
    are cached after the first request; subsequent requests that use
    the same prefix pay only for the new query tokens.
    
    Note: Verify cache_control parameter availability in current SDK docs.
    """
    
    # Format retrieved chunks as a single cached context block
    context_text = "\n\n---\n\n".join([
        f"Source Document {i+1}:\n{chunk}"
        for i, chunk in enumerate(retrieved_chunks)
    ])
    
    response = client.messages.create(
        model="claude-opus-4-8",    # verify current model IDs
        max_tokens=1024,
        system=[
            {
                "type": "text",
                "text": system_prompt,
                "cache_control": {"type": "ephemeral"},  # Cache the system prompt
            }
        ],
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"Retrieved context:\n\n{context_text}",
                        "cache_control": {"type": "ephemeral"},  # Cache context chunks
                    },
                    {
                        "type": "text",
                        "text": f"\nQuery: {user_query}",
                        # No cache_control — this changes per request
                    }
                ]
            }
        ]
    )
    
    # Inspect cache usage in response
    usage = response.usage
    if hasattr(usage, 'cache_creation_input_tokens'):
        print(f"Cache created: {usage.cache_creation_input_tokens} tokens")
    if hasattr(usage, 'cache_read_input_tokens'):
        print(f"Cache read: {usage.cache_read_input_tokens} tokens")
    
    return response.content[0].text

from enum import Enum
from dataclasses import dataclass

class DocumentType(Enum):
    CLINICAL_GUIDELINE = "clinical_guideline"
    FORMULARY = "formulary"
    HOSPITAL_PROTOCOL = "hospital_protocol"
    DRUG_INTERACTION = "drug_interaction"

# Document type → cache TTL policy
CACHE_TTL_POLICY = {
    DocumentType.CLINICAL_GUIDELINE: {
        "ttl_seconds": 7 * 86400,       # 7 days
        "rationale": "Guidelines updated quarterly; short TTL provides safety margin"
    },
    DocumentType.FORMULARY: {
        "ttl_seconds": 86400,           # 1 day
        "rationale": "Formulary changes 5-business-day SLA; 1-day TTL ensures freshness"
    },
    DocumentType.HOSPITAL_PROTOCOL: {
        "ttl_seconds": 3 * 86400,       # 3 days
        "rationale": "Protocols updated infrequently but are safety-critical"
    },
    DocumentType.DRUG_INTERACTION: {
        "ttl_seconds": 4 * 3600,        # 4 hours
        "rationale": "Drug interaction data updated continuously; aggressive TTL"
    },
}