Keep conversations in SQLite with full-text search

One file under XDG_DATA_HOME, 0700 directory and 0600 file. turns.messages
holds what the model is sent for the turn, user/reply/tools what the canvas
draws. FTS5 with remove_diacritics finds città from citta. The query is
quoted as one phrase so FTS syntax typed by the user is matched as text,
and the self-check's integrity-check proves a conversation delete cascades
into the index.

Co-Authored-By: Claude Haiku 5.5 <noreply@anthropic.com>

1 files changed, 126 insertions, 0 deletions

diff --git a/assistant.py b/assistant.py
index b69f199..3bd8b0c 100644
--- a/assistant.py
+++ b/assistant.py
@@ -109,6 +109,101 @@ def log(*a):
     print(time.strftime("%H:%M:%S"), *a, flush=True)
 
 
+# History: one SQLite file, written only by this process. turns.messages is
+# what the model is sent for that turn; user, reply and tools are what the
+# canvas draws. The FTS index follows turns through triggers, and a
+# conversation delete cascades into turns and so into the index.
+SCHEMA = """
+CREATE TABLE IF NOT EXISTS conversations(id INTEGER PRIMARY KEY, title TEXT,
+                                         created INTEGER, updated INTEGER);
+CREATE TABLE IF NOT EXISTS turns(id INTEGER PRIMARY KEY,
+    conv INTEGER REFERENCES conversations ON DELETE CASCADE,
+    seq INTEGER, user TEXT, reply TEXT, tools TEXT, messages TEXT);
+CREATE VIRTUAL TABLE IF NOT EXISTS turns_fts USING fts5(user, reply, content='turns',
+    content_rowid='id', tokenize='unicode61 remove_diacritics 2');
+CREATE TRIGGER IF NOT EXISTS turns_ai AFTER INSERT ON turns BEGIN
+    INSERT INTO turns_fts(rowid, user, reply) VALUES (new.id, new.user, new.reply);
+END;
+CREATE TRIGGER IF NOT EXISTS turns_ad AFTER DELETE ON turns BEGIN
+    INSERT INTO turns_fts(turns_fts, rowid, user, reply)
+        VALUES ('delete', old.id, old.user, old.reply);
+END;
+"""
+
+
+def db_open(path):
+    if path != ":memory:":
+        os.makedirs(os.path.dirname(path), mode=0o700, exist_ok=True)
+    old = os.umask(0o177)                 # 0600, and the WAL files inherit it
+    try:
+        conn = sqlite3.connect(path, check_same_thread=False)
+    finally:
+        os.umask(old)
+    conn.execute("PRAGMA journal_mode=WAL")
+    conn.execute("PRAGMA foreign_keys=ON")
+    conn.executescript(SCHEMA)
+    return conn
+
+
+def save_turn(conn, conv, user, reply, tools, messages):
+    # Returns the conversation id: the row is made at its first finished turn,
+    # so a New that is never used leaves nothing behind.
+    now = int(time.time())
+    with conn:
+        if conv is None:
+            conv = conn.execute("INSERT INTO conversations(title, created, updated) "
+                                "VALUES (?, ?, ?)",
+                                (" ".join(user.split()[:6]), now, now)).lastrowid
+        else:
+            conn.execute("UPDATE conversations SET updated = ? WHERE id = ?", (now, conv))
+        seq = conn.execute("SELECT COALESCE(MAX(seq), 0) + 1 FROM turns WHERE conv = ?",
+                           (conv,)).fetchone()[0]
+        conn.execute("INSERT INTO turns(conv, seq, user, reply, tools, messages) "
+                     "VALUES (?, ?, ?, ?, ?, ?)",
+                     (conv, seq, user, reply, json.dumps(tools), json.dumps(messages)))
+    return conv
+
+
+def list_convs(conn):
+    return [{"id": i, "title": t, "updated": u} for i, t, u in conn.execute(
+        "SELECT id, title, updated FROM conversations ORDER BY updated DESC, id DESC")]
+
+
+def load_conv(conn, conv):
+    # (model history, shown turns), or None for an id that is not there.
+    rows = conn.execute("SELECT user, reply, tools, messages FROM turns WHERE conv = ? "
+                        "ORDER BY seq", (conv,)).fetchall()
+    if not rows:
+        return None
+    return ([json.loads(m) for *_, m in rows],
+            [{"user": u, "reply": r, "tools": json.loads(t)} for u, r, t, _ in rows])
+
+
+def search(conn, q):
+    # The query is one quoted phrase with a prefix star on its last word, so
+    # FTS syntax typed by the user is matched as text and cannot widen it.
+    # ponytail: rank per turn, first hit per conversation; 30 results.
+    q = q.strip()
+    if not q:
+        return []
+    seen, out = set(), []
+    for conv, title, snip in conn.execute(
+            "SELECT t.conv, c.title, snippet(turns_fts, -1, '**', '**', '…', 12) "
+            "FROM turns_fts JOIN turns t ON t.id = turns_fts.rowid "
+            "JOIN conversations c ON c.id = t.conv "
+            "WHERE turns_fts MATCH ? ORDER BY rank LIMIT 300",
+            ('"' + q.replace('"', '""') + '"*',)):
+        if conv not in seen:
+            seen.add(conv)
+            out.append({"id": conv, "title": title, "snippet": snip})
+    return out[:30]
+
+
+def delete_conv(conn, conv):
+    with conn:
+        conn.execute("DELETE FROM conversations WHERE id = ?", (conv,))
+
+
 # The canvas's feed: JSON lines on a Unix socket. With no client connected
 # every event is dropped.
 clients, clients_lock = [], threading.Lock()
