Coverage for src/lilbee/server/handlers/rag.py: 100%

1"""Search, ask, and chat handlers (one-shot and streaming).""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import contextlib 

7import logging 

8import threading 

9from collections.abc import AsyncGenerator, AsyncIterator, Callable 

10from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast 

11 

12from lilbee.app.memory import auto_extract, auto_extract_enabled 

13from lilbee.app.search import clean_result 

14from lilbee.app.services import get_services 

15from lilbee.core.config import cfg 

16from lilbee.core.results import DocumentResult, group 

17from lilbee.data.store import ChunkType, EmbeddingModelMismatchError 

18from lilbee.providers.base import ProviderError, ProviderErrorKind 

19from lilbee.providers.roles import WorkerRole 

20from lilbee.retrieval.query.compaction import ( 

21 compaction_due, 

22 foldable, 

23 history_budget, 

24 prompt_history, 

25) 

26from lilbee.retrieval.query.formatting import ( 

27 StreamingCitationFilter, 

28 cited_subset, 

29 strip_llm_citations, 

30) 

31from lilbee.retrieval.query.searcher import ( 

32 GROUNDED_REFUSAL, 

33 SEARCH_NEEDS_EMBEDDER, 

34 RagContext, 

35) 

36from lilbee.retrieval.reasoning import ( 

37 CAP_NOTICE_TEMPLATE, 

38 REASONING_EXHAUSTED_NOTICE, 

39 CapNotice, 

40 StreamToken, 

41 TagParser, 

42 effective_reasoning_cap, 

43 stream_chat_with_cap, 

44 strip_reasoning, 

45) 

46from lilbee.runtime.progress import SseErrorCode, SseEvent 

47from lilbee.server.chat_completions_api.errors import ( 

48 _BACKEND_FAILURE_MESSAGE, 

49 _INFRASTRUCTURE_KINDS, 

50 CompletionsErrorCode, 

51) 

52from lilbee.server.chat_dispatch.canonical import ( 

53 CanonicalChatRequest, 

54 CanonicalMessage, 

55 ContentBlockDelta, 

56 TextBlock, 

57 TextDelta, 

58) 

59from lilbee.server.chat_dispatch.dispatch import ( 

60 ModelDoesNotSupportToolsError, 

61 ModelNotFoundError, 

62 dispatch_chat, 

63 dispatch_chat_stream, 

64) 

65from lilbee.server.chat_dispatch.reasoning_cap import nudged_request 

66from lilbee.server.handlers.sse import ( 

67 SseErrorCodeValue, 

68 SseStream, 

69 _resolve_generation_options, 

70 classify_load_error, 

71 sse_done, 

72 sse_error, 

73 sse_event, 

74) 

75from lilbee.server.models import ( 

76 AskResponse, 

77 CleanedChunk, 

78 CompactionInfo, 

79 MemoryExtractedEvent, 

80 MemoryExtractedItem, 

81) 

82from lilbee.sessions import SessionNotFoundError, sessions_enabled 

83 

84if TYPE_CHECKING: 

85 from lilbee.core.results import SearchChunk 

86 from lilbee.retrieval.query import ChatMessage 

87 from lilbee.retrieval.query.searcher import Searcher 

88 

89log = logging.getLogger(__name__) 

90 

91 

92# Unmapped kinds surface as their ProviderErrorKind string; shipped clients branch 

93# on it. The kinds that describe the backend keep their code and lose their text 

94# (see _classify_stream_error), so a client can still branch without being handed 

95# engine internals. 

96_STREAM_KIND_CODES: dict[ProviderErrorKind, CompletionsErrorCode] = { 

97 ProviderErrorKind.CONTEXT_OVERFLOW: CompletionsErrorCode.CONTEXT_LENGTH_EXCEEDED, 

98 ProviderErrorKind.NOT_FOUND: CompletionsErrorCode.MODEL_NOT_FOUND, 

99} 

100 

101 

102def _classify_stream_error(exc: BaseException) -> tuple[SseErrorCodeValue | None, str]: 

103 """Return ``(code, user_message)`` for an SSE error event, typed-exception aware.""" 

104 if isinstance(exc, ModelNotFoundError): 

105 return CompletionsErrorCode.MODEL_NOT_FOUND, str(exc) 

106 if isinstance(exc, ModelDoesNotSupportToolsError): 

107 return CompletionsErrorCode.MODEL_DOES_NOT_SUPPORT_TOOLS, str(exc) 

108 if isinstance(exc, ProviderError): 

109 mapped = _STREAM_KIND_CODES.get(exc.kind) 

110 if mapped is not None: 

111 return mapped, str(exc) 

112 code = None if exc.kind is ProviderErrorKind.UNKNOWN else exc.kind 

113 if exc.kind in _INFRASTRUCTURE_KINDS: 

114 # Kinds that describe the backend rather than the request. Their text 

115 # is built at the fleet boundary and carries the dead engine's stderr, 

116 # so it is logged rather than sent, exactly as the completions surface 

117 # already does. Both surfaces answer to the same set. 

118 log.warning("Backend failure on the stream surface: %s", exc) 

119 return code, _BACKEND_FAILURE_MESSAGE 

120 return code, str(exc) 

121 return classify_load_error(str(exc)) 

122 

123 

124async def search( 

125 q: str, top_k: int = 5, chunk_type: ChunkType | None = None 

126) -> list[DocumentResult]: 

127 """Search and return grouped DocumentResults.""" 

128 if not q or not q.strip(): 

129 raise ValueError("query must not be empty") 

130 # search() blocks on retrieval; run it off the event loop so other admitted 

131 # requests stay responsive, matching the sibling ask() handler. 

132 results = await asyncio.to_thread( 

133 get_services().searcher.search, q, top_k=top_k, chunk_type=chunk_type 

134 ) 

135 return group(results) 

136 

137 

138async def ask( 

139 question: str, 

140 top_k: int = 0, 

141 options: dict[str, Any] | None = None, 

142 chunk_type: ChunkType | None = None, 

143) -> AskResponse: 

144 """One-shot RAG answer. Returns answer and sources.""" 

145 if not question or not question.strip(): 

146 raise ValueError("question must not be empty") 

147 opts = _resolve_generation_options(options) 

148 searcher = get_services().searcher 

149 # ask_raw blocks for retrieval plus the whole generation; run it off the 

150 # event loop so other admitted requests stay responsive. 

151 result = await asyncio.to_thread( 

152 searcher.ask_raw, 

153 question, 

154 top_k=top_k, 

155 options=opts, 

156 chunk_type=chunk_type, 

157 ) 

158 # Mirror the streaming ask path: auto-extract memories from a real answer, 

159 # but never from the search-needs-embedder refusal ask_raw returns. 

160 if not searcher.search_unavailable(): 

161 await _store_extracted_memories(question, result.answer) 

162 return AskResponse( 

163 answer=result.answer, 

164 sources=[CleanedChunk(**clean_result(s)) for s in result.sources], 

165 cited_sources=[CleanedChunk(**clean_result(s)) for s in result.cited_sources], 

166 retrieval_query=result.retrieval_query, 

167 ) 

168 

169 

170def _chat_warming_events() -> list[str]: 

171 """One ``warming`` SSE event when the chat server is cold, else nothing. 

