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 dataclasses 

8import logging 

9import threading 

10from collections.abc import AsyncGenerator, AsyncIterator, Callable 

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

12 

13from lilbee.app.memory import auto_extract, auto_extract_enabled 

14from lilbee.app.search import clean_result 

15from lilbee.app.services import get_services 

16from lilbee.core.config import cfg 

17from lilbee.core.results import DocumentResult, group 

18from lilbee.data.store import ChunkType, EmbeddingModelMismatchError 

19from lilbee.providers.base import ProviderError, ProviderErrorKind 

20from lilbee.providers.roles import WorkerRole 

21from lilbee.retrieval.query.compaction import ( 

22 compaction_due, 

23 foldable, 

24 history_budget, 

25 prompt_history, 

26) 

27from lilbee.retrieval.query.formatting import ( 

28 StreamingCitationFilter, 

29 cited_subset, 

30 strip_llm_citations, 

31) 

32from lilbee.retrieval.query.searcher import ( 

33 GROUNDED_REFUSAL, 

34 SEARCH_NEEDS_EMBEDDER, 

35 RagContext, 

36) 

37from lilbee.retrieval.reasoning import ( 

38 CAP_CONTINUATION_PROMPT, 

39 CAP_NOTICE_TEMPLATE, 

40 REASONING_EXHAUSTED_NOTICE, 

41 CapNotice, 

42 StreamToken, 

43 TagParser, 

44 effective_reasoning_cap, 

45 stream_chat_with_cap, 

46 strip_reasoning, 

47) 

48from lilbee.runtime.progress import SseErrorCode, SseEvent 

49from lilbee.server.chat_completions_api.errors import ( 

50 _BACKEND_FAILURE_MESSAGE, 

51 _INFRASTRUCTURE_KINDS, 

52 CompletionsErrorCode, 

53) 

54from lilbee.server.chat_dispatch.canonical import ( 

55 CanonicalChatRequest, 

56 CanonicalMessage, 

57 ContentBlockDelta, 

58 TextBlock, 

59 TextDelta, 

60) 

61from lilbee.server.chat_dispatch.dispatch import ( 

62 ModelDoesNotSupportToolsError, 

63 ModelNotFoundError, 

64 dispatch_chat, 

65 dispatch_chat_stream, 

66) 

67from lilbee.server.handlers.sse import ( 

68 SseErrorCodeValue, 

69 SseStream, 

70 _resolve_generation_options, 

71 classify_load_error, 

72 sse_done, 

73 sse_error, 

74 sse_event, 

75) 

76from lilbee.server.models import ( 

77 AskResponse, 

78 CleanedChunk, 

79 CompactionInfo, 

80 MemoryExtractedEvent, 

81 MemoryExtractedItem, 

82) 

83from lilbee.sessions import SessionNotFoundError, sessions_enabled 

84 

85if TYPE_CHECKING: 

86 from lilbee.core.results import SearchChunk 

87 from lilbee.retrieval.query import ChatMessage 

88 from lilbee.retrieval.query.searcher import Searcher 

89 

90log = logging.getLogger(__name__) 

91 

92 

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

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

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

96# engine internals. 

97_STREAM_KIND_CODES: dict[ProviderErrorKind, CompletionsErrorCode] = { 

98 ProviderErrorKind.CONTEXT_OVERFLOW: CompletionsErrorCode.CONTEXT_LENGTH_EXCEEDED, 

99 ProviderErrorKind.NOT_FOUND: CompletionsErrorCode.MODEL_NOT_FOUND, 

100} 

101 

102 

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

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

105 if isinstance(exc, ModelNotFoundError): 

106 return CompletionsErrorCode.MODEL_NOT_FOUND, str(exc) 

107 if isinstance(exc, ModelDoesNotSupportToolsError): 

108 return CompletionsErrorCode.MODEL_DOES_NOT_SUPPORT_TOOLS, str(exc) 

109 if isinstance(exc, ProviderError): 

110 mapped = _STREAM_KIND_CODES.get(exc.kind) 

111 if mapped is not None: 

112 return mapped, str(exc) 

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

114 if exc.kind in _INFRASTRUCTURE_KINDS: 

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

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

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

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

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

120 return code, _BACKEND_FAILURE_MESSAGE 

121 return code, str(exc) 

