diff options
Diffstat (limited to 'assistant.py')
| -rw-r--r-- | assistant.py | 23 |
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 |