From 283fc9ae0371f4d318a539f021a4b88e05fac190 Mon Sep 17 00:00:00 2001 From: "Danilo M." Date: Fri, 31 Jul 2026 18:36:13 +0200 Subject: 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 --- CHANGELOG.md | 40 ++ README.md | 92 +++- .../specs/2026-07-31-web-search-design.md | 2 +- llamachat/backend.py | 116 +++- llamachat/config.py | 46 ++ llamachat/db.py | 37 +- llamachat/search.py | 210 ++++++++ llamachat/ui.py | 211 +++++++- test_llamachat.py | 589 ++++++++++++++++++++- 9 files changed, 1302 insertions(+), 41 deletions(-) create mode 100644 llamachat/search.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d3542e..7b3aad6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,46 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Web search through a SearXNG instance. The model is offered a `web_search` + tool and decides for itself when a question needs current information; + llamachat runs the query and feeds the results back. Off by default, and + enabled only when both `search_enabled` and `search_url` are set, so an + upgrade never starts talking to the network on its own. +- Searches appear in the transcript as a collapsible block above the + thinking block, listing each query and its sources, with a status line + while a search runs. The queries are visible because a bad answer is + usually a bad query, and showing the sources is what lets a poisoned + result be recognised as one. +- Six configuration keys: `search_enabled`, `search_url`, `search_results`, + `search_snippet_chars`, `search_timeout` and `max_searches`. Searches are + capped per turn, after which the tool is withdrawn and the model must + answer. +- Searches are stored per message in a new `searches` column, so reopening a + conversation still shows what was looked up. They are deliberately kept + out of the full-text index, or web text nobody wrote would compete with + the user's own messages. +- The current date is added to the system prompt when search is on. A model + has no clock and falls back on its training cutoff, which it then writes + into the query itself ("latest kernel ... 2025"), poisoning the results + before they are fetched. It is also told not to date its own queries. +- A search where every engine failed is now reported as a failure naming + the engines, instead of as a successful search that found nothing. Rate + limits and CAPTCHAs are the normal way a self-hosted SearXNG stops + working, and "no results" invited a confident answer from stale training + data. + +### Security + +- Search results are untrusted text entering the model's context. Only + `title`, `url` and the snippet survive, snippets are truncated, all result + text is escaped on display, result links are shown but never fetched + automatically, and a reply forging the search toggle's URL scheme has it + defused the same way the reasoning scheme already was. A poisoned snippet + can still influence what the model says; the sources are shown so that it + can be judged. + ## [0.2.1] - 2026-07-31 ### Fixed diff --git a/README.md b/README.md index 778010f..546fa1b 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,10 @@ persistent process and toggles like a scratchpad from a Hyprland keybind. any conversation can take a one-off custom prompt or none at all. The choice is stored per session, so reopening an old chat restores the prompt it was built with. +- **Web search** through a SearXNG instance, off by default. The model is + offered a `web_search` tool and calls it when a question needs current + information; the queries and their sources show as a collapsible block + above the reply. See [Web search](#web-search). - **Tray icon** for show/hide/quit, hosted by waybar's tray module. Not in this version: RAG or embedding search over history, multi-user @@ -132,6 +136,15 @@ default_prompt = "" attach_ctx_fraction = 0.5 chars_per_token = 3.5 + +# Web search. Both keys are required: a flag with no URL stays off. +search_enabled = false +search_url = "" + +search_results = 5 +search_snippet_chars = 300 +search_timeout = 10 +max_searches = 2 ``` `presets.ini` is read for two things the API does not report: which models @@ -242,7 +255,8 @@ llamachat.py launcher, venv shebang, symlink target llamachat/ __main__.py argument dispatch, daemon, tray, socket event loop config.py config.toml and presets.ini parsing - backend.py httpx streaming client, attachment loading + backend.py httpx streaming client, tool loop, attachment loading + search.py SearXNG queries, tool schema, result sanitising db.py SQLite schema, FTS5 search ui.py the window test_llamachat.py self-checks for everything except the GUI @@ -271,7 +285,7 @@ there is no second thread. ```sql sessions (id, mode, title, model, prompt_name, prompt_custom, created_at, updated_at) -messages (id, session_id, role, content, reasoning, created_at) +messages (id, session_id, role, content, reasoning, searches, created_at) attachments (id, message_id, path, kind, mime, size, sha256, thumb, truncated) messages_fts -- FTS5 external-content table over messages.content ``` @@ -285,6 +299,11 @@ reasons: the FTS index covers `content` only, so thinking text cannot bury real search hits, and only `content` is replayed when continuing a conversation, which is what these APIs expect. +`searches` holds a JSON array of what was looked up for that reply, one entry +per search with its query, results and any error, or NULL when nothing was +searched. Like `reasoning` it stays out of the FTS index, or snippets from +web pages would compete with messages the user actually wrote. + `prompt_name` records which system prompt a conversation was built with, and `prompt_custom` holds the text when that prompt is a one-off rather than a file. Storing the name rather than the resolved text means editing a preset @@ -331,6 +350,75 @@ a preset name, `"none"`, or empty for `default.md`. The choice is recorded per session, so reopening a chat restores the prompt it was built with rather than silently applying whatever is selected now. +### Web search + +Off unless both `search_enabled = true` and a `search_url` are set. A flag +with no URL behind it resolves to disabled at startup, so an upgrade never +starts making network requests on its own. + +The model decides when to search. It is offered an OpenAI-style `web_search` +function, and when it judges a question needs current information it calls +it; llamachat queries SearXNG's JSON API, hands back the results, and the +model answers from them. Nothing needs enabling on the llama-server side: +this is plain tool calling, not llama-server's `--tools` (which runs +server-side tools, none of which search the web) and not `--ui-mcp-proxy` +(a CORS shim that exists for browsers). + +The SearXNG instance must have its JSON output format enabled, which many +installs turn off by default. Without it, queries fail with a message saying +so. + +Searches show in the transcript as a collapsible block above the thinking +block, with a status line while one runs: + +``` +▸ searched: latest stable linux kernel (5 results) +▸ thinking (1441 chars, 38 lines) +Model: +... +``` + +Expanding it lists each result's title, link and snippet. 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. + +`max_searches` caps searches per turn, default 2. On the last round the tool +is withdrawn from the request, which forces the model to answer instead of +searching again. A local 9B will otherwise keep searching when it is unsure. + +Failures do not abort the turn. A timeout, a refused connection, a non-JSON +response or zero results all come back to the model as a tool result saying +what happened, so it answers with that knowledge rather than silently +falling back on stale training data. The block then reads +`▸ search failed: `. + +A search that returns nothing *because every engine failed* counts as a +failure rather than an empty result, and the message names the engines. On +a self-hosted SearXNG the usual way search stops working is engines getting +rate-limited or CAPTCHA'd one by one, and reporting that as "no results" +tells the model the web is empty, which invites exactly the confident wrong +answer this feature exists to prevent. + +When search is on, the current date is appended to the system prompt. A +model has no clock, so it falls back on its training cutoff and writes that +year into the search query itself, poisoning the results before they are +fetched. It is also told not to put a year in a query unless asked. This is +added even when `no system prompt` is selected, since the model has no clock +either way. + +**On trusting results.** Search snippets are attacker-influenced text going +into the model's context, in an app whose output you may paste into a +terminal. Structurally: only `title`, `url` and the snippet are kept and +everything else SearXNG returns is dropped, snippets are truncated to +`search_snippet_chars`, all result text is escaped when displayed, result +links are shown but never fetched automatically, and a reply that forges the +search block's internal URL scheme has it defused. The tool result is also +labelled untrusted for the model, which helps inconsistently. + +None of that stops a poisoned snippet from influencing what the model says. +If a top result asserts something false, a 9B may repeat it. That is why the +sources stay visible rather than being hidden behind the answer. + ### The context meter The bar in the top bar shows how much of the active model's context window diff --git a/docs/superpowers/specs/2026-07-31-web-search-design.md b/docs/superpowers/specs/2026-07-31-web-search-design.md index f7895b8..7e14b2d 100644 --- a/docs/superpowers/specs/2026-07-31-web-search-design.md +++ b/docs/superpowers/specs/2026-07-31-web-search-design.md @@ -1,7 +1,7 @@ # Web search via SearXNG Date: 2026-07-31 -Status: approved, not yet implemented +Status: implemented ## Problem diff --git a/llamachat/backend.py b/llamachat/backend.py index ec3b7d6..d3eb3fa 100644 --- a/llamachat/backend.py +++ b/llamachat/backend.py @@ -24,6 +24,8 @@ from pathlib import Path import httpx from PIL import Image +from . import search + THUMB_SIZE = (128, 128) # Extensions we are willing to read as text even though mimetypes does not @@ -47,6 +49,17 @@ class BackendError(Exception): """Any failure talking to the router, already phrased for the user.""" +@dataclass +class SearchConfig: + """What stream_chat needs to know to offer and run web search.""" + enabled: bool = False + url: str = "" + results: int = 5 + snippet_chars: int = 300 + timeout: int = 10 + max_searches: int = 2 + + @dataclass class Attachment: """A file the user attached, ready for both the API and the database.""" @@ -172,7 +185,9 @@ class Client: raise BackendError(f"Router sent invalid JSON: {exc}") return [m["id"] for m in payload.get("data", []) if "id" in m] - def stream_chat(self, model: str, messages: list[dict]): + def stream_chat( + self, model: str, messages: list[dict], search_cfg: "SearchConfig | None" = None + ): """Yield (kind, text) pairs as the model produces them. `kind` is 'reasoning' for the model's thinking and 'content' for the @@ -180,10 +195,99 @@ class Client: separate `reasoning_content` delta field, which is what lets the UI keep the two apart. + With `search_cfg`, the model is offered a `web_search` tool and this + becomes a loop: a round that ends in a tool call is searched for, the + result appended, and the request re-sent. Two further kinds are + yielded around each search, 'search_start' and 'search_done'. After + `max_searches` rounds the tool is withdrawn, which forces an answer. + Loading a different model makes the router unload the previous one, so the first chunk can take several seconds. That wait happens inside the initial `stream()` call. """ + if search_cfg is None or not search_cfg.enabled: + yield from self._stream_once(model, messages, tools=None) + return + + # The loop mutates its own copy; the caller's list is history. + convo = list(messages) + for round_number in range(search_cfg.max_searches + 1): + last_round = round_number == search_cfg.max_searches + tools = None if last_round else [search.TOOL_SCHEMA] + calls: dict = {} + saw_tool_finish = False + + for kind, piece in self._stream_once(model, convo, tools=tools): + if kind == "tool_calls": + search.accumulate(calls, json.loads(piece)) + elif kind == "tool_finish": + saw_tool_finish = True + else: + yield (kind, piece) + + wanted = [ + call + for call in calls.values() + if call["name"] == search.TOOL_SCHEMA["function"]["name"] + ] + if not (saw_tool_finish and wanted) or last_round: + return + + convo.append( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call["id"], + "type": "function", + "function": { + "name": call["name"], + "arguments": call["arguments"], + }, + } + for call in wanted + ], + } + ) + for call in wanted: + yield from self._run_search(convo, call, search_cfg) + + def _run_search(self, convo: list[dict], call: dict, cfg: "SearchConfig"): + """Perform one search, append its result to `convo`, report both ends. + + Every failure still appends a tool message: the model is waiting on + a result it asked for, and the turn has to complete either way. + """ + query = search.parse_query(call["arguments"]) + if not query: + # Unusable arguments: nothing to search, but the call still needs + # an answer or the next request is malformed. + convo.append( + search.tool_message(call["id"], None, error="Malformed search request") + ) + return + + yield ("search_start", query) + try: + results = search.search( + cfg.url, query, cfg.results, cfg.snippet_chars, cfg.timeout + ) + except search.SearchError as exc: + convo.append(search.tool_message(call["id"], None, error=str(exc))) + yield ( + "search_done", + json.dumps({"query": query, "results": [], "error": str(exc)}), + ) + return + convo.append(search.tool_message(call["id"], results)) + yield ( + "search_done", + json.dumps({"query": query, "results": results, "error": ""}), + ) + + def _stream_once(self, model: str, messages: list[dict], tools: list | None): + """One request/response round, yielding parsed SSE pairs.""" body = { "model": model, "messages": messages, @@ -193,6 +297,8 @@ class Client: # estimate, at no extra request. "stream_options": {"include_usage": True}, } + if tools: + body["tools"] = tools try: with httpx.stream( "POST", @@ -267,6 +373,14 @@ def _parse_sse_line(line: str) -> tuple[str, str] | None: return None delta = choices[0].get("delta") or {} + # Tool-call fragments arrive in their own delta field, alongside a + # finish_reason of 'tool_calls' on the final chunk of the round. + tool_calls = delta.get("tool_calls") + if tool_calls: + return ("tool_calls", json.dumps(tool_calls)) + if choices[0].get("finish_reason") == "tool_calls": + return ("tool_finish", "") + # Both fields can arrive in one delta, and a reasoning delta can be an # empty string. Test for presence, not truthiness, so a chunk carrying # reasoning is never mistaken for a content chunk. diff --git a/llamachat/config.py b/llamachat/config.py index dec2ba6..d7682f0 100644 --- a/llamachat/config.py +++ b/llamachat/config.py @@ -37,6 +37,14 @@ DEFAULTS = { # Named prompt selected for new conversations. Empty means the global # default.md, and "none" means no system prompt at all. "default_prompt": "", + # Web search through a SearXNG instance. Off unless both the flag is set + # and a URL is given, so an upgrade never starts talking to the network. + "search_enabled": False, + "search_url": "", + "search_results": 5, + "search_snippet_chars": 300, + "search_timeout": 10, + "max_searches": 2, } # The global prompt lives here; every other .md beside it is a named preset. @@ -67,6 +75,12 @@ class Config: default_prompt: str prompts_dir: Path state_path: Path + search_enabled: bool + search_url: str + search_results: int + search_snippet_chars: int + search_timeout: int + max_searches: int def _runtime_dir() -> Path: @@ -94,6 +108,11 @@ def load(path: Path = CONFIG_PATH) -> Config: socket = values["socket"] or (_runtime_dir() / "llamachat.sock") db = values["db"] or (_data_dir() / "history.db") + # A flag with no URL behind it is not enabled, it is misconfigured. + # Resolving that here means nothing downstream has to test both. + search_url = str(values["search_url"]).rstrip("/") + search_enabled = bool(values["search_enabled"]) and bool(search_url) + return Config( base_url=str(values["base_url"]).rstrip("/"), presets_path=Path(values["presets"]).expanduser(), @@ -108,6 +127,12 @@ def load(path: Path = CONFIG_PATH) -> Config: # Window layout, remembered between runs. Not user-editable config, # so it sits beside it rather than inside config.toml. state_path=path.parent / "state.ini", + search_enabled=search_enabled, + search_url=search_url, + search_results=int(values["search_results"]), + search_snippet_chars=int(values["search_snippet_chars"]), + search_timeout=int(values["search_timeout"]), + max_searches=int(values["max_searches"]), ) @@ -143,6 +168,27 @@ def write_default(path: Path = CONFIG_PATH) -> Path: '# chars-per-token estimate used to convert ctx-size into characters.\n' f'attach_ctx_fraction = {DEFAULTS["attach_ctx_fraction"]}\n' f'chars_per_token = {DEFAULTS["chars_per_token"]}\n' + '\n' + '# Web search through SearXNG. The model decides when to search; it\n' + '# is offered a web_search tool and calls it when a question needs\n' + '# current information. Both search_enabled and search_url are\n' + '# required: a flag with no URL stays off.\n' + '#\n' + '# The instance must have its JSON API enabled (format=json), which\n' + '# many installs turn off by default.\n' + '#\n' + '# Search results are untrusted text entering the model context.\n' + '# Only title, url and snippet are kept, snippets are truncated, and\n' + '# result links are shown but never fetched automatically.\n' + f'search_enabled = {str(DEFAULTS["search_enabled"]).lower()}\n' + f'search_url = "{DEFAULTS["search_url"]}"\n' + '\n' + '# Results per query, characters kept from each snippet, seconds\n' + '# before a query is abandoned, and searches allowed per turn.\n' + f'search_results = {DEFAULTS["search_results"]}\n' + f'search_snippet_chars = {DEFAULTS["search_snippet_chars"]}\n' + f'search_timeout = {DEFAULTS["search_timeout"]}\n' + f'max_searches = {DEFAULTS["max_searches"]}\n' ) return path diff --git a/llamachat/db.py b/llamachat/db.py index c13f03c..65a2b29 100644 --- a/llamachat/db.py +++ b/llamachat/db.py @@ -38,6 +38,7 @@ CREATE TABLE IF NOT EXISTS messages ( role TEXT NOT NULL, content TEXT NOT NULL, reasoning TEXT NOT NULL DEFAULT '', + searches TEXT, created_at INTEGER NOT NULL ); @@ -101,7 +102,12 @@ class History: def _migrate(self) -> None: """Add columns introduced after a database was first created.""" added = { - "messages": {"reasoning": "TEXT NOT NULL DEFAULT ''"}, + "messages": { + "reasoning": "TEXT NOT NULL DEFAULT ''", + # Nullable rather than defaulted: NULL means "no searches", + # which is exactly what every pre-migration row wants. + "searches": "TEXT", + }, "sessions": { "prompt_name": "TEXT NOT NULL DEFAULT ''", "prompt_custom": "TEXT NOT NULL DEFAULT ''", @@ -202,18 +208,25 @@ class History: return int(cur.lastrowid) def update_message( - self, message_id: int, content: str, reasoning: str | None = None + self, + message_id: int, + content: str, + reasoning: str | None = None, + searches: str | None = None, ) -> None: - if reasoning is None: - self.conn.execute( - "UPDATE messages SET content = ? WHERE id = ?", - (content, message_id), - ) - else: - self.conn.execute( - "UPDATE messages SET content = ?, reasoning = ? WHERE id = ?", - (content, reasoning, message_id), - ) + """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) + values.append(message_id) + self.conn.execute( + f"UPDATE messages SET {', '.join(columns)} WHERE id = ?", values + ) self.conn.commit() def messages(self, session_id: int) -> list[sqlite3.Row]: diff --git a/llamachat/search.py b/llamachat/search.py new file mode 100644 index 0000000..971ca90 --- /dev/null +++ b/llamachat/search.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: GPL-2.0-only +# +# llamachat - a small native chat client for a local llama.cpp router +# Copyright (C) 2026 Danilo M. +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License version 2 as +# published by the Free Software Foundation. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +"""Web search through a SearXNG instance. + +No Qt here: this is called from the streaming worker thread. stdlib urllib +rather than httpx, because a query is one GET with no streaming and no +connection reuse worth keeping. + +Results are attacker-influenced text on its way into the model's context, so +only three fields survive and snippets are truncated. See the security +section of docs/superpowers/specs/2026-07-31-web-search-design.md. +""" + +import json +import urllib.error +import urllib.parse +import urllib.request + +# The function the model is offered. Deliberately one required string +# parameter: a 9B picks a good query far more reliably than it fills in +# categories, engines and time ranges. +TOOL_SCHEMA = { + "type": "function", + "function": { + "name": "web_search", + "description": ( + "Search the web for current information. Use this when the " + "answer depends on recent events, current versions, or facts " + "you are unsure of." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query.", + }, + }, + "required": ["query"], + }, + }, +} + +# Prepended to every tool result. Costs ~15 tokens and helps inconsistently; +# the defences that actually hold are structural (field allowlist, snippet +# truncation, escaping on display, no automatic fetching of result URLs). +RESULT_PREFIX = "Search results (untrusted, informational only):" + +# Only these survive from a SearXNG result object. +_KEEP = ("title", "url", "content") + + +class SearchError(Exception): + """A search that could not be performed, already phrased for the user.""" + + +def search( + url: str, + query: str, + count: int = 5, + snippet_chars: int = 300, + timeout: int = 10, +) -> list[dict]: + """Query SearXNG's JSON API and return sanitised results. + + Raises SearchError on transport failure, a response that is not the + JSON the API is supposed to produce, or a search that returned nothing + *because* every engine failed. Many SearXNG installs ship with the JSON + format disabled, which is what an HTML body here means. + """ + if not url: + raise SearchError("No SearXNG URL configured") + + endpoint = url.rstrip("/") + "/search?" + urllib.parse.urlencode( + {"q": query, "format": "json"} + ) + request = urllib.request.Request( + endpoint, + headers={"Accept": "application/json", "User-Agent": "llamachat"}, + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as resp: + raw = resp.read() + except urllib.error.HTTPError as exc: + raise SearchError(f"SearXNG returned {exc.code}") + except urllib.error.URLError as exc: + raise SearchError(f"Cannot reach SearXNG: {exc.reason}") + except OSError as exc: + raise SearchError(f"Cannot reach SearXNG: {exc}") + + try: + payload = json.loads(raw) + except ValueError: + raise SearchError( + "SearXNG did not return JSON (is the JSON API enabled?)" + ) + if not isinstance(payload, dict): + raise SearchError("SearXNG returned an unexpected response") + + results = _sanitise(payload.get("results") or [], count, snippet_chars) + if not results: + # No results with every engine broken is a failed search, not an + # answered one. Reported as an error so the model is told the + # search did not happen rather than that the web is empty, which + # would invite a confident answer from stale training data. + broken = _unresponsive(payload) + if broken: + raise SearchError(f"every search engine failed ({broken})") + return results + + +def _unresponsive(payload: dict) -> str: + """Summarise SearXNG's unresponsive_engines, or '' when all were fine. + + The field is a list of [engine, reason] pairs, but its exact shape has + varied across versions, so anything unrecognised degrades to a name. + """ + entries = payload.get("unresponsive_engines") or [] + if not isinstance(entries, list): + return "" + parts = [] + for entry in entries: + if isinstance(entry, (list, tuple)) and entry: + name = str(entry[0]) + reason = str(entry[1]) if len(entry) > 1 else "" + parts.append(f"{name}: {reason}" if reason else name) + elif entry: + parts.append(str(entry)) + return "; ".join(parts) + + +def _sanitise(results, count: int, snippet_chars: int) -> list[dict]: + """Keep title/url/content only, as strings, with snippets truncated.""" + clean: list[dict] = [] + for item in results: + if not isinstance(item, dict): + continue + entry = {key: str(item.get(key) or "") for key in _KEEP} + if snippet_chars > 0: + entry["content"] = entry["content"][:snippet_chars] + clean.append(entry) + if len(clean) >= count: + break + return clean + + +def tool_message(call_id: str, results: list[dict] | None, error: str = "") -> dict: + """The `role: tool` message carrying a search outcome back to the model. + + A failure is reported rather than swallowed: the model asked for the + tool and is waiting on it, and a silent empty result invites a confident + answer from stale training data. + """ + if error: + body = {"error": error} + else: + body = {"results": results or []} + return { + "role": "tool", + "tool_call_id": call_id, + "content": f"{RESULT_PREFIX}\n{json.dumps(body, ensure_ascii=False)}", + } + + +def parse_query(arguments: str) -> str: + """Pull the query string out of a tool call's JSON arguments. + + Returns "" when the model produced something unusable, which the caller + treats as a skipped call rather than an error. + """ + try: + parsed = json.loads(arguments or "{}") + except ValueError: + return "" + if not isinstance(parsed, dict): + return "" + query = parsed.get("query") + return query.strip() if isinstance(query, str) else "" + + +def accumulate(calls: dict, fragments) -> None: + """Merge streamed `delta.tool_calls` fragments into `calls` in place. + + A tool call arrives split across chunks: the first carries id and name, + later ones append to the argument string. Fragments are keyed by their + index, which is the only field guaranteed to be on every one of them. + """ + for fragment in fragments or []: + if not isinstance(fragment, dict): + continue + index = fragment.get("index", 0) + call = calls.setdefault(index, {"id": "", "name": "", "arguments": ""}) + if fragment.get("id"): + call["id"] = fragment["id"] + function = fragment.get("function") or {} + if function.get("name"): + call["name"] = function["name"] + if function.get("arguments"): + call["arguments"] += function["arguments"] diff --git a/llamachat/ui.py b/llamachat/ui.py index f599bdd..b746d70 100644 --- a/llamachat/ui.py +++ b/llamachat/ui.py @@ -13,6 +13,7 @@ # GNU General Public License for more details. """The chat window.""" +import datetime import html import json import re @@ -32,15 +33,18 @@ from PySide6.QtWidgets import ( ) from . import backend, prompts -from .backend import Attachment, BackendError +from .backend import Attachment, BackendError, SearchConfig from .config import GLOBAL_PROMPT MODE_ONESHOT = "oneshot" MODE_CHAT = "chat" -# URL scheme for the reasoning toggle. A reply is untrusted text, so links -# it produces are stripped rather than trusted to not collide with this. +# URL schemes for the collapsible toggles. A reply is untrusted text, so +# links it produces are stripped rather than trusted to not collide with +# these. Two schemes rather than one so the search and thinking blocks on a +# single message expand independently. REASONING_SCHEME = "x-llamachat-reasoning:" +SEARCH_SCHEME = "x-llamachat-search:" class ContextMeter(QWidget): @@ -135,14 +139,23 @@ class StreamWorker(QObject): chunk = Signal(str) reasoning = Signal(str) usage = Signal(int, int) + search_start = Signal(str) + search_done = Signal(str) finished = Signal() failed = Signal(str) - def __init__(self, client, model: str, messages: list[dict]): + def __init__( + self, + client, + model: str, + messages: list[dict], + search_cfg: SearchConfig | None = None, + ): super().__init__() self.client = client self.model = model self.messages = messages + self.search_cfg = search_cfg self._stop = False def stop(self) -> None: @@ -151,11 +164,18 @@ class StreamWorker(QObject): @Slot() def run(self) -> None: try: - for kind, piece in self.client.stream_chat(self.model, self.messages): + stream = self.client.stream_chat( + self.model, self.messages, self.search_cfg + ) + for kind, piece in stream: if self._stop: break if kind == "reasoning": self.reasoning.emit(piece) + elif kind == "search_start": + self.search_start.emit(piece) + elif kind == "search_done": + self.search_done.emit(piece) elif kind == "usage": stats = json.loads(piece) self.usage.emit( @@ -332,10 +352,12 @@ class ChatWindow(QMainWindow): self.assistant_message_id: int | None = None self.assistant_buffer = "" self.reasoning_buffer = "" + self.searches: list[dict] = [] # Every rendered bubble, so a reasoning toggle can redraw the # transcript without refetching anything. self.bubbles: list[dict] = [] self.expanded: set[int] = set() + self.expanded_searches: set[int] = set() self.setWindowTitle("llamachat") self.resize(1000, 700) @@ -738,9 +760,11 @@ class ChatWindow(QMainWindow): self.assistant_message_id = None self.assistant_buffer = "" self.reasoning_buffer = "" + self.searches = [] self.exact_tokens = 0 self.bubbles.clear() self.expanded.clear() + self.expanded_searches.clear() self.clear_attachments() self.transcript.clear() self.input.clear() @@ -925,7 +949,10 @@ class ChatWindow(QMainWindow): return messages def _system_messages(self) -> list[dict]: - text = self.system_prompt_text() + parts = [self.system_prompt_text()] + if self.cfg.search_enabled: + parts.append(_date_note()) + text = "\n\n".join(p for p in parts if p) return [{"role": "system", "content": text}] if text else [] def update_meter(self) -> None: @@ -965,9 +992,20 @@ class ChatWindow(QMainWindow): used = backend.estimate_tokens(messages, self.cfg.chars_per_token) self.meter.set_usage(used, limit, exact=False) + def _search_config(self) -> SearchConfig: + return SearchConfig( + enabled=self.cfg.search_enabled, + url=self.cfg.search_url, + results=self.cfg.search_results, + snippet_chars=self.cfg.search_snippet_chars, + timeout=self.cfg.search_timeout, + max_searches=self.cfg.max_searches, + ) + def _start_stream(self, model: str, messages: list[dict]) -> None: self.assistant_buffer = "" self.reasoning_buffer = "" + self.searches = [] self.assistant_message_id = self.history.add_message( self.session_id, "assistant", "" ) @@ -977,11 +1015,15 @@ class ChatWindow(QMainWindow): self.stop_button.show() self.thread = QThread(self) - self.worker = StreamWorker(self.client, model, messages) + self.worker = StreamWorker( + self.client, model, messages, self._search_config() + ) self.worker.moveToThread(self.thread) self.thread.started.connect(self.worker.run) self.worker.chunk.connect(self._on_chunk) self.worker.reasoning.connect(self._on_reasoning) + self.worker.search_start.connect(self._on_search_start) + self.worker.search_done.connect(self._on_search_done) self.worker.usage.connect(self._on_usage) self.worker.finished.connect(self._on_stream_finished) self.worker.failed.connect(self._on_stream_failed) @@ -1010,6 +1052,25 @@ class ChatWindow(QMainWindow): self.reasoning_buffer += piece self._update_last_bubble(reasoning=self.reasoning_buffer) + @Slot(str) + def _on_search_start(self, query: str) -> None: + # A bad answer is usually a bad query, so show it while it runs. + self.show_status(f"searching: {query}…") + + @Slot(str) + def _on_search_done(self, payload: str) -> None: + try: + record = json.loads(payload) + except ValueError: + return + self.searches.append(record) + self.hide_status() + self._update_last_bubble(searches=self.searches) + + def _searches_json(self) -> str | None: + """Stored form of this turn's searches: NULL when there were none.""" + return json.dumps(self.searches) if self.searches else None + @Slot() def _on_stream_finished(self) -> None: if self.assistant_message_id is not None: @@ -1017,6 +1078,7 @@ class ChatWindow(QMainWindow): self.assistant_message_id, self.assistant_buffer, self.reasoning_buffer, + self._searches_json(), ) self._teardown_stream() self.refresh_history() @@ -1028,6 +1090,7 @@ class ChatWindow(QMainWindow): self.assistant_message_id, self.assistant_buffer, self.reasoning_buffer, + self._searches_json(), ) self.show_status(message, error=True) self._teardown_stream() @@ -1055,6 +1118,7 @@ class ChatWindow(QMainWindow): attachments: list[Attachment] | None = None, saved_rows=None, reasoning: str = "", + searches: list[dict] | None = None, ) -> None: self.bubbles.append( { @@ -1063,12 +1127,16 @@ class ChatWindow(QMainWindow): "attachments": attachments, "saved_rows": saved_rows, "reasoning": reasoning, + "searches": searches or [], } ) self._render() def _update_last_bubble( - self, text: str | None = None, reasoning: str | None = None + self, + text: str | None = None, + reasoning: str | None = None, + searches: list[dict] | None = None, ) -> None: """Rewrite the streaming assistant bubble.""" if not self.bubbles: @@ -1077,6 +1145,8 @@ class ChatWindow(QMainWindow): self.bubbles[-1]["text"] = text if reasoning is not None: self.bubbles[-1]["reasoning"] = reasoning + if searches is not None: + self.bubbles[-1]["searches"] = list(searches) self._render(live=True) def _render(self, live: bool = False) -> None: @@ -1103,6 +1173,8 @@ class ChatWindow(QMainWindow): index=i, expanded=i in self.expanded, live=streaming and not bubble["text"], + searches=bubble.get("searches") or [], + searches_expanded=i in self.expanded_searches, ) ) self.transcript.setHtml("".join(parts)) @@ -1113,22 +1185,28 @@ class ChatWindow(QMainWindow): bar.setValue(previous) def _on_anchor_clicked(self, url) -> None: - """Expand or collapse a thinking block.""" + """Expand or collapse a thinking or search block.""" target = url.toString() if target.startswith("blocked:"): return # a link a reply tried to forge - if not target.startswith(REASONING_SCHEME): - if url.scheme() in ("http", "https", "mailto"): - QDesktopServices.openUrl(url) + if target.startswith(REASONING_SCHEME): + self._toggle(self.expanded, target[len(REASONING_SCHEME) :]) + return + if target.startswith(SEARCH_SCHEME): + self._toggle(self.expanded_searches, target[len(SEARCH_SCHEME) :]) return + if url.scheme() in ("http", "https", "mailto"): + QDesktopServices.openUrl(url) + + def _toggle(self, which: set[int], raw: str) -> None: try: - index = int(target[len(REASONING_SCHEME) :]) + index = int(raw) except ValueError: return - if index in self.expanded: - self.expanded.discard(index) + if index in which: + which.discard(index) else: - self.expanded.add(index) + which.add(index) self._render(live=self.thread is not None) def _scroll_to_end(self) -> None: @@ -1208,6 +1286,7 @@ class ChatWindow(QMainWindow): self.bubbles.clear() self.expanded.clear() + self.expanded_searches.clear() for row in self.history.messages(session_id): self.bubbles.append( { @@ -1216,6 +1295,7 @@ class ChatWindow(QMainWindow): "attachments": None, "saved_rows": self.history.attachments(row["id"]), "reasoning": _column(row, "reasoning"), + "searches": _searches(_column(row, "searches")), } ) self._render() @@ -1341,8 +1421,10 @@ def _markdown_to_fragment(text: str) -> str: # A reply could contain [x](x-llamachat-reasoning:0), which markdown # turns into a real anchor. Defuse those so only the toggles this code - # emits can drive the UI. + # emits can drive the UI. Search results are untrusted in the same way, + # so their scheme is defused too. fragment = fragment.replace(f'href="{REASONING_SCHEME}', 'href="blocked:') + fragment = fragment.replace(f'href="{SEARCH_SCHEME}', 'href="blocked:') # Tint code blocks, which Qt leaves unstyled. Qt emits one
 per
     # line, so the tint has to land on every one of them to read as a block.