172 

173 A cold chat server blocks the first token while it loads; the early event 

174 lets the client show a warming state instead of an apparently-dead stream. 

175 """ 

176 if get_services().provider.role_ready(WorkerRole.CHAT): 

177 return [] 

178 log.info("Chat engine cold; streaming a warming notice before the first token.") 

179 return [sse_event(SseEvent.WARMING, {"role": WorkerRole.CHAT.value})] 

180 

181 

182def _put_answer_token( 

183 content: str, 

184 put: Callable[[str | None], None], 

185 cite_filter: StreamingCitationFilter | None, 

186 answer_parts: list[str], 

187) -> None: 

188 """Filter one streamed answer chunk (dropping a model Sources block on 

189 grounded turns), record it, and push it to the SSE queue.""" 

190 token = cite_filter.feed(content) if cite_filter else content 

191 if token: 

192 answer_parts.append(token) 

193 put(sse_event(SseEvent.TOKEN, {"token": token})) 

194 

195 

196def _put_answer_tail( 

197 put: Callable[[str | None], None], 

198 cite_filter: StreamingCitationFilter | None, 

199 answer_parts: list[str], 

200) -> None: 

201 """Release any answer text the filter held back once the stream ends.""" 

202 if cite_filter is None: 

203 return 

204 tail = cite_filter.flush() 

205 if tail: 

206 answer_parts.append(tail) 

207 put(sse_event(SseEvent.TOKEN, {"token": tail})) 

208 

209 

210def _run_llm_stream( 

211 messages: list[ChatMessage], 

212 opts: dict[str, Any] | None, 

213 put: Callable[[str | None], None], 

214 cancel: threading.Event, 

215 error_holder: list[BaseException], 

216 answer_parts: list[str], 

217 cite_filter: StreamingCitationFilter | None, 

218) -> None: 

219 """Forward tokens from the cap-aware chat orchestrator into the SSE queue. 

