Coverage for src/lilbee/server/chat_dispatch/canonical.py: 100%
133 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Protocol-neutral chat request, response, stream-event, and token-count types."""
3from __future__ import annotations
5from dataclasses import dataclass
6from enum import StrEnum
7from typing import Any, Literal
10class StopReason(StrEnum):
11 """Why a canonical chat response ended."""
13 END_TURN = "end_turn"
14 MAX_TOKENS = "max_tokens"
15 TOOL_USE = "tool_use"
18@dataclass(frozen=True)
19class TextBlock:
20 """Plain-text content block."""
22 text: str
23 type: Literal["text"] = "text"
26@dataclass(frozen=True)
27class ToolUseBlock:
28 """Assistant-emitted tool invocation with parsed JSON arguments."""
30 id: str
31 name: str
32 input: dict[str, Any]
33 type: Literal["tool_use"] = "tool_use"
36@dataclass(frozen=True)
37class ToolResultBlock:
38 """Caller-supplied tool result paired to a prior ToolUseBlock by id."""
40 tool_use_id: str
41 content: list[ContentBlock]
42 is_error: bool = False
43 type: Literal["tool_result"] = "tool_result"
46ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock
49@dataclass(frozen=True)
50class CanonicalMessage:
51 """One chat turn; content is always a typed-block list."""
53 role: Literal["user", "assistant", "tool"]
54 content: list[ContentBlock]
56 @classmethod
57 def from_string(
58 cls,
59 *,
60 role: Literal["user", "assistant", "tool"],
61 text: str,
62 ) -> CanonicalMessage:
63 """Build a single-text-block message from a raw string."""
64 return cls(role=role, content=[TextBlock(text=text)])
67@dataclass(frozen=True)
68class CanonicalTool:
69 """Tool definition (JSON-Schema input shape)."""
71 name: str
72 description: str
73 input_schema: dict[str, Any]
76@dataclass(frozen=True)
77class CanonicalToolChoice:
78 """Tool-choice mode; ``tool_name`` is required only when ``mode == "tool"``."""
80 mode: Literal["auto", "any", "none", "tool"]
81 tool_name: str | None = None
83 def __post_init__(self) -> None:
84 # Without this the None reaches the provider as
85 # {"function": {"name": None}}, a malformed tool choice rather than a
86 # rejected request.
87 if self.mode == "tool" and not self.tool_name:
88 raise ValueError('CanonicalToolChoice(mode="tool") requires a tool_name')
91@dataclass(frozen=True)
92class CanonicalChatRequest:
93 """Canonical chat request consumed by the dispatch layer."""
95 model: str
96 messages: list[CanonicalMessage]
97 system: str | None = None
98 tools: list[CanonicalTool] | None = None
99 tool_choice: CanonicalToolChoice | None = None
100 temperature: float | None = None
101 top_p: float | None = None
102 top_k: int | None = None
103 max_tokens: int | None = None
104 seed: int | None = None
105 frequency_penalty: float | None = None
106 presence_penalty: float | None = None
107 stop: list[str] | None = None
108 stream: bool = False
109 # Thinking-template control (chat_template_kwargs.enable_thinking); None
110 # leaves the model's template default in place.
111 think: bool | None = None
112 # Image parts the wire request carried and the translator dropped.
113 dropped_images: int = 0
116IMAGES_ONLY_MESSAGE = (
117 "The request holds only image content, and lilbee cannot pass images to the model. "
118 "Send text with the request."
119)
122# Tells the client how many image parts the server dropped from its request.
123IMAGES_DROPPED_HEADER = "X-Lilbee-Images-Dropped"
126def images_dropped_headers(req: CanonicalChatRequest) -> dict[str, str]:
127 """Response headers that report the images dropped from *req*, or none."""
128 if not req.dropped_images:
129 return {}
130 return {IMAGES_DROPPED_HEADER: str(req.dropped_images)}
133def _blocks_carry_content(blocks: list[ContentBlock]) -> bool:
134 for block in blocks:
135 if isinstance(block, TextBlock) and block.text:
136 return True
137 if isinstance(block, ToolUseBlock):
138 return True
139 if isinstance(block, ToolResultBlock) and _blocks_carry_content(block.content):
140 return True
141 return False
144def require_content_beside_images(req: CanonicalChatRequest) -> CanonicalChatRequest:
145 """Return *req*, or raise ``ValueError`` when dropping its images left nothing to read."""
146 if req.dropped_images and not req.system:
147 carries_content = any(_blocks_carry_content(m.content) for m in req.messages)
148 if not carries_content:
149 raise ValueError(IMAGES_ONLY_MESSAGE)
150 return req
153@dataclass(frozen=True)
154class CanonicalUsage:
155 """Token-count summary for one chat response."""
157 input_tokens: int
158 output_tokens: int
159 # Prompt tokens served from the engine's cache, a subset of input_tokens.
160 cached_input_tokens: int = 0
162 @property
163 def uncached_input_tokens(self) -> int:
164 """Prompt tokens the engine did not serve from its cache."""
165 return self.input_tokens - self.cached_input_tokens
168class TokenCountAccuracy(StrEnum):
169 """Whether a prompt token count was measured on the backend or estimated."""
171 EXACT = "exact"
172 ESTIMATED = "estimated"
175@dataclass(frozen=True)
176class PromptTokenCount:
177 """Tokens a prompt costs, and how the number was arrived at."""
179 tokens: int
180 accuracy: TokenCountAccuracy
183@dataclass(frozen=True)
184class CanonicalResponse:
185 """Canonical non-streaming chat response."""
187 id: str
188 model: str
189 content: list[ContentBlock]
190 stop_reason: StopReason
191 usage: CanonicalUsage
194@dataclass(frozen=True)
195class MessageStart:
196 """Stream prelude carrying the message id and model ref."""
198 id: str
199 model: str
202@dataclass(frozen=True)
203class ContentBlockStart:
204 """Opens a fresh content block at ``index`` with an initial shell."""
206 index: int
207 block: ContentBlock
210@dataclass(frozen=True)
211class TextDelta:
212 """One text-token delta within an open text block."""
214 text: str
217@dataclass(frozen=True)
218class ToolUseDelta:
219 """Accumulating JSON fragment within an open tool-use block."""
221 partial_json: str
224@dataclass(frozen=True)
225class ContentBlockDelta:
226 """Delta payload routed to the content block at ``index``."""
228 index: int
229 delta: TextDelta | ToolUseDelta
232@dataclass(frozen=True)
233class ContentBlockStop:
234 """Closes the content block at ``index``."""
236 index: int
239@dataclass(frozen=True)
240class MessageDelta:
241 """Trailing metadata; either field may carry a value, never both required."""
243 stop_reason: StopReason | None = None
244 usage: CanonicalUsage | None = None
247@dataclass(frozen=True)
248class MessageStop:
249 """Stream terminator."""
252CanonicalStreamEvent = (
253 MessageStart
254 | ContentBlockStart
255 | ContentBlockDelta
256 | ContentBlockStop
257 | MessageDelta
258 | MessageStop
259)