pickle -> sqlite store: D6 async writes, D9 reload race, one-shot pickle migration
This commit is contained in:
@@ -12,6 +12,8 @@ from pathlib import Path
|
||||
from pprint import pformat
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from .persistence import PersistentStore
|
||||
|
||||
|
||||
def pp(*args, **kw):
|
||||
if "width" not in kw:
|
||||
@@ -138,17 +140,16 @@ class AIResponder(AIResponderBase):
|
||||
self.history: List[Dict[str, Any]] = []
|
||||
self.memory: str = "I am an assistant."
|
||||
self.rate_limit_backoff = exponential_backoff()
|
||||
self.history_file: Optional[Path] = None
|
||||
self.memory_file: Optional[Path] = None
|
||||
self.store: Optional[PersistentStore] = None
|
||||
if "history-directory" in self.config:
|
||||
self.history_file = Path(self.config["history-directory"]).expanduser() / f"{self.channel}.dat"
|
||||
if self.history_file.exists():
|
||||
with open(self.history_file, "rb") as fd:
|
||||
self.history = pickle.load(fd)
|
||||
self.memory_file = Path(self.config["history-directory"]).expanduser() / f"{self.channel}.memory"
|
||||
if self.memory_file.exists():
|
||||
with open(self.memory_file, "rb") as fd:
|
||||
self.memory = pickle.load(fd)
|
||||
directory = Path(self.config["history-directory"]).expanduser()
|
||||
self.store = PersistentStore(directory / "bot.db")
|
||||
# Legacy pickles import once, then live on as *.migrated (PER-03)
|
||||
self.store.migrate_pickles(self.channel, directory / f"{self.channel}.dat", directory / f"{self.channel}.memory")
|
||||
self.history = self.store.load_history(self.channel)
|
||||
stored_memory = self.store.load_memory(self.channel)
|
||||
if stored_memory is not None:
|
||||
self.memory = stored_memory
|
||||
logging.info(f"memmory:\n{self.memory}")
|
||||
|
||||
def message(self, message: AIMessage, limit: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
@@ -225,9 +226,6 @@ class AIResponder(AIResponderBase):
|
||||
self.history.append({"role": "user", "content": str(message)})
|
||||
while len(self.history) > limit:
|
||||
self.shrink_history_by_one()
|
||||
if self.history_file is not None:
|
||||
with open(self.history_file, "wb") as fd:
|
||||
pickle.dump(self.history, fd)
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -269,15 +267,17 @@ class AIResponder(AIResponderBase):
|
||||
self.history.append(answer)
|
||||
while len(self.history) > limit:
|
||||
self.shrink_history_by_one()
|
||||
if self.history_file is not None:
|
||||
with open(self.history_file, "wb") as fd:
|
||||
pickle.dump(self.history, fd)
|
||||
|
||||
def update_memory(self, memory) -> None:
|
||||
self.memory = memory
|
||||
if self.memory_file is not None:
|
||||
with open(self.memory_file, "wb") as fd:
|
||||
pickle.dump(self.memory, fd)
|
||||
|
||||
async def _persist_history(self) -> None:
|
||||
if self.store is not None:
|
||||
await asyncio.to_thread(self.store.save_history, self.channel, list(self.history))
|
||||
|
||||
async def _persist_memory(self) -> None:
|
||||
if self.store is not None:
|
||||
await asyncio.to_thread(self.store.save_memory, self.channel, self.memory)
|
||||
|
||||
async def handle_picture(self, response: Dict) -> bool:
|
||||
if not isinstance(response.get("picture"), (type(None), str)):
|
||||
@@ -299,6 +299,7 @@ class AIResponder(AIResponderBase):
|
||||
async def memoize(self, message_user: str, answer_user: str, message: str, answer: str) -> None:
|
||||
self.memory = await self.memory_rewrite(self.memory, message_user, answer_user, message, answer)
|
||||
self.update_memory(self.memory)
|
||||
await self._persist_memory()
|
||||
|
||||
async def memoize_reaction(self, message_user: str, reaction_user: str, operation: str, reaction: str, message: str) -> None:
|
||||
quoted_message = message.replace("\n", "\n> ")
|
||||
@@ -312,6 +313,7 @@ class AIResponder(AIResponderBase):
|
||||
|
||||
# Check if a short path applies, return an empty AIResponse if it does
|
||||
if self.short_path(message, limit):
|
||||
await self._persist_history()
|
||||
return AIResponse(None, False, None, None, None, False, False)
|
||||
|
||||
# Number of retries for sending the message; failed attempts are
|
||||
@@ -347,8 +349,9 @@ class AIResponder(AIResponderBase):
|
||||
answer_message = await self.post_process(message, response)
|
||||
answer["content"] = str(answer_message)
|
||||
|
||||
# Update message history
|
||||
# Update message history; persistence runs off the loop (PER-05)
|
||||
self.update_history(messages[-1], answer, limit, message.historise_question)
|
||||
await self._persist_history()
|
||||
logging.info(f"got this answer:\n{str(answer_message)}")
|
||||
|
||||
# Update memory
|
||||
|
||||
@@ -236,14 +236,27 @@ class FjerkroaBot(commands.Bot):
|
||||
)
|
||||
|
||||
def on_config_file_modified(self, event):
|
||||
if event.src_path == self.config_file:
|
||||
new_config = self.load_config(self.config_file)
|
||||
if repr(new_config) != repr(self.config):
|
||||
logging.info(f"config file {self.config_file} changed, reloading.")
|
||||
self.config = new_config
|
||||
self.airesponder.config = self.config
|
||||
for responder in self.aichannels.values():
|
||||
responder.config = self.config
|
||||
# Runs on the watchdog observer thread — the swap itself is
|
||||
# scheduled onto the event loop so no request reads a
|
||||
# half-swapped config (CFG-04 / D9)
|
||||
if event.src_path != self.config_file:
|
||||
return
|
||||
new_config = self.load_config(self.config_file)
|
||||
if repr(new_config) == repr(self.config):
|
||||
return
|
||||
logging.info(f"config file {self.config_file} changed, reloading.")
|
||||
|
||||
def apply() -> None:
|
||||
self.config = new_config
|
||||
self.airesponder.config = new_config
|
||||
for responder in self.aichannels.values():
|
||||
responder.config = new_config
|
||||
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(apply)
|
||||
except (RuntimeError, AttributeError):
|
||||
# event loop not running yet (startup) — no concurrent readers
|
||||
apply()
|
||||
|
||||
@classmethod
|
||||
def load_config(cls, config_file: str = "config.toml"):
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
"""SQLite persistence for history + memory (SPEC-009).
|
||||
|
||||
One database per deployment, stdlib sqlite3 only (D-010). Callers run
|
||||
the sync methods in a worker thread (`asyncio.to_thread`) so the
|
||||
event loop never blocks on disk (PER-05 / D6).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
|
||||
class PersistentStore:
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
self.db_path = Path(db_path)
|
||||
self._init_db()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
return conn
|
||||
|
||||
def _init_db(self) -> None:
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with closing(self._connect()) as conn, conn:
|
||||
if conn.execute("PRAGMA user_version").fetchone()[0] < SCHEMA_VERSION:
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS history (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL)"
|
||||
)
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS history_channel ON history (channel)")
|
||||
conn.execute("CREATE TABLE IF NOT EXISTS memory (channel TEXT PRIMARY KEY, content TEXT NOT NULL)")
|
||||
conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
|
||||
os.chmod(self.db_path, 0o600) # conversation data (PER-04)
|
||||
|
||||
def load_history(self, channel: str) -> List[Dict[str, Any]]:
|
||||
with closing(self._connect()) as conn:
|
||||
rows = conn.execute("SELECT role, content FROM history WHERE channel = ? ORDER BY id", (channel,)).fetchall()
|
||||
return [{"role": role, "content": content} for role, content in rows]
|
||||
|
||||
def save_history(self, channel: str, history: List[Dict[str, Any]]) -> None:
|
||||
# Full replace per save: histories are small (<= history-limit)
|
||||
# and trims must be reflected (D-011)
|
||||
with closing(self._connect()) as conn, conn:
|
||||
conn.execute("DELETE FROM history WHERE channel = ?", (channel,))
|
||||
conn.executemany(
|
||||
"INSERT INTO history (channel, role, content) VALUES (?, ?, ?)",
|
||||
[(channel, str(entry.get("role", "user")), str(entry.get("content", ""))) for entry in history],
|
||||
)
|
||||
|
||||
def load_memory(self, channel: str) -> Optional[str]:
|
||||
with closing(self._connect()) as conn:
|
||||
row = conn.execute("SELECT content FROM memory WHERE channel = ?", (channel,)).fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
def save_memory(self, channel: str, content: str) -> None:
|
||||
with closing(self._connect()) as conn, conn:
|
||||
conn.execute(
|
||||
"INSERT INTO memory (channel, content) VALUES (?, ?) ON CONFLICT(channel) DO UPDATE SET content = excluded.content",
|
||||
(channel, content),
|
||||
)
|
||||
|
||||
def migrate_pickles(self, channel: str, history_file: Path, memory_file: Path) -> None:
|
||||
"""Import legacy pickles once; rename them *.migrated (PER-03)."""
|
||||
if self.load_history(channel) or self.load_memory(channel) is not None:
|
||||
return
|
||||
if history_file.exists():
|
||||
try:
|
||||
with open(history_file, "rb") as fd:
|
||||
self.save_history(channel, pickle.load(fd))
|
||||
history_file.rename(history_file.with_name(history_file.name + ".migrated"))
|
||||
logging.info(f"migrated legacy history pickle for {channel}")
|
||||
except Exception as err:
|
||||
logging.error(f"failed to migrate history pickle {history_file}: {err!r}")
|
||||
if memory_file.exists():
|
||||
try:
|
||||
with open(memory_file, "rb") as fd:
|
||||
self.save_memory(channel, str(pickle.load(fd)))
|
||||
memory_file.rename(memory_file.with_name(memory_file.name + ".migrated"))
|
||||
logging.info(f"migrated legacy memory pickle for {channel}")
|
||||
except Exception as err:
|
||||
logging.error(f"failed to migrate memory pickle {memory_file}: {err!r}")
|
||||
Reference in New Issue
Block a user