122 return classify_load_error(str(exc)) 

123 

124 

125async def search( 

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

127) -> list[DocumentResult]: 

128 """Search and return grouped DocumentResults.""" 

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

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

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

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

133 results = await asyncio.to_thread( 

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

135 ) 

136 return group(results) 

137 

138 

139async def ask( 

140 question: str, 

141 top_k: int = 0, 

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

143 chunk_type: ChunkType | None = None, 

144) -> AskResponse: 

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

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

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

148 opts = _resolve_generation_options(options) 

149 searcher = get_services().searcher 

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

151 # event loop so other admitted requests stay responsive. 

152 result = await asyncio.to_thread( 

153 searcher.ask_raw, 

154 question, 

155 top_k=top_k, 

156 options=opts, 

157 chunk_type=chunk_type, 

158 ) 

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

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

161 if not searcher.search_unavailable(): 

162 await _store_extracted_memories(question, result.answer) 

163 return AskResponse( 

164 answer=result.answer, 

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

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

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 if _retrieval_off(searcher, top_k): 

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

511 sources: list[SearchChunk] = [] 

512 messages = searcher.direct_messages(question, history) 

513 else: 

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

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

516 pre_answer = searcher.pre_retrieval_answer(question) 

517 if pre_answer is not None: 

518 return AskResponse( 

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

520 ) 

521 rag = searcher.build_rag_context( 

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

523 ) 

524 if rag is None: 

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

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

527 return AskResponse( 

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

529 ) 

530 sources, messages = rag.results, rag.messages 

531 req = _build_canonical_request(messages, options) 

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

533 text = _join_text_blocks(response.content) 

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

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

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

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

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

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

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

541 answer = REASONING_EXHAUSTED_NOTICE 

542 else: 

543 await _store_extracted_memories(question, answer) 

544 return AskResponse( 

545 answer=answer, 

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

547 cited_sources=[ 

548 CleanedChunk(**clean_result(s)) 

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

550 ], 

551 compaction=compaction, 

552 ) 

553 

554 

555def chat_stream( 

556 question: str, 

557 history: list[ChatMessage], 

558 top_k: int | None = None, 

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

560 chunk_type: ChunkType | None = None, 

561 summary: str = "", 

562 session_id: str | None = None, 

563) -> AsyncGenerator[str, None]: 

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

565 return _stream_chat_response( 

566 question, 

567 history=history, 

568 top_k=top_k, 

569 options=options, 

570 chunk_type=chunk_type, 

571 summary=summary, 

572 session_id=session_id, 

573 ) 

574 

575 

576class _StreamResolution(NamedTuple): 

577 """Retrieval outcome for a streaming turn. 

578 

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

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

581 or a clean refusal/error). 

582 """ 

583 

584 sources: list[SearchChunk] 

585 messages: list[ChatMessage] | None 

586 preempt_frames: list[str] 

587 

588 

589class _ChatStreamPlan(NamedTuple): 

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

591 

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

593 refusal or error) and stop. 

594 """ 

595 

596 frames: list[str] 

597 context: RagContext | None 

598 

599 

600def _resolve_chat_stream_context( 

601 searcher: Searcher, 

602 question: str, 

603 history: list[ChatMessage], 

604 top_k: int | None, 

605 chunk_type: ChunkType | None, 

606) -> _ChatStreamPlan: 

607 frames = list(_chat_warming_events()) 

608 if searcher.search_unavailable(): 

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

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

611 frames += [ 

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

613 sse_event(SseEvent.SOURCES, []), 

614 sse_done({}), 

615 ] 

616 return _ChatStreamPlan(frames, None) 

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

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

619 sources, messages, preempt = _resolve_stream_context( 

620 searcher, 

621 question, 

622 history, 

623 top_k, 

624 chunk_type, 

625 retrieval_off=_retrieval_off(searcher, top_k), 

626 ) 

627 frames += preempt 

628 if messages is None: 

629 return _ChatStreamPlan(frames, None) 

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

631 

632 

633async def _stream_chat_response( 

634 question: str, 

635 history: list[ChatMessage], 

636 top_k: int | None, 

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

638 chunk_type: ChunkType | None, 

639 summary: str = "", 

640 session_id: str | None = None, 

641) -> AsyncGenerator[str, None]: 

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

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

644 if isinstance(item, str): 

645 yield item 

646 continue 

647 history, _compaction = item 

648 frames, ctx = await asyncio.to_thread( 

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

650 ) 

651 for frame in frames: 

652 yield frame 

653 if ctx is None: 

654 return 

655 sources, messages = ctx.results, ctx.messages 

656 

657 req = _build_canonical_request(messages, options) 

658 answer_parts: list[str] = [] 

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

660 # need a model-generated Sources block suppressed. 

661 cite_filter = StreamingCitationFilter() if sources else None 

662 try: 

663 async for event in _cap_aware_chat_events(req): 

664 frame = _chat_answer_frame(event, cite_filter, answer_parts) 

665 if frame: 

666 yield frame 

667 except Exception as exc: 

668 raw = str(exc) 

669 code, user_message = _classify_stream_error(exc) 

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

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

672 return 

673 

674 tail_frame = _chat_answer_tail_frame(cite_filter, answer_parts) 

675 if tail_frame: 

676 yield tail_frame 

677 

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

679 yield frame 

680 

681 

682async def _cap_aware_chat_events( 

683 req: CanonicalChatRequest, 

684) -> AsyncIterator[StreamToken | CapNotice]: 

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

686 

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

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

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

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

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

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

693 """ 

