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_results, source_labels 

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 | None = None, chunk_type: ChunkType | None = None 

126) -> list[DocumentResult]: 

127 """Search and return grouped DocumentResults. 

128 

129 An omitted ``top_k`` resolves to ``cfg.top_k``, read live at call time. 

130 """ 

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

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

133 resolved_top_k = top_k if top_k is not None else cfg.top_k 

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

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

136 results = await asyncio.to_thread( 

137 get_services().searcher.search, q, top_k=resolved_top_k, chunk_type=chunk_type 

138 ) 

139 labels = source_labels(results) 

140 return [r.model_copy(update={"label": labels[r.source]}) for r in group(results)] 

141 

142 

143async def ask( 

144 question: str, 

145 top_k: int = 0, 

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

147 chunk_type: ChunkType | None = None, 

148) -> AskResponse: 

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

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

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

152 opts = _resolve_generation_options(options) 

153 searcher = get_services().searcher 

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

155 # event loop so other admitted requests stay responsive. 

156 result = await asyncio.to_thread( 

157 searcher.ask_raw, 

158 question, 

159 top_k=top_k, 

160 options=opts, 

161 chunk_type=chunk_type, 

162 ) 

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

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

165 if not searcher.search_unavailable(): 

166 await _store_extracted_memories(question, result.answer) 

167 return AskResponse( 

168 answer=result.answer, 

169 sources=_cleaned_chunks(result.sources, result.sources), 

170 cited_sources=_cleaned_chunks(result.cited_sources, result.sources), 

171 retrieval_query=result.retrieval_query, 

172 dropped_sources=_cleaned_chunks(result.dropped_sources, result.sources), 

173 ) 

174 

175 

176def _cleaned_chunks( 

177 chunks: list[SearchChunk], answered_from: list[SearchChunk] | None = None 

178) -> list[CleanedChunk]: 

179 """*chunks* as response models, numbered by the retrieved list *answered_from*.""" 

180 return [CleanedChunk(**c) for c in clean_results(chunks, answered_from or [])] 

181 

182 

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

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

185 

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

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

188 """ 

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

190 return [] 

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

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

193 

194 

195def _put_answer_token( 

196 content: str, 

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

198 cite_filter: StreamingCitationFilter | None, 

199 answer_parts: list[str], 

200) -> None: 

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

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

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

204 if token: 

205 answer_parts.append(token) 

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

207 

208 

209def _put_answer_tail( 

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

211 cite_filter: StreamingCitationFilter | None, 

212 answer_parts: list[str], 

213) -> None: 

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

215 if cite_filter is None: 

216 return 

217 tail = cite_filter.flush() 

218 if tail: 

219 answer_parts.append(tail) 

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

221 

222 

223def _run_llm_stream( 

224 messages: list[ChatMessage], 

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

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

227 cancel: threading.Event, 

228 error_holder: list[BaseException], 

229 answer_parts: list[str], 

230 cite_filter: StreamingCitationFilter | None, 

231) -> None: 

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

233 

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

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

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

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

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

239 """ 

240 try: 

241 events = stream_chat_with_cap( 

242 get_services().provider, 

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

244 options=opts, 

245 model=cfg.chat_model, 

246 show_reasoning=cfg.show_reasoning, 

247 cap_chars=effective_reasoning_cap(), 

248 ) 

249 for event in events: 

250 if cancel.is_set(): 

251 events.close() 

252 break 

253 if isinstance(event, CapNotice): 

254 put( 

255 sse_event( 

256 SseEvent.REASONING, 

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

258 ) 

259 ) 

260 elif event.is_reasoning: 

261 if event.content: 

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

263 elif event.content: 

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

265 except Exception as exc: 

266 error_holder.append(exc) 

267 finally: 

268 _put_answer_tail(put, cite_filter, answer_parts) 

269 put(None) 

270 

271 

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

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

274 

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

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

277 A failure is logged with its traceback. 

278 """ 

279 if not answer or not auto_extract_enabled(): 

280 return [] 

281 try: 

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

283 except Exception: 

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

285 return [] 

286 

287 

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

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

290 

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

292 nothing was extracted, so existing consumers are unaffected. 

293 """ 

