aboutsummaryrefslogtreecommitdiffstats
path: root/test_llamachat.py
diff options
context:
space:
mode:
authorDanilo M. <danix@danix.xyz>2026-07-31 18:36:13 +0200
committerDanilo M. <danix@danix.xyz>2026-07-31 18:36:13 +0200
commit283fc9ae0371f4d318a539f021a4b88e05fac190 (patch)
tree0bf03685df84c275cb7ab3af3e20d94781a078cc /test_llamachat.py
parent8bda554af191c5749459cd258c793e4826c3f9c7 (diff)
downloadllamachat-283fc9ae0371f4d318a539f021a4b88e05fac190.tar.gz
llamachat-283fc9ae0371f4d318a539f021a4b88e05fac190.zip
feat: web search via SearXNG
The model is offered a web_search tool and decides when a question needs current information. llamachat runs the query against SearXNG's JSON API, feeds the results back, and the model answers from them. Off unless both search_enabled and search_url are set, so an upgrade never starts talking to the network on its own. The loop lives in backend.stream_chat: a round ending in tool_calls is searched for, the result appended, and the request re-sent. After max_searches rounds the tool is withdrawn, which forces an answer rather than letting an uncertain model search forever. Searches show as a collapsible block above the thinking block, listing each query and its sources. The queries are visible on purpose: when an answer is wrong it is usually the query that was wrong, and without seeing it a bad search and a bad answer look identical. Two failures found while testing against the live stack shaped the design: The model has no clock, so it falls back on its training cutoff and writes that year into the query itself ("latest kernel ... 2025"), poisoning the results before they are fetched. The current date now goes into the system prompt whenever search is on, with an instruction not to date its own queries. A SearXNG whose engines are all rate-limited or CAPTCHA'd returns a valid response with zero results. Reporting that as "no results" tells the model the web is empty and invites a confident answer from stale training data, so a search where every engine failed is now an error naming the engines. Results are attacker-influenced text entering the model's context. Only title, url and the snippet survive, snippets are truncated, result text is escaped on display, result links are never fetched automatically, and a reply forging the search block's URL scheme has it defused as the reasoning scheme already was. None of that stops a poisoned snippet from influencing the answer, which is why the sources stay visible. Adds eight test groups, 23 to 31, all hermetic behind a fake HTTP layer. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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")