694 cap_chars = effective_reasoning_cap() 

695 show = cfg.show_reasoning 

696 answered = False 

697 

698 first_parser = TagParser(show=show) 

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

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

701 yield tok 

702 

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

704 yield CapNotice(cap_chars=cap_chars) 

705 nudged = _nudged_request(req) 

706 cont_parser = TagParser(show=show) 

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

708 answered = answered or bool(tok.content) 

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

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

711 

712 if first_parser.reasoning_chars > 0 and not answered: 

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

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

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

716 

717 

718async def _drive_stream( 

719 stream: AsyncIterator[Any], 

720 parser: TagParser, 

721 cap_chars: int, 

722) -> AsyncIterator[StreamToken]: 

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

724 cap_fired = False 

725 try: 

726 async for event in stream: 

727 text = _text_from_event(event) 

728 if not text: 

729 continue 

730 for tok in parser.feed(text): 

731 if tok.content: 

732 yield tok 

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

734 cap_fired = True 

735 break 

736 finally: 

737 if cap_fired: 

738 await _aclose(stream) 

739 tail = parser.flush() 

740 if tail is not None and tail.content: 

741 yield tail 

742 

743 

744def _nudged_request(req: CanonicalChatRequest) -> CanonicalChatRequest: 

745 """Append the cap-continuation user prompt to *req*'s messages.""" 

746 return dataclasses.replace( 

747 req, 

748 messages=[ 

749 *req.messages, 

750 CanonicalMessage.from_string(role="user", text=CAP_CONTINUATION_PROMPT), 

751 ], 

752 ) 

753 

754 

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

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

757 if not isinstance(stream, AsyncGenerator): 

758 return 

759 with contextlib.suppress(Exception): 

760 await stream.aclose() 

761 

762 

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

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

765 if isinstance(event, CapNotice): 

766 return sse_event( 

767 SseEvent.REASONING, 

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

769 ) 

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

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

772 

773 

774def _chat_answer_frame( 

775 event: StreamToken | CapNotice, 

776 cite_filter: StreamingCitationFilter | None, 

777 answer_parts: list[str], 

778) -> str: 

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

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

781 

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

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

784 treated as a citation source. 

785 """ 

786 is_answer = ( 

787 isinstance(event, StreamToken) 

788 and not event.is_reasoning 

789 and event.content != REASONING_EXHAUSTED_NOTICE 

790 ) 

791 if not is_answer: 

792 return _sse_for_chat_event(event) 

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

794 if cite_filter is None: 

795 answer_parts.append(content) 

796 return _sse_for_chat_event(event) 

797 shown = cite_filter.feed(content) 

798 if not shown: 

799 return "" 

800 answer_parts.append(shown) 

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

802 

803 

804def _chat_answer_tail_frame( 

805 cite_filter: StreamingCitationFilter | None, 

806 answer_parts: list[str], 

807) -> str: 

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

809 if cite_filter is None: 

810 return "" 

811 tail = cite_filter.flush() 

812 if not tail: 

813 return "" 

814 answer_parts.append(tail) 

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