NexusLearnAI / cache_utils.py
ChaseHan's picture
Upload 15 files
833dac3 verified
"""
Utility module for caching LLM responses
"""
import os
import json
import hashlib
from typing import Dict, Any, Optional
from config import CACHE_ENABLED, DEBUG_MODE
# Cache directory
CACHE_DIR = "cache"
os.makedirs(CACHE_DIR, exist_ok=True)
def generate_cache_key(prompt_type: str, params: Dict[str, Any]) -> str:
"""
Generate a cache key based on prompt type and parameters
Args:
prompt_type: Prompt type, such as 'decompose' or 'explain'
params: Prompt parameters
Returns:
Cache key string
"""
# Convert parameters to a standardized JSON string
params_str = json.dumps(params, sort_keys=True, ensure_ascii=False)
# Calculate hash value
hash_obj = hashlib.md5(f"{prompt_type}:{params_str}".encode('utf-8'))
return hash_obj.hexdigest()
def save_to_cache(cache_key: str, data: Dict[str, Any]) -> None:
"""
Save data to cache file
Args:
cache_key: Cache key
data: Data to be cached
"""
cache_path = os.path.join(CACHE_DIR, f"{cache_key}.json")
with open(cache_path, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=2)
def load_from_cache(cache_key: str) -> Optional[Dict[str, Any]]:
"""
Load data from cache
Args:
cache_key: Cache key
Returns:
Cached data, or None if it doesn't exist
"""
cache_path = os.path.join(CACHE_DIR, f"{cache_key}.json")
if not os.path.exists(cache_path):
return None
try:
with open(cache_path, 'r', encoding='utf-8') as f:
return json.load(f)
except (json.JSONDecodeError, IOError):
# Return None if file is corrupted or cannot be read
return None
def cached_llm_call(prompt_type: str, params: Dict[str, Any], call_function) -> Dict[str, Any]:
"""
Cache decorator for LLM calls
Args:
prompt_type: Prompt type
params: Prompt parameters
call_function: Actual function to call LLM
Returns:
LLM response or cached response
"""
# Generate cache key
cache_key = generate_cache_key(prompt_type, params)
# Try to load from cache
cached_result = load_from_cache(cache_key)
if cached_result:
print(f"[Cache] Using cached response: {prompt_type}")
return cached_result
# If not cached, call LLM
result = call_function(params)
# Save to cache
save_to_cache(cache_key, result)
return result
def get_from_cache(cache_key: str) -> Optional[str]:
"""
Get data directly from cache
Args:
cache_key: Cache key (can be a string)
Returns:
Cached data, or None if it doesn't exist
"""
if not CACHE_ENABLED:
return None
# Hash the cache key
hash_obj = hashlib.md5(cache_key.encode('utf-8'))
hashed_key = hash_obj.hexdigest()
cache_path = os.path.join(CACHE_DIR, f"{hashed_key}.json")
if not os.path.exists(cache_path):
return None
try:
with open(cache_path, 'r', encoding='utf-8') as f:
if DEBUG_MODE:
print(f"Loading from cache: {cache_key[:30]}...")
return f.read()
except (IOError):
# Return None if file cannot be read
return None
def save_to_cache(cache_key: str, data: str) -> None:
"""
Save data to cache file
Args:
cache_key: Cache key (can be a string)
data: Data string to cache
"""
if not CACHE_ENABLED:
return
# Hash the cache key
hash_obj = hashlib.md5(cache_key.encode('utf-8'))
hashed_key = hash_obj.hexdigest()
cache_path = os.path.join(CACHE_DIR, f"{hashed_key}.json")
with open(cache_path, 'w', encoding='utf-8') as f:
f.write(data)
if DEBUG_MODE:
print(f"Data cached: {cache_key[:30]}...")