aboutsummaryrefslogtreecommitdiffstats
path: root/test_llamachat.py
diff options
context:
space:
mode:
authorDanilo M. <danix@danix.xyz>2026-07-31 11:03:51 +0200
committerDanilo M. <danix@danix.xyz>2026-07-31 11:03:51 +0200
commit9182c84b8ac9d905a9d553b95155a993338568b8 (patch)
treec1a56b7b6a49099b25e7a2dee4e0f908d5a986c5 /test_llamachat.py
parent325a6a7c12832d8b9fb9b67567e8a1c496dc1330 (diff)
downloadllamachat-9182c84b8ac9d905a9d553b95155a993338568b8.tar.gz
llamachat-9182c84b8ac9d905a9d553b95155a993338568b8.zip
feat: add context meter and file-based system prompts
Two additions to the chat window. A context meter in the top bar shows how much of the active model's window the next request will occupy: system prompt, prior turns, attachments and the current draft. The router does not report a usable context size in router mode (/props returns n_ctx 0), so the limit comes from ctx-size in presets.ini. Streaming requests now ask for usage statistics, which makes the figure exact after each reply at no extra round trip; before that it is an estimate marked with a leading ~. The estimate reads low on models with a reasoning budget, since thinking tokens cannot be known before the reply arrives. The tooltip says so. System prompts are markdown files in the prompts/ directory beside the config, one per file, so they can be edited in an editor and kept in version control. default.md is global, any other file is a named preset that replaces it rather than adding to it, and a conversation may instead carry one-off text or opt out entirely. The choice is stored per session so reopening a chat restores the prompt it was built with. Prompt names are reduced to safe filenames, so a name like ../../etc/passwd cannot write outside the prompts directory, and a name that reduces to nothing is rejected rather than creating a dotfile. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Diffstat (limited to 'test_llamachat.py')
-rwxr-xr-xtest_llamachat.py175
1 files changed, 175 insertions, 0 deletions
diff --git a/test_llamachat.py b/test_llamachat.py
index 4a113e0..7707b7b 100755
--- a/test_llamachat.py
+++ b/test_llamachat.py
@@ -14,6 +14,7 @@
# GNU General Public License for more details.
"""Self-checks for the non-GUI logic. Run: ./test_llamachat.py"""
+import json
import os
import sys
import tempfile
@@ -409,6 +410,175 @@ def test_system_qt_theme_guard():
print("ok system qt theme guard")
+def test_prompt_store():
+ """Prompts are files; a preset replaces the global one."""
+ from llamachat import prompts
+
+ with tempfile.TemporaryDirectory() as tmp:
+ store = prompts.PromptStore(Path(tmp) / "prompts")
+ assert store.names() == []
+
+ store.ensure_default()
+ assert store.names() == ["default"]
+ assert store.path_for("default").exists()
+
+ # ensure_default must not clobber an edited global prompt.
+ store.save("default", "edited global")
+ store.ensure_default()
+ assert store.load("default").text == "edited global"
+
+ store.save("coding", "you write code")
+ store.save("terse", "be brief")
+ # Global first, then presets alphabetically.
+ assert store.names() == ["default", "coding", "terse"]
+
+ # A preset replaces the global prompt rather than adding to it.
+ assert store.resolve("coding") == "you write code"
+ assert store.resolve("") == "edited global"
+ assert store.resolve("default") == "edited global"
+
+ # Sentinels.
+ assert store.resolve(prompts.NONE) == ""
+ assert store.resolve(prompts.CUSTOM, "one off") == "one off"
+
+ # A preset deleted behind our back falls back rather than sending
+ # nothing, which would silently change the model's behaviour.
+ assert store.resolve("missing") == "edited global"
+
+ assert store.delete("coding") is True
+ assert store.names() == ["default", "terse"]
+ # The global prompt is the fallback for everything, so it stays.
+ assert store.delete("default") is False
+ assert "default" in store.names()
+
+ # Names become filenames, so path separators must not escape.
+ assert prompts.safe_name("../../etc/passwd") == "etc-passwd"
+ assert prompts.safe_name("/etc/passwd") == "etc-passwd"
+ assert prompts.safe_name("my prompt!") == "my-prompt"
+ for hostile in ("../escape", "a/b", "/tmp/x"):
+ saved = store.save(hostile, "x")
+ assert saved.path.parent == store.dir, saved.path
+ assert saved.path.resolve().parent == store.dir.resolve()
+
+ # A name that reduces to nothing is rejected rather than writing
+ # a dotfile called ".md".
+ for empty in ("..", "", "....", "///"):
+ assert prompts.safe_name(empty) == ""
+ assert store.load(empty) is None
+ assert store.delete(empty) is False
+ try:
+ store.path_for(empty)
+ except ValueError:
+ pass
+ else:
+ raise AssertionError(f"{empty!r} should be rejected")
+ print("ok prompt store")
+
+
+def test_token_estimate():
+ """The meter's pre-send estimate tracks message size."""
+ small = [{"role": "user", "content": "hi"}]
+ large = [{"role": "user", "content": "word " * 1000}]
+
+ assert backend.estimate_tokens([], 3.5) == 0
+ assert backend.estimate_tokens(small, 3.5) < 20
+ assert backend.estimate_tokens(large, 3.5) > 1000
+
+ # More history means a bigger prompt.
+ grown = large + [{"role": "assistant", "content": "reply " * 500}]
+ assert backend.estimate_tokens(grown, 3.5) > backend.estimate_tokens(
+ large, 3.5
+ )
+
+ # An image costs far more than its text part suggests.
+ with_image = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "what is this"},
+ {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
+ ],
+ }
+ ]
+ assert backend.estimate_tokens(with_image, 3.5) > 500
+
+ # A silly ratio must not divide by zero.
+ assert backend.estimate_tokens(small, 0) > 0
+ print("ok token estimate")
+
+
+def test_usage_parsing():
+ """The final stream chunk carries the exact prompt token count."""
+ line = (
+ 'data: {"choices":[],"usage":{"prompt_tokens":1234,'
+ '"completion_tokens":56,"total_tokens":1290}}'
+ )
+ kind, payload = backend._parse_sse_line(line)
+ assert kind == "usage"
+ stats = json.loads(payload)
+ assert stats["prompt_tokens"] == 1234
+ assert stats["total_tokens"] == 1290
+
+ # A usage-less chunk with empty choices is not mistaken for one.
+ assert backend._parse_sse_line('data: {"choices":[]}') is None
+ print("ok usage parsing")
+
+
+def test_session_prompt_storage():
+ """A conversation remembers the prompt it was built with."""
+ from llamachat import prompts
+
+ with tempfile.TemporaryDirectory() as tmp:
+ history = db.History(Path(tmp) / "p.db")
+
+ sid = history.create_session(
+ "chat", "m", "t", prompt_name="coding", prompt_custom=""
+ )
+ row = history.get_session(sid)
+ assert row["prompt_name"] == "coding"
+ assert row["prompt_custom"] == ""
+
+ history.set_prompt(sid, prompts.CUSTOM, "just this once")
+ row = history.get_session(sid)
+ assert row["prompt_name"] == prompts.CUSTOM
+ assert row["prompt_custom"] == "just this once"
+
+ # Defaults keep older call sites working.
+ plain = history.create_session("oneshot", "m", "t")
+ assert history.get_session(plain)["prompt_name"] == ""
+ history.close()
+ print("ok session prompt storage")
+
+
+def test_prompt_column_migration():
+ """A database predating the prompt columns must still open."""
+ 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);"
+ )
+ conn.commit()
+ conn.close()
+
+ history = db.History(path)
+ row = history.get_session(1)
+ assert row["title"] == "old"
+ assert row["prompt_name"] == ""
+ assert row["prompt_custom"] == ""
+ history.set_prompt(1, "coding")
+ assert history.get_session(1)["prompt_name"] == "coding"
+ history.close()
+ print("ok prompt column migration")
+
+
def test_version_matches_changelog():
"""The package version must be the newest release in the changelog."""
import re
@@ -506,6 +676,11 @@ if __name__ == "__main__":
test_reasoning_storage()
test_migration_adds_reasoning()
test_user_content()
+ test_prompt_store()
+ test_token_estimate()
+ test_usage_parsing()
+ test_session_prompt_storage()
+ test_prompt_column_migration()
test_markdown_rendering()
test_system_qt_theme_guard()
test_version_matches_changelog()