220 

221 Answer tokens (not reasoning) are also accumulated into *answer_parts* so the 

222 caller can feed the finished answer to auto-extraction. When *cite_filter* is 

223 set (grounded turns), answer tokens pass through it so a model-generated 

224 ``Sources:`` block never reaches the client alongside the authoritative 

225 SOURCES event; ungrounded turns pass ``None`` and stream verbatim. 

226 """ 

227 try: 

228 events = stream_chat_with_cap( 

229 get_services().provider, 

230 cast("list[dict[str, Any]]", messages), 

231 options=opts, 

232 model=cfg.chat_model, 

233 show_reasoning=cfg.show_reasoning, 

234 cap_chars=effective_reasoning_cap(), 

235 ) 

236 for event in events: 

237 if cancel.is_set(): 

238 events.close() 

239 break 

240 if isinstance(event, CapNotice): 

241 put( 

242 sse_event( 

243 SseEvent.REASONING, 

244 {"token": CAP_NOTICE_TEMPLATE.format(chars=event.cap_chars)}, 

245 ) 

246 ) 

247 elif event.is_reasoning: 

248 if event.content: 

249 put(sse_event(SseEvent.REASONING, {"token": event.content})) 

250 elif event.content: 

251 _put_answer_token(event.content, put, cite_filter, answer_parts) 

252 except Exception as exc: 

253 error_holder.append(exc) 

254 finally: 

255 _put_answer_tail(put, cite_filter, answer_parts) 

256 put(None) 

257 

258 

259async def _store_extracted_memories(question: str, answer: str) -> list[Any]: 

260 """Run the auto-extraction LLM pass off the event loop and return stored memories. 

261 

262 Returns an empty list (no-op) when the answer is empty or auto-extraction is 

263 off, so one-shot and streaming callers share one extraction path. 

264 """ 

265 if not answer or not auto_extract_enabled(): 

266 return [] 

267 return await asyncio.to_thread(auto_extract, question, answer) 

268 

269 

270async def _emit_extracted_memories(question: str, answer: str) -> AsyncGenerator[str, None]: 

271 """Yield a ``memory_extracted`` SSE event if the turn auto-saved any memories. 

272 

273 Silent (yields nothing) when the answer is empty, auto-extraction is off, or 

274 nothing was extracted, so existing consumers are unaffected. 

275 """ 

276 stored = await _store_extracted_memories(question, answer) 

277 if not stored: 

278 return 

279 event = MemoryExtractedEvent( 

280 count=len(stored), 

281 items=[MemoryExtractedItem(id=m.id, kind=m.kind, text=m.text) for m in stored], 

282 ) 

283 yield sse_event(SseEvent.MEMORY_EXTRACTED, event.model_dump(mode="json")) 

284 

285 

286def _mismatch_detail(exc: EmbeddingModelMismatchError) -> str | None: 

287 """The index's persisted embedder when dims match, so a client can offer to 

288 adopt it; None when they don't match and adoption wouldn't help.""" 

289 return exc.persisted_model if exc.dims_match else None 

290 

291 

292async def _emit_sources_and_memories( 

293 question: str, 

294 answer_parts: list[str], 

295 sources: list[SearchChunk], 

296) -> AsyncGenerator[str, None]: 

297 """Emit the trailing SOURCES event, ``done``, and any memory-extracted event. 

298 

299 SOURCES carries the cited subset (what the answer referenced), falling back to 

300 the full retrieved set when the answer cited nothing, mirroring 

301 ``Searcher.ask_stream``. Auto-extraction trails ``done`` so clients that stop 

302 at ``done`` are unaffected; the memories are stored regardless. 

