| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025 |
- """Tests for the LLM client and chat session (no network)."""
- import pytest
- from stt_daemon.assistant.llm import (
- BriefEngine,
- ChatSession,
- ConversationEngine,
- LlmClient,
- LlmError,
- LlmResponse,
- )
- from stt_daemon.config import LlmConfig
- def _make_client(responder, cfg=None):
- cfg = cfg or LlmConfig()
- calls: list[dict] = []
- def opener(payload, timeout):
- calls.append(payload)
- return responder(payload)
- return LlmClient(cfg, opener=opener), calls
- def _reply(content="", reasoning=None, tool_calls=None, finish="stop"):
- message = {"role": "assistant", "content": content}
- if reasoning is not None:
- message["reasoning_content"] = reasoning
- if tool_calls:
- message["tool_calls"] = tool_calls
- return {"choices": [{"message": message, "finish_reason": finish}]}
- def test_reasoning_content_is_ignored():
- client, _ = _make_client(
- lambda _p: _reply(content="It is noon.", reasoning="thinking hard about it")
- )
- response = client.chat([{"role": "user", "content": "hi"}])
- assert response.content == "It is noon."
- assert "thinking hard" not in response.content
- def test_cache_prompt_and_endpoint_are_sent():
- client, calls = _make_client(lambda _p: _reply(content="ok"))
- client.chat([{"role": "user", "content": "hi"}])
- assert calls[0]["body"]["cache_prompt"] is True
- assert calls[0]["body"]["stream"] is False
- assert calls[0]["__url__"] == LlmConfig().endpoint
- def test_tools_sent_only_when_provided():
- client, calls = _make_client(lambda _p: _reply(content="ok"))
- client.chat([{"role": "user", "content": "hi"}])
- assert "tools" not in calls[0]["body"]
- client.chat(
- [{"role": "user", "content": "hi"}],
- tools=[{"type": "function", "function": {"name": "get_datetime"}}],
- )
- assert calls[1]["body"]["tools"][0]["function"]["name"] == "get_datetime"
- assert calls[1]["body"]["tool_choice"] == "auto"
- def test_id_slot_sent_only_when_slot_id_provided():
- client, calls = _make_client(lambda _p: _reply(content="ok"))
- client.chat([{"role": "user", "content": "hi"}])
- assert "id_slot" not in calls[0]["body"]
- client.chat([{"role": "user", "content": "hi"}], slot_id=1)
- assert calls[1]["body"]["id_slot"] == 1
- def test_slot_and_on_io_never_passed_to_old_fake_when_none():
- seen: list[dict] = []
- class _OldFake:
- def chat(self, messages, tools=None, max_tokens=None, timeout_s=None):
- seen.append(
- {
- "messages": messages,
- "tools": tools,
- "max_tokens": max_tokens,
- "timeout_s": timeout_s,
- }
- )
- return {"ok": True}
- fake = _OldFake()
- slot_id = None
- on_io = None
- kwargs: dict = {}
- if slot_id is not None:
- kwargs["slot_id"] = slot_id
- if on_io is not None:
- kwargs["on_io"] = on_io
- fake.chat([{"role": "user", "content": "hi"}], **kwargs)
- assert seen[0] == {
- "messages": [{"role": "user", "content": "hi"}],
- "tools": None,
- "max_tokens": None,
- "timeout_s": None,
- }
- def test_tool_calls_are_parsed():
- client, _ = _make_client(
- lambda _p: _reply(
- tool_calls=[
- {
- "id": "1",
- "type": "function",
- "function": {"name": "get_datetime", "arguments": "{}"},
- }
- ],
- finish="tool_calls",
- )
- )
- response = client.chat([{"role": "user", "content": "time?"}], tools=[{"type": "function"}])
- assert len(response.tool_calls) == 1
- assert response.tool_calls[0].name == "get_datetime"
- assert response.finish_reason == "tool_calls"
- def test_missing_choices_raises():
- client, _ = _make_client(lambda _p: {"choices": []})
- with pytest.raises(LlmError):
- client.chat([{"role": "user", "content": "hi"}])
- def test_chat_session_appends_and_keeps_prefix():
- client, calls = _make_client(lambda _p: _reply(content="hello"))
- session = ChatSession(LlmConfig(), client)
- assert session.send("hi") == "hello"
- session.send("again")
- assert [m["role"] for m in session.history] == [
- "user",
- "assistant",
- "user",
- "assistant",
- ]
- first = calls[0]["body"]["messages"]
- second = calls[1]["body"]["messages"]
- assert second[: len(first)] == first # append-only, prefix preserved
- assert second[0]["content"] == LlmConfig().system_prompt_chat
- def test_chat_session_trims_in_blocks_and_keeps_system_prompt():
- cfg = LlmConfig(
- max_context_tokens=1330,
- max_tokens=1024,
- compaction_safety_margin=256,
- trim_block_turns=1,
- compaction_enabled=False,
- ) # prompt_budget == 50
- client, calls = _make_client(lambda _p: _reply(content="ok"))
- session = ChatSession(cfg, client)
- for _ in range(10):
- session.send("x" * 40)
- assert len(session.history) <= 6
- assert calls[-1]["body"]["messages"][0]["content"] == cfg.system_prompt_chat
- def test_chat_session_reset_forgets_history():
- client, _ = _make_client(lambda _p: _reply(content="ok"))
- session = ChatSession(LlmConfig(), client)
- session.send("hi")
- session.reset()
- assert session.history == []
- def test_chat_session_drops_reply_when_reset_mid_flight():
- holder: list[ChatSession] = []
- def opener(_payload, _timeout):
- holder[0].reset()
- return _reply(content="late reply")
- holder.append(ChatSession(LlmConfig(), LlmClient(LlmConfig(), opener=opener)))
- assert holder[0].send("hi") == ""
- assert holder[0].history == []
- # ---------------------------------------------------------------------------
- # Phase C tests
- # ---------------------------------------------------------------------------
- def _compact_cfg(
- *,
- budget_chars=200,
- trigger_ratio=0.8,
- keep_turns=1,
- compaction_max_tokens=64,
- compaction_timeout_s=5.0,
- compaction_input_tokens=8192,
- compaction_enabled=True,
- trim_block_turns=1,
- ):
- """Return a small LlmConfig whose prompt_budget == budget_chars // 4 tokens.
- prompt_budget = max_context_tokens - max_tokens - compaction_safety_margin
- We want budget_chars // 4 tokens, so:
- max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
- """
- max_tokens = 16
- compaction_safety_margin = 0
- max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
- return LlmConfig(
- max_context_tokens=max_context_tokens,
- max_tokens=max_tokens,
- compaction_safety_margin=compaction_safety_margin,
- compaction_enabled=compaction_enabled,
- compaction_trigger_ratio=trigger_ratio,
- compaction_keep_turns=keep_turns,
- compaction_max_tokens=compaction_max_tokens,
- compaction_timeout_s=compaction_timeout_s,
- compaction_input_tokens=compaction_input_tokens,
- trim_block_turns=trim_block_turns,
- system_prompt_chat="SYS",
- system_prompt_command="CMD",
- compaction_system_prompt="COMPACT",
- )
- def test_compaction_fires_and_window_shrinks():
- """Crossing the trigger compacts: window shrinks, summary present in next prompt."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- # prompt_budget = 200//4 = 50 tokens; trigger at 40 tokens (0.8 * 50)
- summary_text = "summary of old turns"
- call_index = 0
- def responder(payload):
- nonlocal call_index
- msgs = payload["body"]["messages"]
- call_index += 1
- # Summarizer calls have the compaction system prompt.
- if msgs[0]["content"] == "COMPACT":
- return _reply(content=summary_text)
- return _reply(content="ok")
- client, calls = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- # Send enough turns to exceed the trigger (each "x"*40 = 40 chars = 10 tokens).
- # After 4 turns (user+assistant pairs = 8 messages * 10 tokens = 80 tokens > 40 trigger).
- for _ in range(4):
- session.send("x" * 40)
- # After compaction the window must be smaller than the full history.
- assert len(session.history) <= 4 # keep_turns=1 → keep=2 messages
- # The summary must be set.
- assert session.summary != ""
- # The next chat call must include the summary as a second system message,
- # with the chat system prompt FIRST.
- session.send("after compaction")
- chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
- last_chat = chat_calls[-1]["body"]["messages"]
- assert last_chat[0]["role"] == "system"
- assert last_chat[0]["content"] == "SYS"
- assert last_chat[1]["role"] == "system"
- assert "Earlier conversation summary:" in last_chat[1]["content"]
- assert summary_text in last_chat[1]["content"]
- def test_summary_never_appended_to_history():
- """The summary text must not appear as a history entry."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- summary_text = "SUMMARY_SENTINEL"
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- return _reply(content=summary_text)
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- for _ in range(4):
- session.send("x" * 40)
- for msg in session.history:
- assert summary_text not in msg["content"]
- def test_summarizer_failure_falls_back_to_trim():
- """On summarizer failure the session falls back to _trim() and stays within budget."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- return _reply(content="", finish="stop") # empty → failure
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- for _ in range(4):
- session.send("x" * 40)
- # Summary must remain empty (fallback used trim, not compaction).
- assert session.summary == ""
- # History must be bounded within the budget.
- assert session._estimate_tokens() <= cfg.prompt_budget
- def test_summarizer_length_finish_reason_falls_back_to_trim():
- """finish_reason == 'length' on the summarizer is treated as failure."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- return _reply(content="partial summary", finish="length")
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- for _ in range(4):
- session.send("x" * 40)
- assert session.summary == ""
- assert session._estimate_tokens() <= cfg.prompt_budget
- def test_summarizer_llmerror_falls_back_to_trim():
- """LlmError from the summarizer falls back to _trim()."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- raise LlmError("network error")
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- for _ in range(4):
- session.send("x" * 40)
- assert session.summary == ""
- assert session._estimate_tokens() <= cfg.prompt_budget
- def test_reset_clears_summary():
- """reset() clears both history and summary."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- summary_text = "old summary"
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- return _reply(content=summary_text)
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- for _ in range(4):
- session.send("x" * 40)
- assert session.summary != ""
- session.reset()
- assert session.history == []
- assert session.summary == ""
- def test_compaction_disabled_never_calls_summarizer():
- """compaction_enabled=False never calls the summarizer and still trims."""
- cfg = _compact_cfg(
- budget_chars=200,
- trigger_ratio=0.8,
- keep_turns=1,
- compaction_enabled=False,
- )
- summarizer_called = False
- def responder(payload):
- nonlocal summarizer_called
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- summarizer_called = True
- return _reply(content="should not happen")
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- for _ in range(6):
- session.send("x" * 40)
- assert not summarizer_called
- assert session.summary == ""
- # History must still be bounded (trim ran).
- assert session._estimate_tokens() <= cfg.prompt_budget
- def test_nothing_to_evict_skips_summarizer():
- """When len(history) <= keep, compaction is skipped and trim runs instead."""
- # keep_turns=10 → keep=20 messages; we only send 2 turns (4 messages).
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=10)
- summarizer_called = False
- def responder(payload):
- nonlocal summarizer_called
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- summarizer_called = True
- return _reply(content="should not happen")
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- # Push over the trigger but keep history shorter than keep=20.
- for _ in range(4):
- session.send("x" * 40)
- assert not summarizer_called
- def test_chunk_and_fold_large_evicted():
- """Large evicted list triggers chunk-and-fold (more than one summarizer call).
- budget_chars=160 → prompt_budget=40 tokens, trigger at 32 tokens (0.8 * 40).
- Each turn: user='x'*40 (40 chars) + assistant='ok' (2 chars) = ~10 tokens.
- Compaction fires at turn 4 with 6 evicted messages.
- compaction_input_tokens=1 → char_limit=4 chars per chunk, so each 40-char
- or 2-char message is its own chunk → 6 summarizer calls for 6 chunks.
- """
- cfg = _compact_cfg(
- budget_chars=160,
- trigger_ratio=0.8,
- keep_turns=1,
- compaction_input_tokens=1, # 1 token * 4 = 4 chars per chunk → each msg is a chunk
- )
- summarizer_calls = []
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- summarizer_calls.append(payload)
- return _reply(content="folded")
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- # 4 turns → 8 messages → fires compaction with 6 evicted (keep=2).
- for _ in range(4):
- session.send("x" * 40)
- # With compaction_input_tokens=1 (4 chars/chunk) and 40-char messages,
- # each message is its own chunk → multiple summarizer calls.
- assert len(summarizer_calls) > 1
- def test_chat_per_request_max_tokens_override():
- """LlmClient.chat sends per-request max_tokens when given, cfg default otherwise."""
- cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
- client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)
- # Default: uses cfg.max_tokens.
- client.chat([{"role": "user", "content": "hi"}])
- assert calls[0]["body"]["max_tokens"] == 1024
- # Override: uses the supplied value.
- client.chat([{"role": "user", "content": "hi"}], max_tokens=128)
- assert calls[1]["body"]["max_tokens"] == 128
- def test_chat_per_request_timeout_override():
- """LlmClient.chat uses per-request timeout_s when given, cfg default otherwise."""
- cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
- recorded_timeouts: list[float] = []
- def opener(payload, timeout):
- recorded_timeouts.append(timeout)
- return _reply(content="ok")
- client = LlmClient(cfg, opener=opener)
- # Default timeout.
- client.chat([{"role": "user", "content": "hi"}])
- assert recorded_timeouts[0] == 60.0
- # Per-request override.
- client.chat([{"role": "user", "content": "hi"}], timeout_s=5.0)
- assert recorded_timeouts[1] == 5.0
- def test_summarizer_uses_compaction_max_tokens_and_timeout():
- """The summarizer call uses compaction_max_tokens and compaction_timeout_s."""
- cfg = _compact_cfg(
- budget_chars=200,
- trigger_ratio=0.8,
- keep_turns=1,
- compaction_max_tokens=42,
- compaction_timeout_s=7.0,
- )
- recorded: list[dict] = []
- def opener(payload, timeout):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- recorded.append({"max_tokens": payload["body"]["max_tokens"], "timeout": timeout})
- return _reply(content="summary")
- return _reply(content="ok")
- client = LlmClient(cfg, opener=opener)
- session = ChatSession(cfg, client)
- for _ in range(4):
- session.send("x" * 40)
- assert len(recorded) >= 1
- assert recorded[0]["max_tokens"] == 42
- assert recorded[0]["timeout"] == 7.0
- def test_notify_called_on_compaction():
- """notify('compacted', turns=N) is called after a successful compaction."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- notifications: list[tuple] = []
- def notify(kind, **fields):
- notifications.append((kind, fields))
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- return _reply(content="summary")
- return _reply(content="ok")
- client = LlmClient(cfg, opener=lambda p, t: responder(p))
- session = ChatSession(cfg, client, notify=notify)
- for _ in range(4):
- session.send("x" * 40)
- assert any(kind == "compacted" for kind, _ in notifications)
- compacted = [(k, f) for k, f in notifications if k == "compacted"]
- assert compacted[0][1]["turns"] > 0
- def test_notify_not_called_when_none():
- """ChatSession works without a notify callback (default None)."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- return _reply(content="summary")
- return _reply(content="ok")
- client = LlmClient(cfg, opener=lambda p, t: responder(p))
- session = ChatSession(cfg, client) # no notify
- # Must not raise.
- for _ in range(4):
- session.send("x" * 40)
- def test_system_prompt_is_first_message_after_compaction():
- """After compaction the chat system prompt is always the first message."""
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "COMPACT":
- return _reply(content="summary text")
- return _reply(content="ok")
- client, calls = _make_client(responder, cfg)
- session = ChatSession(cfg, client)
- for _ in range(4):
- session.send("x" * 40)
- # Find the last chat call (system prompt == "SYS", not "COMPACT").
- chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
- assert chat_calls, "no chat calls found after compaction"
- last_msgs = chat_calls[-1]["body"]["messages"]
- assert last_msgs[0]["role"] == "system"
- assert last_msgs[0]["content"] == "SYS"
- def test_chat_session_truncated_empty_reply_notifies_and_is_not_appended():
- """A length-limited empty reply is reported and never stored as an empty turn."""
- notifications: list[tuple] = []
- def notify(kind, **fields):
- notifications.append((kind, fields))
- client, _ = _make_client(lambda _p: _reply(content="", finish="length"))
- session = ChatSession(LlmConfig(), client, notify=notify)
- assert session.send("hi") == ""
- assert session.history == [{"role": "user", "content": "hi"}]
- assert notifications == [("truncated", {"message": "response truncated"})]
- def test_chat_session_truncated_partial_reply_is_appended_and_notified():
- """A truncated but non-empty reply is kept and reported as truncated."""
- notifications: list[tuple] = []
- def notify(kind, **fields):
- notifications.append((kind, fields))
- client, _ = _make_client(lambda _p: _reply(content="partial answer", finish="length"))
- session = ChatSession(LlmConfig(), client, notify=notify)
- assert session.send("hi") == "partial answer"
- assert session.history[-1] == {"role": "assistant", "content": "partial answer"}
- assert any(kind == "truncated" for kind, _ in notifications)
- def test_on_io_receives_request_then_response():
- """chat(..., on_io=cb) reports the request (url+body) then the response payload."""
- client, _ = _make_client(lambda _p: _reply(content="ok"))
- seen: list[dict] = []
- response = client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)
- assert response.content == "ok"
- assert [event["phase"] for event in seen] == ["request", "response"]
- assert seen[0]["url"] == LlmConfig().endpoint
- assert seen[0]["body"]["messages"] == [{"role": "user", "content": "hi"}]
- assert seen[0]["body"]["max_tokens"] == LlmConfig().max_tokens
- assert seen[1]["payload"]["choices"][0]["message"]["content"] == "ok"
- assert seen[1]["error"] is None
- def test_on_io_response_phase_reports_opener_error():
- """An opener raising still yields a response phase with the error set."""
- import urllib.error
- def opener(_payload, _timeout):
- raise urllib.error.URLError("down")
- client = LlmClient(LlmConfig(), opener=opener)
- seen: list[dict] = []
- with pytest.raises(LlmError):
- client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)
- assert [event["phase"] for event in seen] == ["request", "response"]
- assert seen[1]["payload"] is None
- assert "down" in seen[1]["error"]
- def test_raising_on_io_does_not_break_chat():
- """A throwing on_io callback is swallowed and the reply still parses."""
- def bad(_event):
- raise RuntimeError("boom")
- client, _ = _make_client(lambda _p: _reply(content="ok"))
- assert client.chat([{"role": "user", "content": "hi"}], on_io=bad).content == "ok"
- # ---------------------------------------------------------------------------
- # ConversationEngine tests (Todo 2)
- # ---------------------------------------------------------------------------
- def _fake_engine_client(record=None, content="ok", finish="stop"):
- """Old-style fake: accepts no slot_id/on_io keywords, returns LlmResponse."""
- class _Fake:
- def chat(self, messages, max_tokens=None, timeout_s=None):
- if record is not None:
- record.append(
- {
- "messages": messages,
- "max_tokens": max_tokens,
- "timeout_s": timeout_s,
- }
- )
- return LlmResponse(content=content, finish_reason=finish)
- return _Fake()
- def test_engine_custom_system_prompt_is_first():
- client, calls = _make_client(lambda _p: _reply(content="ok"))
- engine = ConversationEngine(LlmConfig(), client, system_prompt="CUSTOM")
- assert engine.send("hi") == "ok"
- assert calls[0]["body"]["messages"][0] == {
- "role": "system",
- "content": "CUSTOM",
- }
- def test_engine_custom_compaction_prompt_is_used():
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- seen_prompts: list[str] = []
- def responder(payload):
- msgs = payload["body"]["messages"]
- if msgs[0]["content"] == "BRIEF-COMPACT":
- seen_prompts.append(msgs[0]["content"])
- return _reply(content="brief summary")
- assert msgs[0]["content"] == "SYS2"
- return _reply(content="ok")
- client, _ = _make_client(responder, cfg)
- engine = ConversationEngine(
- cfg,
- client,
- system_prompt="SYS2",
- compaction_system_prompt="BRIEF-COMPACT",
- )
- for _ in range(4):
- engine.send("x" * 40)
- assert seen_prompts, "summarizer never ran with the custom prompt"
- assert engine.summary == "brief summary"
- def test_engine_slot_id_sent_only_when_set():
- client, calls = _make_client(lambda _p: _reply(content="ok"))
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
- engine.send("hi")
- assert "id_slot" not in calls[0]["body"]
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS", slot_id=3)
- engine.send("hi")
- assert calls[1]["body"]["id_slot"] == 3
- def test_engine_omits_slot_and_on_io_for_old_fake_when_none():
- record: list[dict] = []
- fake = _fake_engine_client(record)
- engine = ConversationEngine(LlmConfig(), fake, system_prompt="SYS")
- assert engine.send("hi") == "ok"
- assert record[0]["messages"][0] == {"role": "system", "content": "SYS"}
- assert record[0]["max_tokens"] is None
- assert record[0]["timeout_s"] is None
- def test_engine_forwards_per_call_on_io():
- client, _ = _make_client(lambda _p: _reply(content="ok"))
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
- seen: list[dict] = []
- assert engine.send("hi", on_io=seen.append) == "ok"
- assert [event["phase"] for event in seen] == ["request", "response"]
- def test_engine_forwards_slot_and_on_io_when_set():
- record: list[dict] = []
- class _NewFake:
- def chat(self, messages, max_tokens=None, timeout_s=None, **kwargs):
- record.append(kwargs)
- return LlmResponse(content="ok", finish_reason="stop")
- def hook(_event):
- pass
- engine = ConversationEngine(
- LlmConfig(), _NewFake(), system_prompt="SYS", slot_id=7
- )
- assert engine.send("hi", on_io=hook) == "ok"
- assert record[0]["slot_id"] == 7
- assert record[0]["on_io"] is hook
- def test_engine_forwards_max_tokens_and_timeout():
- record: list[dict] = []
- fake = _fake_engine_client(record)
- engine = ConversationEngine(
- LlmConfig(), fake, system_prompt="SYS", max_tokens=11, timeout_s=2.5
- )
- engine.send("hi")
- assert record[0]["max_tokens"] == 11
- assert record[0]["timeout_s"] == 2.5
- def test_last_reply_returns_latest_assistant():
- fake = _fake_engine_client(content="first")
- engine = ConversationEngine(LlmConfig(), fake, system_prompt="SYS")
- assert engine.last_reply() == ""
- engine.send("hi")
- assert engine.last_reply() == "first"
- engine.history.append({"role": "user", "content": "again"})
- engine.history.append({"role": "assistant", "content": "second"})
- assert engine.last_reply() == "second"
- def test_finalize_none_restores_history_and_summary_copy():
- notifications: list[tuple] = []
- def notify(kind, **fields):
- notifications.append((kind, fields))
- client, _ = _make_client(lambda _p: _reply(content="raw reply"))
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS", notify=notify)
- engine.history = [
- {"role": "user", "content": "u1"},
- {"role": "assistant", "content": "a1"},
- ]
- engine.summary = "prior summary"
- assert engine.send("new?", finalize=lambda reply, finish: None) == ""
- assert engine.history == [
- {"role": "user", "content": "u1"},
- {"role": "assistant", "content": "a1"},
- ]
- assert engine.summary == "prior summary"
- assert notifications == []
- def test_finalize_transformed_string_is_appended():
- client, _ = _make_client(lambda _p: _reply(content="raw"))
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
- seen_args: list[tuple] = []
- def finalize(reply, finish):
- seen_args.append((reply, finish))
- return "EDITED: " + reply
- assert engine.send("hi", finalize=finalize) == "EDITED: raw"
- assert seen_args == [("raw", "stop")]
- assert engine.history[-1] == {"role": "assistant", "content": "EDITED: raw"}
- assert engine.last_reply() == "EDITED: raw"
- def test_finalize_empty_string_is_appended_not_restored():
- client, _ = _make_client(lambda _p: _reply(content="raw"))
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
- engine.history = [
- {"role": "user", "content": "u1"},
- {"role": "assistant", "content": "a1"},
- ]
- assert engine.send("new?", finalize=lambda reply, finish: "") == ""
- assert engine.history == [
- {"role": "user", "content": "u1"},
- {"role": "assistant", "content": "a1"},
- {"role": "user", "content": "new?"},
- {"role": "assistant", "content": ""},
- ]
- assert engine.history[-1] == {"role": "assistant", "content": ""}
- def test_finalize_reject_after_reset_mid_flight_stays_empty():
- client, _ = _make_client(lambda _p: _reply(content="late"))
- holder: list[ConversationEngine] = []
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
- holder.append(engine)
- engine.summary = "prior summary"
- def finalize(reply, finish):
- holder[0].reset()
- return None
- assert engine.send("hi", finalize=finalize) == ""
- assert engine.history == []
- assert engine.summary == ""
- def test_finalize_snapshot_is_a_copy():
- client, _ = _make_client(lambda _p: _reply(content="raw"))
- engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
- engine.history = [
- {"role": "user", "content": "u1"},
- {"role": "assistant", "content": "a1"},
- ]
- def finalize(reply, finish):
- engine.history.append({"role": "assistant", "content": "POISON"})
- engine.history.append({"role": "user", "content": "POISON"})
- return None
- assert engine.send("new?", finalize=finalize) == ""
- assert engine.history == [
- {"role": "user", "content": "u1"},
- {"role": "assistant", "content": "a1"},
- ]
- def test_summarizer_never_pinned_to_engine_slot():
- cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
- client, calls = _make_client(
- lambda p: (
- _reply(content="s")
- if p["body"]["messages"][0]["content"] == "CUSTOM-COMPACT"
- else _reply(content="ok")
- ),
- cfg,
- )
- engine = ConversationEngine(
- cfg,
- client,
- system_prompt="SYS",
- compaction_system_prompt="CUSTOM-COMPACT",
- slot_id=5,
- )
- for _ in range(4):
- engine.send("x" * 40)
- chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
- compact_calls = [
- c for c in calls if c["body"]["messages"][0]["content"] == "CUSTOM-COMPACT"
- ]
- assert chat_calls and all(c["body"]["id_slot"] == 5 for c in chat_calls)
- assert compact_calls and all("id_slot" not in c["body"] for c in compact_calls)
- def test_chat_session_is_engine_with_chat_prompt():
- client, calls = _make_client(lambda _p: _reply(content="hello"))
- session = ChatSession(LlmConfig(), client)
- assert isinstance(session, ConversationEngine)
- assert session.system == LlmConfig().system_prompt_chat
- assert session.send("hi") == "hello"
- assert calls[0]["body"]["messages"][0]["content"] == LlmConfig().system_prompt_chat
- assert session.last_reply() == "hello"
- def test_brief_engine_pins_single_turn_without_history():
- client, calls = _make_client(lambda _p: _reply(content="BRIEF v1"))
- engine = BriefEngine(LlmConfig(), client, system_prompt="SYS", max_tokens=99)
- assert engine.last_reply() == ""
- assert engine.send("CURRENT BRIEF:\n(none yet)\n\nNEW UTTERANCES:\n- one") == "BRIEF v1"
- assert engine.last_reply() == "BRIEF v1"
- assert len(calls) == 1
- messages = calls[0]["body"]["messages"]
- assert messages[0] == {"role": "system", "content": "SYS"}
- assert len(messages) == 2
- assert messages[1]["content"].startswith("CURRENT BRIEF:")
- assert calls[0]["body"]["max_tokens"] == 99
- def test_brief_engine_second_call_still_sends_two_messages():
- client, calls = _make_client(lambda _p: _reply(content="BRIEF"))
- engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")
- engine.send("first")
- engine.send("second")
- assert len(calls) == 2
- for call in calls:
- assert len(call["body"]["messages"]) == 2
- def test_brief_engine_finalize_none_leaves_pinned_brief():
- client, _ = _make_client(lambda _p: _reply(content="raw"))
- engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")
- engine._brief = "prior"
- assert engine.send("q", finalize=lambda reply, finish: None) == ""
- assert engine.last_reply() == "prior"
- def test_brief_engine_finalize_transform_is_pinned():
- client, _ = _make_client(lambda _p: _reply(content="raw"))
- engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")
- result = engine.send("q", finalize=lambda reply, finish: "CLEANED: " + reply)
- assert result == "CLEANED: raw"
- assert engine.last_reply() == "CLEANED: raw"
- def test_brief_engine_length_truncation_does_not_pin():
- client, _ = _make_client(lambda _p: _reply(content="half", finish="length"))
- engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")
- assert engine.send("q", finalize=lambda reply, finish: reply) == "half"
- assert engine.last_reply() == ""
- def test_brief_engine_reset_drops_pinned_brief():
- client, _ = _make_client(lambda _p: _reply(content="BRIEF"))
- engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")
- engine.send("q")
- assert engine.last_reply() == "BRIEF"
- engine.reset()
- assert engine.last_reply() == ""
- def test_brief_engine_passes_slot_and_io_hook():
- client, calls = _make_client(lambda _p: _reply(content="BRIEF"))
- engine = BriefEngine(LlmConfig(), client, system_prompt="SYS", slot_id=5)
- seen: list[dict] = []
- engine.send("q", on_io=lambda part: seen.append(part))
- assert calls[0]["body"]["id_slot"] == 5
- assert any(part["phase"] == "request" for part in seen)
|