aboutsummaryrefslogtreecommitdiffstats
path: root/test_llamachat.py
diff options
context:
space:
mode:
Diffstat (limited to 'test_llamachat.py')
-rwxr-xr-xtest_llamachat.py589
1 files changed, 583 insertions, 6 deletions
diff --git a/test_llamachat.py b/test_llamachat.py
index f23a576..936a18d 100755
--- a/test_llamachat.py
+++ b/test_llamachat.py
@@ -88,12 +88,20 @@ def test_real_presets():
print("skip real presets (file absent)")
return
presets = config.parse_presets(path)
- assert "gemma-4-12B-it" in presets
- assert presets["gemma-4-12B-it"].vision is True
- # Qwen3.5-9B has its mmproj line commented out with '#'.
- if "Qwen3.5-9B" in presets:
- assert presets["Qwen3.5-9B"].vision is False
- print("ok real presets classification")
+ assert presets, "presets.ini exists but parsed to nothing"
+
+ # Section names track whatever the user currently runs, so assert the
+ # parsing properties rather than a list of names that goes stale on
+ # every rename.
+ for preset in presets.values():
+ assert preset.name
+ assert preset.ctx_size > 0
+ assert isinstance(preset.vision, bool)
+
+ # Whether a given section has vision is the user's choice and changes
+ # when they edit the file; only the commented-out case is a parsing
+ # claim, and PRESETS_SAMPLE covers that hermetically above.
+ print(f"ok real presets classification ({len(presets)} sections)")
def test_fts_query_escaping():
@@ -897,6 +905,567 @@ def test_config_defaults():
print("ok config defaults")
+class _FakeResponse:
+ """Enough of an http.client response for urlopen's context manager."""
+
+ def __init__(self, body: bytes):
+ self._body = body
+
+ def read(self):
+ return self._body
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ return False
+
+
+def _fake_urlopen(body, capture=None):
+ """A urlopen replacement returning `body`, or raising it when an error."""
+
+ def opener(request, timeout=None):
+ if capture is not None:
+ capture.append(request.full_url)
+ if isinstance(body, Exception):
+ raise body
+ return _FakeResponse(body)
+
+ return opener
+
+
+def _with_urlopen(body, capture=None):
+ """Swap search's urlopen for a fake. Returns the original to restore."""
+ import urllib.request
+
+ original = urllib.request.urlopen
+ urllib.request.urlopen = _fake_urlopen(body, capture)
+ return original
+
+
+def test_search_tool_schema():
+ from llamachat import search
+
+ schema = search.TOOL_SCHEMA
+ assert schema["type"] == "function"
+ function = schema["function"]
+ assert function["name"] == "web_search"
+ assert function["description"]
+ params = function["parameters"]
+ assert params["type"] == "object"
+ assert params["required"] == ["query"]
+ assert params["properties"]["query"]["type"] == "string"
+
+ # Disabled search must not put a tools key on the wire at all, or a
+ # model that ignores it still pays for the tokens.
+ sent = []
+ client = backend.Client("http://x")
+ client._stream_once = lambda model, messages, tools: (
+ sent.append(tools) or iter([("content", "hi")])
+ )
+ out = list(client.stream_chat("m", [], backend.SearchConfig(enabled=False)))
+ assert out == [("content", "hi")]
+ assert sent == [None], sent
+
+ # A configured-but-urlless setup resolves to disabled at config load.
+ with tempfile.TemporaryDirectory() as tmp:
+ path = Path(tmp) / "config.toml"
+ path.write_text('search_enabled = true\nsearch_url = ""\n')
+ assert config.load(path).search_enabled is False
+
+ path.write_text(
+ 'search_enabled = true\nsearch_url = "http://searx.local:8888/"\n'
+ )
+ cfg = config.load(path)
+ assert cfg.search_enabled is True
+ assert cfg.search_url == "http://searx.local:8888" # trailing / gone
+ print("ok search tool schema")
+
+
+def test_search_results_sanitising():
+ import urllib.request
+
+ from llamachat import search
+ from llamachat.ui import SEARCH_SCHEME, _markdown_to_fragment, _search_html
+
+ payload = json.dumps(
+ {
+ "results": [
+ {
+ "title": "First",
+ "url": "https://example.com/a",
+ "content": "x" * 500,
+ # Fields the model has no business seeing.
+ "engine": "duckduckgo",
+ "score": 1.5,
+ "positions": [1],
+ },
+ {"title": "Second", "url": "https://example.com/b", "content": "s"},
+ {"title": "Third", "url": "https://example.com/c", "content": "t"},
+ ]
+ }
+ ).encode()
+
+ original = _with_urlopen(payload)
+ try:
+ results = search.search("http://searx", "q", count=2, snippet_chars=100)
+ finally:
+ urllib.request.urlopen = original
+
+ # count caps the list; only three fields survive; snippets truncate.
+ assert len(results) == 2, results
+ assert set(results[0]) == {"title", "url", "content"}, results[0]
+ assert len(results[0]["content"]) == 100
+ assert "engine" not in results[0]
+
+ # Result text is escaped on display, so markup in a snippet stays text.
+ hostile = [
+ {
+ "query": "q",
+ "error": "",
+ "results": [
+ {
+ "title": "<script>alert(1)</script>",
+ "url": "https://example.com/x",
+ "content": f"click [here]({SEARCH_SCHEME}0)",
+ }
+ ],
+ }
+ ]
+ rendered = _search_html(hostile, 0, expanded=True)
+ assert "<script>" not in rendered, rendered
+ assert "&lt;script&gt;" in rendered
+ # The scheme appears exactly once, as this block's own toggle. The copy
+ # inside the snippet stayed literal text rather than becoming an href.
+ assert rendered.count(f'href="{SEARCH_SCHEME}') == 1, rendered
+
+ # A reply forging the scheme as a markdown link gets it defused.
+ forged = _markdown_to_fragment(f"[expand]({SEARCH_SCHEME}0)")
+ assert SEARCH_SCHEME not in forged, forged
+ assert "blocked:" in forged
+ print("ok search result sanitising")
+
+
+def test_tool_call_accumulation():
+ from llamachat import search
+
+ calls: dict = {}
+ # The id and name arrive first, the arguments in pieces after it.
+ search.accumulate(calls, [
+ {"index": 0, "id": "call_1", "function": {"name": "web_search", "arguments": ""}}
+ ])
+ search.accumulate(calls, [{"index": 0, "function": {"arguments": '{"que'}}])
+ search.accumulate(calls, [{"index": 0, "function": {"arguments": 'ry": "kern'}}])
+ search.accumulate(calls, [{"index": 0, "function": {"arguments": 'el"}'}}])
+
+ assert list(calls) == [0]
+ assert calls[0]["id"] == "call_1"
+ assert calls[0]["name"] == "web_search"
+ assert search.parse_query(calls[0]["arguments"]) == "kernel"
+
+ # Two calls in one round stay apart, keyed by index.
+ pair: dict = {}
+ search.accumulate(pair, [
+ {"index": 0, "id": "a", "function": {"name": "web_search", "arguments": '{"query":"one"}'}},
+ {"index": 1, "id": "b", "function": {"name": "web_search", "arguments": '{"query":"two"}'}},
+ ])
+ assert search.parse_query(pair[0]["arguments"]) == "one"
+ assert search.parse_query(pair[1]["arguments"]) == "two"
+
+ # Unusable arguments yield no query rather than an exception.
+ assert search.parse_query("{not json") == ""
+ assert search.parse_query("[]") == ""
+ assert search.parse_query('{"query": 7}') == ""
+ assert search.parse_query("") == ""
+ print("ok tool call accumulation")
+
+
+def _tool_round(query: str = "q"):
+ """One streamed round that asks for a search."""
+ return [
+ (
+ "tool_calls",
+ json.dumps(
+ [
+ {
+ "index": 0,
+ "id": "call_1",
+ "function": {
+ "name": "web_search",
+ "arguments": json.dumps({"query": query}),
+ },
+ }
+ ]
+ ),
+ ),
+ ("tool_finish", ""),
+ ]
+
+
+def test_search_loop_cap():
+ import urllib.request
+
+ from llamachat import search
+
+ sent_tools = []
+ client = backend.Client("http://x")
+
+ def always_tool_calls(model, messages, tools):
+ sent_tools.append(tools)
+ yield from _tool_round()
+
+ client._stream_once = always_tool_calls
+ original = _with_urlopen(json.dumps({"results": []}).encode())
+ try:
+ cfg = backend.SearchConfig(
+ enabled=True, url="http://searx", max_searches=2
+ )
+ list(client.stream_chat("m", [{"role": "user", "content": "hi"}], cfg))
+ finally:
+ urllib.request.urlopen = original
+
+ # max_searches rounds offer the tool, then one final round without it.
+ assert len(sent_tools) == 3, sent_tools
+ assert sent_tools[0] == [search.TOOL_SCHEMA]
+ assert sent_tools[1] == [search.TOOL_SCHEMA]
+ assert sent_tools[-1] is None, "the final round must withdraw the tool"
+
+ # A round that never asks for a tool ends the turn immediately.
+ quiet = []
+
+ def no_tools(model, messages, tools):
+ quiet.append(tools)
+ yield ("content", "done")
+
+ client._stream_once = no_tools
+ out = list(client.stream_chat("m", [], backend.SearchConfig(
+ enabled=True, url="http://searx")))
+ assert out == [("content", "done")]
+ assert len(quiet) == 1, quiet
+ print("ok search loop cap")
+
+
+def test_search_failure_paths():
+ import urllib.error
+ import urllib.request
+
+ from llamachat import search
+
+ cases = [
+ (TimeoutError("timed out"), "Cannot reach SearXNG"),
+ (urllib.error.URLError("Connection refused"), "Cannot reach SearXNG"),
+ (b"<html>not json</html>", "did not return JSON"),
+ ]
+ for body, expected in cases:
+ original = _with_urlopen(body)
+ try:
+ failed = ""
+ try:
+ search.search("http://searx", "q")
+ except search.SearchError as exc:
+ failed = str(exc)
+ finally:
+ urllib.request.urlopen = original
+ assert expected in failed, (body, failed)
+
+ # Zero results with healthy engines is a success with an empty list:
+ # the web really had nothing.
+ original = _with_urlopen(json.dumps({"results": []}).encode())
+ try:
+ assert search.search("http://searx", "q") == []
+ finally:
+ urllib.request.urlopen = original
+
+ # Zero results *because* every engine was rate-limited or CAPTCHA'd is
+ # a failed search. Reporting it as "no results" would tell the model
+ # the web is empty and invite an answer from stale training data.
+ dead = json.dumps(
+ {
+ "results": [],
+ "unresponsive_engines": [
+ ["duckduckgo", "CAPTCHA"],
+ ["brave", "Suspended: too many requests"],
+ ],
+ }
+ ).encode()
+ original = _with_urlopen(dead)
+ try:
+ failed = ""
+ try:
+ search.search("http://searx", "q")
+ except search.SearchError as exc:
+ failed = str(exc)
+ finally:
+ urllib.request.urlopen = original
+ assert "every search engine failed" in failed, failed
+ assert "duckduckgo: CAPTCHA" in failed, failed
+
+ # Engines that failed while others still answered are not an error:
+ # partial results are results.
+ partial = json.dumps(
+ {
+ "results": [{"title": "T", "url": "https://e.com", "content": "c"}],
+ "unresponsive_engines": [["brave", "timeout"]],
+ }
+ ).encode()
+ original = _with_urlopen(partial)
+ try:
+ assert len(search.search("http://searx", "q")) == 1
+ finally:
+ urllib.request.urlopen = original
+
+ # A malformed unresponsive_engines field must not crash the summary.
+ assert search._unresponsive({"unresponsive_engines": "nonsense"}) == ""
+ assert search._unresponsive({"unresponsive_engines": [["solo"]]}) == "solo"
+ assert search._unresponsive({}) == ""
+
+ # An unconfigured URL fails before any request is attempted.
+ try:
+ search.search("", "q")
+ raise AssertionError("empty url must raise")
+ except search.SearchError:
+ pass
+
+ # Every failure still completes the turn and still tells the model.
+ client = backend.Client("http://x")
+ rounds = [_tool_round("kernel"), [("content", "answered anyway")]]
+ seen_messages = []
+
+ def scripted(model, messages, tools):
+ seen_messages.append([dict(m) for m in messages])
+ yield from rounds[min(len(seen_messages) - 1, len(rounds) - 1)]
+
+ client._stream_once = scripted
+ original = _with_urlopen(urllib.error.URLError("Connection refused"))
+ try:
+ cfg = backend.SearchConfig(enabled=True, url="http://searx")
+ out = list(client.stream_chat("m", [{"role": "user", "content": "hi"}], cfg))
+ finally:
+ urllib.request.urlopen = original
+
+ assert ("content", "answered anyway") in out, out
+ starts = [p for k, p in out if k == "search_start"]
+ dones = [json.loads(p) for k, p in out if k == "search_done"]
+ assert starts == ["kernel"], starts
+ assert dones and dones[0]["error"], dones
+ assert dones[0]["results"] == []
+
+ # The second request carries the assistant tool call and a tool reply
+ # naming the failure, so the model knows the search did not happen.
+ second = seen_messages[1]
+ assert second[-2]["role"] == "assistant"
+ assert second[-2]["tool_calls"][0]["id"] == "call_1"
+ tool_msg = second[-1]
+ assert tool_msg["role"] == "tool"
+ assert tool_msg["tool_call_id"] == "call_1"
+ assert search.RESULT_PREFIX in tool_msg["content"]
+ assert "error" in json.loads(
+ tool_msg["content"][len(search.RESULT_PREFIX):].strip()
+ )
+
+ # A malformed tool call is answered rather than left dangling, or the
+ # next request would be rejected for an unanswered call.
+ bad = {"id": "call_9", "name": "web_search", "arguments": "{not json"}
+ convo: list = []
+ emitted = list(client._run_search(convo, bad, backend.SearchConfig(
+ enabled=True, url="http://searx")))
+ assert emitted == [], emitted # nothing searched, nothing displayed
+ assert convo[-1]["role"] == "tool"
+ assert convo[-1]["tool_call_id"] == "call_9"
+ print("ok search failure paths")
+
+
+def test_search_storage():
+ import sqlite3
+
+ with tempfile.TemporaryDirectory() as tmp:
+ history = db.History(Path(tmp) / "s.db")
+ sid = history.create_session("chat", "m", "t")
+ mid = history.add_message(sid, "assistant", "")
+
+ records = [
+ {
+ "query": "latest kernel",
+ "results": [
+ {
+ "title": "Kernel",
+ "url": "https://kernel.org",
+ "content": "snippet",
+ }
+ ],
+ "error": "",
+ }
+ ]
+ history.update_message(mid, "7.1.5", "thinking", json.dumps(records))
+
+ row = history.messages(sid)[0]
+ assert row["content"] == "7.1.5"
+ assert row["reasoning"] == "thinking"
+ assert json.loads(row["searches"]) == records
+
+ # Snippets stay out of the FTS index: web text the user never wrote
+ # must not compete with their own messages.
+ assert history.search("7.1.5") != []
+ assert history.search("snippet") == []
+
+ # Omitting the argument leaves stored searches untouched.
+ history.update_message(mid, "7.1.6")
+ assert json.loads(history.messages(sid)[0]["searches"]) == records
+
+ # A message with no searches stores NULL rather than an empty list.
+ plain = history.add_message(sid, "assistant", "")
+ history.update_message(plain, "no search", "", None)
+ assert history.messages(sid)[1]["searches"] is None
+ history.close()
+
+ # A database predating the column must still open and read as no
+ # searches rather than raising.
+ 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',0);"
+ )
+ conn.commit()
+ conn.close()
+
+ history = db.History(path)
+ from llamachat.ui import _column, _searches
+
+ row = history.messages(1)[0]
+ assert row["searches"] is None
+ assert _searches(_column(row, "searches")) == []
+ history.close()
+ print("ok search storage")
+
+
+def test_date_note():
+ import datetime
+
+ from llamachat import ui
+ from llamachat.ui import _date_note
+
+ fixed = datetime.datetime(2026, 7, 31, 14, 30)
+ note = _date_note(fixed)
+
+ # The real date must be stated, or the model falls back on its cutoff.
+ assert "Friday, 31 July 2026" in note, note
+ # And it must be told not to date its own queries, which is the bug
+ # that poisons results before they are even fetched.
+ assert "year in a search query" in note, note
+
+ # Composed onto the chosen prompt rather than replacing it, so a preset
+ # keeps its instructions.
+ class _Cfg:
+ search_enabled = True
+
+ class _Win:
+ cfg = _Cfg()
+ system_prompt_text = staticmethod(lambda: "Be terse.")
+ _system_messages = ui.ChatWindow._system_messages
+
+ msgs = _Win._system_messages(_Win())
+ assert len(msgs) == 1
+ assert msgs[0]["role"] == "system"
+ assert msgs[0]["content"].startswith("Be terse.")
+ assert "2026" in msgs[0]["content"]
+
+ # With no system prompt selected, the date still goes: the model has no
+ # clock either way.
+ class _Bare(_Win):
+ system_prompt_text = staticmethod(lambda: "")
+
+ bare = _Bare._system_messages(_Bare())
+ assert len(bare) == 1
+ assert "Today's date is" in bare[0]["content"]
+
+ # Search off means no date note and no system message at all.
+ class _Off(_Bare):
+ class cfg:
+ search_enabled = False
+
+ assert _Off._system_messages(_Off()) == []
+ print("ok date note")
+
+
+def test_search_html():
+ from llamachat.ui import SEARCH_SCHEME, _search_html, _searches
+
+ assert _search_html([], 0, False) == ""
+
+ ok = [
+ {
+ "query": "latest kernel",
+ "error": "",
+ "results": [
+ {
+ "title": "Kernel.org",
+ "url": "https://kernel.org",
+ "content": "The Linux Kernel Archives",
+ },
+ {
+ "title": "Wikipedia",
+ "url": "https://en.wikipedia.org/wiki/Linux",
+ "content": "An operating system kernel",
+ },
+ ],
+ }
+ ]
+
+ collapsed = _search_html(ok, 3, expanded=False)
+ assert "▸" in collapsed
+ assert "searched: latest kernel (2 results)" in collapsed
+ assert f'href="{SEARCH_SCHEME}3"' in collapsed
+ # Collapsed shows the summary only, never the sources.
+ assert "kernel.org" not in collapsed
+
+ expanded = _search_html(ok, 3, expanded=True)
+ assert "▾" in expanded
+ assert "Kernel.org" in expanded
+ assert 'href="https://kernel.org"' in expanded
+ assert "The Linux Kernel Archives" in expanded
+
+ # A failure says so, and names the reason.
+ failed = _search_html(
+ [{"query": "kernel", "results": [], "error": "Cannot reach SearXNG"}],
+ 0,
+ expanded=False,
+ )
+ assert "search failed: kernel" in failed
+ assert "Cannot reach SearXNG" in failed
+
+ # A non-web URL is shown but never becomes a clickable anchor.
+ sneaky = _search_html(
+ [
+ {
+ "query": "q",
+ "error": "",
+ "results": [
+ {"title": "T", "url": "file:///etc/passwd", "content": "c"}
+ ],
+ }
+ ],
+ 0,
+ expanded=True,
+ )
+ assert 'href="file:' not in sneaky, sneaky
+ assert "file:///etc/passwd" in sneaky
+
+ # The stored column decodes back into what rendering expects, and junk
+ # in that column degrades to no block rather than raising.
+ assert _searches(json.dumps(ok)) == ok
+ assert _searches("") == []
+ assert _searches("{not json") == []
+ assert _searches('{"query": "not a list"}') == []
+ print("ok search html")
+
+
if __name__ == "__main__":
test_presets()
test_real_presets()
@@ -921,4 +1490,12 @@ if __name__ == "__main__":
test_version_matches_changelog()
test_venv_discovery()
test_config_defaults()
+ test_search_tool_schema()
+ test_search_results_sanitising()
+ test_tool_call_accumulation()
+ test_search_loop_cap()
+ test_search_failure_paths()
+ test_search_storage()
+ test_date_note()
+ test_search_html()
print("\nall checks passed")