Coverage for src/lilbee/server/handlers/__init__.py: 100%
116 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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
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 update_config,
32)
33from lilbee.server.handlers.crawl import crawl_stream
34from lilbee.server.handlers.documents import (
35 delete_documents,
36 get_source_content,
37 list_documents,
38)
39from lilbee.server.handlers.ingest import (
40 add_files_stream,
41 add_uploads_stream,
42 import_stream,
43 sync_stream,
44 validate_add_paths,
45 validate_upload_names,
46)
47from lilbee.server.handlers.models import (
48 TASK_ENDPOINT_PATH,
49 ModelCatalogSection,
50 ModelsResponse,
51 enforce_pull_arch_compat,
52 format_task_mismatch,
53 list_external_models,
54 list_models,
55 models_catalog,
56 models_delete,
57 models_installed,
58 models_pull,
59 models_show,
60 set_chat_model,
61 set_embedding_model,
62 set_reranker_model,
63 set_vision_model,
64)
65from lilbee.server.handlers.rag import (
66 ask,
67 ask_stream,
68 chat,
69 chat_stream,
70 search,
71)
72from lilbee.server.handlers.sse import (
73 SseStream,
74 classify_load_error,
75 sse_done,
76 sse_error,
77 sse_event,
78)
79from lilbee.server.handlers.wiki import (
80 wiki_build_stream,
81 wiki_generate_stream,
82 wiki_synthesize_stream,
83)
84from lilbee.server.models import (
85 GpusResponse,
86 HealthResponse,
87 PlacementResponse,
88 ShutdownResponse,
89 StatusResponse,
90)
92if TYPE_CHECKING:
93 from lilbee.app.placement import GpuInfo, PlacementView
94 from lilbee.providers.base import LLMProvider
96log = logging.getLogger(__name__)
98# How often the warm stream re-snapshots provider state; sub-second so the read
99# bar advances smoothly without busy-spinning.
100_WARM_POLL_INTERVAL_S = 0.25
101# Upper bound on the warm stream; the launcher hands off when this elapses, so a
102# model still loading past it just warms on the client's first call. Generous to
103# cover a cold tensor-split giant off a slow filesystem.
104_WARM_STREAM_TIMEOUT_S = 1800.0
107def _chat_status(
108 provider: LLMProvider,
109) -> tuple[Literal["ready", "loading", "not_started", "error"], str | None]:
110 """Classify the chat engine's readiness for /api/health, with the error reason.
112 ``ready`` once the role serves; ``error`` when warm-up failed (paired with the
113 warm tracker's reason); ``loading`` while a warm is in flight; ``not_started``
114 when nothing is warming and the role isn't up (no chat model planned, or chat
115 is swapped out for its co-tenant; the next chat request loads it).
116 """
117 if provider.role_ready(WorkerRole.CHAT):
118 return "ready", None
119 snapshot = provider.warm_progress()
120 if snapshot is None:
121 return "not_started", None
122 if snapshot.phase is WarmPhase.ERROR:
123 return "error", snapshot.error
124 return "loading", None
127async def health() -> HealthResponse:
128 """Return service health, version, and whether the chat engine is warm."""
129 provider = get_services().provider
130 chat_status, chat_error = _chat_status(provider)
131 return HealthResponse(
132 status="ok",
133 version=get_version(),
134 chat_ready=provider.role_ready(WorkerRole.CHAT),
135 chat_status=chat_status,
136 chat_error=chat_error,
137 chat_ctx=provider.served_chat_ctx(),
138 )
141async def shutdown() -> ShutdownResponse:
142 """Accept an API-requested stop; the route's background task sends SIGTERM.
144 Litestar runs that task only after the response has been handed to the
145 transport, so the signal cannot beat the 202 out and no wall-clock delay
146 has to be guessed. Routing through SIGTERM keeps the fleet teardown and
147 shutdown logging identical however the stop arrives.
148 """
149 log.info("Shutdown requested via the API")
150 return ShutdownResponse(status="shutting_down")
153async def warm_stream() -> AsyncGenerator[str, None]:
154 """Stream chat-model cold-load progress as SSE until the engine is ready.
156 A launcher subscribes to render granular warm feedback. Each
157 :data:`SseEvent.WARM` event carries a :class:`WarmProgress` snapshot; a
158 terminal :data:`SseEvent.DONE` closes the stream once the chat role is ready
159 or has failed, or when the budget elapses (the caller proceeds either way, so
160 a still-loading model just warms on its first call). When nothing is loading
161 because the engine is already warm, a single ready snapshot is emitted.
162 """
163 provider = get_services().provider
164 deadline = time.monotonic() + _WARM_STREAM_TIMEOUT_S
165 while time.monotonic() < deadline:
166 snapshot = provider.warm_progress()
167 if snapshot is None:
168 if provider.role_ready(WorkerRole.CHAT):
169 yield sse_event(SseEvent.WARM, WarmProgress(phase=WarmPhase.READY).model_dump())
170 break
171 yield sse_event(SseEvent.WARM, WarmProgress(phase=WarmPhase.STARTING).model_dump())
172 else:
173 yield sse_event(SseEvent.WARM, snapshot.model_dump())
174 if snapshot.phase in (WarmPhase.READY, WarmPhase.ERROR):
175 break
176 await asyncio.sleep(_WARM_POLL_INTERVAL_S)
177 yield sse_done({})
180_GPU_STATS_INTERVAL_S = 1.0
183async def gpu_stats_stream(
184 devices: Sequence[GpuInfo],
185 interval_s: float = _GPU_STATS_INTERVAL_S,
186 max_ticks: int | None = None,
187) -> AsyncGenerator[str, None]:
188 """Stream live per-GPU utilization + free memory as SSE for the placement view.
190 Devices are resolved by the caller before the stream starts so a ProviderError
191 surfaces as a 503 at route time, not mid-stream. The client keeps the stream
192 open while visible; ``max_ticks`` bounds it for tests. A heartbeat is emitted
193 every ``cfg.sse_heartbeat_interval`` seconds of idle so clients don't time out.
195 The per-vendor probe runs on a worker thread, not here. It is not light: every
196 backend shells out to an SMI tool with a five-second timeout, and the Intel
197 paths sleep and scan /proc on top of that. Driven inline it held the event
198 loop for the whole subprocess on every tick, once per connected client, which
199 stalls chat, search and embedding requests along with it.
200 """
201 from lilbee.cli.tui import messages as msg
202 from lilbee.providers.fleet.gpu_stats import intel_util_hint, probe_gpu_stats_shared
204 last_heartbeat = time.monotonic()
205 tick = 0
206 while max_ticks is None or tick < max_ticks:
207 stats = await asyncio.to_thread(probe_gpu_stats_shared, devices)
208 payload: dict[str, object] = {"gpus": [dataclasses.asdict(s) for s in stats.values()]}
209 hint = intel_util_hint(devices, stats)
210 if hint:
211 payload["notice"] = msg.intel_util_hint_text(hint)
212 yield sse_event(SseEvent.GPU_STATS, payload)
213 tick += 1
214 if max_ticks is None or tick < max_ticks:
215 await asyncio.sleep(interval_s)
216 now = time.monotonic()
217 heartbeat_interval = cfg.sse_heartbeat_interval
218 if heartbeat_interval > 0 and now - last_heartbeat >= heartbeat_interval:
219 last_heartbeat = now
220 yield sse_event(SseEvent.HEARTBEAT, {"ts": time.time()})
223async def status() -> StatusResponse:
224 """Return config, sources, and chunk counts."""
225 raw = gather_status()
226 return StatusResponse(**raw.model_dump(exclude_none=True))
229async def placement() -> PlacementResponse:
230 """Current effective placement."""
231 from lilbee.app.placement import get_placement
233 return await _placement_response_off_loop(get_placement)
236async def placement_preview(spec_json: str | None) -> PlacementResponse:
237 """Preview a candidate spec (or auto when no spec). No persistence."""
238 from lilbee.app.placement import preview_placement
239 from lilbee.providers.fleet.placement_spec import PlacementSpec
241 spec = PlacementSpec.from_json(spec_json) if spec_json else None
242 return await _placement_response_off_loop(lambda: preview_placement(spec))
245async def placement_set(spec_json: str) -> PlacementResponse:
246 """Apply a manual placement spec; persists and rebuilds the fleet."""
247 from lilbee.app.placement import set_placement
248 from lilbee.providers.fleet.placement_spec import PlacementSpec
250 spec = PlacementSpec.from_json(spec_json)
251 return await _placement_response_off_loop(lambda: set_placement(spec))
254async def placement_clear() -> PlacementResponse:
255 """Clear manual placement; returns to the auto planner and rebuilds the fleet."""
256 from lilbee.app.placement import set_placement
258 return await _placement_response_off_loop(lambda: set_placement(None))
261async def _placement_response_off_loop(action: Callable[[], PlacementView]) -> PlacementResponse:
262 """Run a placement action and serialize it off the event loop.
264 Placement actions and the Intel util notice both shell out to GPU probes,
265 so neither may run on the loop.
266 """
267 return await asyncio.to_thread(lambda: _placement_response(action()))
270def _placement_response(view: PlacementView) -> PlacementResponse:
271 """Serialize a placement view with the host-level Intel util notice attached."""
272 resp = PlacementResponse.from_view(view)
273 resp.notice = _intel_notice_text(view.gpus)
274 return resp
277def _intel_notice_text(devices: Sequence[GpuInfo]) -> str | None:
278 """Formatted Intel util fix for the JSON surfaces, or None when util reads fine."""
279 from lilbee.cli.tui import messages as msg
280 from lilbee.providers.fleet.gpu_stats import probe_intel_util_hint
282 hint = probe_intel_util_hint(devices)
283 return msg.intel_util_hint_text(hint) if hint else None
286async def gpus() -> GpusResponse:
287 """Detected GPUs with free/total VRAM, plus the host-level Intel util notice."""
288 from lilbee.app.placement import get_placement
290 def _body() -> GpusResponse:
291 view = get_placement()
292 return GpusResponse(
293 gpus=PlacementResponse.from_view(view).gpus,
294 notice=_intel_notice_text(view.gpus),
295 )
297 return await asyncio.to_thread(_body)
300__all__ = [
301 "TASK_ENDPOINT_PATH",
302 "ModelCatalogSection",
303 "ModelsResponse",
304 "SseStream",
305 "add_files_stream",
306 "add_uploads_stream",
307 "agent_config",
308 "agent_config_index",
309 "ask",
310 "ask_stream",
311 "chat",
312 "chat_stream",
313 "classify_load_error",
314 "crawl_stream",
315 "delete_documents",
316 "enforce_pull_arch_compat",
317 "format_task_mismatch",
318 "get_config",
319 "get_config_defaults",
320 "get_source_content",
321 "gpu_stats_stream",
322 "gpus",
323 "health",
324 "import_stream",
325 "list_documents",
326 "list_external_models",
327 "list_models",
328 "models_catalog",
329 "models_delete",
330 "models_installed",
331 "models_pull",
332 "models_show",
333 "placement",
334 "placement_clear",
335 "placement_preview",
336 "placement_set",
337 "search",
338 "set_chat_model",
339 "set_embedding_model",
340 "set_reranker_model",
341 "set_vision_model",
342 "sse_done",
343 "sse_error",
344 "sse_event",
345 "status",
346 "sync_stream",
347 "update_config",
348 "validate_add_paths",
349 "validate_upload_names",
350 "warm_stream",
351 "wiki_build_stream",
352 "wiki_generate_stream",
353 "wiki_synthesize_stream",
354]