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 dropped_sources=[CleanedChunk(**clean_result(s)) for s in result.dropped_sources], 

168 ) 

169 

170 

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

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

173 

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

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

176 """ 

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

178 return [] 

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

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

181 

182 

183def _put_answer_token( 

184 content: str, 

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

186 cite_filter: StreamingCitationFilter | None, 

187 answer_parts: list[str], 

188) -> None: 

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

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

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

192 if token: 

193 answer_parts.append(token) 

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

195 

196 

197def _put_answer_tail( 

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

199 cite_filter: StreamingCitationFilter | None, 

200 answer_parts: list[str], 

201) -> None: 

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

203 if cite_filter is None: 

204 return 

205 tail = cite_filter.flush() 

206 if tail: 

207 answer_parts.append(tail) 

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

209 

210 

211def _run_llm_stream( 

212 messages: list[ChatMessage], 

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

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

215 cancel: threading.Event, 

216 error_holder: list[BaseException], 

217 answer_parts: list[str], 

218 cite_filter: StreamingCitationFilter | None, 

219) -> None: 

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

221 

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

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

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

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

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

227 """ 

228 try: 

229 events = stream_chat_with_cap( 

230 get_services().provider, 

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

232 options=opts, 

233 model=cfg.chat_model, 

234 show_reasoning=cfg.show_reasoning, 

235 cap_chars=effective_reasoning_cap(), 

236 ) 

237 for event in events: 

238 if cancel.is_set(): 

239 events.close() 

240 break 

241 if isinstance(event, CapNotice): 

242 put( 

243 sse_event( 

244 SseEvent.REASONING, 

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

246 ) 

247 ) 

248 elif event.is_reasoning: 

249 if event.content: 

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

251 elif event.content: 

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

253 except Exception as exc: 

254 error_holder.append(exc) 

255 finally: 

256 _put_answer_tail(put, cite_filter, answer_parts) 

257 put(None) 

258 

259 

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

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

262 

263 The pass is best effort and never raises: an empty list comes back when the 

264 answer is empty, when auto-extraction is off, and when the pass itself fails. 

265 A failure is logged with its traceback. 

266 """ 

267 if not answer or not auto_extract_enabled(): 

268 return [] 

269 try: 

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

271 except Exception: 

272 log.exception("auto-extraction failed; the answer is unaffected") 

273 return [] 

274 

275 

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

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

278 

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

280 nothing was extracted, so existing consumers are unaffected. 

281 """ 

282 stored = await _store_extracted_memories(question, answer) 

283 if not stored: 

284 return 

285 event = MemoryExtractedEvent( 

286 count=len(stored), 

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

288 ) 

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

290 

291 

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

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

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

295 return exc.persisted_model if exc.dims_match else None 

296 

297 

298async def _emit_sources_and_memories( 

299 question: str, 

300 answer_parts: list[str], 

301 sources: list[SearchChunk], 

302) -> AsyncGenerator[str, None]: 

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

304 

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

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

307 ``Searcher.ask_stream``. Recalled memories always ride along: they were in the 

308 prompt whether the answer cited a document or not. Auto-extraction trails 

309 ``done``; it cannot fail the stream, so a client that drains past ``done`` 

310 still keeps the answer. 

311 """ 

312 answer = "".join(answer_parts) 

313 cited = cited_subset(answer, sources) 

314 if cited: 

315 cited_sources = {c.source for c in cited} 

316 memories = [s for s in sources if s.memory_id is not None and s.source not in cited_sources] 

317 source_list = [*cited, *memories] 

318 else: 

319 source_list = sources 

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

321 yield sse_done({}) 

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

323 yield event 

324 

325 

326async def _stream_rag_response( 

327 question: str, 

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

329 top_k: int = 0, 

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

331 chunk_type: ChunkType | None = None, 

332) -> AsyncGenerator[str, None]: 

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

334 

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

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

337 otherwise the answer is grounded in retrieved sources. 

338 """ 

339 yield "" # force generator 

340 

341 for warming in _chat_warming_events(): 

342 yield warming 

343 

344 searcher = get_services().searcher 

345 if searcher.search_unavailable(): 

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

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

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

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

350 yield sse_event(SseEvent.SOURCES, []) 

351 yield sse_done({}) 

352 return 

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

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

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

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

357 _resolve_stream_context, 

358 searcher, 

359 question, 

360 history, 

361 top_k, 

362 chunk_type, 

363 retrieval_off=searcher.skip_retrieval(), 

364 ) 

365 for frame in preempt: 

366 yield frame 

367 if messages is None: 

368 return 

369 

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

371 

372 sse = SseStream() 

373 error_holder: list[BaseException] = [] 

374 answer_parts: list[str] = [] 

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

376 # need a model-generated Sources block suppressed. 

377 cite_filter = StreamingCitationFilter() if results else None 

378 

