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 _with_spent(event: CanonicalStreamEvent, spent: CanonicalUsage | None) -> CanonicalStreamEvent: 

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

145 

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

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

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

149 that event, and those tokens go unreported. 

150 """ 

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

152 return event 

153 return dataclasses.replace( 

154 event, 

155 usage=CanonicalUsage( 

156 input_tokens=event.usage.input_tokens + spent.input_tokens, 

157 output_tokens=event.usage.output_tokens + spent.output_tokens, 

158 ), 

159 ) 

160 

161 

162def cap_aware_chat( 

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

164) -> CanonicalResponse: 

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

166 

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

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

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

170 """ 

171 resp = dispatch_chat(req, canonical_model=canonical_model) 

172 if cap_chars <= 0: 

173 return resp 

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

175 reasoning, answer = split_reasoning(text) 

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

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

178 return resp 

179 

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

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

182 

183 

184def _merged_response( 

185 capped: CanonicalResponse, 

186 continuation: CanonicalResponse, 

187 *, 

188 cap_chars: int, 

189 reasoning: str, 

190) -> CanonicalResponse: 

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

192 

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

194 """ 

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

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

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

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

199 return dataclasses.replace( 

200 continuation, 

201 content=content, 

202 usage=CanonicalUsage( 

203 input_tokens=capped.usage.input_tokens + continuation.usage.input_tokens, 

204 output_tokens=capped.usage.output_tokens + continuation.usage.output_tokens, 

205 ), 

206 )