Coverage for src/lilbee/server/chat_dispatch/dispatch.py: 100%
244 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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 PromptTokenCount,
43 StopReason,
44 TextBlock,
45 TextDelta,
46 TokenCountAccuracy,
47 ToolResultBlock,
48 ToolUseBlock,
49 ToolUseDelta,
50)
51from lilbee.server.chat_dispatch.capability import model_supports_tools
52from lilbee.server.chat_dispatch.tool_args import parse_tool_arguments
54log = logging.getLogger(__name__)
57class ModelNotFoundError(Exception):
58 """Raised when the requested model is not installed or reachable."""
60 def __init__(self, model: str) -> None:
61 self.model = model
62 super().__init__(
63 f"Model {model!r} is not installed. Run 'lilbee model list' to see "
64 f"installed models, or 'lilbee model pull {model}' to download it."
65 )
68class ModelDoesNotSupportToolsError(Exception):
69 """Raised when the request carries tools but the model template cannot use them."""
71 def __init__(self, model: str) -> None:
72 self.model = model
73 super().__init__(
74 f"Model {model!r} does not support tool calls. Pick a chat model "
75 f"with a tool-aware chat template, or remove tools from the request."
76 )
79_FINISH_REASON_TO_STOP: dict[FinishReason, StopReason] = {
80 FinishReason.STOP: StopReason.END_TURN,
81 FinishReason.LENGTH: StopReason.MAX_TOKENS,
82 FinishReason.TOOL_CALLS: StopReason.TOOL_USE,
83 FinishReason.CONTENT_FILTER: StopReason.END_TURN,
84}
86_CanonicalChoiceMode = Literal["auto", "any", "none"]
87_ProviderChoiceMode = Literal["auto", "required", "none"]
89_TOOL_CHOICE_MODES: dict[_CanonicalChoiceMode, _ProviderChoiceMode] = {
90 "auto": "auto",
91 "any": "required",
92 "none": "none",
93}
95# Token allowances for template text the request does not carry. Survey of
96# 47 chat templates across 51 repositories, taken 2026-09-18. Widest value
97# measured beside each constant: preamble 1216, per message 17, tool block
98# 608, per tool 45. See docs/architecture.md for what this means for accuracy.
99# The preamble a template renders around any request, including a system block
100# it substitutes when the request carries none.
101_TEMPLATE_PREAMBLE_TOKENS = 300
102# The role markers and turn delimiters around one wire message.
103_TEMPLATE_MESSAGE_TOKENS = 8
104# The tool-calling instructions a template emits once when tools are present.
105_TEMPLATE_TOOL_BLOCK_TOKENS = 110
106# The wrapper a template renders around one tool schema.
107_TEMPLATE_PER_TOOL_TOKENS = 32
110class _OpenBlockKind(StrEnum):
111 NONE = "none"
112 TEXT = "text"
113 TOOL = "tool"
116def _provider_chat_kwargs(req: CanonicalChatRequest, canonical_model: str) -> dict[str, Any]:
117 """Shared provider.chat keyword arguments for both stream and non-stream paths."""
118 return {
119 "messages": _provider_messages(req),
120 "options": _provider_options(req),
121 "model": canonical_model,
122 "tools": _provider_tools(req.tools),
123 "tool_choice": _provider_tool_choice(req.tool_choice),
124 }
127def _stop_reason_for(result: ChatResult) -> StopReason:
128 """Closing stop reason for a non-streaming result.
130 Tool calls win over the reported finish reason. FinishReason.coerce falls
131 back to STOP for a missing or unknown value, so a provider that returns
132 tool calls without saying so produced tool_use content under end_turn, and
133 a client reading stop_reason decides whether to run the tools. The
134 streaming path already refuses the same downgrade.
135 """
136 if result.tool_calls:
137 return StopReason.TOOL_USE
138 return _FINISH_REASON_TO_STOP.get(result.finish_reason, StopReason.END_TURN)
141def _content_blocks_from_result(result: ChatResult) -> list[ContentBlock]:
142 """Build canonical content blocks from a non-streaming provider result."""
143 content: list[ContentBlock] = []
144 if result.text:
145 content.append(TextBlock(text=result.text))
146 for call in result.tool_calls:
147 content.append(
148 ToolUseBlock(
149 id=call.id or _new_call_id(),
150 name=call.name,
151 input=parse_tool_arguments(call.arguments),
152 )
153 )
154 return content
157def _canonical_usage(usage: TokenUsage) -> CanonicalUsage:
158 """Convert a provider token count, clamping the cached part to the prompt size.
160 The canonical layer promises ``cached_input_tokens <= input_tokens``, and the
161 Anthropic wire subtracts one from the other. The clamp lives here, at the one
162 boundary where provider numbers enter, so no translator can emit a negative.
163 """
164 return CanonicalUsage(
165 input_tokens=usage.prompt_tokens,
166 output_tokens=usage.completion_tokens,
167 cached_input_tokens=min(usage.cached_prompt_tokens, usage.prompt_tokens),
168 )
171def dispatch_chat(
172 req: CanonicalChatRequest, *, canonical_model: str | None = None
173) -> CanonicalResponse:
174 """Run a non-streaming chat request through the provider and return canonical output.
176 Pass *canonical_model* when the caller has already run
177 :func:`preflight_chat_request` (the route does, so the preflight runs once per
178 request); leave it ``None`` to resolve and validate the model here.
179 """
180 if canonical_model is None:
181 canonical_model = preflight_chat_request(req)
182 result = get_services().provider.chat(**_provider_chat_kwargs(req, canonical_model))
183 return CanonicalResponse(
184 id=_new_message_id(),
185 model=canonical_model,
186 content=_content_blocks_from_result(result),
187 stop_reason=_stop_reason_for(result),
188 usage=_canonical_usage(result.usage),
189 )
192async def dispatch_chat_stream(
193 req: CanonicalChatRequest, *, canonical_model: str | None = None
194) -> AsyncIterator[CanonicalStreamEvent]:
195 """Stream a canonical event sequence by translating provider frames on the fly.
197 Pass *canonical_model* when the caller has already run
198 :func:`preflight_chat_request` (the route does, so the preflight runs once per
199 request); leave it ``None`` to resolve and validate the model here.
200 """
201 # The preflight can do blocking HTTP model discovery when its TTL lapses, and
202 # opening the stream can issue a one-time template probe; run both in a thread
203 # so the event loop stays responsive.
204 if canonical_model is None:
205 canonical_model = await asyncio.to_thread(preflight_chat_request, req)
206 stream = await asyncio.to_thread(
207 lambda: get_services().provider.chat(
208 stream=True, **_provider_chat_kwargs(req, canonical_model)
209 )
210 )
211 try:
212 yield MessageStart(id=_new_message_id(), model=canonical_model)
213 state = _StreamState()
214 async for frame in _async_iter_provider_stream(stream):
215 for event in state.feed(frame):
216 yield event
217 for event in state.finish():
218 yield event
219 yield MessageStop()
220 finally:
221 # close() tears down the provider HTTP connection and can block; offload
222 # it like the open and per-frame reads so the event loop stays responsive.
223 await asyncio.to_thread(stream.close)
226async def _async_iter_provider_stream(
227 stream: Iterator[ChatStreamItem],
228) -> AsyncIterator[ChatStreamItem]:
229 """Iterate a provider chat stream without blocking the event loop.
231 ``LLMProvider.chat`` types a streaming result as a ClosableIterator, and
232 every provider in the tree returns a plain sync generator; iterating one
233 inline on the event loop would block, so each ``next()`` runs in a worker
234 thread via ``asyncio.to_thread``.
236 There used to be an async-native branch here for a provider shape that
237 does not exist. It was dead and also wrong: the caller's cleanup is
238 ``await asyncio.to_thread(stream.close)``, which an async-native stream
239 would not satisfy. Adding one means changing the Protocol and that
240 cleanup together, not restoring a branch nothing reaches.
241 """
242 while True:
243 frame = await asyncio.to_thread(_next_or_done, stream)
244 if frame is _STREAM_DONE:
245 return
246 yield frame
249_STREAM_DONE: Any = object()
250"""Sentinel returned by :func:`_next_or_done` to mean ``StopIteration``."""
253def _next_or_done(
254 stream: Iterator[ChatStreamItem],
255) -> ChatStreamItem | Any:
256 """Pull the next frame from *stream*; return ``_STREAM_DONE`` at exhaustion.
258 Raising ``StopIteration`` inside a coroutine becomes ``RuntimeError`` per
259 PEP 479; this helper converts that signal into a sentinel value the async
260 caller can branch on.
261 """
262 try:
263 return next(stream)
264 except StopIteration:
265 return _STREAM_DONE
268class _StreamState:
269 """Tracks open content blocks so deltas land in the right index."""
271 def __init__(self) -> None:
272 self._open: _OpenBlockKind = _OpenBlockKind.NONE
273 self._index: int = -1
274 self._tool_index: int | None = None
275 # Provider tool index -> the (id, name) its first delta carried.
276 # Continuation deltas typically carry neither.
277 self._tool_identity: dict[int, tuple[str, str]] = {}
278 self._stop_reason: StopReason = StopReason.END_TURN
279 self._usage: TokenUsage | None = None
281 def feed(self, frame: ChatStreamItem) -> Iterator[CanonicalStreamEvent]:
282 if isinstance(frame, str):
283 yield from self._feed_text(frame)
284 elif isinstance(frame, TokenUsage):
285 # Terminator-only frame: carries token totals, no content. Stash it
286 # so finish() can attach the counts to the closing MessageDelta.
287 self._usage = frame
288 elif isinstance(frame, StreamFinish):
289 self._feed_finish(frame)
290 else:
291 yield from self._feed_tool(frame)
293 def finish(self) -> Iterator[CanonicalStreamEvent]:
294 if self._open != _OpenBlockKind.NONE:
295 yield ContentBlockStop(index=self._index)
296 self._open = _OpenBlockKind.NONE
297 usage = _canonical_usage(self._usage) if self._usage is not None else None
298 yield MessageDelta(stop_reason=self._stop_reason, usage=usage)
300 def _feed_finish(self, frame: StreamFinish) -> None:
301 # The finish frame sets the closing stop reason (e.g. MAX_TOKENS on a
302 # length truncation). A tool-call stream already settled on TOOL_USE via
303 # the deltas, so never let a trailing finish frame downgrade that.
304 if self._stop_reason is StopReason.TOOL_USE:
305 return
306 self._stop_reason = _FINISH_REASON_TO_STOP.get(frame.reason, StopReason.END_TURN)
308 def _feed_text(self, text: str) -> Iterator[CanonicalStreamEvent]:
309 if self._open != _OpenBlockKind.TEXT:
310 yield from self._close_current()
311 self._index += 1
312 self._open = _OpenBlockKind.TEXT
313 yield ContentBlockStart(index=self._index, block=TextBlock(text=""))
314 yield ContentBlockDelta(index=self._index, delta=TextDelta(text=text))
316 def _feed_tool(self, frame: ToolCallDelta) -> Iterator[CanonicalStreamEvent]:
317 self._stop_reason = StopReason.TOOL_USE
318 is_new_call = self._open != _OpenBlockKind.TOOL or frame.index != self._tool_index
319 if is_new_call:
320 yield from self._close_current()
321 self._index += 1
322 self._open = _OpenBlockKind.TOOL
323 self._tool_index = frame.index
324 yield ContentBlockStart(
325 index=self._index,
326 block=ToolUseBlock(**self._tool_block_fields(frame)),
327 )
328 if frame.arguments_delta is not None:
329 yield ContentBlockDelta(
330 index=self._index,
331 delta=ToolUseDelta(partial_json=frame.arguments_delta),
332 )
334 def _tool_block_fields(self, frame: ToolCallDelta) -> dict[str, Any]:
335 """Identity for the block opening on *frame*, remembered per tool index.
337 A text frame between two argument deltas of one call (streamed
338 reasoning surfaced as text, say) closes the open tool block, so the
339 next delta for the same call has to open a second block. Continuation
340 deltas carry no id and no name, so that block used to get a fresh
341 synthetic id and an empty name, splitting one logical call across two
342 blocks the second of which matched no tool. Reusing the identity the
343 call already announced at least leaves both blocks stitchable by id.
344 """
345 known = self._tool_identity.get(frame.index)
346 identity = (
347 frame.id or (known[0] if known else _new_call_id()),
348 frame.name or (known[1] if known else ""),
349 )
350 self._tool_identity[frame.index] = identity
351 return {"id": identity[0], "name": identity[1], "input": {}}
353 def _close_current(self) -> Iterator[CanonicalStreamEvent]:
354 if self._open != _OpenBlockKind.NONE:
355 yield ContentBlockStop(index=self._index)
356 self._open = _OpenBlockKind.NONE
359def _resolve_canonical_model(model: str) -> str:
360 """Return the canonical ref for *model*, or raise ``ModelNotFoundError``.
362 Consults the cached union of native + remote + frontier refs on
363 Services, so an Ollama-managed model resolves the same way a locally
364 installed GGUF does. A bare ``name:tag`` matches the corresponding
365 ``ollama/<name:tag>`` entry when one exists in the discovered set.
366 """
367 canonical = get_services().known_models.resolve(model)
368 if canonical is None:
369 raise ModelNotFoundError(model)
370 return canonical
373def _ensure_tool_capability(req: CanonicalChatRequest, model: str) -> None:
374 if req.tools and not model_supports_tools(model):
375 raise ModelDoesNotSupportToolsError(model)
378def _ensure_configured_local_model(canonical: str) -> None:
379 """Reject a local-route model that is not the configured chat model.
381 Mirrors the fleet's own configured-model guard (which stays in place as
382 defense in depth for direct provider users) so streaming clients get a
383 clean 400 before headers instead of an SSE error frame mid-stream.
384 """
385 if not parse_model_ref(canonical).is_local or canonical == cfg.chat_model:
386 return
387 raise ProviderError(
388 configured_model_message(WorkerRole.CHAT, cfg.chat_model, canonical),
389 kind=ProviderErrorKind.BAD_REQUEST,
390 )
393def resolve_served_model(req: CanonicalChatRequest) -> str:
394 """Resolve *req*'s model ref and confirm this server serves it.
396 Raises ``ModelNotFoundError`` or a ``BAD_REQUEST`` ``ProviderError`` so the
397 route layer can return a real 4xx status. Returns the canonical model ref.
398 """
399 canonical = _resolve_canonical_model(req.model)
400 _ensure_configured_local_model(canonical)
401 return canonical
404def _log_dropped_images(req: CanonicalChatRequest, model: str) -> None:
405 """Log once when the translator dropped image parts from *req*."""
406 if req.dropped_images:
407 log.warning(
408 "Dropped %d image part(s) from a chat request: lilbee cannot pass images to %s",
409 req.dropped_images,
410 model,
411 )
414def preflight_chat_request(req: CanonicalChatRequest) -> str:
415 """Synchronously validate *req* before any streaming response starts.
417 Adds the tool-capability check to :func:`resolve_served_model`, so a model
418 whose template cannot render tool calls fails with a 4xx instead of burying
419 the failure in an SSE error frame after headers flush. It also logs any
420 image parts the translator dropped.
421 """
422 canonical = resolve_served_model(req)
423 _ensure_tool_capability(req, canonical)
424 _log_dropped_images(req, canonical)
425 return canonical
428def count_request_tokens(req: CanonicalChatRequest, *, canonical_model: str) -> PromptTokenCount:
429 """Tokens the served model prefills for *req*'s prompt, and how they were counted.
431 Hands the provider the arguments a chat call would send, so the count covers
432 the chat template's role markers and tool preamble as well as the content.
433 A backend that cannot render or tokenize falls back to the estimate, and the
434 result carries which of the two answered.
435 """
436 kwargs = _provider_chat_kwargs(req, canonical_model)
437 try:
438 tokens = get_services().provider.count_chat_prompt_tokens(**kwargs)
439 except NotImplementedError:
440 return PromptTokenCount(
441 tokens=_estimate_prompt_tokens(req, wire_messages=kwargs["messages"]),
442 accuracy=TokenCountAccuracy.ESTIMATED,
443 )
444 return PromptTokenCount(tokens=tokens, accuracy=TokenCountAccuracy.EXACT)
447def _provider_messages(req: CanonicalChatRequest) -> list[dict[str, Any]]:
448 """Flatten canonical messages to the OpenAI-shaped wire format the provider speaks."""
449 out: list[dict[str, Any]] = []
450 if req.system is not None:
451 out.append({"role": "system", "content": req.system})
452 for msg in req.messages:
453 out.extend(_translate_message(msg))
454 return out
457def _translate_message(msg: CanonicalMessage) -> list[dict[str, Any]]:
458 text_parts = [b.text for b in msg.content if isinstance(b, TextBlock)]
459 tool_uses = [b for b in msg.content if isinstance(b, ToolUseBlock)]
460 tool_results = [b for b in msg.content if isinstance(b, ToolResultBlock)]
461 text = "".join(text_parts)
463 # One ``tool`` wire-message per result block; tool_call_id pairs it back to
464 # the originating ToolUseBlock. Text blocks in the same canonical message
465 # follow as their own content message rather than being dropped.
466 out: list[dict[str, Any]] = [
467 {
468 "role": "tool",
469 "tool_call_id": block.tool_use_id,
470 "content": _flatten_text(block.content),
471 }
472 for block in tool_results
473 ]
474 if tool_uses:
475 out.append(
476 {
477 "role": msg.role,
478 "content": text,
479 "tool_calls": [
480 {
481 "id": tu.id,
482 "type": "function",
483 "function": {
484 "name": tu.name,
485 "arguments": json.dumps(tu.input),
486 },
487 }
488 for tu in tool_uses
489 ],
490 }
491 )
492 elif text or not tool_results:
493 out.append({"role": msg.role, "content": text})
494 return out
497def _flatten_text(blocks: list[ContentBlock]) -> str:
498 return "".join(b.text for b in blocks if isinstance(b, TextBlock))
501def _provider_tools(
502 tools: list[CanonicalTool] | None,
503) -> list[dict[str, Any]] | None:
504 if not tools:
505 return None
506 return [
507 {
508 "type": "function",
509 "function": {
510 "name": tool.name,
511 "description": tool.description,
512 "parameters": tool.input_schema,
513 },
514 }
515 for tool in tools
516 ]
519def _estimate_prompt_tokens(
520 req: CanonicalChatRequest, *, wire_messages: list[dict[str, Any]]
521) -> int:
522 """Estimate of *req*'s prompt tokens, for a backend with no tokenizer.
524 The request's own text is counted in UTF-8 bytes, which no token encodes
525 fewer than one of. The template text the request does not carry gets the
526 fixed allowances above, which are measured rather than proved: four of the
527 47 surveyed templates substitute more than they cover. A chars-per-token
528 ratio fails differently, reading dense input short.
530 The per-message allowance is charged against *wire_messages*, the messages
531 the provider is sent: one canonical message carrying several tool results
532 becomes one wire message each, and the template renders role markers around
533 every one.
534 """
535 tools = req.tools or []
536 allowance = _TEMPLATE_PREAMBLE_TOKENS + _TEMPLATE_MESSAGE_TOKENS * len(wire_messages)
537 if tools:
538 allowance += _TEMPLATE_TOOL_BLOCK_TOKENS + _TEMPLATE_PER_TOOL_TOKENS * len(tools)
539 return _content_bytes(req) + allowance
542def _content_bytes(req: CanonicalChatRequest) -> int:
543 """UTF-8 bytes of the text *req* itself puts into the rendered prompt."""
544 total = _utf8_len(req.system or "")
545 total += sum(_block_bytes(block) for msg in req.messages for block in msg.content)
546 for tool in req.tools or []:
547 total += _utf8_len(tool.name) + _utf8_len(tool.description)
548 total += _utf8_len(json.dumps(tool.input_schema))
549 return total
552def _block_bytes(block: ContentBlock) -> int:
553 """UTF-8 bytes *block* contributes to the rendered prompt."""
554 if block.type == "text":
555 return _utf8_len(block.text)
556 if block.type == "tool_use":
557 return _utf8_len(block.name) + _utf8_len(json.dumps(block.input))
558 return sum(_block_bytes(inner) for inner in block.content)
561def _utf8_len(text: str) -> int:
562 return len(text.encode("utf-8"))
565def _provider_tool_choice(
566 choice: CanonicalToolChoice | None,
567) -> str | dict[str, Any] | None:
568 if choice is None:
569 return None
570 if choice.mode == "tool":
571 return {"type": "function", "function": {"name": choice.tool_name}}
572 return _TOOL_CHOICE_MODES[choice.mode]
575def _provider_options(req: CanonicalChatRequest) -> dict[str, Any] | None:
576 out: dict[str, Any] = {}
577 if req.temperature is not None:
578 out["temperature"] = req.temperature
579 if req.top_p is not None:
580 out["top_p"] = req.top_p
581 if req.top_k is not None:
582 out["top_k"] = req.top_k
583 if req.max_tokens is not None:
584 out["num_predict"] = req.max_tokens
585 if req.seed is not None:
586 out["seed"] = req.seed
587 if req.frequency_penalty is not None:
588 out["frequency_penalty"] = req.frequency_penalty
589 if req.presence_penalty is not None:
590 out["presence_penalty"] = req.presence_penalty
591 if req.stop is not None:
592 out["stop"] = req.stop
593 if req.think is not None:
594 out["think"] = req.think
595 return out or None
598def _new_call_id() -> str:
599 return f"call_{uuid.uuid4().hex[:24]}"
602def _new_message_id() -> str:
603 return f"msg_{uuid.uuid4().hex[:24]}"