test_llm.py 34 KB

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