Coverage for src/lilbee/retrieval/reasoning.py: 100%
151 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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; the canonical
13 chat-dispatch path mirrors the same filter + cap-nudge behavior over
14 its own async driver.
15- ``effective_reasoning_cap``: resolves the cap from the global config
16 with per-model ``ModelDefaults`` overrides.
17"""
19from __future__ import annotations
21import contextlib
22import re
23from collections.abc import Callable, Generator, Iterator
24from dataclasses import dataclass
25from typing import TYPE_CHECKING, Any, NamedTuple
27from lilbee.core.config import cfg
28from lilbee.providers.base import THINK_CLOSE_TAG, THINK_OPEN_TAG, ClosableIterator
30if TYPE_CHECKING:
31 from lilbee.providers.base import LLMProvider
33_THINK_BLOCK_RE = re.compile(
34 rf"{THINK_OPEN_TAG}[\s\S]*?{THINK_CLOSE_TAG}\s*|{THINK_OPEN_TAG}[\s\S]*$"
35)
36_PROGRESS_TICK_CHARS = 256
37"""Coarseness of the progress callback: fire when reasoning grows by at least this many chars."""
39CAP_CONTINUATION_PROMPT = (
40 "Stop thinking now. Give your final answer directly, without any further <think> blocks."
41)
42"""The user-message nudge appended on the continuation call after the cap fires."""
44CAP_NOTICE_TEMPLATE = "\n[reasoning capped at {chars} chars, asking for a direct answer]\n"
45"""User-visible marker emitted between the truncated reasoning and the continuation answer."""
47REASONING_EXHAUSTED_NOTICE = (
48 "The model spent its whole response budget on reasoning and produced no final "
49 "answer. Try a shorter question, raise the generation token limit, or lower the "
50 "reasoning effort."
51)
52"""Returned in place of an empty answer when reasoning consumed the entire generation.
54Lets a caller tell "the model thought itself to death" apart from a genuine empty
55response, which an empty string alone cannot."""
58@dataclass
59class StreamToken:
60 """A classified token from the stream."""
62 content: str
63 is_reasoning: bool
66@dataclass
67class CapNotice:
68 """Emitted once when the reasoning cap fires, before the continuation stream."""
70 cap_chars: int
73@dataclass
74class TagParser:
75 """Stateful parser that tracks whether we're inside a thinking block."""
77 show: bool
78 buf: str = ""
79 in_thinking: bool = False
80 reasoning_chars: int = 0
82 def feed(self, token: str) -> list[StreamToken]:
83 """Feed a token and return any complete StreamTokens."""
84 self.buf += token
85 result: list[StreamToken] = []
86 while self.buf:
87 emitted = self._process_thinking() if self.in_thinking else self._process_normal()
88 if emitted is None:
89 break
90 if emitted.content:
91 result.append(emitted)
92 return result
94 def flush(self) -> StreamToken | None:
95 """Flush remaining buffer at end of stream."""
96 if not self.buf:
97 return None
98 if self.in_thinking:
99 self.reasoning_chars += len(self.buf)
100 return StreamToken(content=self.buf, is_reasoning=True) if self.show else None
101 return StreamToken(content=self.buf, is_reasoning=False)
103 def _process_thinking(self) -> StreamToken | None:
104 close_idx = self.buf.find(THINK_CLOSE_TAG)
105 if close_idx == -1:
106 if _could_be_partial(THINK_CLOSE_TAG, self.buf):
107 return None
108 content = self.buf
109 self.reasoning_chars += len(content)
110 self.buf = ""
111 return (
112 StreamToken(content=content, is_reasoning=True)
113 if self.show
114 else StreamToken(content="", is_reasoning=True)
115 )
116 thinking_content = self.buf[:close_idx]
117 self.reasoning_chars += len(thinking_content)
118 self.buf = self.buf[close_idx + len(THINK_CLOSE_TAG) :]
119 self.in_thinking = False
120 if thinking_content and self.show:
121 return StreamToken(content=thinking_content, is_reasoning=True)
122 return StreamToken(content="", is_reasoning=True)
124 def _process_normal(self) -> StreamToken | None:
125 open_idx = self.buf.find(THINK_OPEN_TAG)
126 if open_idx == -1:
127 if _could_be_partial(THINK_OPEN_TAG, self.buf):
128 return None
129 content = self.buf
130 self.buf = ""
131 return StreamToken(content=content, is_reasoning=False)
132 before = self.buf[:open_idx]
133 self.buf = self.buf[open_idx + len(THINK_OPEN_TAG) :]
134 self.in_thinking = True
135 return StreamToken(content=before, is_reasoning=False)
138def filter_reasoning(
139 tokens: Iterator[str],
140 *,
141 show: bool,
142 cap_chars: int,
143 on_cap: Callable[[], None] | None = None,
144 on_progress: Callable[[int], None] | None = None,
145) -> Iterator[StreamToken]:
146 """Classify ``<think>...</think>`` tokens and stop when reasoning exceeds the cap.
148 *cap_chars* bounds reasoning content. When exceeded, ``on_cap`` is
149 fired (no payload), the upstream iterator is closed, and iteration
150 stops. The caller decides what to do next via the higher-level
151 ``stream_chat_with_cap`` orchestrator. *on_progress* is fired with
152 the running reasoning-chars count each time it grows by at least 256
153 characters. A non-positive *cap_chars* disables the cap.
154 """
155 parser = TagParser(show=show)
156 last_progress_tick = 0
157 try:
158 for token in tokens:
159 for st in parser.feed(token):
160 if st.content:
161 yield st
162 if (
163 on_progress is not None
164 and parser.reasoning_chars >= last_progress_tick + _PROGRESS_TICK_CHARS
165 ):
166 last_progress_tick = parser.reasoning_chars
167 on_progress(parser.reasoning_chars)
168 if cap_chars > 0 and parser.reasoning_chars > cap_chars:
169 if on_cap is not None:
170 on_cap()
171 return
172 final = parser.flush()
173 if final and final.content:
174 yield final
175 if on_progress is not None and parser.reasoning_chars > last_progress_tick:
176 on_progress(parser.reasoning_chars)
177 finally:
178 _close_iterator(tokens)
181def effective_reasoning_cap() -> int:
182 """Return the active reasoning cap; 0 means unlimited.
184 A per-model ``ModelDefaults.max_reasoning_chars`` value (including
185 ``0`` for "this model is allowed to think forever") beats the global
186 ``cfg.max_reasoning_chars`` setting. A missing (``None``) or negative
187 per-model value falls through to the global: ModelDefaults is an
188 unvalidated dataclass, so a negative is treated as "unset" rather than
189 trusted as a cap.
190 """
191 defaults = cfg.model_defaults
192 override = defaults.max_reasoning_chars if defaults is not None else None
193 return override if isinstance(override, int) and override >= 0 else cfg.max_reasoning_chars
196def stream_chat_with_cap(
197 provider: LLMProvider,
198 messages: list[dict[str, Any]],
199 *,
200 options: dict[str, Any] | None,
201 model: str,
202 show_reasoning: bool,
203 cap_chars: int,
204) -> Generator[StreamToken | CapNotice, None, None]:
205 """Stream chat tokens; on cap-fire, re-issue with a stop-thinking nudge.
207 Yields ``StreamToken`` events for both reasoning and response tokens
208 in the first pass. If reasoning exceeds *cap_chars*, the upstream
209 iterator is closed, a single ``CapNotice`` is yielded, and the
210 continuation stream starts (same messages plus a user message asking
211 the model to answer directly).
213 The continuation is parsed too: a chat template can force-open a
214 ``<think>`` block whatever the nudge asks, and that reasoning must not
215 reach the visible answer. Its cap is disabled (the cap already fired
216 once, and re-capping would cut the answer off), and every continuation
217 token is reported as final-answer text, matching the async HTTP path.
219 A run that reasoned but never produced final-answer text closes with
220 ``REASONING_EXHAUSTED_NOTICE``, so the CLI/TUI/library path ends with an
221 explanation rather than silence -- the same close the HTTP path makes.
222 """
223 cap_fired = False
224 reasoned = False
225 answered = False
227 def _on_cap() -> None:
228 nonlocal cap_fired, reasoned
229 cap_fired = True
230 reasoned = True
232 def _on_reasoning(_chars: int) -> None:
233 nonlocal reasoned
234 reasoned = True
236 first_stream = provider.chat(messages, stream=True, options=options or None, model=model)
237 for token in filter_reasoning(
238 _text_only(first_stream),
239 show=show_reasoning,
240 cap_chars=cap_chars,
241 on_cap=_on_cap,
242 on_progress=_on_reasoning,
243 ):
244 answered = answered or not token.is_reasoning
245 yield token
246 if cap_fired:
247 yield CapNotice(cap_chars=cap_chars)
248 nudged = [*messages, {"role": "user", "content": CAP_CONTINUATION_PROMPT}]
249 second_stream = provider.chat(nudged, stream=True, options=options or None, model=model)
250 try:
251 for token in filter_reasoning(
252 _text_only(second_stream), show=show_reasoning, cap_chars=0
253 ):
254 if token.content:
255 answered = True
256 yield StreamToken(content=token.content, is_reasoning=False)
257 finally:
258 _close_iterator(second_stream)
259 if reasoned and not answered:
260 yield StreamToken(content=REASONING_EXHAUSTED_NOTICE, is_reasoning=False)
263def _text_only(stream: Iterator[Any]) -> Iterator[str]:
264 """Filter a chat stream down to its text deltas.
266 Tool-call deltas (when ``tools`` is passed) and the trailing token-usage
267 frame both ride the same iterator; the RAG / reasoning paths only consume
268 text, so any non-str frame is dropped here rather than crashing the chat.
269 """
270 try:
271 for item in stream:
272 if isinstance(item, str):
273 yield item
274 finally:
275 # Forward close to the source: when a consumer (filter_reasoning on
276 # cap-fire) closes this generator, a plain for-loop would not propagate
277 # GeneratorExit to *stream*, leaking its HTTP connection / in_flight slot.
278 _close_iterator(stream)
281def cap_events_as_stream_tokens(
282 events: Iterator[StreamToken | CapNotice],
283) -> Iterator[StreamToken]:
284 """Render ``CapNotice`` events as user-visible reasoning ``StreamToken``s.
286 Library and CLI surfaces speak ``StreamToken`` only. This helper lets
287 them consume the orchestrator's union output without a per-call
288 isinstance dance for the cap notice.
289 """
290 for event in events:
291 if isinstance(event, CapNotice):
292 yield StreamToken(
293 content=CAP_NOTICE_TEMPLATE.format(chars=event.cap_chars),
294 is_reasoning=True,
295 )
296 elif event.content:
297 yield event
300def _close_iterator(tokens: Iterator[Any]) -> None:
301 """Close *tokens* if it satisfies the ClosableIterator protocol."""
302 if isinstance(tokens, ClosableIterator):
303 with contextlib.suppress(Exception):
304 tokens.close()
307def strip_reasoning(text: str) -> str:
308 """Remove ``<think>...</think>`` blocks from a complete (non-streaming) string."""
309 return _THINK_BLOCK_RE.sub("", text)
312class ReasoningSplit(NamedTuple):
313 """A model reply separated into its reasoning and its visible answer."""
315 reasoning: str
316 answer: str
319def split_reasoning(text: str) -> ReasoningSplit:
320 """Split a complete string into its reasoning and its answer.
322 An unterminated ``<think>`` block means the model never reached an answer, so
323 everything after the tag is reasoning. Surfaces that must not leak lilbee's
324 inline ``<think>`` convention (the OpenAI-compatible API) use this to report
325 reasoning in its own field.
326 """
327 blocks = [
328 match.group(0).strip().removeprefix(THINK_OPEN_TAG).removesuffix(THINK_CLOSE_TAG).strip()
329 for match in _THINK_BLOCK_RE.finditer(text)
330 ]
331 return ReasoningSplit("".join(blocks), strip_reasoning(text))
334def _could_be_partial(tag: str, buf: str) -> bool:
335 """Check if the end of buf could be the start of the given tag."""
336 return any(buf.endswith(tag[:length]) for length in range(1, len(tag)))