NIRAJz's picture
Update utils/cache_manager.py
4d6c43c verified
# cache_manager.py
from diskcache import Cache
from functools import wraps
import hashlib
import json
import os
from config import settings
# Handle Spaces environment - use /tmp if available
def get_cache_dir():
"""Get cache directory that works on Spaces"""
if os.path.exists("/tmp"):
return "/tmp/.cache"
return settings.CACHE_DIR
cache_dir = get_cache_dir()
if not os.path.exists(cache_dir):
os.makedirs(cache_dir, exist_ok=True)
# Initialize cache with error handling
try:
cache = Cache(cache_dir, size_limit=1000000000) # 1GB limit
except Exception as e:
print(f"Cache initialization failed: {e}. Using memory cache fallback.")
# Fallback to a simple dictionary-based cache
cache = {}
def cache_llm_response(expire=86400):
"""Cache decorator for LLM responses with fallback"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
if not settings.CACHE_ENABLED:
return func(*args, **kwargs)
# Create cache key
key_data = {
"func": func.__name__,
"args": str(args),
"kwargs": str(kwargs)
}
key = hashlib.md5(json.dumps(key_data, sort_keys=True).encode()).hexdigest()
# Check cache (with fallback for dictionary cache)
if isinstance(cache, dict):
if key in cache:
return cache[key]
else:
if key in cache:
return cache[key]
# Call function and cache result
result = func(*args, **kwargs)
# Store in cache (with fallback)
if isinstance(cache, dict):
cache[key] = result
else:
try:
cache.set(key, result, expire=expire)
except Exception:
# Fallback to dictionary if disk cache fails
cache[key] = result
return result
return wrapper
return decorator
def clear_cache():
"""Clear all cached responses"""
try:
if isinstance(cache, dict):
cache.clear()
else:
cache.clear()
except Exception as e:
print(f"Cache clearing failed: {e}")
def get_cache_stats():
"""Get cache statistics"""
try:
if isinstance(cache, dict):
return {
"size": sum(len(str(v)) for v in cache.values()),
"count": len(cache),
"enabled": settings.CACHE_ENABLED
}
else:
return {
"size": cache.volume(),
"count": len(cache),
"enabled": settings.CACHE_ENABLED
}
except Exception:
return {
"size": 0,
"count": 0,
"enabled": False
}