Coverage for src/lilbee/server/chat_dispatch/reasoning_cap.py: 100%

1"""Applies ``cfg.max_reasoning_chars`` to the canonical chat surfaces. 

2 

3Counting and the stop-thinking nudge live in :mod:`lilbee.retrieval.reasoning`. 

4The cap notice is written into the still-open ``<think>`` block rather than 

5handled per surface, so every downstream translator presents a capped turn. 

6Under ``inline`` that block is the answer text, so the notice reads there. 

7""" 

8 

9from __future__ import annotations 

10 

11import contextlib 

12import dataclasses 

13from collections.abc import AsyncGenerator, AsyncIterator 

14 

15from lilbee.providers.base import THINK_CLOSE_TAG, THINK_OPEN_TAG 

16from lilbee.retrieval.reasoning import ( 

17 CAP_CONTINUATION_PROMPT, 

18 CAP_NOTICE_TEMPLATE, 

19 TagParser, 

20 split_reasoning, 

21) 

22from lilbee.server.chat_dispatch.canonical import ( 

23 CanonicalChatRequest, 

24 CanonicalMessage, 

25 CanonicalResponse, 

26 CanonicalStreamEvent, 

27 CanonicalUsage, 

28 ContentBlock, 

29 ContentBlockDelta, 

30 ContentBlockStart, 

31 ContentBlockStop, 

32 MessageDelta, 

33 MessageStart, 

34 TextBlock, 

35 TextDelta, 

36 ToolUseBlock, 

37) 

38from lilbee.server.chat_dispatch.dispatch import dispatch_chat, dispatch_chat_stream 

39 

40CHARS_PER_TOKEN = 4 

41"""Approximate chars per token, for reading ``budget_tokens`` as a char cap.""" 

42 

43 

44def budget_capped_chars(cap_chars: int, budget_tokens: int | None) -> int: 

45 """Tighten *cap_chars* with a per-request token budget; never loosen it. 

46 

47 ``0`` means unlimited on both sides. A budget of zero or less is no budget, 

48 not a request for unlimited thinking. 

49 """ 

50 if budget_tokens is None or budget_tokens <= 0: 

51 return cap_chars 

52 budget_chars = budget_tokens * CHARS_PER_TOKEN 

53 if cap_chars <= 0: 

54 return budget_chars 

55 return min(cap_chars, budget_chars) 

56 

57 

58def nudged_request(req: CanonicalChatRequest) -> CanonicalChatRequest: 

59 """Append the cap-continuation user prompt; every other request field is kept.""" 

60 return dataclasses.replace( 

61 req, 

62 messages=[ 

63 *req.messages, 

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

65 ], 

66 ) 

67 

68 

69def _cap_notice(cap_chars: int) -> str: 

70 """The notice written into the thinking block when the cap fires.""" 

71 return CAP_NOTICE_TEMPLATE.format(chars=cap_chars) 

72 

73 

74async def _aclose(stream: AsyncIterator[CanonicalStreamEvent]) -> None: 

75 """Best-effort close; only an async generator has ``aclose``.""" 

76 if isinstance(stream, AsyncGenerator): 

77 with contextlib.suppress(Exception): 

78 await stream.aclose() 

79 

80 

81def _reindexed(event: CanonicalStreamEvent, offset: int) -> CanonicalStreamEvent: 

82 """Shift a continuation event's block index past the first stream's blocks. 

83 

84 Without it the continuation reopens index 0 and translators merge it into 

85 the block the cap just closed. 

86 """ 

87 if isinstance(event, ContentBlockStart | ContentBlockDelta | ContentBlockStop): 

88 return dataclasses.replace(event, index=event.index + offset) 

89 return event 

90 

91 

92async def cap_aware_chat_stream( 

93 stream: AsyncIterator[CanonicalStreamEvent], 

94 req: CanonicalChatRequest, 

95 *, 

96 canonical_model: str, 

97 cap_chars: int, 

98) -> AsyncIterator[CanonicalStreamEvent]: 

99 """Forward *stream*, stopping the reasoning at *cap_chars* and forcing an answer. 

100 

101 On cap-fire: close upstream, write the notice into the open thinking block, 

102 splice in the continuation. The continuation is not capped again. 

103 ``cap_chars <= 0`` forwards the stream verbatim. 

104 """ 

105 parser = TagParser(show=True) 

106 cap_fired = False 

107 max_index = -1 

108 open_index = 0 

109 spent = None 

110 try: 

111 async for event in stream: 