294 stored = await _store_extracted_memories(question, answer) 

295 if not stored: 

296 return 

297 event = MemoryExtractedEvent( 

298 count=len(stored), 

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

300 ) 

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

302 

303 

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

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

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

307 return exc.persisted_model if exc.dims_match else None 

308 

309 

310async def _emit_sources_and_memories( 

311 question: str, 

312 answer_parts: list[str], 

313 sources: list[SearchChunk], 

314) -> AsyncGenerator[str, None]: 

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

316 

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

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

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

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

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

322 still keeps the answer. 

323 """ 

324 answer = "".join(answer_parts) 

325 cited = cited_subset(answer, sources) 

326 if cited: 

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

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

329 source_list = [*cited, *memories] 

330 else: 

331 source_list = sources 

332 yield sse_event(SseEvent.SOURCES, clean_results(source_list, sources)) 

333 yield sse_done({}) 

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

335 yield event 

336 

337 

338async def _stream_rag_response( 

339 question: str, 

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

341 top_k: int = 0, 

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

343 chunk_type: ChunkType | None = None, 

344) -> AsyncGenerator[str, None]: 

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

346 

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

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

349 otherwise the answer is grounded in retrieved sources. 

350 """ 

351 yield "" # force generator 

352 

353 for warming in _chat_warming_events(): 

354 yield warming 

355 

356 searcher = get_services().searcher 

357 if searcher.search_unavailable(): 

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

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

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

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

362 yield sse_event(SseEvent.SOURCES, []) 

363 yield sse_done({}) 

364 return 

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

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

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

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

369 _resolve_stream_context, 

370 searcher, 

371 question, 

372 history, 

373 top_k, 

374 chunk_type, 

375 retrieval_off=searcher.skip_retrieval(), 

376 ) 

377 for frame in preempt: 

378 yield frame 

379 if messages is None: 

380 return 

381 

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

383 

384 sse = SseStream() 

385 error_holder: list[BaseException] = [] 

386 answer_parts: list[str] = [] 

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

388 # need a model-generated Sources block suppressed. 

389 cite_filter = StreamingCitationFilter() if results else None 

390 

391 executor_fut = sse.loop.run_in_executor( 

392 None, 

393 _run_llm_stream, 

394 messages, 

395 opts, 

396 sse.put_threadsafe, 

397 sse.cancel, 

398 error_holder, 

399 answer_parts, 

400 cite_filter, 

401 ) 

402 task = asyncio.ensure_future(executor_fut) 

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

404 yield event 

405 

406 if error_holder: 

407 exc = error_holder[0] 

408 raw = str(exc) 

409 code, user_message = _classify_stream_error(exc) 

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

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

412 sse.cancel.set() 

413 return 

414 

415 # Ensure executor thread has finished before yielding final events 

416 await executor_fut 

417 

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

419 yield event 

420 

421 

422def ask_stream( 

423 question: str, 

424 top_k: int = 0, 

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

426 chunk_type: ChunkType | None = None, 

427) -> AsyncGenerator[str, None]: 

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

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

430 

431 

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

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

434 budget = history_budget(cfg.chat_n_ctx_target) 

435 return bool( 

436 cfg.chat_compaction 

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

438 and foldable(history) 

439 ) 

440 

441 