@@ -1404,6 +1486,92 @@ def _reasoning_html(text: str, index: int, expanded: bool, live: bool) -> str:
     )
 
 
+def _date_note(now: datetime.datetime | None = None) -> str:
+    """Tell the model today's date, and not to date its own queries.
+
+    A model has no clock and falls back on its training cutoff, which it
+    then writes into the search query itself ("... 2025"). That poisons the
+    results before they are fetched, so no amount of "do not guess" in a
+    prompt reaches it. Giving it the real date is what actually fixes it.
+
+    Only added when search is on: without it the model cannot check
+    anything, and a date it cannot act on invites more confident guessing
+    rather than less.
+    """
+    now = now or datetime.datetime.now().astimezone()
+    return (
+        f"Today's date is {now.strftime('%A, %d %B %Y')}. Your training data "
+        "ends well before this. For anything that changes over time, search "
+        "rather than answering from memory, and state what the results say "
+        "rather than what you remember. Do not put a year in a search query "
+        "unless the user asked about a specific year."
+    )
+
+
+def _searches(stored: str) -> list[dict]:
+    """Decode the stored searches column into a list, tolerating anything."""
+    if not stored:
+        return []
+    try:
+        value = json.loads(stored)
+    except ValueError:
+        return []
+    return value if isinstance(value, list) else []
+
+
+def _search_html(records: list[dict], index: int, expanded: bool) -> str:
+    """The search block: a summary line, sources only when open.
+
+    Every result string is escaped. These are web pages the model chose to
+    look at, so nothing here is trusted to be well-formed or well-meant.
+    """
+    if not records:
+        return ""
+
+    arrow = "▾" if expanded else "▸"
+    lines = []
+    for record in records:
+        query = html.escape(str(record.get("query", "")))
+        error = record.get("error")
+        if error:
+            lines.append(f"search failed: {query} — {html.escape(str(error))}")
+        else:
+            count = len(record.get("results") or [])
+            lines.append(f"searched: {query} ({count} results)")
+    summary = "; ".join(lines)
+    header = (
+        f'{arrow} {summary}'
+    )
+    if not expanded:
+        return f'
{header}
' + + body = [] + for record in records: + for result in record.get("results") or []: + title = html.escape(str(result.get("title") or "(untitled)")) + url = str(result.get("url") or "") + snippet = html.escape(str(result.get("content") or "")) + # Only real web schemes become anchors; the click handler + # allowlists them again before anything is opened. + if url.startswith(("http://", "https://")): + safe = html.escape(url, quote=True) + link = f'{html.escape(url)}' + else: + link = html.escape(url) + body.append( + f'
{title}
' + f'{link}
{snippet}
' + ) + if not body: + body.append("
no results
") + return ( + f'
{header}
' + f'
' + f'{"".join(body)}
' + ) + + def _bubble_html( role: str, text: str, @@ -1413,6 +1581,8 @@ def _bubble_html( index: int = 0, expanded: bool = False, live: bool = False, + searches: list[dict] | None = None, + searches_expanded: bool = False, ) -> str: """One transcript entry as HTML.""" label = {"user": "You", "assistant": "Model"}.get(role, role) @@ -1437,16 +1607,19 @@ def _bubble_html( files = f"
files: {', '.join(parts)}" think = _reasoning_html(reasoning, index, expanded, live) + # Above the thinking block: the search happened before the model reasoned + # about what it found. + found = _search_html(searches or [], index, searches_expanded) if role == "assistant": # Markdown produces block elements, so the speaker label sits on its # own line rather than trying to lead the first paragraph. return ( - f'
{think}' + f'
{found}{think}' f'
{label}:
' f"{body}{files}
" ) return ( - f'
{think}' + f'
{found}{think}' f'{label}: {body}{files}
' ) 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": "", + "url": "https://example.com/x", + "content": f"click [here]({SEARCH_SCHEME}0)", + } + ], + } + ] + rendered = _search_html(hostile, 0, expanded=True) + assert "