112 if isinstance(event, MessageDelta) and event.usage is not None: 

113 spent = event.usage 

114 if isinstance(event, ContentBlockStart | ContentBlockDelta | ContentBlockStop): 

115 max_index = max(max_index, event.index) 

116 if isinstance(event, ContentBlockDelta) and isinstance(event.delta, TextDelta): 

117 open_index = event.index 

118 parser.feed(event.delta.text) 

119 yield event 

120 # Only cap while still inside <think>; a closed block already answered. 

121 if cap_chars > 0 and parser.in_thinking and parser.reasoning_chars > cap_chars: 

122 cap_fired = True 

123 break 

124 finally: 

125 if cap_fired: 

126 await _aclose(stream) 

127 

128 if not cap_fired: 

129 return 

130 

131 yield ContentBlockDelta( 

132 index=open_index, 

133 delta=TextDelta(text=f"{_cap_notice(cap_chars)}{THINK_CLOSE_TAG}"), 

134 ) 

135 offset = max_index + 1 

136 async for event in dispatch_chat_stream(nudged_request(req), canonical_model=canonical_model): 

137 # The message already started; a second prelude would restart it. 

138 if isinstance(event, MessageStart): 

139 continue 

140 yield _reindexed(_with_spent(event, spent), offset) 

141 

142 

143def _summed(first: CanonicalUsage, second: CanonicalUsage) -> CanonicalUsage: 

144 """Field-wise total of two turns' token counts.""" 

145 return CanonicalUsage( 

146 input_tokens=first.input_tokens + second.input_tokens, 

147 output_tokens=first.output_tokens + second.output_tokens, 

148 cached_input_tokens=first.cached_input_tokens + second.cached_input_tokens, 

149 ) 

150 

151 

152def _with_spent(event: CanonicalStreamEvent, spent: CanonicalUsage | None) -> CanonicalStreamEvent: 

153 """Add the capped call's tokens to the continuation's usage. 

154 

155 The caller paid for both generations, so reporting only the continuation's 

156 hides the reasoning the cap just stopped. A provider that reports usage 

157 once at the end has none to add: the cap closes the first stream before 

158 that event, and those tokens go unreported. 

159 """ 

160 if spent is None or not isinstance(event, MessageDelta) or event.usage is None: 

161 return event 

162 return dataclasses.replace(event, usage=_summed(event.usage, spent)) 

163 

164 

165def cap_aware_chat( 

166 req: CanonicalChatRequest, *, canonical_model: str, cap_chars: int 

167) -> CanonicalResponse: 

168 """Run a non-streaming chat call, re-issuing a turn that only reasoned. 

169 

170 One finished result arrives, so the first call's reasoning cannot be stopped 

171 mid-flight. A turn over the cap with no answer and no tool call is re-issued 

172 with the nudge; a turn that answered is kept. 

173 """ 

174 resp = dispatch_chat(req, canonical_model=canonical_model) 

175 if cap_chars <= 0: 

176 return resp 

177 text = "".join(b.text for b in resp.content if isinstance(b, TextBlock)) 

178 reasoning, answer = split_reasoning(text) 

179 tool_uses = [b for b in resp.content if isinstance(b, ToolUseBlock)] 

180 if len(reasoning) <= cap_chars or answer.strip() or tool_uses: 

181 return resp 

182 

183 continuation = dispatch_chat(nudged_request(req), canonical_model=canonical_model) 

184 return _merged_response(resp, continuation, cap_chars=cap_chars, reasoning=reasoning) 

185 

186 

187def _merged_response( 

188 capped: CanonicalResponse, 

189 continuation: CanonicalResponse, 

190 *, 

191 cap_chars: int, 

192 reasoning: str, 

193) -> CanonicalResponse: 

194 """Fold a capped turn and its continuation into one response. 

195 

196 Reasoning is truncated to the cap; both calls' token counts are summed. 

197 """ 

198 kept = f"{THINK_OPEN_TAG}{reasoning[:cap_chars]}{_cap_notice(cap_chars)}{THINK_CLOSE_TAG}" 

199 tail = "".join(b.text for b in continuation.content if isinstance(b, TextBlock)) 

200 content: list[ContentBlock] = [TextBlock(text=f"{kept}{tail}")] 

201 content.extend(b for b in continuation.content if isinstance(b, ToolUseBlock)) 

202 return dataclasses.replace( 

203 continuation, 

204 content=content, 

205 usage=_summed(capped.usage, continuation.usage), 

206 )