303 """ 

304 answer = "".join(answer_parts) 

305 cited = cited_subset(answer, sources) 

306 source_list = cited if cited else sources 

307 yield sse_event(SseEvent.SOURCES, [clean_result(s) for s in source_list]) 

308 yield sse_done({}) 

309 async for event in _emit_extracted_memories(question, answer): 

310 yield event 

311 

312 

313async def _stream_rag_response( 

314 question: str, 

315 history: list[ChatMessage] | None = None, 

316 top_k: int = 0, 

317 options: dict[str, Any] | None = None, 

318 chunk_type: ChunkType | None = None, 

319) -> AsyncGenerator[str, None]: 

320 """SSE streaming for the ask (search) endpoint. 

321 

322 Mirrors ``Searcher.ask_stream`` so streaming, one-shot, and CLI ask agree: 

323 search mode with no embedder refuses cleanly, chat mode answers ungrounded, 

324 otherwise the answer is grounded in retrieved sources. 

325 """ 

326 yield "" # force generator 

327 

328 for warming in _chat_warming_events(): 

329 yield warming 

330 

331 searcher = get_services().searcher 

332 if searcher.search_unavailable(): 

333 # Search needs an embedder to ground. Mirror Searcher.ask_stream by 

334 # returning the refusal as a normal answer token (not an SSE error) so the 

335 # streaming, one-shot, and CLI ask paths all surface it the same way. 

336 yield sse_event(SseEvent.TOKEN, {"token": SEARCH_NEEDS_EMBEDDER}) 

337 yield sse_event(SseEvent.SOURCES, []) 

338 yield sse_done({}) 

339 return 

340 # Retrieval embeds, searches, reranks, and can spend an LLM call expanding 

341 # the query. On the loop it stalls every other admitted request for the whole 

342 # turn; the non-streaming siblings already thread the same work. 

343 results, messages, preempt = await asyncio.to_thread( 

344 _resolve_stream_context, 

345 searcher, 

346 question, 

347 history, 

348 top_k, 

349 chunk_type, 

350 retrieval_off=searcher.skip_retrieval(), 

351 ) 

352 for frame in preempt: 

353 yield frame 

354 if messages is None: 

355 return 

356 

357 opts = _resolve_generation_options(options) or cfg.generation_options() 

358 

359 sse = SseStream() 

360 error_holder: list[BaseException] = [] 

361 answer_parts: list[str] = [] 

362 # Only grounded turns append an authoritative SOURCES event, so only they 

363 # need a model-generated Sources block suppressed. 

364 cite_filter = StreamingCitationFilter() if results else None 

365 

366 executor_fut = sse.loop.run_in_executor( 

367 None, 

368 _run_llm_stream, 

369 messages, 

370 opts, 

371 sse.put_threadsafe, 

372 sse.cancel, 

373 error_holder, 

374 answer_parts, 

375 cite_filter, 

376 ) 

377 task = asyncio.ensure_future(executor_fut) 

378 async for event in sse.drain(task, "RAG stream"): 

379 yield event 

380 

381 if error_holder: 

382 exc = error_holder[0] 

383 raw = str(exc) 

384 code, user_message = _classify_stream_error(exc) 

385 log.warning("Stream error: %s", raw) 

386 yield sse_error(user_message, code=code, detail=raw if code else None) 

387 sse.cancel.set() 

388 return 

389 

390 # Ensure executor thread has finished before yielding final events 

391 await executor_fut 

392 

393 async for event in _emit_sources_and_memories(question, answer_parts, results): 

394 yield event 

395 

396 

397def ask_stream( 

398 question: str, 

399 top_k: int = 0, 

400 options: dict[str, Any] | None = None, 

401 chunk_type: ChunkType | None = None, 

402) -> AsyncGenerator[str, None]: 

403 """Yield SSE events: token, sources, done.""" 

404 return _stream_rag_response(question, top_k=top_k, options=options, chunk_type=chunk_type) 

405 

406 

407def _compaction_pending(history: list[ChatMessage], summary: str) -> bool: 

408 """Whether this turn will fold turns into notes before answering.""" 

409 budget = history_budget(cfg.chat_n_ctx_target) 

410 return bool( 

411 cfg.chat_compaction 

412 and compaction_due(history, summary, max_tokens=budget) 

413 and foldable(history) 

414 ) 

415 

416 

417def _manage_history( 

418 history: list[ChatMessage], 

419 summary: str, 

420 on_batch: Callable[[int, int], None] | None = None, 

421) -> tuple[list[ChatMessage], CompactionInfo | None]: 

422 """Apply the TUI's pre-turn context discipline to an HTTP conversation.""" 

