Coverage for src/lilbee/server/chat_completions_api/translate.py: 100%
168 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Translation between OpenAI chat-completions models and the canonical types."""
3from __future__ import annotations
5import json
6import time
7from collections.abc import AsyncIterator
8from typing import Literal, assert_never
10from lilbee.core.config.enums import ReasoningMode
11from lilbee.retrieval.reasoning import StreamToken, TagParser, split_reasoning
12from lilbee.server.chat_completions_api.models import (
13 CompletionsImageContent,
14 CompletionsMessage,
15 CompletionsNamedToolChoice,
16 CompletionsRequest,
17 CompletionsResponse,
18 CompletionsResponseChoice,
19 CompletionsResponseMessage,
20 CompletionsResponseToolCall,
21 CompletionsResponseToolCallFunction,
22 CompletionsStreamChoice,
23 CompletionsStreamChunk,
24 CompletionsStreamDelta,
25 CompletionsStreamToolCall,
26 CompletionsStreamToolCallFunction,
27 CompletionsTextContent,
28 CompletionsTool,
29 CompletionsUsage,
30 FinishReason,
31 PromptTokensDetails,
32 ToolChoiceMode,
33)
34from lilbee.server.chat_dispatch.canonical import (
35 CanonicalChatRequest,
36 CanonicalMessage,
37 CanonicalResponse,
38 CanonicalStreamEvent,
39 CanonicalTool,
40 CanonicalToolChoice,
41 CanonicalUsage,
42 ContentBlock,
43 ContentBlockDelta,
44 ContentBlockStart,
45 ContentBlockStop,
46 MessageDelta,
47 MessageStart,
48 MessageStop,
49 StopReason,
50 TextBlock,
51 TextDelta,
52 ToolResultBlock,
53 ToolUseBlock,
54 ToolUseDelta,
55 require_content_beside_images,
56)
57from lilbee.server.chat_dispatch.tool_args import parse_tool_arguments
59_TOOL_CHOICE_MODES: dict[ToolChoiceMode, Literal["auto", "any", "none"]] = {
60 ToolChoiceMode.AUTO: "auto",
61 ToolChoiceMode.NONE: "none",
62 ToolChoiceMode.REQUIRED: "any",
63}
65_STOP_REASON_TO_FINISH: dict[StopReason, FinishReason] = {
66 StopReason.END_TURN: FinishReason.STOP,
67 StopReason.MAX_TOKENS: FinishReason.LENGTH,
68 StopReason.TOOL_USE: FinishReason.TOOL_CALLS,
69}
71_TOOL_CALL_ID_REQUIRED = "A tool message requires tool_call_id naming the call it answers."
74def completions_to_canonical_request(
75 request: CompletionsRequest, *, mode: ReasoningMode = ReasoningMode.SEPARATE
76) -> CanonicalChatRequest:
77 """Translate a validated ``CompletionsRequest`` to the canonical request."""
78 system_parts: list[str] = []
79 messages: list[CanonicalMessage] = []
80 for msg in request.messages:
81 if msg.role == "system":
82 system_parts.append(_system_text(msg))
83 continue
84 messages.append(_message_from_request(msg))
86 req = CanonicalChatRequest(
87 model=request.model,
88 messages=messages,
89 dropped_images=_count_images(request.messages),
90 system="\n\n".join(system_parts) if system_parts else None,
91 tools=_tools_from_request(request.tools),
92 tool_choice=_tool_choice_from_request(request.tool_choice),
93 temperature=request.temperature,
94 top_p=request.top_p,
95 top_k=request.top_k,
96 max_tokens=request.max_tokens,
97 seed=request.seed,
98 frequency_penalty=request.frequency_penalty,
99 presence_penalty=request.presence_penalty,
100 stop=_stop_from_request(request.stop),
101 stream=request.stream,
102 # OFF asks the template to skip thinking; the other modes only change
103 # presentation, so the template default stands.
104 think=False if mode is ReasoningMode.OFF else None,
105 )
106 return require_content_beside_images(req)
109def _count_images(messages: list[CompletionsMessage]) -> int:
110 """Image parts across *messages*, whatever the role."""
111 return sum(
112 isinstance(part, CompletionsImageContent)
113 for msg in messages
114 if isinstance(msg.content, list)
115 for part in msg.content
116 )
119def canonical_to_completions_response(
120 resp: CanonicalResponse, *, response_id: str, mode: ReasoningMode = ReasoningMode.SEPARATE
121) -> CompletionsResponse:
122 """Translate a canonical chat response to the OpenAI ``chat.completion`` model."""
123 text_parts = [b.text for b in resp.content if isinstance(b, TextBlock)]
124 tool_calls = [_response_tool_call(b) for b in resp.content if isinstance(b, ToolUseBlock)]
125 # lilbee carries a reasoning model's thinking inline as <think>...</think>; the
126 # OpenAI surface reports it in its own field so agents render a clean answer.
127 # INLINE streams the thinking as ordinary content for clients that never
128 # render reasoning_content -- with the tag markers stripped, because a
129 # client that ignores reasoning_content renders raw <think> text literally.
130 # OFF keeps the split, because a template that ignores enable_thinking
131 # still thinks.
132 if mode is ReasoningMode.INLINE:
133 reasoning, answer = split_reasoning("".join(text_parts))
134 if reasoning:
135 answer = f"{reasoning}\n\n{answer}" if answer else reasoning
136 reasoning = ""
137 else:
138 reasoning, answer = split_reasoning("".join(text_parts))
139 content: str | None = answer if answer or not tool_calls else None
141 total = resp.usage.input_tokens + resp.usage.output_tokens
142 return CompletionsResponse(
143 id=response_id,
144 created=int(time.time()),
145 model=resp.model,
146 choices=[
147 CompletionsResponseChoice(
148 index=0,
149 message=CompletionsResponseMessage(
150 content=content,
151 reasoning_content=reasoning or None,
152 tool_calls=tool_calls or None,
153 ),
154 finish_reason=_STOP_REASON_TO_FINISH[resp.stop_reason],
155 )
156 ],
157 usage=CompletionsUsage(
158 prompt_tokens=resp.usage.input_tokens,
159 completion_tokens=resp.usage.output_tokens,
160 total_tokens=total,
161 prompt_tokens_details=PromptTokensDetails(cached_tokens=resp.usage.cached_input_tokens),
162 ),
163 )
166class _StreamMapper:
167 """Per-stream state for the canonical-to-OpenAI chunk converter."""
169 def __init__(self, *, mode: ReasoningMode = ReasoningMode.SEPARATE) -> None:
170 self._role_emitted = False
171 self._tool_index_for_block: dict[int, int] = {}
172 self._next_tool_index = 0
173 # INLINE routes thinking into content instead of reasoning_content, with
174 # the <think> markers stripped: a client that ignores reasoning_content
175 # renders raw tags literally, which is exactly what INLINE exists to
176 # avoid. The same stateful parse handles a tag split across deltas.
177 self._inline = mode is ReasoningMode.INLINE
178 # Splits lilbee's inline <think> text into its own delta field. Stateful
179 # because a tag can arrive split across deltas.
180 self._reasoning = TagParser(show=True)
182 def block_start(self, event: ContentBlockStart) -> CompletionsStreamDelta | None:
183 # The first delta must carry role:assistant for OpenAI-SDK accumulation,
184 # whether the response opens with text or a tool call.
185 role: Literal["assistant"] | None = None
186 if not self._role_emitted:
187 self._role_emitted = True
188 role = "assistant"
189 if isinstance(event.block, TextBlock):
190 return CompletionsStreamDelta(role=role) if role is not None else None
191 if isinstance(event.block, ToolUseBlock):
192 tool_index = self._next_tool_index
193 self._tool_index_for_block[event.index] = tool_index
194 self._next_tool_index += 1
195 return CompletionsStreamDelta(
196 role=role,
197 tool_calls=[_tool_call_open(tool_index, event.block.id, event.block.name)],
198 )
199 return None
201 def block_delta(self, event: ContentBlockDelta) -> CompletionsStreamDelta | None:
202 if isinstance(event.delta, TextDelta):
203 return self._text_delta(self._reasoning.feed(event.delta.text))
204 if isinstance(event.delta, ToolUseDelta):
205 # A delta for a block we never saw start is a provider quirk, not a
206 # server fault; a bare subscript turned it into a KeyError that
207 # surfaced as a stream-level internal error.
208 tool_index = self._tool_index_for_block.get(event.index)
209 if tool_index is None:
210 return None
211 return CompletionsStreamDelta(
212 tool_calls=[_tool_call_args(tool_index, event.delta.partial_json)],
213 )
214 return None
216 def block_stop(self) -> CompletionsStreamDelta | None:
217 """Emit whatever the reasoning splitter still holds (a partial or unclosed tag)."""
218 remaining = self._reasoning.flush()
219 return self._text_delta([remaining] if remaining else [])
221 def _text_delta(self, tokens: list[StreamToken]) -> CompletionsStreamDelta | None:
222 """One delta carrying the parsed text: split fields, or tag-free content."""
223 if self._inline:
224 # Arrival order preserved; the parser already dropped the markers.
225 text = "".join(t.content for t in tokens)
226 return CompletionsStreamDelta(content=text) if text else None
227 reasoning = "".join(t.content for t in tokens if t.is_reasoning)
228 answer = "".join(t.content for t in tokens if not t.is_reasoning)
229 if not reasoning and not answer:
230 return None
231 return CompletionsStreamDelta(
232 content=answer or None,
233 reasoning_content=reasoning or None,
234 )
237def _tool_call_open(index: int, call_id: str, name: str) -> CompletionsStreamToolCall:
238 return CompletionsStreamToolCall(
239 index=index,
240 id=call_id,
241 type="function",
242 function=CompletionsStreamToolCallFunction(name=name, arguments=""),
243 )
246def _tool_call_args(index: int, partial_json: str) -> CompletionsStreamToolCall:
247 return CompletionsStreamToolCall(
248 index=index,
249 function=CompletionsStreamToolCallFunction(arguments=partial_json),
250 )
253def _finish_reason_for(event: MessageDelta) -> FinishReason:
254 if event.stop_reason is None:
255 return FinishReason.STOP
256 return _STOP_REASON_TO_FINISH[event.stop_reason]
259async def canonical_stream_to_completions_chunks(
260 events: AsyncIterator[CanonicalStreamEvent],
261 *,
262 model: str,
263 response_id: str,
264 include_usage: bool = False,
265 mode: ReasoningMode = ReasoningMode.SEPARATE,
266) -> AsyncIterator[CompletionsStreamChunk]:
267 """Turn canonical stream events into ``CompletionsStreamChunk`` instances.
269 The trailing usage-only chunk is emitted only when *include_usage* is set,
270 matching OpenAI's ``stream_options.include_usage`` contract.
271 """
272 mapper = _StreamMapper(mode=mode)
273 # One timestamp for the whole completion. OpenAI holds created constant
274 # across a stream; recomputing it per chunk reported several creation times
275 # for one completion, and clients order or dedupe on it.
276 created = int(time.time())
277 async for event in events:
278 for chunk in _chunks_for_event(
279 event,
280 mapper,
281 model=model,
282 response_id=response_id,
283 include_usage=include_usage,
284 created=created,
285 ):
286 yield chunk
289def _chunks_for_event(
290 event: CanonicalStreamEvent,
291 mapper: _StreamMapper,
292 *,
293 model: str,
294 response_id: str,
295 include_usage: bool,
296 created: int,
297) -> list[CompletionsStreamChunk]:
298 """The OpenAI chunks one canonical event translates to; empty for a no-op event."""
299 if isinstance(event, ContentBlockStart):
300 return _maybe_chunk(model, response_id, mapper.block_start(event), created=created)
301 if isinstance(event, ContentBlockDelta):
302 return _maybe_chunk(model, response_id, mapper.block_delta(event), created=created)
303 if isinstance(event, ContentBlockStop):
304 # Closing a text block flushes any text the reasoning splitter still
305 # buffers (an unclosed <think>, or a tag that never completed).
306 return _maybe_chunk(model, response_id, mapper.block_stop(), created=created)
307 if isinstance(event, MessageDelta):
308 return _message_delta_chunks(
309 event,
310 model=model,
311 response_id=response_id,
312 include_usage=include_usage,
313 created=created,
314 )
315 if isinstance(event, MessageStart | MessageStop):
316 # OpenAI's wire format has no equivalent: MessageStart carries metadata
317 # already encoded in the chunk header, and MessageStop is replaced by the
318 # final chunk's finish_reason.
319 return []
320 # Unreachable: the branches above exhaust CanonicalStreamEvent. This makes a new
321 # event type a type error rather than a silently dropped frame.
322 assert_never(event) # pragma: no cover
325def _message_delta_chunks(
326 event: MessageDelta, *, model: str, response_id: str, include_usage: bool, created: int
327) -> list[CompletionsStreamChunk]:
328 """The finish chunk, plus the usage-only chunk when the client asked for it."""
329 chunks = [
330 _chunk(
331 model,
332 response_id,
333 CompletionsStreamDelta(),
334 finish_reason=_finish_reason_for(event),
335 created=created,
336 )
337 ]
338 if include_usage:
339 # OpenAI's contract sends the usage-only chunk unconditionally when
340 # include_usage is set; a client blocking on it must not hang because the
341 # provider streamed no usage frame.
342 usage = event.usage or CanonicalUsage(input_tokens=0, output_tokens=0)
343 chunks.append(_usage_chunk(model, response_id, usage, created=created))
344 return chunks
347def _maybe_chunk(
348 model: str, response_id: str, delta: CompletionsStreamDelta | None, *, created: int
349) -> list[CompletionsStreamChunk]:
350 """Wrap a delta in a chunk, or nothing when the event produced no delta."""
351 return [] if delta is None else [_chunk(model, response_id, delta, created=created)]
354def _chunk(
355 model: str,
356 response_id: str,
357 delta: CompletionsStreamDelta,
358 *,
359 finish_reason: FinishReason | None = None,
360 created: int,
361) -> CompletionsStreamChunk:
362 return CompletionsStreamChunk(
363 id=response_id,
364 created=created,
365 model=model,
366 choices=[CompletionsStreamChoice(index=0, delta=delta, finish_reason=finish_reason)],
367 )
370def _usage_chunk(
371 model: str, response_id: str, usage: CanonicalUsage, *, created: int
372) -> CompletionsStreamChunk:
373 """Final include_usage chunk: empty choices, populated usage totals."""
374 total = usage.input_tokens + usage.output_tokens
375 return CompletionsStreamChunk(
376 id=response_id,
377 created=created,
378 model=model,
379 choices=[],
380 usage=CompletionsUsage(
381 prompt_tokens=usage.input_tokens,
382 completion_tokens=usage.output_tokens,
383 total_tokens=total,
384 prompt_tokens_details=PromptTokensDetails(cached_tokens=usage.cached_input_tokens),
385 ),
386 )
389def _system_text(msg: CompletionsMessage) -> str:
390 """Flatten a system message to its text parts."""
391 if isinstance(msg.content, str):
392 return msg.content
393 if isinstance(msg.content, list):
394 return "".join(
395 part.text for part in msg.content if isinstance(part, CompletionsTextContent)
396 )
397 return ""
400def _message_from_request(msg: CompletionsMessage) -> CanonicalMessage:
401 role = msg.role
402 if role == "system":
403 raise ValueError("system messages should be extracted by the caller")
404 if role == "tool":
405 if not msg.tool_call_id:
406 # Substituting "" produced a tool result no tool call can pair with,
407 # which the provider sees as a result for a call it never made.
408 raise ValueError(_TOOL_CALL_ID_REQUIRED)
409 return CanonicalMessage(
410 role="tool",
411 content=[
412 ToolResultBlock(
413 tool_use_id=msg.tool_call_id,
414 content=_tool_result_content(msg.content),
415 )
416 ],
417 )
419 blocks: list[ContentBlock] = list(_content_blocks(msg.content))
420 for call in msg.tool_calls or []:
421 blocks.append(
422 ToolUseBlock(
423 id=call.id,
424 name=call.function.name,
425 input=parse_tool_arguments(call.function.arguments),
426 )
427 )
428 return CanonicalMessage(role=role, content=blocks)
431def _content_blocks(content: str | list | None) -> list[ContentBlock]:
432 if content is None or content == "":
433 return []
434 if isinstance(content, str):
435 return [TextBlock(text=content)]
436 blocks: list[ContentBlock] = []
437 for part in content:
438 if isinstance(part, CompletionsTextContent):
439 blocks.append(TextBlock(text=part.text))
440 return blocks
443def _tool_result_content(content: str | list | None) -> list[ContentBlock]:
444 if isinstance(content, str):
445 return [TextBlock(text=content)]
446 if isinstance(content, list):
447 return _content_blocks(content)
448 return [TextBlock(text="" if content is None else str(content))]
451def _response_tool_call(block: ToolUseBlock) -> CompletionsResponseToolCall:
452 return CompletionsResponseToolCall(
453 id=block.id,
454 function=CompletionsResponseToolCallFunction(
455 name=block.name, arguments=json.dumps(block.input)
456 ),
457 )
460def _tools_from_request(tools: list[CompletionsTool] | None) -> list[CanonicalTool] | None:
461 if not tools:
462 return None
463 return [
464 CanonicalTool(
465 name=tool.function.name,
466 description=tool.function.description or "",
467 input_schema=tool.function.parameters,
468 )
469 for tool in tools
470 ]
473def _tool_choice_from_request(
474 choice: ToolChoiceMode | CompletionsNamedToolChoice | None,
475) -> CanonicalToolChoice | None:
476 if choice is None:
477 return None
478 if isinstance(choice, ToolChoiceMode):
479 return CanonicalToolChoice(mode=_TOOL_CHOICE_MODES[choice])
480 return CanonicalToolChoice(mode="tool", tool_name=choice.function.name)
483def _stop_from_request(stop: str | list[str] | None) -> list[str] | None:
484 if stop is None:
485 return None
486 if isinstance(stop, str):
487 return [stop]
488 return list(stop)