diff options
Diffstat (limited to 'test_llamachat.py')
| -rwxr-xr-x | test_llamachat.py | 328 |
1 files changed, 318 insertions, 10 deletions
diff --git a/test_llamachat.py b/test_llamachat.py index 3f1f678..74b0d0c 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -208,24 +208,24 @@ def test_classify(): def test_sse_parsing(): line = 'data: {"choices":[{"delta":{"content":"hi"}}]}' - assert backend._parse_sse_line(line) == ("content", "hi") + assert backend._parse_sse_line(line) == [("content", "hi")] # Reasoning arrives in its own field, which is what lets the UI keep # thinking and reply apart without parsing <think> tags. think = 'data: {"choices":[{"delta":{"reasoning_content":"hmm"}}]}' - assert backend._parse_sse_line(think) == ("reasoning", "hmm") + assert backend._parse_sse_line(think) == [("reasoning", "hmm")] assert backend._parse_sse_line("data: [DONE]") == backend.DONE - assert backend._parse_sse_line(": keepalive") is None - assert backend._parse_sse_line("") is None - assert backend._parse_sse_line("data: {bad json") is None + assert backend._parse_sse_line(": keepalive") == [] + assert backend._parse_sse_line("") == [] + assert backend._parse_sse_line("data: {bad json") == [] # An opening delta of {'role': 'assistant', 'content': None} carries # nothing to render and must not be mistaken for end-of-stream. opening = backend._parse_sse_line( 'data: {"choices":[{"delta":{"role":"assistant","content":null}}]}' ) - assert opening is None + assert opening == [], opening assert opening != backend.DONE # A delta carrying both fields is thinking, not reply. Classifying it as @@ -234,14 +234,30 @@ def test_sse_parsing(): 'data: {"choices":[{"delta":' '{"reasoning_content":"still thinking","content":""}}]}' ) - assert both == ("reasoning", "still thinking"), both + assert both == [("reasoning", "still thinking")], both # An empty reasoning delta renders nothing and must not fall through to # the content branch. empty = backend._parse_sse_line( 'data: {"choices":[{"delta":{"reasoning_content":""}}]}' ) - assert empty is None, empty + assert empty == [], empty + + # The final chunk of a tool round carries finish_reason and the usage + # block together; both must survive, or the tool call is dropped. + combined = backend._parse_sse_line( + 'data: {"choices":[{"index":0,"delta":{"content":"",' + '"reasoning_content":null},"finish_reason":"tool_calls"}],' + '"usage":{"prompt_tokens":388,"completion_tokens":74,' + '"total_tokens":462}}' + ) + assert combined == [ + ( + "usage", + '{"prompt_tokens": 388, "completion_tokens": 74, "total_tokens": 462}', + ), + ("tool_finish", ""), + ], combined print("ok sse parsing") @@ -573,14 +589,16 @@ def test_usage_parsing(): 'data: {"choices":[],"usage":{"prompt_tokens":1234,' '"completion_tokens":56,"total_tokens":1290}}' ) - kind, payload = backend._parse_sse_line(line) + events = backend._parse_sse_line(line) + assert [k for k, _ in events] == ["usage"], events + kind, payload = events[0] assert kind == "usage" stats = json.loads(payload) assert stats["prompt_tokens"] == 1234 assert stats["total_tokens"] == 1290 # A usage-less chunk with empty choices is not mistaken for one. - assert backend._parse_sse_line('data: {"choices":[]}') is None + assert backend._parse_sse_line('data: {"choices":[]}') == [] print("ok usage parsing") @@ -928,6 +946,7 @@ def test_provider_parsing(): "price_in": 0.6, "price_out": 0.9, "thinking_budget": 8192, + "replay_reasoning": True, }, } } @@ -941,6 +960,9 @@ def test_provider_parsing(): assert parsed["together"].price_in == 0.6 assert parsed["together"].price_out == 0.9 assert parsed["together"].thinking_budget == 8192 + assert parsed["together"].replay_reasoning is True + # Off by default: replaying reasoning costs context and input tokens. + assert parsed["local"].replay_reasoning is False # An old config: bare base_url, no providers table at all. legacy = providers.parse({"base_url": "http://localhost:8181"}) @@ -1415,6 +1437,32 @@ def test_config_providers(): print("ok config provider table") +def test_replays_reasoning(): + """reasoning replay is per-provider, gated on search being on.""" + from llamachat import providers + + table = providers.parse( + { + "providers": { + "local": {"base_url": "http://localhost:8181"}, + "deepseek": {"base_url": "https://api.deepseek.com", + "replay_reasoning": True}, + "siliconflow": {"base_url": "https://api.siliconflow.com"}, + } + } + ) + + # Only the opted-in provider replays, and only when tools are sent. + assert providers.replays_reasoning("deepseek:m", table, True) is True + assert providers.replays_reasoning("deepseek:m", table, False) is False + assert providers.replays_reasoning("siliconflow:m", table, True) is False + assert providers.replays_reasoning("localmodel", table, True) is False + + # An unknown prefix resolves to local, which never replays. + assert providers.replays_reasoning("unconfigured:m", table, True) is False + print("ok replays reasoning") + + def test_models_store(): """models.ini round-trips per-model metadata and the cancel record.""" from llamachat import models @@ -1829,6 +1877,123 @@ def test_token_column_migration(): print("ok token column migration") +def test_usage_columns(): + """Provider, raw usage and the GUI estimate persist per reply.""" + with tempfile.TemporaryDirectory() as tmp: + history = db.History(Path(tmp) / "u.db") + sid = history.create_session("chat", "m", "t") + mid = history.add_message(sid, "assistant", "") + usage = '{"prompt_tokens":100,"completion_tokens":20,"total_tokens":120}' + history.update_message( + mid, "done", + prompt_tokens=100, completion_tokens=20, + model="siliconflow:zai-org/GLM-5.2", provider="siliconflow", + usage_json=usage, reported_cost_usd=0.000212, + ) + row = history.messages(sid)[0] + assert row["provider"] == "siliconflow" + assert row["usage_json"] == usage + assert row["reported_cost_usd"] == 0.000212 + + # None means keep, exactly like the token columns. + history.update_message(mid, "edited") + kept = history.messages(sid)[0] + assert kept["provider"] == "siliconflow" + assert kept["usage_json"] == usage + assert kept["reported_cost_usd"] == 0.000212 + history.close() + print("ok usage columns") + + +def test_usage_accumulation(): + """Every round of a searched turn reaches usage_all, not only the last. + + Each search round is its own billed API call, so dropping the earlier + rounds' usage would under-report what a searched turn cost. + """ + from llamachat import ui + + class FakeClient: + def client_for(self, model): + return self + + def wire_name(self, model): + return model + + def stream_chat(self, wire, messages, search_cfg=None): + # One usage block per round: a tool round, then the final answer. + yield ("usage", '{"prompt_tokens": 100, "completion_tokens": 5,' + ' "total_tokens": 105}') + yield ("usage", '{"prompt_tokens": 140, "completion_tokens": 60,' + ' "total_tokens": 200}') + + worker = ui.StreamWorker(FakeClient(), "m", []) + got: list[str] = [] + worker.usage_all.connect(got.append) + worker.run() + + assert len(got) == 1, got + blocks = json.loads(got[0]) + assert [b["prompt_tokens"] for b in blocks] == [100, 140] + assert [b["completion_tokens"] for b in blocks] == [5, 60] + + # A turn with no usage chunk at all (no include_usage support) yields + # an empty array, which the window stores as NULL. + class SilentClient(FakeClient): + def stream_chat(self, wire, messages, search_cfg=None): + yield ("content", "hi") + return + + worker2 = ui.StreamWorker(SilentClient(), "m", []) + got2: list[str] = [] + worker2.usage_all.connect(got2.append) + worker2.run() + assert got2 == ["[]"], got2 + print("ok usage accumulation") + + +def test_usage_column_migration(): + """A pre-provider database opens, and the new columns read as NULL.""" + import sqlite3 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "old.db" + conn = sqlite3.connect(path) + conn.executescript( + "CREATE TABLE sessions (id INTEGER PRIMARY KEY, mode TEXT," + " title TEXT, model TEXT, created_at INTEGER, updated_at INTEGER);" + "CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id INTEGER," + " role TEXT NOT NULL, content TEXT NOT NULL," + " prompt_tokens INTEGER, completion_tokens INTEGER, model TEXT," + " created_at INTEGER NOT NULL);" + "INSERT INTO sessions VALUES (1,'chat','old','m',0,0);" + "INSERT INTO messages VALUES (1,1,'assistant','older reply'," + "1200,340,'m',0);" + ) + conn.commit() + conn.close() + + history = db.History(path) + # The pre-migration row reads as unknown, never as an empty string + # or a fabricated zero. + old = history.messages(1)[0] + assert old["provider"] is None + assert old["usage_json"] is None + assert old["reported_cost_usd"] is None + + mid = history.add_message(1, "assistant", "") + history.update_message( + mid, "new reply", prompt_tokens=10, completion_tokens=5, + model="q:r", provider="q", usage_json='{"total_tokens":15}', + ) + fresh = history.messages(1)[1] + assert fresh["provider"] == "q" + assert fresh["usage_json"] == '{"total_tokens":15}' + assert fresh["reported_cost_usd"] is None + history.close() + print("ok usage column migration") + + def test_client_auth_header(): """A client with a key sends Bearer auth; one without sends no header.""" sent = {} @@ -1908,6 +2073,33 @@ def test_client_auth_header(): print("ok client authorization header") +def test_debug_log(): + """The diagnostic log is off unless $LLAMACHAT_DEBUG_LOG names a path.""" + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "debug.log" + saved = os.environ.get("LLAMACHAT_DEBUG_LOG") + try: + # Unset -> a no-op that creates nothing to rotate. + os.environ.pop("LLAMACHAT_DEBUG_LOG", None) + backend._debug_log("must not be written") + assert not path.exists() + + # Set -> appends timestamped records. + os.environ["LLAMACHAT_DEBUG_LOG"] = str(path) + backend._debug_log("first") + backend._debug_log("second") + lines = path.read_text(encoding="utf-8").strip().splitlines() + assert len(lines) == 2, lines + assert lines[0].startswith("[") and "first" in lines[0] + assert "second" in lines[1] + finally: + if saved is None: + os.environ.pop("LLAMACHAT_DEBUG_LOG", None) + else: + os.environ["LLAMACHAT_DEBUG_LOG"] = saved + print("ok debug log") + + def test_stream_body_for_cloud(): """Cloud providers get max_tokens and thinking_budget; local does not.""" captured = {} @@ -2507,6 +2699,115 @@ def test_search_failure_paths(): print("ok search failure paths") +def test_reasoning_content_preserved(): + """A tool-call round replays its reasoning_content verbatim. + + DeepSeek V3.2+/V4 and GLM-4.7+ on SiliconFlow emit chain-of-thought as + reasoning_content and require it sent back unchanged on the assistant + tool-call message. Dropping it breaks their multi-step tool flow, which + surfaces as a turn that stops at the thinking with no answer. + """ + import urllib.request + + from llamachat import search + + client = backend.Client("http://x") + seen_messages = [] + + def scripted(model, messages, tools): + seen_messages.append([dict(m) for m in messages]) + if len(seen_messages) == 1: + yield ("reasoning", "I should check the current version.") + yield from _tool_round("latest kernel") + else: + yield ("content", "answered") + + client._stream_once = scripted + original = _with_urlopen(json.dumps({"results": []}).encode()) + try: + cfg = backend.SearchConfig( + enabled=True, url="http://searx", max_searches=1 + ) + list(client.stream_chat("m", [{"role": "user", "content": "hi"}], cfg)) + finally: + urllib.request.urlopen = original + + # The second request's assistant message must carry the thinking back. + second = seen_messages[1] + assistant = next(m for m in second if m["role"] == "assistant" and m.get("tool_calls")) + assert assistant["reasoning_content"] == "I should check the current version.", assistant + + # The reasoning still reaches the UI unchanged. + assert len(seen_messages) == 2, seen_messages + print("ok reasoning content preserved") + + +def test_tool_call_finish_with_usage(): + """finish_reason riding the usage chunk still triggers the search. + + DeepSeek's include_usage stream puts finish_reason: "tool_calls" on the + same SSE line as the usage block. The old single-event parser returned + that line as "usage" and never saw the finish, so the tool call was + dropped and the turn ended at the thinking with no search. + """ + import httpx + import urllib.request + + def sse(obj): + return "data: " + json.dumps(obj) + + lines = [ + sse({"choices": [{"delta": {"reasoning_content": "hmm"}}]}), + sse({"choices": [{"delta": {"tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query":"kernel"}'}, + } + ]}}]}), + sse({ + "choices": [{ + "index": 0, + "delta": {"content": "", "reasoning_content": None}, + "finish_reason": "tool_calls", + }], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }), + "data: [DONE]", + ] + + class _Stream: + status_code = 200 + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def iter_lines(self): + return iter(lines) + + def _fake_stream(method, url, json=None, timeout=None, headers=None): + return _Stream() + + client = backend.Client("http://x") + original = _with_urlopen(json.dumps({"results": []}).encode()) + try: + with _patched(httpx, "stream", _fake_stream): + cfg = backend.SearchConfig(enabled=True, url="http://searx", max_searches=1) + out = list(client.stream_chat("m", [{"role": "user", "content": "hi"}], cfg)) + finally: + urllib.request.urlopen = original + + kinds = [k for k, _ in out] + assert "search_start" in kinds, kinds + assert "search_done" in kinds, kinds + assert ("reasoning", "hmm") in out, out + print("ok tool call finish with usage") + + def test_search_storage(): import sqlite3 @@ -2877,10 +3178,15 @@ if __name__ == "__main__": test_model_ids_and_filtering() test_key_resolution() test_config_providers() + test_replays_reasoning() test_models_store() test_metadata_and_cost() test_token_column_migration() + test_usage_columns() + test_usage_accumulation() + test_usage_column_migration() test_client_auth_header() + test_debug_log() test_stream_body_for_cloud() test_multi_client() test_model_dialog_values() @@ -2890,6 +3196,8 @@ if __name__ == "__main__": test_tool_call_accumulation() test_search_loop_cap() test_search_failure_paths() + test_reasoning_content_preserved() + test_tool_call_finish_with_usage() test_search_storage() test_date_note() test_search_html() |