Coverage for src/lilbee/server/chat_dispatch/dispatch.py: 100%
240 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +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 preflight_chat_request(req: CanonicalChatRequest) -> str:
405 """Synchronously validate *req* before any streaming response starts.
407 Adds the tool-capability check to :func:`resolve_served_model`, so a model
408 whose template cannot render tool calls fails with a 4xx instead of burying
409 the failure in an SSE error frame after headers flush.
410 """
411 canonical = resolve_served_model(req)
412 _ensure_tool_capability(req, canonical)
413 return canonical
416def count_request_tokens(req: CanonicalChatRequest, *, canonical_model: str) -> PromptTokenCount:
417 """Tokens the served model prefills for *req*'s prompt, and how they were counted.
419 Hands the provider the arguments a chat call would send, so the count covers
420 the chat template's role markers and tool preamble as well as the content.
421 A backend that cannot render or tokenize falls back to the estimate, and the
422 result carries which of the two answered.
423 """
424 kwargs = _provider_chat_kwargs(req, canonical_model)
425 try:
426 tokens = get_services().provider.count_chat_prompt_tokens(**kwargs)
427 except NotImplementedError:
428 return PromptTokenCount(
429 tokens=_estimate_prompt_tokens(req, wire_messages=kwargs["messages"]),
430 accuracy=TokenCountAccuracy.ESTIMATED,
431 )
432 return PromptTokenCount(tokens=tokens, accuracy=TokenCountAccuracy.EXACT)
435def _provider_messages(req: CanonicalChatRequest) -> list[dict[str, Any]]:
436 """Flatten canonical messages to the OpenAI-shaped wire format the provider speaks."""
437 out: list[dict[str, Any]] = []
438 if req.system is not None:
439 out.append({"role": "system", "content": req.system})
440 for msg in req.messages:
441 out.extend(_translate_message(msg))
442 return out
445def _translate_message(msg: CanonicalMessage) -> list[dict[str, Any]]:
446 text_parts = [b.text for b in msg.content if isinstance(b, TextBlock)]
447 tool_uses = [b for b in msg.content if isinstance(b, ToolUseBlock)]
448 tool_results = [b for b in msg.content if isinstance(b, ToolResultBlock)]
449 text = "".join(text_parts)
451 # One ``tool`` wire-message per result block; tool_call_id pairs it back to
452 # the originating ToolUseBlock. Text blocks in the same canonical message
453 # follow as their own content message rather than being dropped.
454 out: list[dict[str, Any]] = [
455 {
456 "role": "tool",
457 "tool_call_id": block.tool_use_id,
458 "content": _flatten_text(block.content),
459 }
460 for block in tool_results
461 ]
462 if tool_uses:
463 out.append(
464 {
465 "role": msg.role,
466 "content": text,
467 "tool_calls": [
468 {
469 "id": tu.id,
470 "type": "function",
471 "function": {
472 "name": tu.name,
473 "arguments": json.dumps(tu.input),
474 },
475 }
476 for tu in tool_uses
477 ],
478 }
479 )
480 elif text or not tool_results:
481 out.append({"role": msg.role, "content": text})
482 return out
485def _flatten_text(blocks: list[ContentBlock]) -> str:
486 return "".join(b.text for b in blocks if isinstance(b, TextBlock))
489def _provider_tools(
490 tools: list[CanonicalTool] | None,
491) -> list[dict[str, Any]] | None:
492 if not tools:
493 return None
494 return [
495 {
496 "type": "function",
497 "function": {
498 "name": tool.name,
499 "description": tool.description,
500 "parameters": tool.input_schema,
501 },
502 }
503 for tool in tools
504 ]
507def _estimate_prompt_tokens(
508 req: CanonicalChatRequest, *, wire_messages: list[dict[str, Any]]
509) -> int:
510 """Estimate of *req*'s prompt tokens, for a backend with no tokenizer.
512 The request's own text is counted in UTF-8 bytes, which no token encodes
513 fewer than one of. The template text the request does not carry gets the
514 fixed allowances above, which are measured rather than proved: four of the
515 47 surveyed templates substitute more than they cover. A chars-per-token
516 ratio fails differently, reading dense input short.
518 The per-message allowance is charged against *wire_messages*, the messages
519 the provider is sent: one canonical message carrying several tool results
520 becomes one wire message each, and the template renders role markers around
521 every one.
522 """
523 tools = req.tools or []
524 allowance = _TEMPLATE_PREAMBLE_TOKENS + _TEMPLATE_MESSAGE_TOKENS * len(wire_messages)
525 if tools:
526 allowance += _TEMPLATE_TOOL_BLOCK_TOKENS + _TEMPLATE_PER_TOOL_TOKENS * len(tools)
527 return _content_bytes(req) + allowance
530def _content_bytes(req: CanonicalChatRequest) -> int:
531 """UTF-8 bytes of the text *req* itself puts into the rendered prompt."""
532 total = _utf8_len(req.system or "")
533 total += sum(_block_bytes(block) for msg in req.messages for block in msg.content)
534 for tool in req.tools or []:
535 total += _utf8_len(tool.name) + _utf8_len(tool.description)
536 total += _utf8_len(json.dumps(tool.input_schema))
537 return total
540def _block_bytes(block: ContentBlock) -> int:
541 """UTF-8 bytes *block* contributes to the rendered prompt."""
542 if block.type == "text":
543 return _utf8_len(block.text)
544 if block.type == "tool_use":
545 return _utf8_len(block.name) + _utf8_len(json.dumps(block.input))
546 return sum(_block_bytes(inner) for inner in block.content)
549def _utf8_len(text: str) -> int:
550 return len(text.encode("utf-8"))
553def _provider_tool_choice(
554 choice: CanonicalToolChoice | None,
555) -> str | dict[str, Any] | None:
556 if choice is None:
557 return None
558 if choice.mode == "tool":
559 return {"type": "function", "function": {"name": choice.tool_name}}
560 return _TOOL_CHOICE_MODES[choice.mode]
563def _provider_options(req: CanonicalChatRequest) -> dict[str, Any] | None:
564 out: dict[str, Any] = {}
565 if req.temperature is not None:
566 out["temperature"] = req.temperature
567 if req.top_p is not None:
568 out["top_p"] = req.top_p
569 if req.top_k is not None:
570 out["top_k"] = req.top_k
571 if req.max_tokens is not None:
572 out["num_predict"] = req.max_tokens
573 if req.seed is not None:
574 out["seed"] = req.seed
575 if req.frequency_penalty is not None:
576 out["frequency_penalty"] = req.frequency_penalty
577 if req.presence_penalty is not None:
578 out["presence_penalty"] = req.presence_penalty
579 if req.stop is not None:
580 out["stop"] = req.stop
581 if req.think is not None:
582 out["think"] = req.think
583 return out or None
586def _new_call_id() -> str:
587 return f"call_{uuid.uuid4().hex[:24]}"
590def _new_message_id() -> str:
591 return f"msg_{uuid.uuid4().hex[:24]}"