442def _manage_history( 

443 history: list[ChatMessage], 

444 summary: str, 

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

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

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

448 budget = history_budget(cfg.chat_n_ctx_target) 

449 info: CompactionInfo | None = None 

450 if _compaction_pending(history, summary): 

451 dropped = foldable(history) 

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

453 history = history[len(dropped) :] 

454 summary = result.summary 

455 info = CompactionInfo( 

456 summary=result.summary, condensed=result.condensed, stranded=result.stranded 

457 ) 

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

459 

460 

461async def _context_management_frames( 

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

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

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

465 

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

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

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

469 """ 

470 if not _compaction_pending(history, summary): 

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

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

473 return 

474 

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

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

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

478 yield sse_event(SseEvent.COMPACTING, {}) 

479 stream = SseStream() 

480 

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

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

483 

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

485 try: 

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

487 finally: 

488 stream.put_threadsafe(None) 

489 

490 task = asyncio.ensure_future(_condense()) 

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

492 yield frame 

493 try: 

494 managed_history, compaction = await task 

495 except Exception: 

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

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

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

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

500 budget = history_budget(cfg.chat_n_ctx_target) 

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

502 return 

503 if compaction is not None: 

504 _persist_summary(session_id, compaction) 

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

506 yield (managed_history, compaction) 

507 

508 

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

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

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

512 return 

513 with contextlib.suppress(SessionNotFoundError): 

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

515 

516 

517async def chat( 

518 question: str, 

519 history: list[ChatMessage], 

520 top_k: int | None = None, 

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

522 chunk_type: ChunkType | None = None, 

523 summary: str = "", 

524 session_id: str | None = None, 

525) -> AskResponse: 

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

527 searcher = get_services().searcher 

528 if searcher.search_unavailable(): 

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

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

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

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

533 _persist_summary(session_id, compaction) 

534 retrieval_query: str | None = None 

535 dropped: list[SearchChunk] = [] 

536 if _retrieval_off(searcher, top_k): 

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

538 sources: list[SearchChunk] = [] 

539 messages = searcher.direct_messages(question, history) 

540 else: 

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

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

543 pre_answer = searcher.pre_retrieval_answer(question) 

544 if pre_answer is not None: 

545 return AskResponse( 

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

547 ) 

548 rag = searcher.build_rag_context( 

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

550 ) 

551 if rag is None: 

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

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

554 return AskResponse( 

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

556 ) 

557 sources, messages = rag.results, rag.messages 

558 retrieval_query = rag.retrieval_query 

559 dropped = rag.dropped or [] 

560 req = _build_canonical_request(messages, options) 

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

562 text = _join_text_blocks(response.content) 

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

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

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

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

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

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

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

570 answer = REASONING_EXHAUSTED_NOTICE 

571 else: 

572 await _store_extracted_memories(question, answer) 

573 return AskResponse( 

574 answer=answer, 

575 sources=_cleaned_chunks(sources, sources), 

576 cited_sources=_cleaned_chunks(cited_subset(strip_llm_citations(answer), sources), sources), 

577 compaction=compaction, 

578 retrieval_query=retrieval_query, 

579 dropped_sources=_cleaned_chunks(dropped, sources), 

580 ) 

581 

582 

583def chat_stream( 

584 question: str, 

585 history: list[ChatMessage], 

586 top_k: int | None = None, 

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

588 chunk_type: ChunkType | None = None, 

589 summary: str = "", 

590 session_id: str | None = None, 

591) -> AsyncGenerator[str, None]: 

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

593 return _stream_chat_response( 

594 question, 

595 history=history, 

596 top_k=top_k, 

597 options=options, 

598 chunk_type=chunk_type, 

599 summary=summary, 

600 session_id=session_id, 

601 ) 

602 

603 

604class _StreamResolution(NamedTuple): 

605 """Retrieval outcome for a streaming turn. 

606 

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

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

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

610 """ 

611 

612 sources: list[SearchChunk] 

613 messages: list[ChatMessage] | None 

614 preempt_frames: list[str] 

615 

616 

617class _ChatStreamPlan(NamedTuple): 

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

619 

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

621 refusal or error) and stop. 

622 """ 

623 

624 frames: list[str] 

625 context: RagContext | None 

626 

627 

628def _resolve_chat_stream_context( 

629 searcher: Searcher, 

630 question: str, 

631 history: list[ChatMessage], 

632 top_k: int | None, 

633 chunk_type: ChunkType | None, 

634) -> _ChatStreamPlan: 

635 frames = list(_chat_warming_events()) 

636 if searcher.search_unavailable(): 

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

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

639 frames += [ 

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

641 sse_event(SseEvent.SOURCES, []), 

642 sse_done({}), 

643 ] 

644 return _ChatStreamPlan(frames, None) 

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

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

647 sources, messages, preempt = _resolve_stream_context( 

648 searcher, 

649 question, 

650 history, 

651 top_k, 

652 chunk_type, 

653 retrieval_off=_retrieval_off(searcher, top_k), 

654 ) 

655 frames += preempt 

656 if messages is None: 

657 return _ChatStreamPlan(frames, None) 

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

659 

660 

661async def _stream_chat_response( 

662 question: str, 

663 history: list[ChatMessage], 

664 top_k: int | None, 

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

666 chunk_type: ChunkType | None, 

667 summary: str = "", 

668 session_id: str | None = None, 

669) -> AsyncGenerator[str, None]: 

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

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

672 if isinstance(item, str): 

673 yield item 

674 continue 

675 history, _compaction = item 

676 frames, ctx = await asyncio.to_thread( 

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

678 ) 

679 for frame in frames: 

680 yield frame 

681 if ctx is None: 

682 return 

683 sources, messages = ctx.results, ctx.messages 

684 

685 req = _build_canonical_request(messages, options) 

686 answer_parts: list[str] = [] 

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

688 # need a model-generated Sources block suppressed. 

689 cite_filter = StreamingCitationFilter() if sources else None 

690 try: 

691 async for event in _cap_aware_chat_events(req): 

692 frame = _chat_answer_frame(event, cite_filter, answer_parts) 

693 if frame: 

694 yield frame 

695 except Exception as exc: 

696 raw = str(exc) 

697 code, user_message = _classify_stream_error(exc) 

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

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

700 return 

701 

702 tail_frame = _chat_answer_tail_frame(cite_filter, answer_parts) 

703 if tail_frame: 

704 yield tail_frame 

705 

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

707 yield frame 

708 

709 

710async def _cap_aware_chat_events( 

711 req: CanonicalChatRequest, 

712) -> AsyncIterator[StreamToken | CapNotice]: 

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

714 

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

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

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

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

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

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

721 """ 

722 cap_chars = effective_reasoning_cap() 

723 show = cfg.show_reasoning 

724 answered = False 

725 

726 first_parser = TagParser(show=show) 

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

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

729 yield tok 

730 

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

732 yield CapNotice(cap_chars=cap_chars) 

733 nudged = nudged_request(req) 

734 cont_parser = TagParser(show=show) 

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

736 answered = answered or bool(tok.content) 

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

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

739 

740 if first_parser.reasoning_chars > 0 and not answered: 

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

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

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

744 

745 

746async def _drive_stream( 

747 stream: AsyncIterator[Any], 

748 parser: TagParser, 

749 cap_chars: int, 

750) -> AsyncIterator[StreamToken]: 

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

752 cap_fired = False 

753 try: 

754 async for event in stream: 

755 text = _text_from_event(event) 

756 if not text: 

757 continue 

758 for tok in parser.feed(text): 

759 if tok.content: 

760 yield tok 

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

762 cap_fired = True 

763 break 

764 finally: 

765 if cap_fired: 

766 await _aclose(stream) 

767 tail = parser.flush() 

768 if tail is not None and tail.content: 

769 yield tail 

770 

771 

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

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

774 if not isinstance(stream, AsyncGenerator): 

775 return 

776 with contextlib.suppress(Exception): 

777 await stream.aclose() 

778 

779 

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

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

782 if isinstance(event, CapNotice): 

783 return sse_event( 

784 SseEvent.REASONING, 

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

786 ) 

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

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

789 

790 

791def _chat_answer_frame( 

792 event: StreamToken | CapNotice, 

793 cite_filter: StreamingCitationFilter | None, 

794 answer_parts: list[str], 

795) -> str: 

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

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

798 

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

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

801 treated as a citation source. 

802 """ 

803 is_answer = ( 

804 isinstance(event, StreamToken) 

805 and not event.is_reasoning 

806 and event.content != REASONING_EXHAUSTED_NOTICE 

807 ) 

808 if not is_answer: 

809 return _sse_for_chat_event(event)