Tool call caching system implemented with SQLite for multi-process safety
| 16 | logger = logging.getLogger(__name__) |
| 17 | |
| 18 | class ToolCache: |
| 19 | """ |
| 20 | Tool call caching system implemented with SQLite for multi-process safety |
| 21 | """ |
| 22 | |
| 23 | def __init__(self, cache_dir: str = "./cache", ttl_hours: int = 0, enabled: bool = False, |
| 24 | server_whitelist: Optional[list] = None): |
| 25 | """ |
| 26 | Initialize the cache system |
| 27 | |
| 28 | Args: |
| 29 | cache_dir: Cache directory path |
| 30 | ttl_hours: Cache time-to-live in hours, 0 means permanent cache |
| 31 | enabled: Whether caching is enabled |
| 32 | server_whitelist: Server whitelist, only cache tool calls from these servers (None or empty list means cache all) |
| 33 | """ |
| 34 | self.enabled = enabled |
| 35 | if not self.enabled: |
| 36 | logger.info("Tool cache is disabled") |
| 37 | return |
| 38 | |
| 39 | self.cache_dir = Path(cache_dir) |
| 40 | self.cache_dir.mkdir(parents=True, exist_ok=True) |
| 41 | self.db_path = self.cache_dir / "tool_cache.db" |
| 42 | self.ttl_seconds = ttl_hours * 3600 if ttl_hours > 0 else 0 |
| 43 | self.server_whitelist = server_whitelist or [] |
| 44 | |
| 45 | # Thread-local storage, one connection per thread |
| 46 | self.local = threading.local() |
| 47 | |
| 48 | # Initialize database |
| 49 | self._init_db() |
| 50 | |
| 51 | whitelist_msg = f" with whitelist: {self.server_whitelist}" if self.server_whitelist else "" |
| 52 | logger.info(f"Tool cache initialized at {self.db_path} with TTL={ttl_hours} hours{whitelist_msg}") |
| 53 | |
| 54 | def _get_connection(self) -> sqlite3.Connection: |
| 55 | """Get thread-local database connection""" |
| 56 | if not hasattr(self.local, 'conn'): |
| 57 | self.local.conn = sqlite3.connect(str(self.db_path), timeout=30.0) |
| 58 | # Enable WAL mode for better concurrency |
| 59 | self.local.conn.execute('PRAGMA journal_mode=WAL') |
| 60 | self.local.conn.execute('PRAGMA synchronous=NORMAL') |
| 61 | return self.local.conn |
| 62 | |
| 63 | def _init_db(self): |
| 64 | """Initialize SQLite database""" |
| 65 | conn = sqlite3.connect(str(self.db_path), timeout=30.0) |
| 66 | try: |
| 67 | conn.execute('PRAGMA journal_mode=WAL') |
| 68 | conn.execute(''' |
| 69 | CREATE TABLE IF NOT EXISTS cache ( |
| 70 | cache_key TEXT PRIMARY KEY, |
| 71 | server_name TEXT NOT NULL, |
| 72 | tool_name TEXT NOT NULL, |
| 73 | params TEXT NOT NULL, |
| 74 | result TEXT NOT NULL, |
| 75 | timestamp REAL NOT NULL, |