Managing and Reducing Costs › Strategy 3: Cache Aggressively

from google import genai
import hashlib
import time
from datetime import datetime, timedelta

gemini_client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

class CachingAgent:
    """Agent with intelligent caching to minimize API calls."""
    
    def __init__(self):
        self.client = gemini_client
        
        # Short-term cache: exact query matches
        self.exact_cache = {}
        
        # Long-term cache: common queries that rarely change
        self.persistent_cache = {
            "what is python": "Python is a high-level programming language...",
            "what is machine learning": "Machine learning is a subset of AI...",
            # Preload common queries
        }
        
        # Cache metadata
        self.cache_stats = {
            "hits": 0,
            "misses": 0,
            "api_calls": 0,
            "cost_saved": 0.0
        }
    
    def _hash_query(self, query):
        """Create cache key from query."""
        return hashlib.md5(query.lower().strip().encode()).hexdigest()
    
    def _estimate_cost_saved(self, query):
        """Estimate cost saved by cache hit."""
        # Rough estimate: 100 input tokens + 200 output tokens
        input_cost = (100 / 1_000_000) * 0.40
        output_cost = (200 / 1_000_000) * 1.20
        return input_cost + output_cost
    
    def respond(self, query):
        """Get response, using cache when possible."""
        cache_key = self._hash_query(query)
        query_normalized = query.lower().strip()
        
        # Check persistent cache first (common queries)
        if query_normalized in self.persistent_cache:
            self.cache_stats["hits"] += 1
            self.cache_stats["cost_saved"] += self._estimate_cost_saved(query)
            return self.persistent_cache[query_normalized], "persistent_cache"
        
        # Check exact cache (recent queries)
        if cache_key in self.exact_cache:
            self.cache_stats["hits"] += 1
            self.cache_stats["cost_saved"] += self._estimate_cost_saved(query)
            return self.exact_cache[cache_key], "exact_cache"
        
        # Cache miss: call the model
        self.cache_stats["misses"] += 1
        self.cache_stats["api_calls"] += 1
        
        response = self.client.models.generate_content(
            model="gemini-2.5-flash",
            contents=query
        )
        result = response.text
        
        # Store in exact cache
        self.exact_cache[cache_key] = result
        
        # If cache is getting large, prune old entries
        if len(self.exact_cache) > 1000:
            # Keep only the most recent 500
            keys_to_remove = list(self.exact_cache.keys())[:-500]
            for key in keys_to_remove:
                del self.exact_cache[key]
        
        return result, "api_call"
    
    def get_cache_stats(self):
        """Get caching performance metrics."""
        total_requests = self.cache_stats["hits"] + self.cache_stats["misses"]
        hit_rate = (self.cache_stats["hits"] / total_requests * 100) if total_requests > 0 else 0
        
        return {
            "total_requests": total_requests,
            "cache_hits": self.cache_stats["hits"],
            "cache_misses": self.cache_stats["misses"],
            "hit_rate": hit_rate,
            "api_calls": self.cache_stats["api_calls"],
            "estimated_cost_saved": self.cache_stats["cost_saved"]
        }

## Test the caching agent
agent = CachingAgent()

## Simulate user queries (with some repetition)
queries = [
    "What is Python?",
    "What is machine learning?",
    "How do I learn programming?",
    "What is Python?",  # Duplicate
    "What is machine learning?",  # Duplicate
    "How do I learn programming?",  # Duplicate
    "What are data structures?",
    "What is Python?",  # Duplicate again
]

for query in queries:
    result, source = agent.respond(query)
    print(f"Q: {query}")
    print(f"Source: {source}")
    print()

## Show cache performance
stats = agent.get_cache_stats()
print("Cache Performance:")
print(f"Total requests: {stats['total_requests']}")
print(f"Cache hits: {stats['cache_hits']} ({stats['hit_rate']:.1f}%)")
print(f"API calls: {stats['api_calls']}")
print(f"Estimated cost saved: ${stats['estimated_cost_saved']:.4f}")