test_llm.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947
  1. """Tests for the LLM client and chat session (no network)."""
  2. import pytest
  3. from stt_daemon.assistant.llm import (
  4. ChatSession,
  5. ConversationEngine,
  6. LlmClient,
  7. LlmError,
  8. LlmResponse,
  9. )
  10. from stt_daemon.config import LlmConfig
  11. def _make_client(responder, cfg=None):
  12. cfg = cfg or LlmConfig()
  13. calls: list[dict] = []
  14. def opener(payload, timeout):
  15. calls.append(payload)
  16. return responder(payload)
  17. return LlmClient(cfg, opener=opener), calls
  18. def _reply(content="", reasoning=None, tool_calls=None, finish="stop"):
  19. message = {"role": "assistant", "content": content}
  20. if reasoning is not None:
  21. message["reasoning_content"] = reasoning
  22. if tool_calls:
  23. message["tool_calls"] = tool_calls
  24. return {"choices": [{"message": message, "finish_reason": finish}]}
  25. def test_reasoning_content_is_ignored():
  26. client, _ = _make_client(
  27. lambda _p: _reply(content="It is noon.", reasoning="thinking hard about it")
  28. )
  29. response = client.chat([{"role": "user", "content": "hi"}])
  30. assert response.content == "It is noon."
  31. assert "thinking hard" not in response.content
  32. def test_cache_prompt_and_endpoint_are_sent():
  33. client, calls = _make_client(lambda _p: _reply(content="ok"))
  34. client.chat([{"role": "user", "content": "hi"}])
  35. assert calls[0]["body"]["cache_prompt"] is True
  36. assert calls[0]["body"]["stream"] is False
  37. assert calls[0]["__url__"] == LlmConfig().endpoint
  38. def test_tools_sent_only_when_provided():
  39. client, calls = _make_client(lambda _p: _reply(content="ok"))
  40. client.chat([{"role": "user", "content": "hi"}])
  41. assert "tools" not in calls[0]["body"]
  42. client.chat(
  43. [{"role": "user", "content": "hi"}],
  44. tools=[{"type": "function", "function": {"name": "get_datetime"}}],
  45. )
  46. assert calls[1]["body"]["tools"][0]["function"]["name"] == "get_datetime"
  47. assert calls[1]["body"]["tool_choice"] == "auto"
  48. def test_id_slot_sent_only_when_slot_id_provided():
  49. client, calls = _make_client(lambda _p: _reply(content="ok"))
  50. client.chat([{"role": "user", "content": "hi"}])
  51. assert "id_slot" not in calls[0]["body"]
  52. client.chat([{"role": "user", "content": "hi"}], slot_id=1)
  53. assert calls[1]["body"]["id_slot"] == 1
  54. def test_slot_and_on_io_never_passed_to_old_fake_when_none():
  55. seen: list[dict] = []
  56. class _OldFake:
  57. def chat(self, messages, tools=None, max_tokens=None, timeout_s=None):
  58. seen.append(
  59. {
  60. "messages": messages,
  61. "tools": tools,
  62. "max_tokens": max_tokens,
  63. "timeout_s": timeout_s,
  64. }
  65. )
  66. return {"ok": True}
  67. fake = _OldFake()
  68. slot_id = None
  69. on_io = None
  70. kwargs: dict = {}
  71. if slot_id is not None:
  72. kwargs["slot_id"] = slot_id
  73. if on_io is not None:
  74. kwargs["on_io"] = on_io
  75. fake.chat([{"role": "user", "content": "hi"}], **kwargs)
  76. assert seen[0] == {
  77. "messages": [{"role": "user", "content": "hi"}],
  78. "tools": None,
  79. "max_tokens": None,
  80. "timeout_s": None,
  81. }
  82. def test_tool_calls_are_parsed():
  83. client, _ = _make_client(
  84. lambda _p: _reply(
  85. tool_calls=[
  86. {
  87. "id": "1",
  88. "type": "function",
  89. "function": {"name": "get_datetime", "arguments": "{}"},
  90. }
  91. ],
  92. finish="tool_calls",
  93. )
  94. )
  95. response = client.chat([{"role": "user", "content": "time?"}], tools=[{"type": "function"}])
  96. assert len(response.tool_calls) == 1
  97. assert response.tool_calls[0].name == "get_datetime"
  98. assert response.finish_reason == "tool_calls"
  99. def test_missing_choices_raises():
  100. client, _ = _make_client(lambda _p: {"choices": []})
  101. with pytest.raises(LlmError):
  102. client.chat([{"role": "user", "content": "hi"}])
  103. def test_chat_session_appends_and_keeps_prefix():
  104. client, calls = _make_client(lambda _p: _reply(content="hello"))
  105. session = ChatSession(LlmConfig(), client)
  106. assert session.send("hi") == "hello"
  107. session.send("again")
  108. assert [m["role"] for m in session.history] == [
  109. "user",
  110. "assistant",
  111. "user",
  112. "assistant",
  113. ]
  114. first = calls[0]["body"]["messages"]
  115. second = calls[1]["body"]["messages"]
  116. assert second[: len(first)] == first # append-only, prefix preserved
  117. assert second[0]["content"] == LlmConfig().system_prompt_chat
  118. def test_chat_session_trims_in_blocks_and_keeps_system_prompt():
  119. cfg = LlmConfig(
  120. max_context_tokens=1330,
  121. max_tokens=1024,
  122. compaction_safety_margin=256,
  123. trim_block_turns=1,
  124. compaction_enabled=False,
  125. ) # prompt_budget == 50
  126. client, calls = _make_client(lambda _p: _reply(content="ok"))
  127. session = ChatSession(cfg, client)
  128. for _ in range(10):
  129. session.send("x" * 40)
  130. assert len(session.history) <= 6
  131. assert calls[-1]["body"]["messages"][0]["content"] == cfg.system_prompt_chat
  132. def test_chat_session_reset_forgets_history():
  133. client, _ = _make_client(lambda _p: _reply(content="ok"))
  134. session = ChatSession(LlmConfig(), client)
  135. session.send("hi")
  136. session.reset()
  137. assert session.history == []
  138. def test_chat_session_drops_reply_when_reset_mid_flight():
  139. holder: list[ChatSession] = []
  140. def opener(_payload, _timeout):
  141. holder[0].reset()
  142. return _reply(content="late reply")
  143. holder.append(ChatSession(LlmConfig(), LlmClient(LlmConfig(), opener=opener)))
  144. assert holder[0].send("hi") == ""
  145. assert holder[0].history == []
  146. # ---------------------------------------------------------------------------
  147. # Phase C tests
  148. # ---------------------------------------------------------------------------
  149. def _compact_cfg(
  150. *,
  151. budget_chars=200,
  152. trigger_ratio=0.8,
  153. keep_turns=1,
  154. compaction_max_tokens=64,
  155. compaction_timeout_s=5.0,
  156. compaction_input_tokens=8192,
  157. compaction_enabled=True,
  158. trim_block_turns=1,
  159. ):
  160. """Return a small LlmConfig whose prompt_budget == budget_chars // 4 tokens.
  161. prompt_budget = max_context_tokens - max_tokens - compaction_safety_margin
  162. We want budget_chars // 4 tokens, so:
  163. max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
  164. """
  165. max_tokens = 16
  166. compaction_safety_margin = 0
  167. max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
  168. return LlmConfig(
  169. max_context_tokens=max_context_tokens,
  170. max_tokens=max_tokens,
  171. compaction_safety_margin=compaction_safety_margin,
  172. compaction_enabled=compaction_enabled,
  173. compaction_trigger_ratio=trigger_ratio,
  174. compaction_keep_turns=keep_turns,
  175. compaction_max_tokens=compaction_max_tokens,
  176. compaction_timeout_s=compaction_timeout_s,
  177. compaction_input_tokens=compaction_input_tokens,
  178. trim_block_turns=trim_block_turns,
  179. system_prompt_chat="SYS",
  180. system_prompt_command="CMD",
  181. compaction_system_prompt="COMPACT",
  182. )
  183. def test_compaction_fires_and_window_shrinks():
  184. """Crossing the trigger compacts: window shrinks, summary present in next prompt."""
  185. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  186. # prompt_budget = 200//4 = 50 tokens; trigger at 40 tokens (0.8 * 50)
  187. summary_text = "summary of old turns"
  188. call_index = 0
  189. def responder(payload):
  190. nonlocal call_index
  191. msgs = payload["body"]["messages"]
  192. call_index += 1
  193. # Summarizer calls have the compaction system prompt.
  194. if msgs[0]["content"] == "COMPACT":
  195. return _reply(content=summary_text)
  196. return _reply(content="ok")
  197. client, calls = _make_client(responder, cfg)
  198. session = ChatSession(cfg, client)
  199. # Send enough turns to exceed the trigger (each "x"*40 = 40 chars = 10 tokens).
  200. # After 4 turns (user+assistant pairs = 8 messages * 10 tokens = 80 tokens > 40 trigger).
  201. for _ in range(4):
  202. session.send("x" * 40)
  203. # After compaction the window must be smaller than the full history.
  204. assert len(session.history) <= 4 # keep_turns=1 → keep=2 messages
  205. # The summary must be set.
  206. assert session.summary != ""
  207. # The next chat call must include the summary as a second system message,
  208. # with the chat system prompt FIRST.
  209. session.send("after compaction")
  210. chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
  211. last_chat = chat_calls[-1]["body"]["messages"]
  212. assert last_chat[0]["role"] == "system"
  213. assert last_chat[0]["content"] == "SYS"
  214. assert last_chat[1]["role"] == "system"
  215. assert "Earlier conversation summary:" in last_chat[1]["content"]
  216. assert summary_text in last_chat[1]["content"]
  217. def test_summary_never_appended_to_history():
  218. """The summary text must not appear as a history entry."""
  219. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  220. summary_text = "SUMMARY_SENTINEL"
  221. def responder(payload):
  222. msgs = payload["body"]["messages"]
  223. if msgs[0]["content"] == "COMPACT":
  224. return _reply(content=summary_text)
  225. return _reply(content="ok")
  226. client, _ = _make_client(responder, cfg)
  227. session = ChatSession(cfg, client)
  228. for _ in range(4):
  229. session.send("x" * 40)
  230. for msg in session.history:
  231. assert summary_text not in msg["content"]
  232. def test_summarizer_failure_falls_back_to_trim():
  233. """On summarizer failure the session falls back to _trim() and stays within budget."""
  234. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  235. def responder(payload):
  236. msgs = payload["body"]["messages"]
  237. if msgs[0]["content"] == "COMPACT":
  238. return _reply(content="", finish="stop") # empty → failure
  239. return _reply(content="ok")
  240. client, _ = _make_client(responder, cfg)
  241. session = ChatSession(cfg, client)
  242. for _ in range(4):
  243. session.send("x" * 40)
  244. # Summary must remain empty (fallback used trim, not compaction).
  245. assert session.summary == ""
  246. # History must be bounded within the budget.
  247. assert session._estimate_tokens() <= cfg.prompt_budget
  248. def test_summarizer_length_finish_reason_falls_back_to_trim():
  249. """finish_reason == 'length' on the summarizer is treated as failure."""
  250. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  251. def responder(payload):
  252. msgs = payload["body"]["messages"]
  253. if msgs[0]["content"] == "COMPACT":
  254. return _reply(content="partial summary", finish="length")
  255. return _reply(content="ok")
  256. client, _ = _make_client(responder, cfg)
  257. session = ChatSession(cfg, client)
  258. for _ in range(4):
  259. session.send("x" * 40)
  260. assert session.summary == ""
  261. assert session._estimate_tokens() <= cfg.prompt_budget
  262. def test_summarizer_llmerror_falls_back_to_trim():
  263. """LlmError from the summarizer falls back to _trim()."""
  264. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  265. def responder(payload):
  266. msgs = payload["body"]["messages"]
  267. if msgs[0]["content"] == "COMPACT":
  268. raise LlmError("network error")
  269. return _reply(content="ok")
  270. client, _ = _make_client(responder, cfg)
  271. session = ChatSession(cfg, client)
  272. for _ in range(4):
  273. session.send("x" * 40)
  274. assert session.summary == ""
  275. assert session._estimate_tokens() <= cfg.prompt_budget
  276. def test_reset_clears_summary():
  277. """reset() clears both history and summary."""
  278. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  279. summary_text = "old summary"
  280. def responder(payload):
  281. msgs = payload["body"]["messages"]
  282. if msgs[0]["content"] == "COMPACT":
  283. return _reply(content=summary_text)
  284. return _reply(content="ok")
  285. client, _ = _make_client(responder, cfg)
  286. session = ChatSession(cfg, client)
  287. for _ in range(4):
  288. session.send("x" * 40)
  289. assert session.summary != ""
  290. session.reset()
  291. assert session.history == []
  292. assert session.summary == ""
  293. def test_compaction_disabled_never_calls_summarizer():
  294. """compaction_enabled=False never calls the summarizer and still trims."""
  295. cfg = _compact_cfg(
  296. budget_chars=200,
  297. trigger_ratio=0.8,
  298. keep_turns=1,
  299. compaction_enabled=False,
  300. )
  301. summarizer_called = False
  302. def responder(payload):
  303. nonlocal summarizer_called
  304. msgs = payload["body"]["messages"]
  305. if msgs[0]["content"] == "COMPACT":
  306. summarizer_called = True
  307. return _reply(content="should not happen")
  308. return _reply(content="ok")
  309. client, _ = _make_client(responder, cfg)
  310. session = ChatSession(cfg, client)
  311. for _ in range(6):
  312. session.send("x" * 40)
  313. assert not summarizer_called
  314. assert session.summary == ""
  315. # History must still be bounded (trim ran).
  316. assert session._estimate_tokens() <= cfg.prompt_budget
  317. def test_nothing_to_evict_skips_summarizer():
  318. """When len(history) <= keep, compaction is skipped and trim runs instead."""
  319. # keep_turns=10 → keep=20 messages; we only send 2 turns (4 messages).
  320. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=10)
  321. summarizer_called = False
  322. def responder(payload):
  323. nonlocal summarizer_called
  324. msgs = payload["body"]["messages"]
  325. if msgs[0]["content"] == "COMPACT":
  326. summarizer_called = True
  327. return _reply(content="should not happen")
  328. return _reply(content="ok")
  329. client, _ = _make_client(responder, cfg)
  330. session = ChatSession(cfg, client)
  331. # Push over the trigger but keep history shorter than keep=20.
  332. for _ in range(4):
  333. session.send("x" * 40)
  334. assert not summarizer_called
  335. def test_chunk_and_fold_large_evicted():
  336. """Large evicted list triggers chunk-and-fold (more than one summarizer call).
  337. budget_chars=160 → prompt_budget=40 tokens, trigger at 32 tokens (0.8 * 40).
  338. Each turn: user='x'*40 (40 chars) + assistant='ok' (2 chars) = ~10 tokens.
  339. Compaction fires at turn 4 with 6 evicted messages.
  340. compaction_input_tokens=1 → char_limit=4 chars per chunk, so each 40-char
  341. or 2-char message is its own chunk → 6 summarizer calls for 6 chunks.
  342. """
  343. cfg = _compact_cfg(
  344. budget_chars=160,
  345. trigger_ratio=0.8,
  346. keep_turns=1,
  347. compaction_input_tokens=1, # 1 token * 4 = 4 chars per chunk → each msg is a chunk
  348. )
  349. summarizer_calls = []
  350. def responder(payload):
  351. msgs = payload["body"]["messages"]
  352. if msgs[0]["content"] == "COMPACT":
  353. summarizer_calls.append(payload)
  354. return _reply(content="folded")
  355. return _reply(content="ok")
  356. client, _ = _make_client(responder, cfg)
  357. session = ChatSession(cfg, client)
  358. # 4 turns → 8 messages → fires compaction with 6 evicted (keep=2).
  359. for _ in range(4):
  360. session.send("x" * 40)
  361. # With compaction_input_tokens=1 (4 chars/chunk) and 40-char messages,
  362. # each message is its own chunk → multiple summarizer calls.
  363. assert len(summarizer_calls) > 1
  364. def test_chat_per_request_max_tokens_override():
  365. """LlmClient.chat sends per-request max_tokens when given, cfg default otherwise."""
  366. cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
  367. client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)
  368. # Default: uses cfg.max_tokens.
  369. client.chat([{"role": "user", "content": "hi"}])
  370. assert calls[0]["body"]["max_tokens"] == 1024
  371. # Override: uses the supplied value.
  372. client.chat([{"role": "user", "content": "hi"}], max_tokens=128)
  373. assert calls[1]["body"]["max_tokens"] == 128
  374. def test_chat_per_request_timeout_override():
  375. """LlmClient.chat uses per-request timeout_s when given, cfg default otherwise."""
  376. cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
  377. recorded_timeouts: list[float] = []
  378. def opener(payload, timeout):
  379. recorded_timeouts.append(timeout)
  380. return _reply(content="ok")
  381. client = LlmClient(cfg, opener=opener)
  382. # Default timeout.
  383. client.chat([{"role": "user", "content": "hi"}])
  384. assert recorded_timeouts[0] == 60.0
  385. # Per-request override.
  386. client.chat([{"role": "user", "content": "hi"}], timeout_s=5.0)
  387. assert recorded_timeouts[1] == 5.0
  388. def test_summarizer_uses_compaction_max_tokens_and_timeout():
  389. """The summarizer call uses compaction_max_tokens and compaction_timeout_s."""
  390. cfg = _compact_cfg(
  391. budget_chars=200,
  392. trigger_ratio=0.8,
  393. keep_turns=1,
  394. compaction_max_tokens=42,
  395. compaction_timeout_s=7.0,
  396. )
  397. recorded: list[dict] = []
  398. def opener(payload, timeout):
  399. msgs = payload["body"]["messages"]
  400. if msgs[0]["content"] == "COMPACT":
  401. recorded.append({"max_tokens": payload["body"]["max_tokens"], "timeout": timeout})
  402. return _reply(content="summary")
  403. return _reply(content="ok")
  404. client = LlmClient(cfg, opener=opener)
  405. session = ChatSession(cfg, client)
  406. for _ in range(4):
  407. session.send("x" * 40)
  408. assert len(recorded) >= 1
  409. assert recorded[0]["max_tokens"] == 42
  410. assert recorded[0]["timeout"] == 7.0
  411. def test_notify_called_on_compaction():
  412. """notify('compacted', turns=N) is called after a successful compaction."""
  413. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  414. notifications: list[tuple] = []
  415. def notify(kind, **fields):
  416. notifications.append((kind, fields))
  417. def responder(payload):
  418. msgs = payload["body"]["messages"]
  419. if msgs[0]["content"] == "COMPACT":
  420. return _reply(content="summary")
  421. return _reply(content="ok")
  422. client = LlmClient(cfg, opener=lambda p, t: responder(p))
  423. session = ChatSession(cfg, client, notify=notify)
  424. for _ in range(4):
  425. session.send("x" * 40)
  426. assert any(kind == "compacted" for kind, _ in notifications)
  427. compacted = [(k, f) for k, f in notifications if k == "compacted"]
  428. assert compacted[0][1]["turns"] > 0
  429. def test_notify_not_called_when_none():
  430. """ChatSession works without a notify callback (default None)."""
  431. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  432. def responder(payload):
  433. msgs = payload["body"]["messages"]
  434. if msgs[0]["content"] == "COMPACT":
  435. return _reply(content="summary")
  436. return _reply(content="ok")
  437. client = LlmClient(cfg, opener=lambda p, t: responder(p))
  438. session = ChatSession(cfg, client) # no notify
  439. # Must not raise.
  440. for _ in range(4):
  441. session.send("x" * 40)
  442. def test_system_prompt_is_first_message_after_compaction():
  443. """After compaction the chat system prompt is always the first message."""
  444. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  445. def responder(payload):
  446. msgs = payload["body"]["messages"]
  447. if msgs[0]["content"] == "COMPACT":
  448. return _reply(content="summary text")
  449. return _reply(content="ok")
  450. client, calls = _make_client(responder, cfg)
  451. session = ChatSession(cfg, client)
  452. for _ in range(4):
  453. session.send("x" * 40)
  454. # Find the last chat call (system prompt == "SYS", not "COMPACT").
  455. chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
  456. assert chat_calls, "no chat calls found after compaction"
  457. last_msgs = chat_calls[-1]["body"]["messages"]
  458. assert last_msgs[0]["role"] == "system"
  459. assert last_msgs[0]["content"] == "SYS"
  460. def test_chat_session_truncated_empty_reply_notifies_and_is_not_appended():
  461. """A length-limited empty reply is reported and never stored as an empty turn."""
  462. notifications: list[tuple] = []
  463. def notify(kind, **fields):
  464. notifications.append((kind, fields))
  465. client, _ = _make_client(lambda _p: _reply(content="", finish="length"))
  466. session = ChatSession(LlmConfig(), client, notify=notify)
  467. assert session.send("hi") == ""
  468. assert session.history == [{"role": "user", "content": "hi"}]
  469. assert notifications == [("truncated", {"message": "response truncated"})]
  470. def test_chat_session_truncated_partial_reply_is_appended_and_notified():
  471. """A truncated but non-empty reply is kept and reported as truncated."""
  472. notifications: list[tuple] = []
  473. def notify(kind, **fields):
  474. notifications.append((kind, fields))
  475. client, _ = _make_client(lambda _p: _reply(content="partial answer", finish="length"))
  476. session = ChatSession(LlmConfig(), client, notify=notify)
  477. assert session.send("hi") == "partial answer"
  478. assert session.history[-1] == {"role": "assistant", "content": "partial answer"}
  479. assert any(kind == "truncated" for kind, _ in notifications)
  480. def test_on_io_receives_request_then_response():
  481. """chat(..., on_io=cb) reports the request (url+body) then the response payload."""
  482. client, _ = _make_client(lambda _p: _reply(content="ok"))
  483. seen: list[dict] = []
  484. response = client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)
  485. assert response.content == "ok"
  486. assert [event["phase"] for event in seen] == ["request", "response"]
  487. assert seen[0]["url"] == LlmConfig().endpoint
  488. assert seen[0]["body"]["messages"] == [{"role": "user", "content": "hi"}]
  489. assert seen[0]["body"]["max_tokens"] == LlmConfig().max_tokens
  490. assert seen[1]["payload"]["choices"][0]["message"]["content"] == "ok"
  491. assert seen[1]["error"] is None
  492. def test_on_io_response_phase_reports_opener_error():
  493. """An opener raising still yields a response phase with the error set."""
  494. import urllib.error
  495. def opener(_payload, _timeout):
  496. raise urllib.error.URLError("down")
  497. client = LlmClient(LlmConfig(), opener=opener)
  498. seen: list[dict] = []
  499. with pytest.raises(LlmError):
  500. client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)
  501. assert [event["phase"] for event in seen] == ["request", "response"]
  502. assert seen[1]["payload"] is None
  503. assert "down" in seen[1]["error"]
  504. def test_raising_on_io_does_not_break_chat():
  505. """A throwing on_io callback is swallowed and the reply still parses."""
  506. def bad(_event):
  507. raise RuntimeError("boom")
  508. client, _ = _make_client(lambda _p: _reply(content="ok"))
  509. assert client.chat([{"role": "user", "content": "hi"}], on_io=bad).content == "ok"
  510. # ---------------------------------------------------------------------------
  511. # ConversationEngine tests (Todo 2)
  512. # ---------------------------------------------------------------------------
  513. def _fake_engine_client(record=None, content="ok", finish="stop"):
  514. """Old-style fake: accepts no slot_id/on_io keywords, returns LlmResponse."""
  515. class _Fake:
  516. def chat(self, messages, max_tokens=None, timeout_s=None):
  517. if record is not None:
  518. record.append(
  519. {
  520. "messages": messages,
  521. "max_tokens": max_tokens,
  522. "timeout_s": timeout_s,
  523. }
  524. )
  525. return LlmResponse(content=content, finish_reason=finish)
  526. return _Fake()
  527. def test_engine_custom_system_prompt_is_first():
  528. client, calls = _make_client(lambda _p: _reply(content="ok"))
  529. engine = ConversationEngine(LlmConfig(), client, system_prompt="CUSTOM")
  530. assert engine.send("hi") == "ok"
  531. assert calls[0]["body"]["messages"][0] == {
  532. "role": "system",
  533. "content": "CUSTOM",
  534. }
  535. def test_engine_custom_compaction_prompt_is_used():
  536. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  537. seen_prompts: list[str] = []
  538. def responder(payload):
  539. msgs = payload["body"]["messages"]
  540. if msgs[0]["content"] == "BRIEF-COMPACT":
  541. seen_prompts.append(msgs[0]["content"])
  542. return _reply(content="brief summary")
  543. assert msgs[0]["content"] == "SYS2"
  544. return _reply(content="ok")
  545. client, _ = _make_client(responder, cfg)
  546. engine = ConversationEngine(
  547. cfg,
  548. client,
  549. system_prompt="SYS2",
  550. compaction_system_prompt="BRIEF-COMPACT",
  551. )
  552. for _ in range(4):
  553. engine.send("x" * 40)
  554. assert seen_prompts, "summarizer never ran with the custom prompt"
  555. assert engine.summary == "brief summary"
  556. def test_engine_slot_id_sent_only_when_set():
  557. client, calls = _make_client(lambda _p: _reply(content="ok"))
  558. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
  559. engine.send("hi")
  560. assert "id_slot" not in calls[0]["body"]
  561. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS", slot_id=3)
  562. engine.send("hi")
  563. assert calls[1]["body"]["id_slot"] == 3
  564. def test_engine_omits_slot_and_on_io_for_old_fake_when_none():
  565. record: list[dict] = []
  566. fake = _fake_engine_client(record)
  567. engine = ConversationEngine(LlmConfig(), fake, system_prompt="SYS")
  568. assert engine.send("hi") == "ok"
  569. assert record[0]["messages"][0] == {"role": "system", "content": "SYS"}
  570. assert record[0]["max_tokens"] is None
  571. assert record[0]["timeout_s"] is None
  572. def test_engine_forwards_per_call_on_io():
  573. client, _ = _make_client(lambda _p: _reply(content="ok"))
  574. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
  575. seen: list[dict] = []
  576. assert engine.send("hi", on_io=seen.append) == "ok"
  577. assert [event["phase"] for event in seen] == ["request", "response"]
  578. def test_engine_forwards_slot_and_on_io_when_set():
  579. record: list[dict] = []
  580. class _NewFake:
  581. def chat(self, messages, max_tokens=None, timeout_s=None, **kwargs):
  582. record.append(kwargs)
  583. return LlmResponse(content="ok", finish_reason="stop")
  584. def hook(_event):
  585. pass
  586. engine = ConversationEngine(
  587. LlmConfig(), _NewFake(), system_prompt="SYS", slot_id=7
  588. )
  589. assert engine.send("hi", on_io=hook) == "ok"
  590. assert record[0]["slot_id"] == 7
  591. assert record[0]["on_io"] is hook
  592. def test_engine_forwards_max_tokens_and_timeout():
  593. record: list[dict] = []
  594. fake = _fake_engine_client(record)
  595. engine = ConversationEngine(
  596. LlmConfig(), fake, system_prompt="SYS", max_tokens=11, timeout_s=2.5
  597. )
  598. engine.send("hi")
  599. assert record[0]["max_tokens"] == 11
  600. assert record[0]["timeout_s"] == 2.5
  601. def test_last_reply_returns_latest_assistant():
  602. fake = _fake_engine_client(content="first")
  603. engine = ConversationEngine(LlmConfig(), fake, system_prompt="SYS")
  604. assert engine.last_reply() == ""
  605. engine.send("hi")
  606. assert engine.last_reply() == "first"
  607. engine.history.append({"role": "user", "content": "again"})
  608. engine.history.append({"role": "assistant", "content": "second"})
  609. assert engine.last_reply() == "second"
  610. def test_finalize_none_restores_history_and_summary_copy():
  611. notifications: list[tuple] = []
  612. def notify(kind, **fields):
  613. notifications.append((kind, fields))
  614. client, _ = _make_client(lambda _p: _reply(content="raw reply"))
  615. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS", notify=notify)
  616. engine.history = [
  617. {"role": "user", "content": "u1"},
  618. {"role": "assistant", "content": "a1"},
  619. ]
  620. engine.summary = "prior summary"
  621. assert engine.send("new?", finalize=lambda reply, finish: None) == ""
  622. assert engine.history == [
  623. {"role": "user", "content": "u1"},
  624. {"role": "assistant", "content": "a1"},
  625. ]
  626. assert engine.summary == "prior summary"
  627. assert notifications == []
  628. def test_finalize_transformed_string_is_appended():
  629. client, _ = _make_client(lambda _p: _reply(content="raw"))
  630. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
  631. seen_args: list[tuple] = []
  632. def finalize(reply, finish):
  633. seen_args.append((reply, finish))
  634. return "EDITED: " + reply
  635. assert engine.send("hi", finalize=finalize) == "EDITED: raw"
  636. assert seen_args == [("raw", "stop")]
  637. assert engine.history[-1] == {"role": "assistant", "content": "EDITED: raw"}
  638. assert engine.last_reply() == "EDITED: raw"
  639. def test_finalize_empty_string_is_appended_not_restored():
  640. client, _ = _make_client(lambda _p: _reply(content="raw"))
  641. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
  642. engine.history = [
  643. {"role": "user", "content": "u1"},
  644. {"role": "assistant", "content": "a1"},
  645. ]
  646. assert engine.send("new?", finalize=lambda reply, finish: "") == ""
  647. assert engine.history == [
  648. {"role": "user", "content": "u1"},
  649. {"role": "assistant", "content": "a1"},
  650. {"role": "user", "content": "new?"},
  651. {"role": "assistant", "content": ""},
  652. ]
  653. assert engine.history[-1] == {"role": "assistant", "content": ""}
  654. def test_finalize_reject_after_reset_mid_flight_stays_empty():
  655. client, _ = _make_client(lambda _p: _reply(content="late"))
  656. holder: list[ConversationEngine] = []
  657. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
  658. holder.append(engine)
  659. engine.summary = "prior summary"
  660. def finalize(reply, finish):
  661. holder[0].reset()
  662. return None
  663. assert engine.send("hi", finalize=finalize) == ""
  664. assert engine.history == []
  665. assert engine.summary == ""
  666. def test_finalize_snapshot_is_a_copy():
  667. client, _ = _make_client(lambda _p: _reply(content="raw"))
  668. engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
  669. engine.history = [
  670. {"role": "user", "content": "u1"},
  671. {"role": "assistant", "content": "a1"},
  672. ]
  673. def finalize(reply, finish):
  674. engine.history.append({"role": "assistant", "content": "POISON"})
  675. engine.history.append({"role": "user", "content": "POISON"})
  676. return None
  677. assert engine.send("new?", finalize=finalize) == ""
  678. assert engine.history == [
  679. {"role": "user", "content": "u1"},
  680. {"role": "assistant", "content": "a1"},
  681. ]
  682. def test_summarizer_never_pinned_to_engine_slot():
  683. cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
  684. client, calls = _make_client(
  685. lambda p: (
  686. _reply(content="s")
  687. if p["body"]["messages"][0]["content"] == "CUSTOM-COMPACT"
  688. else _reply(content="ok")
  689. ),
  690. cfg,
  691. )
  692. engine = ConversationEngine(
  693. cfg,
  694. client,
  695. system_prompt="SYS",
  696. compaction_system_prompt="CUSTOM-COMPACT",
  697. slot_id=5,
  698. )
  699. for _ in range(4):
  700. engine.send("x" * 40)
  701. chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
  702. compact_calls = [
  703. c for c in calls if c["body"]["messages"][0]["content"] == "CUSTOM-COMPACT"
  704. ]
  705. assert chat_calls and all(c["body"]["id_slot"] == 5 for c in chat_calls)
  706. assert compact_calls and all("id_slot" not in c["body"] for c in compact_calls)
  707. def test_chat_session_is_engine_with_chat_prompt():
  708. client, calls = _make_client(lambda _p: _reply(content="hello"))
  709. session = ChatSession(LlmConfig(), client)
  710. assert isinstance(session, ConversationEngine)
  711. assert session.system == LlmConfig().system_prompt_chat
  712. assert session.send("hi") == "hello"
  713. assert calls[0]["body"]["messages"][0]["content"] == LlmConfig().system_prompt_chat
  714. assert session.last_reply() == "hello"