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