@@ -332,6 +427,37 @@ def test():
     assert calls == {0: {"id": "c1", "name": "weather", "arguments": '{"place":"Rome"}'},
                      1: {"id": "c2", "name": "now", "arguments": "{}"}}
 
+    # Storage.
+    db = db_open(":memory:")
+    a = save_turn(db, None, "che tempo fa in città oggi pomeriggio?", "Piove.", ["weather"],
+                  [{"role": "user", "content": "q"}])
+    assert save_turn(db, a, "e domani?", "Sole.", [], [{"role": "user", "content": "q2"}]) == a
+    b = save_turn(db, None, "how do I list files", "Use `ls`.", [], [])
+    assert [x["id"] for x in list_convs(db)] == [b, a]
+    assert list_convs(db)[1]["title"] == "che tempo fa in città oggi"   # first six words
+    hist, shown = load_conv(db, a)
+    assert hist == [[{"role": "user", "content": "q"}], [{"role": "user", "content": "q2"}]]
+    assert shown == [{"user": "che tempo fa in città oggi pomeriggio?", "reply": "Piove.",
+                      "tools": ["weather"]},
+                     {"user": "e domani?", "reply": "Sole.", "tools": []}]
+    assert load_conv(db, 999) is None
+    hits = search(db, "citta")                       # no accent typed, accented word found
+    assert [h["id"] for h in hits] == [a] and "**città**" in hits[0]["snippet"]
+    assert [h["id"] for h in search(db, "cit")] == [a]                 # prefix, as you type
+    assert [h["id"] for h in search(db, "piove")] == [a]               # replies are indexed
+    for q in ('"', "*", 'a"b', "NOT", "x OR y", "ls)", ""):            # FTS syntax is text
+        search(db, q)
+    delete_conv(db, a)
+    assert load_conv(db, a) is None and search(db, "citta") == []
+    # Raises if the index still holds the deleted turns: the cascade must
+    # have fired the delete trigger.
+    db.execute("INSERT INTO turns_fts(turns_fts, rank) VALUES ('integrity-check', 1)")
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "sub", "h.db")
+        db_open(path).close()
+        assert os.stat(os.path.dirname(path)).st_mode & 0o777 == 0o700
+        assert os.stat(path).st_mode & 0o777 == 0o600
+
     a, b = socket.socketpair()
     a.setblocking(False)
     clients.append(a)