379 executor_fut = sse.loop.run_in_executor( 

380 None, 

381 _run_llm_stream, 

382 messages, 

383 opts, 

384 sse.put_threadsafe, 

385 sse.cancel, 

386 error_holder, 

387 answer_parts, 

388 cite_filter, 

389 ) 

390 task = asyncio.ensure_future(executor_fut) 

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

392 yield event 

393 

394 if error_holder: 

395 exc = error_holder[0] 

396 raw = str(exc) 

397 code, user_message = _classify_stream_error(exc) 

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

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

400 sse.cancel.set() 

401 return 

402 

403 # Ensure executor thread has finished before yielding final events 

404 await executor_fut 

405 

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

407 yield event 

408 

409 

410def ask_stream( 

411 question: str, 

412 top_k: int = 0, 

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

414 chunk_type: ChunkType | None = None, 

415) -> AsyncGenerator[str, None]: 

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

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

418 

419 

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

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

422 budget = history_budget(cfg.chat_n_ctx_target) 

423 return bool( 

424 cfg.chat_compaction 

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

426 and foldable(history) 

427 ) 

428 

429 

430def _manage_history( 

431 history: list[ChatMessage], 

432 summary: str, 

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

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

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

436 budget = history_budget(cfg.chat_n_ctx_target) 

437 info: CompactionInfo | None = None 

438 if _compaction_pending(history, summary): 

439 dropped = foldable(history) 

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

441 history = history[len(dropped) :] 

442 summary = result.summary 

443 info = CompactionInfo( 

444 summary=result.summary, condensed=result.condensed, stranded=result.stranded 

445 ) 

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

447 

448 

449async def _context_management_frames( 

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

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

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

453 

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

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

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

457 """ 

458 if not _compaction_pending(history, summary): 

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

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

461 return 

462 

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

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

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

466 yield sse_event(SseEvent.COMPACTING, {}) 

467 stream = SseStream() 

468 

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

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

471 

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

473 try: 

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

475 finally: 

476 stream.put_threadsafe(None) 

477 

478 task = asyncio.ensure_future(_condense()) 

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

480 yield frame 

481 try: 

482 managed_history, compaction = await task 

483 except Exception: 

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

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

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

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

488 budget = history_budget(cfg.chat_n_ctx_target) 

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

490 return 

491 if compaction is not None: 

492 _persist_summary(session_id, compaction) 

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

494 yield (managed_history, compaction) 

495 

496 

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

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

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

500 return 

501 with contextlib.suppress(SessionNotFoundError): 

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

503 

504 

505async def chat( 

506 question: str, 

507 history: list[ChatMessage], 

508 top_k: int | None = None, 

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

510 chunk_type: ChunkType | None = None, 

511 summary: str = "", 

512 session_id: str | None = None, 

513) -> AskResponse: 

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

515 searcher = get_services().searcher 

516 if searcher.search_unavailable(): 

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

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

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

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

521 _persist_summary(session_id, compaction) 

522 retrieval_query: str | None = None 

523 dropped: list[SearchChunk] = [] 

524 if _retrieval_off(searcher, top_k): 

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

526 sources: list[SearchChunk] = [] 

527 messages = searcher.direct_messages(question, history) 

528 else: 

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

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

531 pre_answer = searcher.pre_retrieval_answer(question) 

532 if pre_answer is not None: 

533 return AskResponse( 

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

535 ) 

536 rag = searcher.build_rag_context( 

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

538 ) 

539 if rag is None: 

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

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

542 return AskResponse( 

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

544 ) 

545 sources, messages = rag.results, rag.messages 

546 retrieval_query = rag.retrieval_query 

547 dropped = rag.dropped or [] 

548 req = _build_canonical_request(messages, options) 

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

550 text = _join_text_blocks(response.content) 

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

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

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

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

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

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

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

558 answer = REASONING_EXHAUSTED_NOTICE 

559 else: 

560 await _store_extracted_memories(question, answer) 

561 return AskResponse( 

562 answer=answer, 

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

564 cited_sources=[ 

565 CleanedChunk(**clean_result(s)) 

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

567 ], 

568 compaction=compaction, 

569 retrieval_query=retrieval_query, 

570 dropped_sources=[CleanedChunk(**clean_result(s)) for s in dropped], 

571 ) 

572 

573 

574def chat_stream( 

575 question: str, 

576 history: list[ChatMessage], 

577 top_k: int | None = None, 

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

579 chunk_type: ChunkType | None = None, 

580 summary: str = "", 

581 session_id: str | None = None, 

582) -> AsyncGenerator[str, None]: 

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

584 return _stream_chat_response( 

585 question, 

586 history=history, 

587 top_k=top_k, 

588 options=options, 

589 chunk_type=chunk_type, 

590 summary=summary, 

591 session_id=session_id, 

592 ) 

593 

594 

595class _StreamResolution(NamedTuple): 

596 """Retrieval outcome for a streaming turn. 

597 

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

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

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

601 """ 

602 

603 sources: list[SearchChunk] 

604 messages: list[ChatMessage] | None 

605 preempt_frames: list[str] 

606 

607 

608class _ChatStreamPlan(NamedTuple): 

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

610 

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

612 refusal or error) and stop. 

613 """ 

614 

615 frames: list[str] 

616 context: RagContext | None 

617 

618 

619def _resolve_chat_stream_context( 

620 searcher: Searcher, 

621 question: str, 

622 history: list[ChatMessage], 

623 top_k: int | None, 

624 chunk_type: ChunkType | None, 

625) -> _ChatStreamPlan: 

626 frames = list(_chat_warming_events()) 

627 if searcher.search_unavailable(): 

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

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

630 frames += [ 

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

632 sse_event(SseEvent.SOURCES, []), 

633 sse_done({}), 

634 ] 

635 return _ChatStreamPlan(frames, None) 

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

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

638 sources, messages, preempt = _resolve_stream_context( 

639 searcher, 

640 question, 

641 history, 

642 top_k, 

643 chunk_type, 

644 retrieval_off=_retrieval_off(searcher, top_k), 

645 ) 

646 frames += preempt 

647 if messages is None: 

648 return _ChatStreamPlan(frames, None) 

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

650 

651 

652async def _stream_chat_response( 

653 question: str, 

654 history: list[ChatMessage], 

655 top_k: int | None, 

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

657 chunk_type: ChunkType | None, 

658 summary: str = "", 

659 session_id: str | None = None, 

660) -> AsyncGenerator[str, None]: 

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

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

663 if isinstance(item, str): 

664 yield item 

665 continue 

666 history, _compaction = item 

667 frames, ctx = await asyncio.to_thread( 

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

669 ) 

670 for frame in frames: 

671 yield frame 

672 if ctx is None: 

673 return 

674 sources, messages = ctx.results, ctx.messages 

675 

676 req = _build_canonical_request(messages, options) 

677 answer_parts: list[str] = [] 

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

679 # need a model-generated Sources block suppressed. 

680 cite_filter = StreamingCitationFilter() if sources else None 

681 try: 

682 async for event in _cap_aware_chat_events(req): 

683 frame = _chat_answer_frame(event, cite_filter, answer_parts) 

684 if frame: 

685 yield frame 

686 except Exception as exc: 

687 raw = str(exc) 

688 code, user_message = _classify_stream_error(exc) 

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

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

691 return 

692 

693 tail_frame = _chat_answer_tail_frame(cite_filter, answer_parts) 

694 if tail_frame: 

695 yield tail_frame 

696 

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

698 yield frame 

699 

700 

701async def _cap_aware_chat_events( 

702 req: CanonicalChatRequest, 

703) -> AsyncIterator[StreamToken | CapNotice]: 

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

705 

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

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

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

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

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

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

712 """ 

713 cap_chars = effective_reasoning_cap() 

714 show = cfg.show_reasoning 

715 answered = False 

716 

717 first_parser = TagParser(show=show) 

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

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

720 yield tok 

721 

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

723 yield CapNotice(cap_chars=cap_chars) 

724 nudged = nudged_request(req) 

725 cont_parser = TagParser(show=show) 

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

727 answered = answered or bool(tok.content) 

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

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

730 

731 if first_parser.reasoning_chars > 0 and not answered: 

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

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

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

735 

736 

737async def _drive_stream( 

738 stream: AsyncIterator[Any], 

739 parser: TagParser, 

740 cap_chars: int, 

741) -> AsyncIterator[StreamToken]: 

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

743 cap_fired = False 

744 try: 

745 async for event in stream: 

746 text = _text_from_event(event) 

747 if not text: 

748 continue 

749 for tok in parser.feed(text): 

750 if tok.content: 

751 yield tok 

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

753 cap_fired = True 

754 break 

755 finally: 

756 if cap_fired: 

757 await _aclose(stream) 

758 tail = parser.flush() 

759 if tail is not None and tail.content: 

760 yield tail 

761 

762 

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

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

765 if not isinstance(stream, AsyncGenerator): 

766 return 

767 with contextlib.suppress(Exception): 

768 await stream.aclose() 

769 

770 

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

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

773 if isinstance(event, CapNotice): 

774 return sse_event( 

775 SseEvent.REASONING, 

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

777 ) 

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

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

780 

781 

782def _chat_answer_frame( 

783 event: StreamToken | CapNotice, 

784 cite_filter: StreamingCitationFilter | None, 

785 answer_parts: list[str], 

786) -> str: 

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

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

789 

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

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

792 treated as a citation source. 

793 """ 

794 is_answer = ( 

795 isinstance(event, StreamToken) 

796 and not event.is_reasoning 

797 and event.content != REASONING_EXHAUSTED_NOTICE 

798 ) 

799 if not is_answer: 

800 return _sse_for_chat_event(event) 

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

802 if cite_filter is None: 

803 answer_parts.append(content) 

804 return _sse_for_chat_event(event) 

805 shown = cite_filter.feed(content) 

806 if not shown: 

807 return "" 

808 answer_parts.append(shown) 

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