Coverage for src/lilbee/server/chat_completions_api/routes.py: 100%
147 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""HTTP routes for ``/v1/models`` and ``/v1/chat/completions``."""
3from __future__ import annotations
5import asyncio
6import json
7import logging
8import time
9import uuid
10from collections.abc import AsyncGenerator
11from datetime import datetime
13from litestar import Request, Router, get, post
14from litestar.background_tasks import BackgroundTask
15from litestar.exceptions import NotAuthorizedException, ValidationException
16from litestar.response import Response, Stream
18from lilbee.app.models import with_configured_chat
19from lilbee.app.services import get_services
20from lilbee.catalog.types import ModelTask
21from lilbee.core.config import cfg
22from lilbee.core.config.enums import ReasoningMode
23from lilbee.providers.model_ref import default_first
24from lilbee.retrieval.reasoning import effective_reasoning_cap
25from lilbee.server.auth import auth_checked_in_handler, session_manager
26from lilbee.server.chat_completions_api.errors import (
27 CompletionsErrorCode,
28 classify_provider_error,
29 completions_error_body,
30)
31from lilbee.server.chat_completions_api.models import (
32 CompletionsRequest,
33 CompletionsResponse,
34 ModelEntry,
35 ModelsListResponse,
36)
37from lilbee.server.chat_completions_api.streaming import encode_completions_sse
38from lilbee.server.chat_completions_api.translate import (
39 canonical_stream_to_completions_chunks,
40 canonical_to_completions_response,
41 completions_to_canonical_request,
42)
43from lilbee.server.chat_dispatch.canonical import CanonicalChatRequest, images_dropped_headers
44from lilbee.server.chat_dispatch.concurrency import (
45 ChatBusyError,
46 ChatSlotGuard,
47 acquire_chat_slot_or_busy,
48)
49from lilbee.server.chat_dispatch.dispatch import (
50 dispatch_chat_stream,
51 preflight_chat_request,
52)
53from lilbee.server.chat_dispatch.reasoning_cap import cap_aware_chat, cap_aware_chat_stream
54from lilbee.server.handlers.sse import SSE_MEDIA_TYPE
55from lilbee.server.validation_format import format_validation
57log = logging.getLogger(__name__)
59_MULTI_CHOICE_MESSAGE = "lilbee serves one choice per request; set n to 1 or omit it."
62@get("/v1/models")
63@auth_checked_in_handler
64async def list_models_endpoint(request: Request) -> Response:
65 """Return installed chat models in the ``/v1/models`` shape, the configured one leading."""
66 auth_error = _auth_failure(request)
67 if auth_error is not None:
68 return auth_error
70 # list_installed walks the model filesystem and served_chat_ctx may probe the
71 # engine; run both off the event loop like the sibling chat-completions route.
72 payload = await asyncio.to_thread(_build_models_list_payload)
73 # exclude_none like the sibling completions responses: context_window is
74 # deliberately None for every model but the active one, and emitting it as
75 # null adds a field OpenAI clients do not expect on a model entry.
76 return Response(payload.model_dump(exclude_none=True), media_type="application/json")
79def _build_models_list_payload() -> ModelsListResponse:
80 """Synchronous /v1/models body: blocking registry walk + engine ctx probe."""
81 services = get_services()
82 # The served shape applies to the active chat model; advertise it so a
83 # client trims history to fit and an agent harness reads the real
84 # concurrency instead of assuming the configured target.
85 served_ctx = services.provider.served_chat_ctx()
86 served_slots = services.provider.served_chat_slots()
87 installed = {m.ref: m for m in services.registry.list_installed() if m.task == ModelTask.CHAT}
88 # A remote or path-configured chat model has no registry entry but is still
89 # listed, configured model first (the launcher's picker order).
90 listed = with_configured_chat(sorted(installed), cfg.chat_model)
91 refs = default_first(listed, cfg.chat_model)
92 # A ref without a registry entry carries the newest native timestamp so a
93 # client sorting by created desc does not bury the model lilbee serves.
94 fallback_created = max((_parse_created(m.downloaded_at) for m in installed.values()), default=0)
95 return ModelsListResponse(
96 data=[
97 ModelEntry(
98 id=ref,
99 created=_parse_created(installed[ref].downloaded_at)
100 if ref in installed
101 else fallback_created,
102 context_window=served_ctx if ref == cfg.chat_model else None,
103 slots=served_slots if ref == cfg.chat_model else None,
104 )
105 for ref in refs
106 ]
107 )
110async def _auth_before_request(request: Request) -> Response | None:
111 """Reject an unauthenticated caller before Litestar parses the body.
113 The endpoint is marked @auth_checked_in_handler so AuthMiddleware defers to
114 the handler, which answers a bad token with the OpenAI error envelope
115 rather than Litestar's 401 shape. But binding the body as a handler
116 parameter made Litestar parse and pydantic-validate the whole payload
117 first, so an unauthenticated caller got request_max_body_size worth of
118 JSON processed on every request, and a malformed body was answered with a
119 400 naming the failing fields instead of ever reaching the 401. Litestar
120 runs before_request ahead of kwargs resolution and short-circuits on a
121 returned value, so this is the one place the check can sit.
122 """
123 return _auth_failure(request)
126@post("/v1/chat/completions", status_code=200, before_request=_auth_before_request)
127@auth_checked_in_handler
128async def chat_completions_endpoint(
129 request: Request, data: CompletionsRequest
130) -> Response | Stream:
131 """``/v1/chat/completions`` (stream + non-stream + tools).
133 ``stream: true`` switches the 200 response from JSON to an SSE stream of
134 chat.completion.chunk frames. The request body picks the arm, which OpenAPI
135 cannot express, so the schema declares the JSON default and this contract
136 matches OpenAI's own.
137 """
138 rejection = _reject_before_dispatch(request, data)
139 if rejection is not None:
140 return rejection
142 # The request field wins over the completions_reasoning setting, so a
143 # client that can send extra body fields picks its mode per call.
144 mode = data.reasoning if data.reasoning is not None else cfg.completions_reasoning
145 try:
146 req = completions_to_canonical_request(data, mode=mode)
147 except ValueError as exc:
148 # Request shape is wire-valid but carries something we can't translate
149 # (e.g. image content). Surface as 400 instead of a generic 500.
150 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, str(exc))
152 preflight = await _preflight_resolved_model(req)
153 if isinstance(preflight, Response):
154 return preflight
155 resolved_model = preflight
157 try:
158 await acquire_chat_slot_or_busy(get_services().provider.max_concurrent_chats())
159 except ChatBusyError:
160 return _error_response(
161 429,
162 CompletionsErrorCode.RATE_LIMIT_EXCEEDED,
163 "Backend is busy. Retry in a moment.",
164 headers={"Retry-After": "1"},
165 )
167 guard = ChatSlotGuard()
168 if req.stream:
169 include_usage = bool(data.stream_options and data.stream_options.include_usage)
170 # The after-send hook frees the slot when a disconnect lands before the
171 # generator's first iteration (its finally never runs in that case).
172 return Stream(
173 _gated_completions_stream(
174 req, guard, model=resolved_model, include_usage=include_usage, mode=mode
175 ),
176 media_type=SSE_MEDIA_TYPE,
177 background=BackgroundTask(guard.release),
178 headers=images_dropped_headers(req),
179 )
180 return await _run_non_stream(req, guard, canonical_model=resolved_model, mode=mode)
183def _reject_before_dispatch(request: Request, data: CompletionsRequest) -> Response | None:
184 """Multi-choice and unsupported-param checks before any dispatch work.
186 Returns a 4xx Response to short-circuit, or None to proceed. Unmapped OpenAI
187 params (``response_format``, ``logprobs``, and any other unknown field) are
188 accepted but logged at debug so a client learns they had no effect. Auth is
189 not checked here: it runs in the before_request hook, ahead of body parsing.
190 """
191 if data.n is not None and data.n > 1:
192 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, _MULTI_CHOICE_MESSAGE)
193 if data.model_extra:
194 log.debug("chat/completions ignoring unsupported params: %s", sorted(data.model_extra))
195 return None
198async def _preflight_resolved_model(req: CanonicalChatRequest) -> str | Response:
199 """Validate *req* before any streaming response starts, returning the resolved model.
201 A 4xx body here is reachable by any OpenAI-compatible client; once a
202 Stream is returned the headers are flushed at 200 and downstream errors
203 can only travel via SSE frames which not every client surfaces cleanly.
204 Runs the preflight in a thread: a lapsed model-discovery TTL makes it do
205 blocking HTTP probes that must not stall the event loop. Returns the
206 canonical resolved model string when *req* is fit to dispatch, so the
207 streaming response echoes the same model the non-streaming path does.
208 """
209 try:
210 return await asyncio.to_thread(preflight_chat_request, req)
211 except Exception as exc:
212 classified = classify_provider_error(exc)
213 if classified is None:
214 # Mirror _run_non_stream: an unclassified failure still rides the
215 # OpenAI error envelope, not a bare framework 500.
216 return _internal_error_response()
217 return _error_response(classified.http_status, classified.code, classified.message)
220_INTERNAL_ERROR_MESSAGE = "Internal server error. Check the server logs for details."
223def _internal_error_response() -> Response:
224 """Log and return the generic internal_error 500 envelope."""
225 log.exception("chat_completions_endpoint failed")
226 return _error_response(500, CompletionsErrorCode.INTERNAL_ERROR, _INTERNAL_ERROR_MESSAGE)
229async def _run_non_stream(
230 req: CanonicalChatRequest,
231 guard: ChatSlotGuard,
232 *,
233 canonical_model: str,
234 mode: ReasoningMode = ReasoningMode.SEPARATE,
235) -> Response:
236 """Dispatch a non-streaming chat call, translating errors to the wire envelope."""
237 try:
238 # cap_aware_chat blocks for the whole generation; run it off the event loop
239 # so a slow chat does not stall other admitted requests. The preflight
240 # already resolved the model, so hand it in to avoid re-running it.
241 resp = await asyncio.to_thread(
242 cap_aware_chat,
243 req,
244 canonical_model=canonical_model,
245 cap_chars=effective_reasoning_cap(),
246 )
247 except Exception as exc:
248 classified = classify_provider_error(exc)
249 if classified is None:
250 return _internal_error_response()
251 return _error_response(classified.http_status, classified.code, classified.message)
252 finally:
253 await guard.release()
254 body: CompletionsResponse = canonical_to_completions_response(
255 resp, response_id=_response_id(), mode=mode
256 )
257 return Response(
258 body.model_dump(exclude_none=True),
259 media_type="application/json",
260 headers=images_dropped_headers(req),
261 )
264async def _gated_completions_stream(
265 req: CanonicalChatRequest,
266 guard: ChatSlotGuard,
267 *,
268 model: str,
269 include_usage: bool = False,
270 mode: ReasoningMode = ReasoningMode.SEPARATE,
271) -> AsyncGenerator[bytes, None]:
272 """Drive ``cap_aware_chat_stream`` -> translate -> SSE-encode, freeing the slot on exit.
274 Pre-flight errors (unknown model, tools-against-non-tool-model) are
275 surfaced as a single SSE ``data:`` frame carrying the error
276 envelope, then ``[DONE]``. The chat slot is released in ``finally`` so
277 natural completion, exception, and client disconnect (GeneratorExit)
278 all unwind cleanly; a disconnect before the first iteration is covered
279 by the route's after-send release of the same guard.
280 """
281 # One id for the whole stream, error frame included: the frame is shaped as a
282 # real chunk so SDK clients parse it, and those accumulate by chunk id, so a
283 # fresh id made the error look like part of a different completion.
284 response_id = _response_id()
285 try:
286 try:
287 events = cap_aware_chat_stream(
288 dispatch_chat_stream(req, canonical_model=model),
289 req,
290 canonical_model=model,
291 cap_chars=effective_reasoning_cap(),
292 )
293 chunks = canonical_stream_to_completions_chunks(
294 events,
295 model=model,
296 response_id=response_id,
297 include_usage=include_usage,
298 mode=mode,
299 )
300 async for frame in encode_completions_sse(chunks):
301 yield frame
302 except Exception as exc:
303 classified = classify_provider_error(exc)
304 if classified is None:
305 log.exception("chat_completions stream failed")
306 yield _sse_error_frame(
307 CompletionsErrorCode.INTERNAL_ERROR,
308 _INTERNAL_ERROR_MESSAGE,
309 model=model,
310 response_id=response_id,
311 )
312 else:
313 yield _sse_error_frame(
314 classified.code, classified.message, model=model, response_id=response_id
315 )
316 finally:
317 await guard.release()
320def _sse_error_frame(
321 code: CompletionsErrorCode,
322 message: str,
323 *,
324 model: str = "",
325 response_id: str | None = None,
326) -> bytes:
327 """SSE frame carrying a mid-stream error in OpenAI's chunk-shaped wire format.
329 OpenAI-SDK clients only parse ``chat.completion.chunk``-shaped frames, so the
330 error rides a real chunk (empty delta, inline ``error`` field) followed by
331 ``[DONE]`` rather than a bare error frame. ``finish_reason`` stays null, as
332 in OpenAI's non-final chunks: a concrete reason like ``"length"`` would tell
333 clients the answer was merely truncated, and some auto-continue on it.
334 """
335 body = completions_error_body(code, message)
336 chunk: dict[str, object] = {
337 "id": response_id or _response_id(),
338 "object": "chat.completion.chunk",
339 "created": int(time.time()),
340 "model": model,
341 "choices": [{"index": 0, "delta": {}, "finish_reason": None}],
342 "error": body["error"],
343 }
344 payload = json.dumps(chunk, separators=(",", ":"))
345 return f"data: {payload}\n\ndata: [DONE]\n\n".encode()
348def _error_response(
349 status: int,
350 code: CompletionsErrorCode,
351 message: str,
352 *,
353 headers: dict[str, str] | None = None,
354) -> Response:
355 return Response(
356 completions_error_body(code, message),
357 status_code=status,
358 headers=headers or {},
359 media_type="application/json",
360 )
363def _auth_failure(request: Request) -> Response | None:
364 """Return a 401 error response if the bearer token is missing/wrong, else None.
366 validate() fails closed by raising when auth is uninitialized; surface that
367 as the same OpenAI 401 envelope rather than letting it escape as a 500.
368 """
369 auth_header = request.headers.get("authorization", "")
370 try:
371 authorized = session_manager.validate(auth_header)
372 except NotAuthorizedException:
373 authorized = False
374 if authorized:
375 return None
376 return _error_response(401, CompletionsErrorCode.INVALID_API_KEY, "Missing or invalid API key.")
379def _validation_exception_handler(_: Request, exc: ValidationException) -> Response:
380 """Wrap Litestar's body-parse failures in the OpenAI error envelope."""
381 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, format_validation(exc))
384def _parse_created(downloaded_at: str | None) -> int:
385 """Best-effort ISO-8601 to Unix-timestamp conversion; zero on failure."""
386 if not downloaded_at:
387 return 0
388 try:
389 return int(datetime.fromisoformat(downloaded_at).timestamp())
390 except ValueError:
391 return 0
394def _response_id() -> str:
395 """OpenAI-style ``chatcmpl-*`` id."""
396 return f"chatcmpl-{uuid.uuid4().hex[:24]}"
399completions_router = Router(
400 path="/",
401 route_handlers=[list_models_endpoint, chat_completions_endpoint],
402 exception_handlers={ValidationException: _validation_exception_handler},
403)