Coverage for src/lilbee/server/handlers/__init__.py: 100%
118 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Framework-agnostic route handlers for the lilbee HTTP server.
3Every public function is a plain async callable; no framework imports.
4Return types are dicts (JSON responses), lists, or async generators of SSE strings.
6Handlers are grouped by concern (sse, rag, models, ingest, config, documents,
7crawl) under sibling submodules. The names re-exported below are the public
8API consumed by ``server/routes/*.py``.
9"""
11from __future__ import annotations
13import asyncio
14import dataclasses
15import logging
16import time
17from collections.abc import AsyncGenerator, Callable, Sequence
18from typing import TYPE_CHECKING, Literal
20from lilbee.app.services import get_services
21from lilbee.app.status import gather_status
22from lilbee.app.version import get_version
23from lilbee.core.config import cfg
24from lilbee.providers.roles import WorkerRole
25from lilbee.providers.warm_progress import WarmPhase, WarmProgress, is_active_warm
26from lilbee.runtime.progress import SseEvent
27from lilbee.server.handlers.agent_config import agent_config, agent_config_index
28from lilbee.server.handlers.config import (
29 get_config,
30 get_config_defaults,
31 get_config_schema,
32 get_config_sources,
33 reset_config,
34 update_config,
35)
36from lilbee.server.handlers.crawl import crawl_stream
37from lilbee.server.handlers.documents import (
38 delete_documents,
39 get_source_content,
40 list_documents,
41)
42from lilbee.server.handlers.ingest import (
43 add_files_stream,
44 add_uploads_stream,
45 import_stream,
46 parse_ocr_mode,
47 sync_stream,
48 validate_add_paths,
49 validate_upload_names,
50)
51from lilbee.server.handlers.models import (
52 TASK_ENDPOINT_PATH,
53 ModelCatalogSection,
54 ModelsResponse,
55 enforce_pull_arch_compat,
56 format_task_mismatch,
57 list_external_models,
58 list_models,
59 models_catalog,
60 models_delete,
61 models_installed,
62 models_pull,
63 models_show,
64 set_chat_model,
65 set_embedding_model,
66 set_reranker_model,
67 set_vision_model,
68)
69from lilbee.server.handlers.rag import (
70 ask,
71 ask_stream,
72 chat,
73 chat_stream,
74 search,
75)
76from lilbee.server.handlers.sse import (
77 SseStream,
78 classify_load_error,
79 sse_done,
80 sse_error,
81 sse_event,
82)
83from lilbee.server.handlers.wiki import (
84 wiki_build_stream,
85 wiki_generate_stream,
86 wiki_synthesize_stream,
87)
88from lilbee.server.models import (
89 GpusResponse,
90 HealthResponse,
91 PlacementResponse,
92 ShutdownResponse,
93 StatusResponse,
94)
96if TYPE_CHECKING:
97 from lilbee.app.placement import GpuInfo, PlacementView
98 from lilbee.providers.base import LLMProvider
100log = logging.getLogger(__name__)
102# How often the warm stream re-snapshots provider state; sub-second so the read
103# bar advances smoothly without busy-spinning.
104_WARM_POLL_INTERVAL_S = 0.25
105# Upper bound on the warm stream; the launcher hands off when this elapses, so a
106# model still loading past it just warms on the client's first call. Generous to
107# cover a cold tensor-split giant off a slow filesystem.
108_WARM_STREAM_TIMEOUT_S = 1800.0
111def _chat_status(
112 provider: LLMProvider,
113) -> tuple[Literal["ready", "loading", "not_started", "error"], str | None]:
114 """Classify the chat engine's readiness for /api/health, with the error reason.
116 ``ready`` once the role serves; ``error`` when warm-up failed (paired with the
117 warm tracker's reason); ``loading`` while a warm is in flight; ``not_started``
118 when nothing is warming and the role isn't up (no chat model planned, or chat
119 is swapped out for its co-tenant; the next chat request loads it).
120 """
121 if provider.role_ready(WorkerRole.CHAT):
122 return "ready", None
123 snapshot = provider.warm_progress()
124 if snapshot is not None and snapshot.phase is WarmPhase.ERROR:
125 return "error", snapshot.error
126 if is_active_warm(snapshot):
127 return "loading", None
128 return "not_started", None
131async def health() -> HealthResponse:
132 """Return service health, version, and whether the chat engine is warm."""
133 services = get_services()
134 provider = services.provider
135 chat_status, chat_error = _chat_status(provider)
136 prefill = provider.chat_prefill_progress()
137 return HealthResponse(
138 status="ok",
139 version=get_version(),
140 chat_ready=provider.role_ready(WorkerRole.CHAT),
141 chat_status=chat_status,
142 chat_error=chat_error,
143 chat_ctx=provider.served_chat_ctx(),
144 chat_slots=provider.served_chat_slots(),
145 chat_prefill_processed=prefill[0] if prefill else None,
146 chat_prefill_total=prefill[1] if prefill else None,
147 embed_token_cap=provider.embed_token_cap(),
148 warnings=[*services.store.health_warnings(), *provider.health_warnings()],
149 upgrade=services.store.upgrade_progress(),
150 )
153async def shutdown() -> ShutdownResponse:
154 """Accept an API-requested stop; the route's background task stops the server.
156 Litestar runs that task only after the response has been handed to the
157 transport, so the stop cannot beat the 202 out and no wall-clock delay
158 has to be guessed. Stopping through the serving loop keeps teardown
159 ordered; without a loop the task falls back to SIGTERM.
160 """
161 log.info("Shutdown requested via the API")
162 return ShutdownResponse(status="shutting_down")
165async def warm_stream() -> AsyncGenerator[str, None]:
166 """Stream chat-model cold-load progress as SSE until the engine is ready.
168 A launcher subscribes to render granular warm feedback. Each
169 :data:`SseEvent.WARM` event carries a :class:`WarmProgress` snapshot; a
170 terminal :data:`SseEvent.DONE` closes the stream once the chat role is ready
171 or has failed, or when the budget elapses (the caller proceeds either way, so
172 a still-loading model just warms on its first call). When nothing is loading
173 because the engine is already warm, a single ready snapshot is emitted.
174 """
175 provider = get_services().provider
176 deadline = time.monotonic() + _WARM_STREAM_TIMEOUT_S
177 while time.monotonic() < deadline:
178 # Same first question as /api/health: a served role is ready whatever the
179 # local tracker last recorded (a remote chat role leaves a stale local error).
180 if provider.role_ready(WorkerRole.CHAT):
181 yield sse_event(SseEvent.WARM, WarmProgress(phase=WarmPhase.READY).model_dump())
182 break
183 snapshot = provider.warm_progress()
184 if snapshot is None:
185 yield sse_event(SseEvent.WARM, WarmProgress(phase=WarmPhase.STARTING).model_dump())
186 else:
187 yield sse_event(SseEvent.WARM, snapshot.model_dump())
188 if snapshot.phase is WarmPhase.ERROR:
189 break
190 await asyncio.sleep(_WARM_POLL_INTERVAL_S)
191 yield sse_done({})
194_GPU_STATS_INTERVAL_S = 1.0
197async def gpu_stats_stream(
198 devices: Sequence[GpuInfo],
199 interval_s: float = _GPU_STATS_INTERVAL_S,
200 max_ticks: int | None = None,
201) -> AsyncGenerator[str, None]:
202 """Stream live per-GPU utilization + free memory as SSE for the placement view.
204 Devices are resolved by the caller before the stream starts so a ProviderError
205 surfaces as a 503 at route time, not mid-stream. The client keeps the stream
206 open while visible; ``max_ticks`` bounds it for tests. A heartbeat is emitted
207 every ``cfg.sse_heartbeat_interval`` seconds of idle so clients don't time out.
209 The per-vendor probe runs on a worker thread, not here. It is not light: every
210 backend shells out to an SMI tool with a five-second timeout, and the Intel
211 paths sleep and scan /proc on top of that. Driven inline it held the event
212 loop for the whole subprocess on every tick, once per connected client, which
213 stalls chat, search and embedding requests along with it.
214 """
215 from lilbee.cli.tui import messages as msg
216 from lilbee.providers.fleet.gpu_stats import intel_util_hint, probe_gpu_stats_shared
218 last_heartbeat = time.monotonic()
219 tick = 0
220 while max_ticks is None or tick < max_ticks:
221 stats = await asyncio.to_thread(probe_gpu_stats_shared, devices)
222 payload: dict[str, object] = {"gpus": [dataclasses.asdict(s) for s in stats.values()]}
223 hint = intel_util_hint(devices, stats)
224 if hint:
225 payload["notice"] = msg.intel_util_hint_text(hint)
226 yield sse_event(SseEvent.GPU_STATS, payload)
227 tick += 1
228 if max_ticks is None or tick < max_ticks:
229 await asyncio.sleep(interval_s)
230 now = time.monotonic()
231 heartbeat_interval = cfg.sse_heartbeat_interval
232 if heartbeat_interval > 0 and now - last_heartbeat >= heartbeat_interval:
233 last_heartbeat = now
234 yield sse_event(SseEvent.HEARTBEAT, {"ts": time.time()})
237async def status() -> StatusResponse:
238 """Return config, sources, and chunk counts."""
239 raw = gather_status()
240 return StatusResponse(**raw.model_dump(exclude_none=True))
243async def placement() -> PlacementResponse:
244 """Current effective placement."""
245 from lilbee.app.placement import get_placement
247 return await _placement_response_off_loop(get_placement)
250async def placement_preview(spec_json: str | None) -> PlacementResponse:
251 """Preview a candidate spec (or auto when no spec). No persistence."""
252 from lilbee.app.placement import preview_placement
253 from lilbee.providers.fleet.placement_spec import PlacementSpec
255 spec = PlacementSpec.from_json(spec_json) if spec_json else None
256 return await _placement_response_off_loop(lambda: preview_placement(spec))
259async def placement_set(spec_json: str) -> PlacementResponse:
260 """Apply a manual placement spec; persists and rebuilds the fleet."""
261 from lilbee.app.placement import set_placement
262 from lilbee.providers.fleet.placement_spec import PlacementSpec
264 spec = PlacementSpec.from_json(spec_json)
265 return await _placement_response_off_loop(lambda: set_placement(spec))
268async def placement_clear() -> PlacementResponse:
269 """Clear manual placement; returns to the auto planner and rebuilds the fleet."""
270 from lilbee.app.placement import set_placement
272 return await _placement_response_off_loop(lambda: set_placement(None))
275async def _placement_response_off_loop(action: Callable[[], PlacementView]) -> PlacementResponse:
276 """Run a placement action and serialize it off the event loop.
278 Placement actions and the Intel util notice both shell out to GPU probes,
279 so neither may run on the loop.
280 """
281 return await asyncio.to_thread(lambda: _placement_response(action()))
284def _placement_response(view: PlacementView) -> PlacementResponse:
285 """Serialize a placement view with the host-level Intel util notice attached."""
286 resp = PlacementResponse.from_view(view)
287 resp.notice = _intel_notice_text(view.gpus)
288 return resp
291def _intel_notice_text(devices: Sequence[GpuInfo]) -> str | None:
292 """Formatted Intel util fix for the JSON surfaces, or None when util reads fine."""
293 from lilbee.cli.tui import messages as msg
294 from lilbee.providers.fleet.gpu_stats import probe_intel_util_hint
296 hint = probe_intel_util_hint(devices)
297 return msg.intel_util_hint_text(hint) if hint else None
300async def gpus() -> GpusResponse:
301 """Detected GPUs with free/total VRAM, plus the host-level Intel util notice."""
302 from lilbee.app.placement import get_placement
304 def _body() -> GpusResponse:
305 view = get_placement()
306 return GpusResponse(
307 gpus=PlacementResponse.from_view(view).gpus,
308 notice=_intel_notice_text(view.gpus),
309 )
311 return await asyncio.to_thread(_body)
314__all__ = [
315 "TASK_ENDPOINT_PATH",
316 "ModelCatalogSection",
317 "ModelsResponse",
318 "SseStream",
319 "add_files_stream",
320 "add_uploads_stream",
321 "agent_config",
322 "agent_config_index",
323 "ask",
324 "ask_stream",
325 "chat",
326 "chat_stream",
327 "classify_load_error",
328 "crawl_stream",
329 "delete_documents",
330 "enforce_pull_arch_compat",
331 "format_task_mismatch",
332 "get_config",
333 "get_config_defaults",
334 "get_config_schema",
335 "get_config_sources",
336 "get_source_content",
337 "gpu_stats_stream",
338 "gpus",
339 "health",
340 "import_stream",
341 "list_documents",
342 "list_external_models",
343 "list_models",
344 "models_catalog",
345 "models_delete",
346 "models_installed",
347 "models_pull",
348 "models_show",
349 "parse_ocr_mode",
350 "placement",
351 "placement_clear",
352 "placement_preview",
353 "placement_set",
354 "reset_config",
355 "search",
356 "set_chat_model",
357 "set_embedding_model",
358 "set_reranker_model",
359 "set_vision_model",
360 "sse_done",
361 "sse_error",
362 "sse_event",
363 "status",
364 "sync_stream",
365 "update_config",
366 "validate_add_paths",
367 "validate_upload_names",
368 "warm_stream",
369 "wiki_build_stream",
370 "wiki_generate_stream",
371 "wiki_synthesize_stream",
372]