diff options
Diffstat (limited to 'test_llamachat.py')
| -rwxr-xr-x | test_llamachat.py | 161 |
1 files changed, 156 insertions, 5 deletions
diff --git a/test_llamachat.py b/test_llamachat.py index a766382..a06d737 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -967,11 +967,11 @@ def test_search_tool_schema(): assert out == [("content", "hi")] assert sent == [None], sent - # One search per turn by default: a follow-up call arrives as literal - # <tool_call> XML rather than a structured call, and a turn that ends on - # one shows the user an empty reply. - assert config.DEFAULTS["max_searches"] == 1 - assert config.load(Path("/nonexistent/config.toml")).max_searches == 1 + # Two searches per turn, which a follow-up call now survives: the final + # round's tool result says the tool is gone, and quantized KV cache, the + # other half of the empty-reply bug, is a server-side setting. + assert config.DEFAULTS["max_searches"] == 2 + assert config.load(Path("/nonexistent/config.toml")).max_searches == 2 # A configured-but-urlless setup resolves to disabled at config load. with tempfile.TemporaryDirectory() as tmp: @@ -1472,6 +1472,152 @@ def test_search_html(): print("ok search html") +def test_final_round_note(): + """The last search result tells the model it has no searches left. + + Withdrawing the tool schema is invisible to the model: it tries a second + search anyway and emits it as literal <tool_call> text, or trails off + inside the thinking block, either way leaving the reply empty. Saying so + in the tool result is what actually ends the turn with an answer. The + note cannot be a trailing system message: this model's chat template + raises "System message must be at the beginning". + """ + from llamachat import search + + plain = search.tool_message("call_1", [{"title": "T", "url": "u", "content": "c"}]) + final = search.tool_message( + "call_1", [{"title": "T", "url": "u", "content": "c"}], last=True + ) + assert plain["role"] == "tool" and final["role"] == "tool" + assert search.NO_MORE_SEARCHES not in plain["content"] + assert search.NO_MORE_SEARCHES in final["content"] + # The results themselves survive the note being appended. + for message in (plain, final): + assert '"title": "T"' in message["content"] or '"title":"T"' in message["content"] + # A failed final search still gets the note, or the model retries. + failed = search.tool_message("call_1", None, error="boom", last=True) + assert "boom" in failed["content"] + assert search.NO_MORE_SEARCHES in failed["content"] + print("ok final round note") + + +def test_title_cleaning(): + clean = backend.clean_title + + assert clean("Kernel build failure") == "Kernel build failure" + # Quotes, trailing punctuation and a leading label all come off. + assert clean('"Kernel build failure"') == "Kernel build failure" + assert clean("Title: Kernel build failure.") == "Kernel build failure" + assert clean("**Kernel build failure**") == "Kernel build failure" + # Thinking that leaks into the body is dropped, answer kept. + assert clean("<think>hmm, short</think>Kernel build") == "Kernel build" + # An unclosed think block leaves nothing usable rather than a stray tag. + assert clean("<think>still reasoning") == "" + # Chatter before the title: the last non-empty line wins. + assert clean("Sure, here you go:\nKernel build failure") == "Kernel build failure" + # Too long is cut on a word boundary, not mid-word. + long = clean("word " * 40) + assert len(long) <= 60, long + assert not long.endswith("wor"), long + # Nothing usable yields nothing, so the caller keeps its fallback. + assert clean("") == "" + assert clean(" \n\n ") == "" + assert clean('""') == "" + print("ok title cleaning") + + +def test_title_request(): + """Client.complete posts a non-streaming request and returns the text.""" + seen = {} + + class FakeResponse: + status_code = 200 + text = "" + + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"message": {"content": " A Title "}}]} + + def fake_post(url, json=None, timeout=None): + seen["url"] = url + seen["body"] = json + return FakeResponse() + + import httpx + + original = httpx.post + httpx.post = fake_post + try: + client = backend.Client("http://router:8181") + out = client.complete("m", [{"role": "user", "content": "hi"}], max_tokens=32) + finally: + httpx.post = original + + assert out == " A Title ", out + assert seen["url"] == "http://router:8181/v1/chat/completions" + assert seen["body"]["model"] == "m" + assert seen["body"]["stream"] is False + assert seen["body"]["max_tokens"] == 32 + # No tools are offered: a title turn must not trigger a search. + assert "tools" not in seen["body"] + # Thinking off, or a reasoning model spends the budget and returns "". + assert seen["body"]["chat_template_kwargs"] == {"enable_thinking": False} + + # A reply with no choices is empty rather than an exception. + class Empty(FakeResponse): + def json(self): + return {} + + httpx.post = lambda url, json=None, timeout=None: Empty() + try: + assert backend.Client("http://x").complete("m", []) == "" + finally: + httpx.post = original + print("ok title request") + + +def test_needs_title(): + """Which finished turns are worth asking the model to name. + + The trigger is a session still carrying its placeholder title, not the + turn number: a first reply that comes back empty (the follow-up + <tool_call> bug) must not forfeit titling for the whole session, and a + one-shot window keeps appending to the same session until New is + pressed. + """ + needs = backend.needs_title + q = "how do I build a custom kernel on slackware?" + + assert needs(q, q, "Fetch the sources.") is True + # Nothing to title from: the empty-reply turn is skipped, not consumed. + assert needs(q, q, "") is False + assert needs(q, q, " \n ") is False + # Already named by the model, so leave it alone. + assert needs("Building a custom kernel", q, "Fetch the sources.") is False + # A long question is stored truncated; that still counts as untouched. + long_q = "x" * 200 + assert needs(long_q[:60], long_q, "an answer") is True + # A title the user or model set that happens to be short is not a + # placeholder, even though it is under the cut-off. + assert needs("Kernels", long_q, "an answer") is False + print("ok needs title") + + +def test_title_prompt(): + """The title request carries the exchange and asks for a short title.""" + messages = backend.title_messages("how do I build a kernel?", "Run make.") + assert messages[-1]["role"] == "user" + blob = json.dumps(messages) + assert "how do I build a kernel?" in blob + assert "Run make." in blob + # A huge exchange is trimmed so titling never costs a full context. + big = backend.title_messages("x" * 10000, "y" * 10000) + assert len(json.dumps(big)) < 6000, len(json.dumps(big)) + print("ok title prompt") + + if __name__ == "__main__": test_presets() test_real_presets() @@ -1504,4 +1650,9 @@ if __name__ == "__main__": test_search_storage() test_date_note() test_search_html() + test_final_round_note() + test_title_cleaning() + test_title_request() + test_needs_title() + test_title_prompt() print("\nall checks passed") |