diff options
| -rw-r--r-- | llamachat/db.py | 43 | ||||
| -rwxr-xr-x | test_llamachat.py | 97 |
2 files changed, 127 insertions, 13 deletions
diff --git a/llamachat/db.py b/llamachat/db.py index 65a2b29..4465adc 100644 --- a/llamachat/db.py +++ b/llamachat/db.py @@ -33,13 +33,16 @@ CREATE TABLE IF NOT EXISTS sessions ( ); CREATE TABLE IF NOT EXISTS messages ( - id INTEGER PRIMARY KEY, - session_id INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, - role TEXT NOT NULL, - content TEXT NOT NULL, - reasoning TEXT NOT NULL DEFAULT '', - searches TEXT, - created_at INTEGER NOT NULL + id INTEGER PRIMARY KEY, + session_id INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + role TEXT NOT NULL, + content TEXT NOT NULL, + reasoning TEXT NOT NULL DEFAULT '', + searches TEXT, + prompt_tokens INTEGER, + completion_tokens INTEGER, + model TEXT, + created_at INTEGER NOT NULL ); CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id); @@ -107,6 +110,13 @@ class History: # Nullable rather than defaulted: NULL means "no searches", # which is exactly what every pre-migration row wants. "searches": "TEXT", + # Nullable for the same reason: an old row has no counts, + # and zero would read as a reply that cost nothing. + "prompt_tokens": "INTEGER", + "completion_tokens": "INTEGER", + # The model on `sessions` is the current one, which prices a + # switched conversation wrongly. Record what actually replied. + "model": "TEXT", }, "sessions": { "prompt_name": "TEXT NOT NULL DEFAULT ''", @@ -213,16 +223,23 @@ class History: content: str, reasoning: str | None = None, searches: str | None = None, + prompt_tokens: int | None = None, + completion_tokens: int | None = None, + model: str | None = None, ) -> None: """Fill in a streamed reply. Omitted fields keep their stored value.""" columns = ["content = ?"] values: list = [content] - if reasoning is not None: - columns.append("reasoning = ?") - values.append(reasoning) - if searches is not None: - columns.append("searches = ?") - values.append(searches) + for column, value in ( + ("reasoning", reasoning), + ("searches", searches), + ("prompt_tokens", prompt_tokens), + ("completion_tokens", completion_tokens), + ("model", model), + ): + if value is not None: + columns.append(f"{column} = ?") + values.append(value) values.append(message_id) self.conn.execute( f"UPDATE messages SET {', '.join(columns)} WHERE id = ?", values diff --git a/test_llamachat.py b/test_llamachat.py index 5f96a80..4d68cc6 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -1724,6 +1724,102 @@ def test_metadata_and_cost(): print("ok metadata resolution and cost") +def test_token_column_migration(): + """A pre-token database opens, and new rows record counts and model.""" + import sqlite3 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "old.db" + conn = sqlite3.connect(path) + conn.executescript( + "CREATE TABLE sessions (id INTEGER PRIMARY KEY, mode TEXT," + " title TEXT, model TEXT, created_at INTEGER, updated_at INTEGER);" + "CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id INTEGER," + " role TEXT NOT NULL, content TEXT NOT NULL," + " created_at INTEGER NOT NULL);" + "INSERT INTO sessions VALUES (1,'chat','old','m',0,0);" + "INSERT INTO messages VALUES (1,1,'assistant','older reply',0);" + ) + conn.commit() + conn.close() + + history = db.History(path) + rows = history.messages(1) + # The pre-migration row survives and reads as unknown, not as zero. + assert rows[0]["content"] == "older reply" + assert rows[0]["prompt_tokens"] is None + assert rows[0]["completion_tokens"] is None + assert rows[0]["model"] is None + + mid = history.add_message(1, "assistant", "") + history.update_message( + mid, "new reply", + prompt_tokens=1200, completion_tokens=340, + model="together:Qwen/Qwen2.5", + ) + fresh = history.messages(1)[1] + assert fresh["prompt_tokens"] == 1200 + assert fresh["completion_tokens"] == 340 + assert fresh["model"] == "together:Qwen/Qwen2.5" + + # Omitting them leaves stored values alone, as with reasoning. + history.update_message(mid, "edited") + kept = history.messages(1)[1] + assert kept["content"] == "edited" + assert kept["prompt_tokens"] == 1200 + assert kept["model"] == "together:Qwen/Qwen2.5" + history.close() + + # A fresh database must end up with the same columns as a migrated one, + # or cost would read back on one path and raise on the other. This + # passes via SCHEMA or via _migrate() indifferently, which is the point: + # both paths run on every open and either one alone suffices. + from llamachat import models, providers + + with tempfile.TemporaryDirectory() as tmp: + fresh_db = db.History(Path(tmp) / "new.db") + columns = [ + row["name"] + for row in fresh_db.conn.execute("PRAGMA table_info(messages)") + ] + assert {"prompt_tokens", "completion_tokens", "model"} <= set(columns) + # SCHEMA puts created_at last while ALTER TABLE appends after it, so + # the two paths hold the same columns in a different order. That is + # tolerable only because every read goes by name: a positional read + # of a message row would be right on one path and wrong on the other. + assert columns[-1] == "created_at" + sid = fresh_db.create_session("chat", "m", "t") + mid = fresh_db.add_message(sid, "assistant", "") + fresh_db.update_message( + mid, "hi", prompt_tokens=7, completion_tokens=3, + model="together:Qwen/Qwen2.5", + ) + row = fresh_db.messages(sid)[0] + # A stored row prices straight through models.message_cost(), which + # proves the three column names are exactly the ones it reads. + table = providers.parse( + { + "providers": { + "together": { + "base_url": "https://api.example.org", + "api_key": "env:X", + "price_in": 1_000_000.0, + "price_out": 1_000_000.0, + } + } + } + ) + store = models.ModelStore(Path(tmp) / "models.ini") + assert models.conversation_cost([row], table, store) == 10.0 + + # Reopening runs _migrate() again over columns that already exist. + fresh_db.close() + again = db.History(Path(tmp) / "new.db") + assert again.messages(sid)[0]["prompt_tokens"] == 7 + again.close() + print("ok token column migration") + + class _FakeResponse: """Enough of an http.client response for urlopen's context manager.""" @@ -2469,6 +2565,7 @@ if __name__ == "__main__": test_config_providers() test_models_store() test_metadata_and_cost() + test_token_column_migration() test_search_tool_schema() test_search_results_sanitising() test_tool_call_accumulation() |