Coverage for src/lilbee/server/chat_completions_api/translate.py: 100%
160 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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.retrieval.reasoning import StreamToken, TagParser, split_reasoning
11from lilbee.server.chat_completions_api.models import (
12 CompletionsImageContent,
13 CompletionsMessage,
14 CompletionsNamedToolChoice,
15 CompletionsRequest,
16 CompletionsResponse,
17 CompletionsResponseChoice,
18 CompletionsResponseMessage,
19 CompletionsResponseToolCall,
20 CompletionsResponseToolCallFunction,
21 CompletionsStreamChoice,
22 CompletionsStreamChunk,
23 CompletionsStreamDelta,
24 CompletionsStreamToolCall,
25 CompletionsStreamToolCallFunction,
26 CompletionsTextContent,
27 CompletionsTool,
28 CompletionsUsage,
29 FinishReason,
30 ToolChoiceMode,
31)
32from lilbee.server.chat_dispatch.canonical import (
33 CanonicalChatRequest,
34 CanonicalMessage,
35 CanonicalResponse,
36 CanonicalStreamEvent,
37 CanonicalTool,
38 CanonicalToolChoice,
39 CanonicalUsage,
40 ContentBlock,
41 ContentBlockDelta,
42 ContentBlockStart,
43 ContentBlockStop,
44 MessageDelta,
45 MessageStart,
46 MessageStop,
47 StopReason,
48 TextBlock,
49 TextDelta,
50 ToolResultBlock,
51 ToolUseBlock,
52 ToolUseDelta,
53)
54from lilbee.server.chat_dispatch.tool_args import parse_tool_arguments
56_TOOL_CHOICE_MODES: dict[ToolChoiceMode, Literal["auto", "any", "none"]] = {
57 ToolChoiceMode.AUTO: "auto",
58 ToolChoiceMode.NONE: "none",
59 ToolChoiceMode.REQUIRED: "any",
60}
62_STOP_REASON_TO_FINISH: dict[StopReason, FinishReason] = {
63 StopReason.END_TURN: FinishReason.STOP,
64 StopReason.MAX_TOKENS: FinishReason.LENGTH,
65 StopReason.TOOL_USE: FinishReason.TOOL_CALLS,
66}
68_TOOL_CALL_ID_REQUIRED = "A tool message requires tool_call_id naming the call it answers."
69_IMAGE_CONTENT_UNSUPPORTED = (
70 "Image content is not supported by /v1/chat/completions yet. Send a text-only request."
71)
74def completions_to_canonical_request(request: CompletionsRequest) -> CanonicalChatRequest:
75 """Translate a validated ``CompletionsRequest`` to the canonical request."""
76 system_parts: list[str] = []
77 messages: list[CanonicalMessage] = []
78 for msg in request.messages:
79 if msg.role == "system":
80 system_parts.append(_system_text(msg))
81 continue
82 messages.append(_message_from_request(msg))
84 return CanonicalChatRequest(
85 model=request.model,
86 messages=messages,
87 system="\n\n".join(system_parts) if system_parts else None,
88 tools=_tools_from_request(request.tools),
89 tool_choice=_tool_choice_from_request(request.tool_choice),
90 temperature=request.temperature,
91 top_p=request.top_p,
92 top_k=request.top_k,
93 max_tokens=request.max_tokens,
94 seed=request.seed,
95 frequency_penalty=request.frequency_penalty,
96 presence_penalty=request.presence_penalty,
97 stop=_stop_from_request(request.stop),
98 stream=request.stream,
99 )
102def canonical_to_completions_response(
103 resp: CanonicalResponse, *, response_id: str
104) -> CompletionsResponse:
105 """Translate a canonical chat response to the OpenAI ``chat.completion`` model."""
106 text_parts = [b.text for b in resp.content if isinstance(b, TextBlock)]
107 tool_calls = [_response_tool_call(b) for b in resp.content if isinstance(b, ToolUseBlock)]
108 # lilbee carries a reasoning model's thinking inline as <think>...</think>; the
109 # OpenAI surface reports it in its own field so agents render a clean answer.
110 reasoning, answer = split_reasoning("".join(text_parts))
111 content: str | None = answer if answer or not tool_calls else None
113 total = resp.usage.input_tokens + resp.usage.output_tokens
114 return CompletionsResponse(
115 id=response_id,
116 created=int(time.time()),
117 model=resp.model,
118 choices=[
119 CompletionsResponseChoice(
120 index=0,
121 message=CompletionsResponseMessage(
122 content=content,
123 reasoning_content=reasoning or None,
124 tool_calls=tool_calls or None,
125 ),
126 finish_reason=_STOP_REASON_TO_FINISH[resp.stop_reason],
127 )
128 ],
129 usage=CompletionsUsage(
130 prompt_tokens=resp.usage.input_tokens,
131 completion_tokens=resp.usage.output_tokens,
132 total_tokens=total,
133 ),
134 )
137class _StreamMapper:
138 """Per-stream state for the canonical-to-OpenAI chunk converter."""
140 def __init__(self) -> None:
141 self._role_emitted = False
142 self._tool_index_for_block: dict[int, int] = {}
143 self._next_tool_index = 0
144 # Splits lilbee's inline <think> text into its own delta field. Stateful
145 # because a tag can arrive split across deltas.
146 self._reasoning = TagParser(show=True)
148 def block_start(self, event: ContentBlockStart) -> CompletionsStreamDelta | None:
149 # The first delta must carry role:assistant for OpenAI-SDK accumulation,
150 # whether the response opens with text or a tool call.
151 role: Literal["assistant"] | None = None
152 if not self._role_emitted:
153 self._role_emitted = True
154 role = "assistant"
155 if isinstance(event.block, TextBlock):
156 return CompletionsStreamDelta(role=role) if role is not None else None
157 if isinstance(event.block, ToolUseBlock):
158 tool_index = self._next_tool_index
159 self._tool_index_for_block[event.index] = tool_index
160 self._next_tool_index += 1
161 return CompletionsStreamDelta(
162 role=role,
163 tool_calls=[_tool_call_open(tool_index, event.block.id, event.block.name)],
164 )
165 return None
167 def block_delta(self, event: ContentBlockDelta) -> CompletionsStreamDelta | None:
168 if isinstance(event.delta, TextDelta):
169 return self._text_delta(self._reasoning.feed(event.delta.text))
170 if isinstance(event.delta, ToolUseDelta):
171 # A delta for a block we never saw start is a provider quirk, not a
172 # server fault; a bare subscript turned it into a KeyError that
173 # surfaced as a stream-level internal error.
174 tool_index = self._tool_index_for_block.get(event.index)
175 if tool_index is None:
176 return None
177 return CompletionsStreamDelta(
178 tool_calls=[_tool_call_args(tool_index, event.delta.partial_json)],
179 )
180 return None
182 def block_stop(self) -> CompletionsStreamDelta | None:
183 """Emit whatever the reasoning splitter still holds (a partial or unclosed tag)."""
184 remaining = self._reasoning.flush()
185 return self._text_delta([remaining] if remaining else [])
187 def _text_delta(self, tokens: list[StreamToken]) -> CompletionsStreamDelta | None:
188 """One delta carrying the reasoning and answer text split out of *tokens*."""
189 reasoning = "".join(t.content for t in tokens if t.is_reasoning)
190 answer = "".join(t.content for t in tokens if not t.is_reasoning)
191 if not reasoning and not answer:
192 return None
193 return CompletionsStreamDelta(
194 content=answer or None,
195 reasoning_content=reasoning or None,
196 )
199def _tool_call_open(index: int, call_id: str, name: str) -> CompletionsStreamToolCall:
200 return CompletionsStreamToolCall(
201 index=index,
202 id=call_id,
203 type="function",
204 function=CompletionsStreamToolCallFunction(name=name, arguments=""),
205 )
208def _tool_call_args(index: int, partial_json: str) -> CompletionsStreamToolCall:
209 return CompletionsStreamToolCall(
210 index=index,
211 function=CompletionsStreamToolCallFunction(arguments=partial_json),
212 )
215def _finish_reason_for(event: MessageDelta) -> FinishReason:
216 if event.stop_reason is None:
217 return FinishReason.STOP
218 return _STOP_REASON_TO_FINISH[event.stop_reason]
221async def canonical_stream_to_completions_chunks(
222 events: AsyncIterator[CanonicalStreamEvent],
223 *,
224 model: str,
225 response_id: str,
226 include_usage: bool = False,
227) -> AsyncIterator[CompletionsStreamChunk]:
228 """Turn canonical stream events into ``CompletionsStreamChunk`` instances.
230 The trailing usage-only chunk is emitted only when *include_usage* is set,
231 matching OpenAI's ``stream_options.include_usage`` contract.
232 """
233 mapper = _StreamMapper()
234 # One timestamp for the whole completion. OpenAI holds created constant
235 # across a stream; recomputing it per chunk reported several creation times
236 # for one completion, and clients order or dedupe on it.
237 created = int(time.time())
238 async for event in events:
239 for chunk in _chunks_for_event(
240 event,
241 mapper,
242 model=model,
243 response_id=response_id,
244 include_usage=include_usage,
245 created=created,
246 ):
247 yield chunk
250def _chunks_for_event(
251 event: CanonicalStreamEvent,
252 mapper: _StreamMapper,
253 *,
254 model: str,
255 response_id: str,
256 include_usage: bool,
257 created: int,
258) -> list[CompletionsStreamChunk]:
259 """The OpenAI chunks one canonical event translates to; empty for a no-op event."""
260 if isinstance(event, ContentBlockStart):
261 return _maybe_chunk(model, response_id, mapper.block_start(event), created=created)
262 if isinstance(event, ContentBlockDelta):
263 return _maybe_chunk(model, response_id, mapper.block_delta(event), created=created)
264 if isinstance(event, ContentBlockStop):
265 # Closing a text block flushes any text the reasoning splitter still
266 # buffers (an unclosed <think>, or a tag that never completed).
267 return _maybe_chunk(model, response_id, mapper.block_stop(), created=created)
268 if isinstance(event, MessageDelta):
269 return _message_delta_chunks(
270 event,
271 model=model,
272 response_id=response_id,
273 include_usage=include_usage,
274 created=created,
275 )
276 if isinstance(event, MessageStart | MessageStop):
277 # OpenAI's wire format has no equivalent: MessageStart carries metadata
278 # already encoded in the chunk header, and MessageStop is replaced by the
279 # final chunk's finish_reason.
280 return []
281 # Unreachable: the branches above exhaust CanonicalStreamEvent. This makes a new
282 # event type a type error rather than a silently dropped frame.
283 assert_never(event) # pragma: no cover
286def _message_delta_chunks(
287 event: MessageDelta, *, model: str, response_id: str, include_usage: bool, created: int
288) -> list[CompletionsStreamChunk]:
289 """The finish chunk, plus the usage-only chunk when the client asked for it."""
290 chunks = [
291 _chunk(
292 model,
293 response_id,
294 CompletionsStreamDelta(),
295 finish_reason=_finish_reason_for(event),
296 created=created,
297 )
298 ]
299 if include_usage:
300 # OpenAI's contract sends the usage-only chunk unconditionally when
301 # include_usage is set; a client blocking on it must not hang because the
302 # provider streamed no usage frame.
303 usage = event.usage or CanonicalUsage(input_tokens=0, output_tokens=0)
304 chunks.append(_usage_chunk(model, response_id, usage, created=created))
305 return chunks
308def _maybe_chunk(
309 model: str, response_id: str, delta: CompletionsStreamDelta | None, *, created: int
310) -> list[CompletionsStreamChunk]:
311 """Wrap a delta in a chunk, or nothing when the event produced no delta."""
312 return [] if delta is None else [_chunk(model, response_id, delta, created=created)]
315def _chunk(
316 model: str,
317 response_id: str,
318 delta: CompletionsStreamDelta,
319 *,
320 finish_reason: FinishReason | None = None,
321 created: int,
322) -> CompletionsStreamChunk:
323 return CompletionsStreamChunk(
324 id=response_id,
325 created=created,
326 model=model,
327 choices=[CompletionsStreamChoice(index=0, delta=delta, finish_reason=finish_reason)],
328 )
331def _usage_chunk(
332 model: str, response_id: str, usage: CanonicalUsage, *, created: int
333) -> CompletionsStreamChunk:
334 """Final include_usage chunk: empty choices, populated usage totals."""
335 total = usage.input_tokens + usage.output_tokens
336 return CompletionsStreamChunk(
337 id=response_id,
338 created=created,
339 model=model,
340 choices=[],
341 usage=CompletionsUsage(
342 prompt_tokens=usage.input_tokens,
343 completion_tokens=usage.output_tokens,
344 total_tokens=total,
345 ),
346 )
349def _system_text(msg: CompletionsMessage) -> str:
350 """Flatten a system message to text, rejecting image parts.
352 The same image part is a 400 in a user message, so silently dropping it
353 here answered as though the request had been honoured.
354 """
355 if isinstance(msg.content, str):
356 return msg.content
357 if isinstance(msg.content, list):
358 if any(isinstance(part, CompletionsImageContent) for part in msg.content):
359 raise ValueError(_IMAGE_CONTENT_UNSUPPORTED)
360 return "".join(
361 part.text for part in msg.content if isinstance(part, CompletionsTextContent)
362 )
363 return ""
366def _message_from_request(msg: CompletionsMessage) -> CanonicalMessage:
367 role = msg.role
368 if role == "system":
369 raise ValueError("system messages should be extracted by the caller")
370 if role == "tool":
371 if not msg.tool_call_id:
372 # Substituting "" produced a tool result no tool call can pair with,
373 # which the provider sees as a result for a call it never made.
374 raise ValueError(_TOOL_CALL_ID_REQUIRED)
375 return CanonicalMessage(
376 role="tool",
377 content=[
378 ToolResultBlock(
379 tool_use_id=msg.tool_call_id,
380 content=_tool_result_content(msg.content),
381 )
382 ],
383 )
385 blocks: list[ContentBlock] = list(_content_blocks(msg.content))
386 for call in msg.tool_calls or []:
387 blocks.append(
388 ToolUseBlock(
389 id=call.id,
390 name=call.function.name,
391 input=parse_tool_arguments(call.function.arguments),
392 )
393 )
394 return CanonicalMessage(role=role, content=blocks)
397def _content_blocks(content: str | list | None) -> list[ContentBlock]:
398 if content is None or content == "":
399 return []
400 if isinstance(content, str):
401 return [TextBlock(text=content)]
402 blocks: list[ContentBlock] = []
403 for part in content:
404 if isinstance(part, CompletionsTextContent):
405 blocks.append(TextBlock(text=part.text))
406 elif isinstance(part, CompletionsImageContent):
407 raise ValueError(_IMAGE_CONTENT_UNSUPPORTED)
408 return blocks
411def _tool_result_content(content: str | list | None) -> list[ContentBlock]:
412 if isinstance(content, str):
413 return [TextBlock(text=content)]
414 if isinstance(content, list):
415 return _content_blocks(content)
416 return [TextBlock(text="" if content is None else str(content))]
419def _response_tool_call(block: ToolUseBlock) -> CompletionsResponseToolCall:
420 return CompletionsResponseToolCall(
421 id=block.id,
422 function=CompletionsResponseToolCallFunction(
423 name=block.name, arguments=json.dumps(block.input)
424 ),
425 )
428def _tools_from_request(tools: list[CompletionsTool] | None) -> list[CanonicalTool] | None:
429 if not tools:
430 return None
431 return [
432 CanonicalTool(
433 name=tool.function.name,
434 description=tool.function.description or "",
435 input_schema=tool.function.parameters,
436 )
437 for tool in tools
438 ]
441def _tool_choice_from_request(
442 choice: ToolChoiceMode | CompletionsNamedToolChoice | None,
443) -> CanonicalToolChoice | None:
444 if choice is None:
445 return None
446 if isinstance(choice, ToolChoiceMode):
447 return CanonicalToolChoice(mode=_TOOL_CHOICE_MODES[choice])
448 return CanonicalToolChoice(mode="tool", tool_name=choice.function.name)
451def _stop_from_request(stop: str | list[str] | None) -> list[str] | None:
452 if stop is None:
453 return None
454 if isinstance(stop, str):
455 return [stop]
456 return list(stop)