Spaces:
Sleeping
Sleeping
| import sqlite3 | |
| import json | |
| import os | |
| DB_PATH = "askmynotes_cache.db" | |
| def init_db(): | |
| """Initialize the SQLite database and create the cache table.""" | |
| with sqlite3.connect(DB_PATH) as conn: | |
| conn.execute(""" | |
| CREATE TABLE IF NOT EXISTS query_cache ( | |
| query TEXT PRIMARY KEY, | |
| tool TEXT, | |
| answer TEXT, | |
| chunks TEXT | |
| ) | |
| """) | |
| conn.commit() | |
| # Initialize DB on module load | |
| init_db() | |
| def get_cached_response(query: str): | |
| """Retrieve a cached response for a given query.""" | |
| with sqlite3.connect(DB_PATH) as conn: | |
| cursor = conn.execute( | |
| "SELECT tool, answer, chunks FROM query_cache WHERE query = ?", | |
| (query,) | |
| ) | |
| row = cursor.fetchone() | |
| if row: | |
| tool, answer, chunks_json = row | |
| return tool, answer, json.loads(chunks_json) | |
| return None | |
| def set_cached_response(query: str, tool: str, answer: str, chunks: list): | |
| """Store a response in the cache.""" | |
| with sqlite3.connect(DB_PATH) as conn: | |
| conn.execute( | |
| "INSERT OR REPLACE INTO query_cache (query, tool, answer, chunks) VALUES (?, ?, ?, ?)", | |
| (query, tool, answer, json.dumps(chunks)) | |
| ) | |
| conn.commit() | |
| def clear_cache(): | |
| """Clear all cached queries. Useful when a new document is uploaded.""" | |
| with sqlite3.connect(DB_PATH) as conn: | |
| conn.execute("DELETE FROM query_cache") | |
| conn.commit() | |