423 budget = history_budget(cfg.chat_n_ctx_target) 

424 info: CompactionInfo | None = None 

425 if _compaction_pending(history, summary): 

426 dropped = foldable(history) 

427 result = get_services().searcher.summarize_history(dropped, summary, on_batch=on_batch) 

428 history = history[len(dropped) :] 

429 summary = result.summary 

430 info = CompactionInfo( 

431 summary=result.summary, condensed=result.condensed, stranded=result.stranded 

432 ) 

433 return prompt_history(history, summary, max_tokens=budget), info 

434 

435 

436async def _context_management_frames( 

437 history: list[ChatMessage], summary: str, session_id: str | None 

438) -> AsyncGenerator[str | tuple[list[ChatMessage], CompactionInfo | None], None]: 

439 """Manage this turn's history off-loop, yielding its SSE frames, then the result. 

440 

441 Yields the ``compacting`` announcement, per-batch progress frames, and the 

442 closing ``compaction`` frame (str items), persisting a fresh summary along 

443 the way; the ``(history, info)`` tuple arrives exactly once, last. 

444 """ 

445 if not _compaction_pending(history, summary): 

446 # Windowing only: pure arithmetic over the messages, no model call to wait on. 

447 yield await asyncio.to_thread(_manage_history, history, summary) 

448 return 

449 

450 # Condensing blocks this turn on model calls; announce it like warming, then 

451 # relay per-batch progress. SseStream carries the heartbeats and the 

452 # client-disconnect cancellation every other streaming endpoint here gets. 

453 yield sse_event(SseEvent.COMPACTING, {}) 

454 stream = SseStream() 

455 

456 def _on_batch(batch: int, total: int) -> None: 

457 stream.put_threadsafe(sse_event(SseEvent.COMPACTING, {"batch": batch, "batches": total})) 

458 

459 async def _condense() -> tuple[list[ChatMessage], CompactionInfo | None]: 

460 try: 

461 return await asyncio.to_thread(_manage_history, history, summary, _on_batch) 

462 finally: 

463 stream.put_threadsafe(None) 

464 

465 task = asyncio.ensure_future(_condense()) 

466 async for frame in stream.drain(task, "Compaction stream"): 

467 yield frame 

468 try: 

469 managed_history, compaction = await task 

470 except Exception: 

471 # Condensing is best-effort: a failed fold degrades to plain windowing, 

472 # which is what this turn would have done with compaction off. Losing the 

473 # summary costs context; failing the turn costs the user their answer. 

474 log.warning("Compaction failed; falling back to windowing", exc_info=True) 

475 budget = history_budget(cfg.chat_n_ctx_target) 

476 yield (prompt_history(history, summary, max_tokens=budget), None) 

477 return 

478 if compaction is not None: 

479 _persist_summary(session_id, compaction) 

480 yield sse_event(SseEvent.COMPACTION, compaction.model_dump()) 

481 yield (managed_history, compaction) 

482 

483 

484def _persist_summary(session_id: str | None, info: CompactionInfo | None) -> None: 

485 """Store fresh notes on the session; a session deleted mid-chat is tolerated.""" 

486 if info is None or not session_id or not info.summary or not sessions_enabled(): 

487 return 

488 with contextlib.suppress(SessionNotFoundError): 

489 get_services().session_store.set_summary(session_id, info.summary) 

490 

491 

492async def chat( 

493 question: str, 

494 history: list[ChatMessage], 

495 top_k: int | None = None, 

496 options: dict[str, Any] | None = None, 

497 chunk_type: ChunkType | None = None, 

498 summary: str = "", 

499 session_id: str | None = None, 

500) -> AskResponse: 

501 """Chat with history. Returns answer and sources via canonical dispatch.""" 

502 searcher = get_services().searcher 

503 if searcher.search_unavailable(): 

504 # Search mode with no embedder can't ground; refuse cleanly with the same 

505 # message ask returns instead of silently answering off-corpus. 

