|
import hashlib
|
|
import json
|
|
import sqlite3
|
|
import time
|
|
from typing import Any, Optional, Set
|
|
|
|
|
|
class ToolCache:
|
|
DETERMINISTIC_TOOLS: Set[str] = {
|
|
"read_file",
|
|
"list_directory",
|
|
"get_current_directory",
|
|
"db_get",
|
|
"db_query",
|
|
"index_directory",
|
|
"http_fetch",
|
|
"web_search",
|
|
"web_search_news",
|
|
"search_knowledge",
|
|
"get_knowledge_entry",
|
|
"get_knowledge_by_category",
|
|
"get_knowledge_statistics",
|
|
}
|
|
|
|
def __init__(self, db_path: str, ttl_seconds: int = 300):
|
|
self.db_path = db_path
|
|
self.ttl_seconds = ttl_seconds
|
|
self._initialize_cache()
|
|
|
|
def _initialize_cache(self):
|
|
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"\n CREATE TABLE IF NOT EXISTS tool_cache (\n cache_key TEXT PRIMARY KEY,\n tool_name TEXT NOT NULL,\n result_data TEXT NOT NULL,\n created_at INTEGER NOT NULL,\n expires_at INTEGER NOT NULL,\n hit_count INTEGER DEFAULT 0\n )\n "
|
|
)
|
|
cursor.execute(
|
|
"\n CREATE INDEX IF NOT EXISTS idx_tool_expires ON tool_cache(expires_at)\n "
|
|
)
|
|
cursor.execute(
|
|
"\n CREATE INDEX IF NOT EXISTS idx_tool_name ON tool_cache(tool_name)\n "
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def _generate_cache_key(self, tool_name: str, arguments: dict) -> str:
|
|
cache_data = {"tool": tool_name, "args": arguments}
|
|
serialized = json.dumps(cache_data, sort_keys=True)
|
|
return hashlib.sha256(serialized.encode()).hexdigest()
|
|
|
|
def is_cacheable(self, tool_name: str) -> bool:
|
|
return tool_name in self.DETERMINISTIC_TOOLS
|
|
|
|
def get(self, tool_name: str, arguments: dict) -> Optional[Any]:
|
|
if not self.is_cacheable(tool_name):
|
|
return None
|
|
cache_key = self._generate_cache_key(tool_name, arguments)
|
|
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
|
cursor = conn.cursor()
|
|
current_time = int(time.time())
|
|
cursor.execute(
|
|
"\n SELECT result_data, hit_count FROM tool_cache\n WHERE cache_key = ? AND expires_at > ?\n ",
|
|
(cache_key, current_time),
|
|
)
|
|
row = cursor.fetchone()
|
|
if row:
|
|
cursor.execute(
|
|
"\n UPDATE tool_cache SET hit_count = hit_count + 1\n WHERE cache_key = ?\n ",
|
|
(cache_key,),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
return json.loads(row[0])
|
|
conn.close()
|
|
return None
|
|
|
|
def set(self, tool_name: str, arguments: dict, result: Any):
|
|
if not self.is_cacheable(tool_name):
|
|
return
|
|
cache_key = self._generate_cache_key(tool_name, arguments)
|
|
current_time = int(time.time())
|
|
expires_at = current_time + self.ttl_seconds
|
|
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"\n INSERT OR REPLACE INTO tool_cache\n (cache_key, tool_name, result_data, created_at, expires_at, hit_count)\n VALUES (?, ?, ?, ?, ?, 0)\n ",
|
|
(cache_key, tool_name, json.dumps(result), current_time, expires_at),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def clear_expired(self):
|
|
current_time = int(time.time())
|
|
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM tool_cache WHERE expires_at <= ?", (current_time,))
|
|
deleted_count = cursor.rowcount
|
|
conn.commit()
|
|
conn.close()
|
|
return deleted_count
|
|
|
|
def clear_all(self):
|
|
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM tool_cache")
|
|
deleted_count = cursor.rowcount
|
|
conn.commit()
|
|
conn.close()
|
|
return deleted_count
|
|
|
|
def get_statistics(self) -> dict:
|
|
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT COUNT(*) FROM tool_cache")
|
|
total_entries = cursor.fetchone()[0]
|
|
current_time = int(time.time())
|
|
cursor.execute("SELECT COUNT(*) FROM tool_cache WHERE expires_at > ?", (current_time,))
|
|
valid_entries = cursor.fetchone()[0]
|
|
cursor.execute(
|
|
"SELECT SUM(hit_count) FROM tool_cache WHERE expires_at > ?", (current_time,)
|
|
)
|
|
total_hits = cursor.fetchone()[0] or 0
|
|
cursor.execute(
|
|
"\n SELECT tool_name, COUNT(*), SUM(hit_count)\n FROM tool_cache\n WHERE expires_at > ?\n GROUP BY tool_name\n ",
|
|
(current_time,),
|
|
)
|
|
tool_stats = {}
|
|
for row in cursor.fetchall():
|
|
tool_stats[row[0]] = {"cached_entries": row[1], "total_hits": row[2] or 0}
|
|
conn.close()
|
|
return {
|
|
"total_entries": total_entries,
|
|
"valid_entries": valid_entries,
|
|
"expired_entries": total_entries - valid_entries,
|
|
"total_cache_hits": total_hits,
|
|
"by_tool": tool_stats,
|
|
}
|