wehub-resource-sync
4a19d70af1
Lint with Ruff / ruff (push) Has been cancelled
MCP Server Tests / live-mcp-tests (push) Has been cancelled
Tests / unit-tests (push) Has been cancelled
Tests / database-integration-tests (push) Has been cancelled
CodeQL Advanced / Analyze (actions) (push) Has been cancelled
CodeQL Advanced / Analyze (python) (push) Has been cancelled
Server Tests / live-server-tests (push) Has been cancelled
Pyright Type Check / pyright (push) Has been cancelled
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""
|
|
Copyright 2024, Zep Software, Inc.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import sqlite3
|
|
import typing
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LLMCache:
|
|
"""Simple SQLite + JSON cache for LLM responses.
|
|
|
|
Replaces diskcache to avoid unsafe pickle deserialization (CVE in diskcache <= 5.6.3).
|
|
Only stores JSON-serializable data.
|
|
"""
|
|
|
|
def __init__(self, directory: str):
|
|
os.makedirs(directory, exist_ok=True)
|
|
db_path = os.path.join(directory, 'cache.db')
|
|
self._conn = sqlite3.connect(db_path, check_same_thread=False)
|
|
self._conn.execute('CREATE TABLE IF NOT EXISTS cache (key TEXT PRIMARY KEY, value TEXT)')
|
|
self._conn.commit()
|
|
|
|
def get(self, key: str) -> dict[str, typing.Any] | None:
|
|
row = self._conn.execute('SELECT value FROM cache WHERE key = ?', (key,)).fetchone()
|
|
if row is None:
|
|
return None
|
|
try:
|
|
return json.loads(row[0])
|
|
except json.JSONDecodeError:
|
|
logger.warning(f'Corrupted cache entry for key {key}, ignoring')
|
|
return None
|
|
|
|
def set(self, key: str, value: dict[str, typing.Any]) -> None:
|
|
try:
|
|
serialized = json.dumps(value)
|
|
except TypeError:
|
|
logger.warning(f'Non-JSON-serializable cache value for key {key}, skipping')
|
|
return
|
|
self._conn.execute(
|
|
'INSERT OR REPLACE INTO cache (key, value) VALUES (?, ?)',
|
|
(key, serialized),
|
|
)
|
|
self._conn.commit()
|
|
|
|
def close(self) -> None:
|
|
self._conn.close()
|
|
|
|
def __del__(self) -> None:
|
|
with contextlib.suppress(Exception):
|
|
self._conn.close()
|