506 return AskResponse(answer=SEARCH_NEEDS_EMBEDDER, sources=[], cited_sources=[]) 

507 history, compaction = await asyncio.to_thread(_manage_history, history, summary) 

508 _persist_summary(session_id, compaction) 

509 retrieval_query: str | None = None 

510 if _retrieval_off(searcher, top_k): 

511 # Chat-only mode or an explicit top_k:0 pure-LLM call. 

512 sources: list[SearchChunk] = [] 

513 messages = searcher.direct_messages(question, history) 

514 else: 

515 # Grounded turn: the searcher's own pre-retrieval ladder (empty 

516 # library, count routing, memory-awareness) so surfaces cannot drift. 

517 pre_answer = searcher.pre_retrieval_answer(question) 

518 if pre_answer is not None: 

519 return AskResponse( 

520 answer=pre_answer, sources=[], cited_sources=[], compaction=compaction 

521 ) 

522 rag = searcher.build_rag_context( 

523 question, top_k=top_k or 0, history=history, chunk_type=chunk_type 

524 ) 

525 if rag is None: 

526 # Refuse like every sibling surface; the old fallback silently 

527 # answered off-corpus with nothing telling the caller so. 

528 return AskResponse( 

529 answer=GROUNDED_REFUSAL, sources=[], cited_sources=[], compaction=compaction 

530 ) 

531 sources, messages = rag.results, rag.messages 

532 retrieval_query = rag.retrieval_query 

533 req = _build_canonical_request(messages, options) 

534 response = await asyncio.to_thread(dispatch_chat, req) 

535 text = _join_text_blocks(response.content) 

536 answer = text if cfg.show_reasoning else strip_reasoning(text) 

537 if not answer.strip() and text.strip(): 

538 # The model emitted only reasoning (stripped to nothing) and no final 

539 # answer. Surface that distinctly instead of a silent empty string the 

540 # caller can't tell apart from a legitimate empty response (bb-cpu). The 

541 # synthetic notice is not an answer, so -- like the search-needs-embedder 

542 # refusal -- it doesn't seed memory. 

543 answer = REASONING_EXHAUSTED_NOTICE 

544 else: 

545 await _store_extracted_memories(question, answer) 

546 return AskResponse( 

547 answer=answer, 

548 sources=[CleanedChunk(**clean_result(s)) for s in sources], 

549 cited_sources=[ 

550 CleanedChunk(**clean_result(s)) 

551 for s in cited_subset(strip_llm_citations(answer), sources) 

552 ], 

553 compaction=compaction, 

554 retrieval_query=retrieval_query, 

555 ) 

556 

557 

558def chat_stream( 

559 question: str, 

560 history: list[ChatMessage], 

561 top_k: int | None = None, 

562 options: dict[str, Any] | None = None, 

563 chunk_type: ChunkType | None = None, 

564 summary: str = "", 

565 session_id: str | None = None, 

566) -> AsyncGenerator[str, None]: 

567 """Stream RAG chat tokens through canonical dispatch as token/sources/done events.""" 

568 return _stream_chat_response( 

569 question, 

570 history=history, 

571 top_k=top_k, 

572 options=options, 

573 chunk_type=chunk_type, 

574 summary=summary, 

575 session_id=session_id, 

576 ) 

577 

578 

579class _StreamResolution(NamedTuple): 

580 """Retrieval outcome for a streaming turn. 

581 

582 ``preempt_frames`` are emitted verbatim before anything else; ``messages`` 

583 of ``None`` means the stream ends after them (a direct exact-scan answer 

584 or a clean refusal/error), otherwise they lead the answer. 

585 """ 

586 

587 sources: list[SearchChunk] 

588 messages: list[ChatMessage] | None 

589 preempt_frames: list[str] 

590 

591 

592class _ChatStreamPlan(NamedTuple): 

593 """Leading SSE frames plus the grounded context for a chat stream. 

594 

595 A ``None`` context means the turn can't proceed: emit the frames (a clean 

596 refusal or error) and stop. 

597 """ 

598 

599 frames: list[str] 

600 context: RagContext | None 

601 

602 

603def _resolve_chat_stream_context( 

604 searcher: Searcher, 

605 question: str, 

606 history: list[ChatMessage], 

607 top_k: int | None, 

608 chunk_type: ChunkType | None, 

609) -> _ChatStreamPlan: 

