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