Coverage for src/lilbee/server/chat_dispatch/dispatch.py: 100%
201 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Canonical chat dispatch: canonical request to provider call to canonical response."""
3from __future__ import annotations
5import asyncio
6import json
7import logging
8import uuid
9from collections.abc import AsyncIterator, Iterator
10from enum import StrEnum
11from typing import Any, Literal
13from lilbee.app.services import get_services
14from lilbee.core.config import cfg
15from lilbee.providers.base import (
16 ChatResult,
17 ChatStreamItem,
18 FinishReason,
19 ProviderError,
20 ProviderErrorKind,
21 StreamFinish,
22 TokenUsage,
23 ToolCallDelta,
24)
25from lilbee.providers.model_ref import parse_model_ref
26from lilbee.providers.roles import WorkerRole, configured_model_message
27from lilbee.server.chat_dispatch.canonical import (
28 CanonicalChatRequest,
29 CanonicalMessage,
30 CanonicalResponse,
31 CanonicalStreamEvent,
32 CanonicalTool,
33 CanonicalToolChoice,
34 CanonicalUsage,
35 ContentBlock,
36 ContentBlockDelta,
37 ContentBlockStart,
38 ContentBlockStop,
39 MessageDelta,
40 MessageStart,
41 MessageStop,
42 StopReason,
43 TextBlock,
44 TextDelta,
45 ToolResultBlock,
46 ToolUseBlock,
47 ToolUseDelta,
48)
49from lilbee.server.chat_dispatch.capability import model_supports_tools
50from lilbee.server.chat_dispatch.tool_args import parse_tool_arguments
52log = logging.getLogger(__name__)
55class ModelNotFoundError(Exception):
56 """Raised when the requested model is not installed or reachable."""
58 def __init__(self, model: str) -> None:
59 self.model = model
60 super().__init__(
61 f"Model {model!r} is not installed. Run 'lilbee model list' to see "
62 f"installed models, or 'lilbee model pull {model}' to download it."
63 )
66class ModelDoesNotSupportToolsError(Exception):
67 """Raised when the request carries tools but the model template cannot use them."""
69 def __init__(self, model: str) -> None:
70 self.model = model
71 super().__init__(
72 f"Model {model!r} does not support tool calls. Pick a chat model "
73 f"with a tool-aware chat template, or remove tools from the request."
74 )
77_FINISH_REASON_TO_STOP: dict[FinishReason, StopReason] = {
78 FinishReason.STOP: StopReason.END_TURN,
79 FinishReason.LENGTH: StopReason.MAX_TOKENS,
80 FinishReason.TOOL_CALLS: StopReason.TOOL_USE,
81 FinishReason.CONTENT_FILTER: StopReason.END_TURN,
82}
84_CanonicalChoiceMode = Literal["auto", "any", "none"]
85_ProviderChoiceMode = Literal["auto", "required", "none"]
87_TOOL_CHOICE_MODES: dict[_CanonicalChoiceMode, _ProviderChoiceMode] = {
88 "auto": "auto",
89 "any": "required",
90 "none": "none",
91}
94class _OpenBlockKind(StrEnum):
95 NONE = "none"
96 TEXT = "text"
97 TOOL = "tool"
100def _provider_chat_kwargs(req: CanonicalChatRequest, canonical_model: str) -> dict[str, Any]:
101 """Shared provider.chat keyword arguments for both stream and non-stream paths."""
102 return {
103 "messages": _provider_messages(req),
104 "options": _provider_options(req),
105 "model": canonical_model,
106 "tools": _provider_tools(req.tools),
107 "tool_choice": _provider_tool_choice(req.tool_choice),
108 }
111def _stop_reason_for(result: ChatResult) -> StopReason:
112 """Closing stop reason for a non-streaming result.
114 Tool calls win over the reported finish reason. FinishReason.coerce falls
115 back to STOP for a missing or unknown value, so a provider that returns
116 tool calls without saying so produced tool_use content under end_turn, and
117 a client reading stop_reason decides whether to run the tools. The
118 streaming path already refuses the same downgrade.
119 """
120 if result.tool_calls:
121 return StopReason.TOOL_USE
122 return _FINISH_REASON_TO_STOP.get(result.finish_reason, StopReason.END_TURN)
125def _content_blocks_from_result(result: ChatResult) -> list[ContentBlock]:
126 """Build canonical content blocks from a non-streaming provider result."""
127 content: list[ContentBlock] = []
128 if result.text:
129 content.append(TextBlock(text=result.text))
130 for call in result.tool_calls:
131 content.append(
132 ToolUseBlock(
133 id=call.id or _new_call_id(),
134 name=call.name,
135 input=parse_tool_arguments(call.arguments),
136 )
137 )
138 return content
141def dispatch_chat(
142 req: CanonicalChatRequest, *, canonical_model: str | None = None
143) -> CanonicalResponse:
144 """Run a non-streaming chat request through the provider and return canonical output.
146 Pass *canonical_model* when the caller has already run
147 :func:`preflight_chat_request` (the route does, so the preflight runs once per
148 request); leave it ``None`` to resolve and validate the model here.
149 """
150 if canonical_model is None:
151 canonical_model = preflight_chat_request(req)
152 result = get_services().provider.chat(**_provider_chat_kwargs(req, canonical_model))
153 return CanonicalResponse(
154 id=_new_message_id(),
155 model=canonical_model,
156 content=_content_blocks_from_result(result),
157 stop_reason=_stop_reason_for(result),
158 usage=CanonicalUsage(
159 input_tokens=result.usage.prompt_tokens,
160 output_tokens=result.usage.completion_tokens,
161 ),
162 )
165async def dispatch_chat_stream(
166 req: CanonicalChatRequest, *, canonical_model: str | None = None
167) -> AsyncIterator[CanonicalStreamEvent]:
168 """Stream a canonical event sequence by translating provider frames on the fly.
170 Pass *canonical_model* when the caller has already run
171 :func:`preflight_chat_request` (the route does, so the preflight runs once per
172 request); leave it ``None`` to resolve and validate the model here.
173 """
174 # The preflight can do blocking HTTP model discovery when its TTL lapses, and
175 # opening the stream can issue a one-time template probe; run both in a thread
176 # so the event loop stays responsive.
177 if canonical_model is None:
178 canonical_model = await asyncio.to_thread(preflight_chat_request, req)
179 stream = await asyncio.to_thread(
180 lambda: get_services().provider.chat(
181 stream=True, **_provider_chat_kwargs(req, canonical_model)
182 )
183 )
184 try:
185 yield MessageStart(id=_new_message_id(), model=canonical_model)
186 state = _StreamState()
187 async for frame in _async_iter_provider_stream(stream):
188 for event in state.feed(frame):
189 yield event
190 for event in state.finish():
191 yield event
192 yield MessageStop()
193 finally:
194 # close() tears down the provider HTTP connection and can block; offload
195 # it like the open and per-frame reads so the event loop stays responsive.
196 await asyncio.to_thread(stream.close)
199async def _async_iter_provider_stream(
200 stream: Iterator[ChatStreamItem],
201) -> AsyncIterator[ChatStreamItem]:
202 """Iterate a provider chat stream without blocking the event loop.
204 ``LLMProvider.chat`` types a streaming result as a ClosableIterator, and
205 every provider in the tree returns a plain sync generator; iterating one
206 inline on the event loop would block, so each ``next()`` runs in a worker
207 thread via ``asyncio.to_thread``.
209 There used to be an async-native branch here for a provider shape that
210 does not exist. It was dead and also wrong: the caller's cleanup is
211 ``await asyncio.to_thread(stream.close)``, which an async-native stream
212 would not satisfy. Adding one means changing the Protocol and that
213 cleanup together, not restoring a branch nothing reaches.
214 """
215 while True:
216 frame = await asyncio.to_thread(_next_or_done, stream)
217 if frame is _STREAM_DONE:
218 return
219 yield frame
222_STREAM_DONE: Any = object()
223"""Sentinel returned by :func:`_next_or_done` to mean ``StopIteration``."""
226def _next_or_done(
227 stream: Iterator[ChatStreamItem],
228) -> ChatStreamItem | Any:
229 """Pull the next frame from *stream*; return ``_STREAM_DONE`` at exhaustion.
231 Raising ``StopIteration`` inside a coroutine becomes ``RuntimeError`` per
232 PEP 479; this helper converts that signal into a sentinel value the async
233 caller can branch on.
234 """
235 try:
236 return next(stream)
237 except StopIteration:
238 return _STREAM_DONE
241class _StreamState:
242 """Tracks open content blocks so deltas land in the right index."""
244 def __init__(self) -> None:
245 self._open: _OpenBlockKind = _OpenBlockKind.NONE
246 self._index: int = -1
247 self._tool_index: int | None = None
248 # Provider tool index -> the (id, name) its first delta carried.
249 # Continuation deltas typically carry neither.
250 self._tool_identity: dict[int, tuple[str, str]] = {}
251 self._stop_reason: StopReason = StopReason.END_TURN
252 self._usage: TokenUsage | None = None
254 def feed(self, frame: ChatStreamItem) -> Iterator[CanonicalStreamEvent]:
255 if isinstance(frame, str):
256 yield from self._feed_text(frame)
257 elif isinstance(frame, TokenUsage):
258 # Terminator-only frame: carries token totals, no content. Stash it
259 # so finish() can attach the counts to the closing MessageDelta.
260 self._usage = frame
261 elif isinstance(frame, StreamFinish):
262 self._feed_finish(frame)
263 else:
264 yield from self._feed_tool(frame)
266 def finish(self) -> Iterator[CanonicalStreamEvent]:
267 if self._open != _OpenBlockKind.NONE:
268 yield ContentBlockStop(index=self._index)
269 self._open = _OpenBlockKind.NONE
270 usage = (
271 CanonicalUsage(
272 input_tokens=self._usage.prompt_tokens,
273 output_tokens=self._usage.completion_tokens,
274 )
275 if self._usage is not None
276 else None
277 )
278 yield MessageDelta(stop_reason=self._stop_reason, usage=usage)
280 def _feed_finish(self, frame: StreamFinish) -> None:
281 # The finish frame sets the closing stop reason (e.g. MAX_TOKENS on a
282 # length truncation). A tool-call stream already settled on TOOL_USE via
283 # the deltas, so never let a trailing finish frame downgrade that.
284 if self._stop_reason is StopReason.TOOL_USE:
285 return
286 self._stop_reason = _FINISH_REASON_TO_STOP.get(frame.reason, StopReason.END_TURN)
288 def _feed_text(self, text: str) -> Iterator[CanonicalStreamEvent]:
289 if self._open != _OpenBlockKind.TEXT:
290 yield from self._close_current()
291 self._index += 1
292 self._open = _OpenBlockKind.TEXT
293 yield ContentBlockStart(index=self._index, block=TextBlock(text=""))
294 yield ContentBlockDelta(index=self._index, delta=TextDelta(text=text))
296 def _feed_tool(self, frame: ToolCallDelta) -> Iterator[CanonicalStreamEvent]:
297 self._stop_reason = StopReason.TOOL_USE
298 is_new_call = self._open != _OpenBlockKind.TOOL or frame.index != self._tool_index
299 if is_new_call:
300 yield from self._close_current()
301 self._index += 1
302 self._open = _OpenBlockKind.TOOL
303 self._tool_index = frame.index
304 yield ContentBlockStart(
305 index=self._index,
306 block=ToolUseBlock(**self._tool_block_fields(frame)),
307 )
308 if frame.arguments_delta is not None:
309 yield ContentBlockDelta(
310 index=self._index,
311 delta=ToolUseDelta(partial_json=frame.arguments_delta),
312 )
314 def _tool_block_fields(self, frame: ToolCallDelta) -> dict[str, Any]:
315 """Identity for the block opening on *frame*, remembered per tool index.
317 A text frame between two argument deltas of one call (streamed
318 reasoning surfaced as text, say) closes the open tool block, so the
319 next delta for the same call has to open a second block. Continuation
320 deltas carry no id and no name, so that block used to get a fresh
321 synthetic id and an empty name, splitting one logical call across two
322 blocks the second of which matched no tool. Reusing the identity the
323 call already announced at least leaves both blocks stitchable by id.
324 """
325 known = self._tool_identity.get(frame.index)
326 identity = (
327 frame.id or (known[0] if known else _new_call_id()),
328 frame.name or (known[1] if known else ""),
329 )
330 self._tool_identity[frame.index] = identity
331 return {"id": identity[0], "name": identity[1], "input": {}}
333 def _close_current(self) -> Iterator[CanonicalStreamEvent]:
334 if self._open != _OpenBlockKind.NONE:
335 yield ContentBlockStop(index=self._index)
336 self._open = _OpenBlockKind.NONE
339def _resolve_canonical_model(model: str) -> str:
340 """Return the canonical ref for *model*, or raise ``ModelNotFoundError``.
342 Consults the cached union of native + remote + frontier refs on
343 Services, so an Ollama-managed model resolves the same way a locally
344 installed GGUF does. A bare ``name:tag`` matches the corresponding
345 ``ollama/<name:tag>`` entry when one exists in the discovered set.
346 """
347 canonical = get_services().known_models.resolve(model)
348 if canonical is None:
349 raise ModelNotFoundError(model)
350 return canonical
353def _ensure_tool_capability(req: CanonicalChatRequest, model: str) -> None:
354 if req.tools and not model_supports_tools(model):
355 raise ModelDoesNotSupportToolsError(model)
358def _ensure_configured_local_model(canonical: str) -> None:
359 """Reject a local-route model that is not the configured chat model.
361 Mirrors the fleet's own configured-model guard (which stays in place as
362 defense in depth for direct provider users) so streaming clients get a
363 clean 400 before headers instead of an SSE error frame mid-stream.
364 """
365 if not parse_model_ref(canonical).is_local or canonical == cfg.chat_model:
366 return
367 raise ProviderError(
368 configured_model_message(WorkerRole.CHAT, cfg.chat_model, canonical),
369 kind=ProviderErrorKind.BAD_REQUEST,
370 )
373def preflight_chat_request(req: CanonicalChatRequest) -> str:
374 """Synchronously validate *req* before any streaming response starts.
376 Raises ``ModelNotFoundError``, ``ModelDoesNotSupportToolsError``, or a
377 ``BAD_REQUEST`` ``ProviderError`` so the route layer can return a real
378 4xx HTTP status instead of burying the failure in an SSE error frame
379 after headers flush. Returns the resolved canonical model ref.
380 """
381 canonical = _resolve_canonical_model(req.model)
382 _ensure_configured_local_model(canonical)
383 _ensure_tool_capability(req, canonical)
384 return canonical
387def _provider_messages(req: CanonicalChatRequest) -> list[dict[str, Any]]:
388 """Flatten canonical messages to the OpenAI-shaped wire format the provider speaks."""
389 out: list[dict[str, Any]] = []
390 if req.system is not None:
391 out.append({"role": "system", "content": req.system})
392 for msg in req.messages:
393 out.extend(_translate_message(msg))
394 return out
397def _translate_message(msg: CanonicalMessage) -> list[dict[str, Any]]:
398 text_parts = [b.text for b in msg.content if isinstance(b, TextBlock)]
399 tool_uses = [b for b in msg.content if isinstance(b, ToolUseBlock)]
400 tool_results = [b for b in msg.content if isinstance(b, ToolResultBlock)]
401 text = "".join(text_parts)
403 # One ``tool`` wire-message per result block; tool_call_id pairs it back to
404 # the originating ToolUseBlock. Text blocks in the same canonical message
405 # follow as their own content message rather than being dropped.
406 out: list[dict[str, Any]] = [
407 {
408 "role": "tool",
409 "tool_call_id": block.tool_use_id,
410 "content": _flatten_text(block.content),
411 }
412 for block in tool_results
413 ]
414 if tool_uses:
415 out.append(
416 {
417 "role": msg.role,
418 "content": text,
419 "tool_calls": [
420 {
421 "id": tu.id,
422 "type": "function",
423 "function": {
424 "name": tu.name,
425 "arguments": json.dumps(tu.input),
426 },
427 }
428 for tu in tool_uses
429 ],
430 }
431 )
432 elif text or not tool_results:
433 out.append({"role": msg.role, "content": text})
434 return out
437def _flatten_text(blocks: list[ContentBlock]) -> str:
438 return "".join(b.text for b in blocks if isinstance(b, TextBlock))
441def _provider_tools(
442 tools: list[CanonicalTool] | None,
443) -> list[dict[str, Any]] | None:
444 if not tools:
445 return None
446 return [
447 {
448 "type": "function",
449 "function": {
450 "name": tool.name,
451 "description": tool.description,
452 "parameters": tool.input_schema,
453 },
454 }
455 for tool in tools
456 ]
459def _provider_tool_choice(
460 choice: CanonicalToolChoice | None,
461) -> str | dict[str, Any] | None:
462 if choice is None:
463 return None
464 if choice.mode == "tool":
465 return {"type": "function", "function": {"name": choice.tool_name}}
466 return _TOOL_CHOICE_MODES[choice.mode]
469def _provider_options(req: CanonicalChatRequest) -> dict[str, Any] | None:
470 out: dict[str, Any] = {}
471 if req.temperature is not None:
472 out["temperature"] = req.temperature
473 if req.top_p is not None:
474 out["top_p"] = req.top_p
475 if req.top_k is not None:
476 out["top_k"] = req.top_k
477 if req.max_tokens is not None:
478 out["num_predict"] = req.max_tokens
479 if req.seed is not None:
480 out["seed"] = req.seed
481 if req.frequency_penalty is not None:
482 out["frequency_penalty"] = req.frequency_penalty
483 if req.presence_penalty is not None:
484 out["presence_penalty"] = req.presence_penalty
485 if req.stop is not None:
486 out["stop"] = req.stop
487 return out or None
490def _new_call_id() -> str:
491 return f"call_{uuid.uuid4().hex[:24]}"
494def _new_message_id() -> str:
495 return f"msg_{uuid.uuid4().hex[:24]}"