history: keep the model a conversation is asked with

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

1 files changed, 19 insertions, 4 deletions

diff --git a/assistant.py b/assistant.py
index 7f31add..630f5f8 100644
--- a/assistant.py
+++ b/assistant.py
@@ -186,10 +186,14 @@ def db_open(path):
     # Added with attachments; a database from before them gains the column.
     if "files" not in {r[1] for r in conn.execute("PRAGMA table_info(turns)")}:
         conn.execute("ALTER TABLE turns ADD COLUMN files TEXT")
+    # Added with cloud models: the model a conversation is asked with, NULL
+    # for whatever the router has loaded.
+    if "model" not in {r[1] for r in conn.execute("PRAGMA table_info(conversations)")}:
+        conn.execute("ALTER TABLE conversations ADD COLUMN model TEXT")
     return conn
 
 
-def save_turn(conn, conv, user, reply, tools, messages, files=()):
+def save_turn(conn, conv, user, reply, tools, messages, files=(), model=None):
     # Returns the conversation id: the row is made at its first finished turn,
     # so a New that is never used leaves nothing behind. A question that is
     # only files is titled by their names.
@@ -197,10 +201,11 @@ def save_turn(conn, conv, user, reply, tools, messages, files=()):
     title = " ".join(user.split()[:6]) or ", ".join(f["name"] for f in files)
     with conn:
         if conv is None:
-            conv = conn.execute("INSERT INTO conversations(title, created, updated) "
-                                "VALUES (?, ?, ?)", (title, now, now)).lastrowid
+            conv = conn.execute("INSERT INTO conversations(title, created, updated, model) "
+                                "VALUES (?, ?, ?, ?)", (title, now, now, model)).lastrowid
         else:
-            conn.execute("UPDATE conversations SET updated = ? WHERE id = ?", (now, conv))
+            conn.execute("UPDATE conversations SET updated = ?, model = ? WHERE id = ?",
+                         (now, model, 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, files) "
@@ -1271,6 +1276,14 @@ def test():
                       "tools": ["weather"], "files": []},
                      {"user": "e domani?", "reply": "Sole.", "tools": [], "files": []}]
     assert load_conv(db, 999) is None
+    m = save_turn(db, None, "q", "r", [], [], model="fake:x")
+    pick = lambda: db.execute("SELECT model FROM conversations WHERE id = ?", (m,)).fetchone()[0]
+    assert pick() == "fake:x"
+    save_turn(db, m, "q", "r", [], [], model="local")
+    assert pick() == "local"
+    save_turn(db, m, "q", "r", [], [])
+    assert pick() is None                            # an ask with no choice: the loaded one
+    delete_conv(db, m)
     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
@@ -1415,6 +1428,8 @@ def test():
         old.close()
         assert "files" in {r[1] for r in db_open(os.path.join(d, "h.db"))
                            .execute("PRAGMA table_info(turns)")}
+        assert "model" in {r[1] for r in db_open(os.path.join(d, "h.db"))
+                           .execute("PRAGMA table_info(conversations)")}
 
     # Protocol, against a fake llama-server.
     global DB, TOOLS, TOOL_SCHEMA