| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947 |
- """Tests for the LLM client and chat session (no network)."""
- import pytest
- from stt_daemon.assistant.llm import (
- 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"
|