610 frames = list(_chat_warming_events()) 

611 if searcher.search_unavailable(): 

612 # Search mode with no embedder can't ground; refuse cleanly with the same 

613 # token the ask stream emits instead of silently answering off-corpus. 

614 frames += [ 

615 sse_event(SseEvent.TOKEN, {"token": SEARCH_NEEDS_EMBEDDER}), 

616 sse_event(SseEvent.SOURCES, []), 

617 sse_done({}), 

618 ] 

619 return _ChatStreamPlan(frames, None) 

620 # Retrieval itself is resolved by the shared helper, so the chat stream 

621 # routes empty libraries and count questions exactly like the ask stream. 

622 sources, messages, preempt = _resolve_stream_context( 

623 searcher, 

624 question, 

625 history, 

626 top_k, 

627 chunk_type, 

628 retrieval_off=_retrieval_off(searcher, top_k), 

629 ) 

630 frames += preempt 

631 if messages is None: 

632 return _ChatStreamPlan(frames, None) 

633 return _ChatStreamPlan(frames, RagContext(sources, messages)) 

634 

635 

636async def _stream_chat_response( 

637 question: str, 

638 history: list[ChatMessage], 

639 top_k: int | None, 

640 options: dict[str, Any] | None, 

641 chunk_type: ChunkType | None, 

642 summary: str = "", 

643 session_id: str | None = None, 

644) -> AsyncGenerator[str, None]: 

645 """Drive ``dispatch_chat_stream`` and emit reasoning/token/sources/done SSE events.""" 

646 async for item in _context_management_frames(history, summary, session_id): 

647 if isinstance(item, str): 

648 yield item 

649 continue 

650 history, _compaction = item 

651 frames, ctx = await asyncio.to_thread( 

652 _resolve_chat_stream_context, get_services().searcher, question, history, top_k, chunk_type 

653 ) 

654 for frame in frames: 

655 yield frame 

656 if ctx is None: 

657 return 

658 sources, messages = ctx.results, ctx.messages 

659 

660 req = _build_canonical_request(messages, options) 

661 answer_parts: list[str] = [] 

662 # Only grounded turns append an authoritative SOURCES event, so only they 

663 # need a model-generated Sources block suppressed. 

664 cite_filter = StreamingCitationFilter() if sources else None 

665 try: 

666 async for event in _cap_aware_chat_events(req): 

667 frame = _chat_answer_frame(event, cite_filter, answer_parts) 

668 if frame: 

669 yield frame 

670 except Exception as exc: 

671 raw = str(exc) 

672 code, user_message = _classify_stream_error(exc) 

673 log.warning("Stream error: %s", raw) 

674 yield sse_error(user_message, code=code, detail=raw if code else None) 

675 return 

676 

677 tail_frame = _chat_answer_tail_frame(cite_filter, answer_parts) 

678 if tail_frame: 

679 yield tail_frame 

680 

681 async for frame in _emit_sources_and_memories(question, answer_parts, sources): 

682 yield frame 

683 

684 

685async def _cap_aware_chat_events( 

686 req: CanonicalChatRequest, 

687) -> AsyncIterator[StreamToken | CapNotice]: 

688 """Run ``dispatch_chat_stream``, split reasoning, and re-issue on cap-fire. 

689 

690 Mirrors :func:`stream_chat_with_cap` but consumes the canonical async 

691 stream. ``CapNotice`` is yielded once between the truncated reasoning 

692 and the continuation answer; ``StreamToken`` carries the 

693 reasoning-vs-response split for downstream SSE shaping. When reasoning 

694 runs but no final answer follows, a closing ``StreamToken`` carrying 

695 ``REASONING_EXHAUSTED_NOTICE`` is yielded so the run isn't silent (bb-cpu). 

696 """ 

697 cap_chars = effective_reasoning_cap() 

698 show = cfg.show_reasoning 

699 answered = False 

700 

701 first_parser = TagParser(show=show) 

702 async for tok in _drive_stream(dispatch_chat_stream(req), first_parser, cap_chars): 

703 answered = answered or (not tok.is_reasoning and bool(tok.content)) 

704 yield tok 

705 

706 if cap_chars > 0 and first_parser.reasoning_chars > cap_chars: 

707 yield CapNotice(cap_chars=cap_chars) 

