Coverage for src/lilbee/retrieval/reasoning.py: 100%
203 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""Reasoning token filter and cap-aware chat orchestrator.
3Reasoning models (Qwen3, DeepSeek-R1) wrap their thinking process in
4``<think>...</think>`` tags. This module provides:
6- ``filter_reasoning``: a stateful streaming filter that classifies
7 tokens as reasoning vs response and signals when reasoning exceeds a
8 caller-supplied cap.
9- ``stream_chat_with_cap``: the high-level orchestrator. Wraps a
10 provider call with the filter; when the cap fires, re-issues the
11 chat with a "stop thinking, answer directly" nudge. The ask/search
12 streaming path and CLI/TUI consume it directly. Two other drivers
13 mirror the same filter + cap-nudge behavior over async streams:
14 the RAG handler's, which needs the reasoning/answer split for its
15 SSE channels, and
16 :mod:`lilbee.server.chat_dispatch.reasoning_cap`, which keeps the
17 canonical event stream intact so the chat surfaces still carry tool
18 calls.
19- ``effective_reasoning_cap``: resolves the cap from the global config
20 with per-model ``ModelDefaults`` overrides.
21"""
23from __future__ import annotations
25import contextlib
26import re
27from collections.abc import Callable, Generator, Iterator
28from dataclasses import dataclass
29from typing import TYPE_CHECKING, Any, NamedTuple
31from lilbee.core.config import cfg
32from lilbee.providers.base import THINK_CLOSE_TAG, THINK_OPEN_TAG, ClosableIterator
34if TYPE_CHECKING:
35 from lilbee.providers.base import LLMProvider
37_THINK_BLOCK_RE = re.compile(
38 rf"{THINK_OPEN_TAG}[\s\S]*?{THINK_CLOSE_TAG}\s*|{THINK_OPEN_TAG}[\s\S]*$"
39)
40_PROGRESS_TICK_CHARS = 256
41"""Coarseness of the progress callback: fire when reasoning grows by at least this many chars."""
43CAP_CONTINUATION_PROMPT = (
44 "Stop thinking now. Give your final answer directly, without any further <think> blocks."
45)
46"""The user-message nudge appended on the continuation call after the cap fires."""
48CAP_NOTICE_TEMPLATE = "\n[reasoning capped at {chars} chars, asking for a direct answer]\n"
49"""User-visible marker emitted between the truncated reasoning and the continuation answer."""
51REASONING_EXHAUSTED_NOTICE = (
52 "The model spent its whole response budget on reasoning and produced no final "
53 "answer. Try a shorter question, raise the generation token limit, or lower the "
54 "reasoning effort."
55)
56"""Returned in place of an empty answer when reasoning consumed the entire generation.
58Lets a caller tell "the model thought itself to death" apart from a genuine empty
59response, which an empty string alone cannot."""
61PSEUDO_THINKING_TAGS = ("anthropic_thinking", "anti_codeblock", "thinking")
62"""Tag names some models emit as literal reply-initial planning blocks in plain text."""
65@dataclass
66class StreamToken:
67 """A classified token from the stream."""
69 content: str
70 is_reasoning: bool
73@dataclass
74class CapNotice:
75 """Emitted once when the reasoning cap fires, before the continuation stream."""
77 cap_chars: int
80@dataclass
81class RetrievalNotice:
82 """Emitted once before the first token when retrieval ran on a rewrite of the question."""
84 query: str
87@dataclass
88class TagParser:
89 """Stateful parser that tracks whether we're inside a thinking block."""
91 show: bool
92 buf: str = ""
93 in_thinking: bool = False
94 reasoning_chars: int = 0
96 def feed(self, token: str) -> list[StreamToken]:
97 """Feed a token and return any complete StreamTokens."""
98 self.buf += token
99 result: list[StreamToken] = []
100 while self.buf:
101 emitted = self._process_thinking() if self.in_thinking else self._process_normal()
102 if emitted is None:
103 break
104 if emitted.content:
105 result.append(emitted)
106 return result
108 def flush(self) -> StreamToken | None:
109 """Flush remaining buffer at end of stream."""
110 if not self.buf:
111 return None
112 if self.in_thinking:
113 self.reasoning_chars += len(self.buf)
114 return StreamToken(content=self.buf, is_reasoning=True) if self.show else None
115 return StreamToken(content=self.buf, is_reasoning=False)
117 def _process_thinking(self) -> StreamToken | None:
118 close_idx = self.buf.find(THINK_CLOSE_TAG)
119 if close_idx == -1:
120 if _could_be_partial(THINK_CLOSE_TAG, self.buf):
121 return None
122 content = self.buf
123 self.reasoning_chars += len(content)
124 self.buf = ""
125 return (
126 StreamToken(content=content, is_reasoning=True)
127 if self.show
128 else StreamToken(content="", is_reasoning=True)
129 )
130 thinking_content = self.buf[:close_idx]
131 self.reasoning_chars += len(thinking_content)
132 self.buf = self.buf[close_idx + len(THINK_CLOSE_TAG) :]
133 self.in_thinking = False
134 if thinking_content and self.show:
135 return StreamToken(content=thinking_content, is_reasoning=True)
136 return StreamToken(content="", is_reasoning=True)
138 def _process_normal(self) -> StreamToken | None:
139 open_idx = self.buf.find(THINK_OPEN_TAG)
140 if open_idx == -1:
141 if _could_be_partial(THINK_OPEN_TAG, self.buf):
142 return None
143 content = self.buf
144 self.buf = ""
145 return StreamToken(content=content, is_reasoning=False)
146 before = self.buf[:open_idx]
147 self.buf = self.buf[open_idx + len(THINK_OPEN_TAG) :]
148 self.in_thinking = True
149 return StreamToken(content=before, is_reasoning=False)
152@dataclass
153class PseudoThinkingNormalizer:
154 """Rewrites a reply-initial pseudo-thinking tag pair to the ``<think>`` tags.
156 Some models under context pressure open a reply with a literal planning tag
157 from ``PSEUDO_THINKING_TAGS`` as ordinary text, outside any reasoning
158 channel. Rewriting the tag pair to ``<think>``/``</think>`` lets
159 ``TagParser`` treat the block as reasoning. Only a reply-initial tag is
160 rewritten; once the reply starts with anything else, text passes through
161 verbatim, so mid-reply XML/HTML is never altered.
162 """
164 buf: str = ""
165 decided: bool = False
166 close_tag: str | None = None
168 def feed(self, text: str) -> str:
169 """Feed a chunk and return the text decided so far."""
170 if self.decided and self.close_tag is None:
171 return text
172 self.buf += text
173 out = "" if self.decided else self._resolve_initial()
174 if self.close_tag is not None:
175 out += self._scan_close(self.close_tag)
176 return out
178 def flush(self) -> str:
179 """Emit the held buffer at end of reply; later feeds pass through."""
180 rest, self.buf = self.buf, ""
181 self.decided = True
182 self.close_tag = None
183 return rest
185 def _resolve_initial(self) -> str:
186 """Decide whether the reply opens with a pseudo tag; ``""`` while undecidable."""
187 stripped = self.buf.lstrip()
188 for name in PSEUDO_THINKING_TAGS:
189 open_tag = f"<{name}>"
190 if stripped.startswith(open_tag):
191 self.decided = True
192 self.close_tag = f"</{name}>"
193 self.buf = stripped[len(open_tag) :]
194 return THINK_OPEN_TAG
195 if any(f"<{name}>".startswith(stripped) for name in PSEUDO_THINKING_TAGS):
196 return ""
197 self.decided = True
198 rest, self.buf = self.buf, ""
199 return rest
201 def _scan_close(self, close_tag: str) -> str:
202 """Pass block content through, rewriting *close_tag* when it arrives."""
203 close_idx = self.buf.find(close_tag)
204 if close_idx == -1:
205 if _could_be_partial(close_tag, self.buf):
206 return ""
207 content, self.buf = self.buf, ""
208 return content
209 content = self.buf[:close_idx]
210 rest = self.buf[close_idx + len(close_tag) :]
211 self.buf = ""
212 self.close_tag = None
213 return f"{content}{THINK_CLOSE_TAG}{rest}"
216def normalize_pseudo_thinking(text: str) -> str:
217 """Rewrite a reply-initial pseudo-thinking block in a complete string."""
218 normalizer = PseudoThinkingNormalizer()
219 return normalizer.feed(text) + normalizer.flush()
222def filter_reasoning(
223 tokens: Iterator[str],
224 *,
225 show: bool,
226 cap_chars: int,
227 on_cap: Callable[[], None] | None = None,
228 on_progress: Callable[[int], None] | None = None,
229) -> Iterator[StreamToken]:
230 """Classify ``<think>...</think>`` tokens and stop when reasoning exceeds the cap.
232 *cap_chars* bounds reasoning content. When exceeded, ``on_cap`` is
233 fired (no payload), the upstream iterator is closed, and iteration
234 stops. The caller decides what to do next via the higher-level
235 ``stream_chat_with_cap`` orchestrator. *on_progress* is fired with
236 the running reasoning-chars count each time it grows by at least 256
237 characters. A non-positive *cap_chars* disables the cap.
238 """
239 parser = TagParser(show=show)
240 last_progress_tick = 0
241 try:
242 for token in tokens:
243 for st in parser.feed(token):
244 if st.content:
245 yield st
246 if (
247 on_progress is not None
248 and parser.reasoning_chars >= last_progress_tick + _PROGRESS_TICK_CHARS
249 ):
250 last_progress_tick = parser.reasoning_chars
251 on_progress(parser.reasoning_chars)
252 if cap_chars > 0 and parser.reasoning_chars > cap_chars:
253 if on_cap is not None:
254 on_cap()
255 return
256 final = parser.flush()
257 if final and final.content:
258 yield final
259 if on_progress is not None and parser.reasoning_chars > last_progress_tick:
260 on_progress(parser.reasoning_chars)
261 finally:
262 _close_iterator(tokens)
265def effective_reasoning_cap() -> int:
266 """Return the active reasoning cap; 0 means unlimited.
268 A per-model ``ModelDefaults.max_reasoning_chars`` value (including
269 ``0`` for "this model is allowed to think forever") beats the global
270 ``cfg.max_reasoning_chars`` setting. A missing (``None``) or negative
271 per-model value falls through to the global: ModelDefaults is an
272 unvalidated dataclass, so a negative is treated as "unset" rather than
273 trusted as a cap.
274 """
275 defaults = cfg.model_defaults
276 override = defaults.max_reasoning_chars if defaults is not None else None
277 return override if isinstance(override, int) and override >= 0 else cfg.max_reasoning_chars
280def stream_chat_with_cap(
281 provider: LLMProvider,
282 messages: list[dict[str, Any]],
283 *,
284 options: dict[str, Any] | None,
285 model: str,
286 show_reasoning: bool,
287 cap_chars: int,
288) -> Generator[StreamToken | CapNotice, None, None]:
289 """Stream chat tokens; on cap-fire, re-issue with a stop-thinking nudge.
291 Yields ``StreamToken`` events for both reasoning and response tokens
292 in the first pass. If reasoning exceeds *cap_chars*, the upstream
293 iterator is closed, a single ``CapNotice`` is yielded, and the
294 continuation stream starts (same messages plus a user message asking
295 the model to answer directly).
297 The continuation is parsed too: a chat template can force-open a
298 ``<think>`` block whatever the nudge asks, and that reasoning must not
299 reach the visible answer. Its cap is disabled (the cap already fired
300 once, and re-capping would cut the answer off), and every continuation
301 token is reported as final-answer text, matching the async HTTP path.
303 A run that reasoned but never produced final-answer text closes with
304 ``REASONING_EXHAUSTED_NOTICE``, so the CLI/TUI/library path ends with an
305 explanation rather than silence -- the same close the HTTP path makes.
306 """
307 cap_fired = False
308 reasoned = False
309 answered = False
311 def _on_cap() -> None:
312 nonlocal cap_fired, reasoned
313 cap_fired = True
314 reasoned = True
316 def _on_reasoning(_chars: int) -> None:
317 nonlocal reasoned
318 reasoned = True
320 first_stream = provider.chat(messages, stream=True, options=options or None, model=model)
321 for token in filter_reasoning(
322 _text_only(first_stream),
323 show=show_reasoning,
324 cap_chars=cap_chars,
325 on_cap=_on_cap,
326 on_progress=_on_reasoning,
327 ):
328 answered = answered or not token.is_reasoning
329 yield token
330 if cap_fired:
331 yield CapNotice(cap_chars=cap_chars)
332 nudged = [*messages, {"role": "user", "content": CAP_CONTINUATION_PROMPT}]
333 second_stream = provider.chat(nudged, stream=True, options=options or None, model=model)
334 try:
335 for token in filter_reasoning(
336 _text_only(second_stream), show=show_reasoning, cap_chars=0
337 ):
338 if token.content:
339 answered = True
340 yield StreamToken(content=token.content, is_reasoning=False)
341 finally:
342 _close_iterator(second_stream)
343 if reasoned and not answered:
344 yield StreamToken(content=REASONING_EXHAUSTED_NOTICE, is_reasoning=False)
347def _text_only(stream: Iterator[Any]) -> Iterator[str]:
348 """Filter a chat stream down to its text deltas.
350 Tool-call deltas (when ``tools`` is passed) and the trailing token-usage
351 frame both ride the same iterator; the RAG / reasoning paths only consume
352 text, so any non-str frame is dropped here rather than crashing the chat.
353 """
354 try:
355 for item in stream:
356 if isinstance(item, str):
357 yield item
358 finally:
359 # Forward close to the source: when a consumer (filter_reasoning on
360 # cap-fire) closes this generator, a plain for-loop would not propagate
361 # GeneratorExit to *stream*, leaking its HTTP connection / in_flight slot.
362 _close_iterator(stream)
365def cap_events_as_stream_tokens(
366 events: Iterator[StreamToken | CapNotice],
367) -> Iterator[StreamToken]:
368 """Render ``CapNotice`` events as user-visible reasoning ``StreamToken``s.
370 Library and CLI surfaces speak ``StreamToken`` only. This helper lets
371 them consume the orchestrator's union output without a per-call
372 isinstance dance for the cap notice.
373 """
374 for event in events:
375 if isinstance(event, CapNotice):
376 yield StreamToken(
377 content=CAP_NOTICE_TEMPLATE.format(chars=event.cap_chars),
378 is_reasoning=True,
379 )
380 elif event.content:
381 yield event
384def _close_iterator(tokens: Iterator[Any]) -> None:
385 """Close *tokens* if it satisfies the ClosableIterator protocol."""
386 if isinstance(tokens, ClosableIterator):
387 with contextlib.suppress(Exception):
388 tokens.close()
391def strip_reasoning(text: str) -> str:
392 """Remove ``<think>...</think>`` blocks from a complete (non-streaming) string."""
393 return _THINK_BLOCK_RE.sub("", text)
396class ReasoningSplit(NamedTuple):
397 """A model reply separated into its reasoning and its visible answer."""
399 reasoning: str
400 answer: str
403def split_reasoning(text: str) -> ReasoningSplit:
404 """Split a complete string into its reasoning and its answer.
406 An unterminated ``<think>`` block means the model never reached an answer, so
407 everything after the tag is reasoning. Surfaces that must not leak lilbee's
408 inline ``<think>`` convention (the OpenAI-compatible API) use this to report
409 reasoning in its own field.
410 """
411 blocks = [
412 match.group(0).strip().removeprefix(THINK_OPEN_TAG).removesuffix(THINK_CLOSE_TAG).strip()
413 for match in _THINK_BLOCK_RE.finditer(text)
414 ]
415 return ReasoningSplit("".join(blocks), strip_reasoning(text))
418def _could_be_partial(tag: str, buf: str) -> bool:
419 """Check if the end of buf could be the start of the given tag."""
420 return any(buf.endswith(tag[:length]) for length in range(1, len(tag)))