708 nudged = nudged_request(req) 

709 cont_parser = TagParser(show=show) 

710 async for tok in _drive_stream(dispatch_chat_stream(nudged), cont_parser, cap_chars=0): 

711 answered = answered or bool(tok.content) 

712 # Continuation tokens are always treated as final-answer text. 

713 yield StreamToken(content=tok.content, is_reasoning=False) 

714 

715 if first_parser.reasoning_chars > 0 and not answered: 

716 # The model spent its budget reasoning and produced no final answer; 

717 # a distinct notice tells a reasoning-only run apart from a completed one. 

718 yield StreamToken(content=REASONING_EXHAUSTED_NOTICE, is_reasoning=False) 

719 

720 

721async def _drive_stream( 

722 stream: AsyncIterator[Any], 

723 parser: TagParser, 

724 cap_chars: int, 

725) -> AsyncIterator[StreamToken]: 

726 """Feed *stream* through *parser*; yield ``StreamToken``s; stop on cap-fire.""" 

727 cap_fired = False 

728 try: 

729 async for event in stream: 

730 text = _text_from_event(event) 

731 if not text: 

732 continue 

733 for tok in parser.feed(text): 

734 if tok.content: 

735 yield tok 

736 if cap_chars > 0 and parser.reasoning_chars > cap_chars: 

737 cap_fired = True 

738 break 

739 finally: 

740 if cap_fired: 

741 await _aclose(stream) 

742 tail = parser.flush() 

743 if tail is not None and tail.content: 

744 yield tail 

745 

746 

747async def _aclose(stream: AsyncIterator[Any]) -> None: 

748 """Best-effort close for async-generator-shaped streams.""" 

749 if not isinstance(stream, AsyncGenerator): 

750 return 

751 with contextlib.suppress(Exception): 

752 await stream.aclose() 

753 

754 

755def _sse_for_chat_event(event: StreamToken | CapNotice) -> str: 

756 """Render one orchestrator event as an SSE frame with the right channel.""" 

757 if isinstance(event, CapNotice): 

758 return sse_event( 

759 SseEvent.REASONING, 

760 {"token": CAP_NOTICE_TEMPLATE.format(chars=event.cap_chars)}, 

761 ) 

762 kind = SseEvent.REASONING if event.is_reasoning else SseEvent.TOKEN 

763 return sse_event(kind, {"token": event.content}) 

764 

765 

766def _chat_answer_frame( 

767 event: StreamToken | CapNotice, 

768 cite_filter: StreamingCitationFilter | None, 

769 answer_parts: list[str], 

770) -> str: 

771 """Render the SSE frame for one chat event and record answer text, dropping a 

772 model Sources block on grounded turns. Returns '' when nothing should emit. 

773 

774 The reasoning-exhausted notice streams to the client but is not a real 

775 answer, so it is left out of *answer_parts*: it seeds no memory and is not 

776 treated as a citation source. 

777 """ 

778 is_answer = ( 

779 isinstance(event, StreamToken) 

780 and not event.is_reasoning 

781 and event.content != REASONING_EXHAUSTED_NOTICE 

782 ) 

783 if not is_answer: 

784 return _sse_for_chat_event(event) 

785 content = cast("StreamToken", event).content 

786 if cite_filter is None: 

787 answer_parts.append(content) 

788 return _sse_for_chat_event(event) 

789 shown = cite_filter.feed(content) 

790 if not shown: 

791 return "" 

792 answer_parts.append(shown) 

793 return sse_event(SseEvent.TOKEN, {"token": shown}) 

794 

795 

796def _chat_answer_tail_frame( 

797 cite_filter: StreamingCitationFilter | None, 

798 answer_parts: list[str], 

799) -> str: 

800 """SSE frame releasing any answer text the filter held back, or '' if none.""" 

801 if cite_filter is None: 

802 return "" 

803 tail = cite_filter.flush() 

804 if not tail: 

805 return "" 

806 answer_parts.append(tail) 

807 return sse_event(SseEvent.TOKEN, {"token": tail}) 

808 

809 

810def _text_from_event(event: Any) -> str: 

811 """Return the text payload of a canonical event, or '' if not a text delta.""" 

812 if isinstance(event, ContentBlockDelta) and isinstance(event.delta, TextDelta):