Coverage for src/lilbee/server/handlers/__init__.py: 100%
117 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +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 prefill = provider.chat_prefill_progress()
132 return HealthResponse(
133 status="ok",
134 version=get_version(),
135 chat_ready=provider.role_ready(WorkerRole.CHAT),
136 chat_status=chat_status,
137 chat_error=chat_error,
138 chat_ctx=provider.served_chat_ctx(),
139 chat_slots=provider.served_chat_slots(),
140 chat_prefill_processed=prefill[0] if prefill else None,
141 chat_prefill_total=prefill[1] if prefill else None,
142 )
145async def shutdown() -> ShutdownResponse:
146 """Accept an API-requested stop; the route's background task sends SIGTERM.
148 Litestar runs that task only after the response has been handed to the
149 transport, so the signal cannot beat the 202 out and no wall-clock delay
150 has to be guessed. Routing through SIGTERM keeps the fleet teardown and
151 shutdown logging identical however the stop arrives.
152 """
153 log.info("Shutdown requested via the API")
154 return ShutdownResponse(status="shutting_down")
157async def warm_stream() -> AsyncGenerator[str, None]:
158 """Stream chat-model cold-load progress as SSE until the engine is ready.
160 A launcher subscribes to render granular warm feedback. Each
161 :data:`SseEvent.WARM` event carries a :class:`WarmProgress` snapshot; a
162 terminal :data:`SseEvent.DONE` closes the stream once the chat role is ready
163 or has failed, or when the budget elapses (the caller proceeds either way, so
164 a still-loading model just warms on its first call). When nothing is loading
165 because the engine is already warm, a single ready snapshot is emitted.
166 """
167 provider = get_services().provider
168 deadline = time.monotonic() + _WARM_STREAM_TIMEOUT_S
169 while time.monotonic() < deadline:
170 snapshot = provider.warm_progress()
171 if snapshot is None:
172 if provider.role_ready(WorkerRole.CHAT):
173 yield sse_event(SseEvent.WARM, WarmProgress(phase=WarmPhase.READY).model_dump())
174 break
175 yield sse_event(SseEvent.WARM, WarmProgress(phase=WarmPhase.STARTING).model_dump())
176 else:
177 yield sse_event(SseEvent.WARM, snapshot.model_dump())
178 if snapshot.phase in (WarmPhase.READY, WarmPhase.ERROR):
179 break
180 await asyncio.sleep(_WARM_POLL_INTERVAL_S)
181 yield sse_done({})
184_GPU_STATS_INTERVAL_S = 1.0
187async def gpu_stats_stream(
188 devices: Sequence[GpuInfo],
189 interval_s: float = _GPU_STATS_INTERVAL_S,
190 max_ticks: int | None = None,
191) -> AsyncGenerator[str, None]:
192 """Stream live per-GPU utilization + free memory as SSE for the placement view.
194 Devices are resolved by the caller before the stream starts so a ProviderError
195 surfaces as a 503 at route time, not mid-stream. The client keeps the stream
196 open while visible; ``max_ticks`` bounds it for tests. A heartbeat is emitted
197 every ``cfg.sse_heartbeat_interval`` seconds of idle so clients don't time out.
199 The per-vendor probe runs on a worker thread, not here. It is not light: every
200 backend shells out to an SMI tool with a five-second timeout, and the Intel
201 paths sleep and scan /proc on top of that. Driven inline it held the event
202 loop for the whole subprocess on every tick, once per connected client, which
203 stalls chat, search and embedding requests along with it.
204 """
205 from lilbee.cli.tui import messages as msg
206 from lilbee.providers.fleet.gpu_stats import intel_util_hint, probe_gpu_stats_shared
208 last_heartbeat = time.monotonic()
209 tick = 0
210 while max_ticks is None or tick < max_ticks:
211 stats = await asyncio.to_thread(probe_gpu_stats_shared, devices)
212 payload: dict[str, object] = {"gpus": [dataclasses.asdict(s) for s in stats.values()]}
213 hint = intel_util_hint(devices, stats)
214 if hint:
215 payload["notice"] = msg.intel_util_hint_text(hint)
216 yield sse_event(SseEvent.GPU_STATS, payload)
217 tick += 1
218 if max_ticks is None or tick < max_ticks:
219 await asyncio.sleep(interval_s)
220 now = time.monotonic()
221 heartbeat_interval = cfg.sse_heartbeat_interval
222 if heartbeat_interval > 0 and now - last_heartbeat >= heartbeat_interval:
223 last_heartbeat = now
224 yield sse_event(SseEvent.HEARTBEAT, {"ts": time.time()})
227async def status() -> StatusResponse:
228 """Return config, sources, and chunk counts."""
229 raw = gather_status()
230 return StatusResponse(**raw.model_dump(exclude_none=True))
233async def placement() -> PlacementResponse:
234 """Current effective placement."""
235 from lilbee.app.placement import get_placement
237 return await _placement_response_off_loop(get_placement)
240async def placement_preview(spec_json: str | None) -> PlacementResponse:
241 """Preview a candidate spec (or auto when no spec). No persistence."""
242 from lilbee.app.placement import preview_placement
243 from lilbee.providers.fleet.placement_spec import PlacementSpec
245 spec = PlacementSpec.from_json(spec_json) if spec_json else None
246 return await _placement_response_off_loop(lambda: preview_placement(spec))
249async def placement_set(spec_json: str) -> PlacementResponse:
250 """Apply a manual placement spec; persists and rebuilds the fleet."""
251 from lilbee.app.placement import set_placement
252 from lilbee.providers.fleet.placement_spec import PlacementSpec
254 spec = PlacementSpec.from_json(spec_json)
255 return await _placement_response_off_loop(lambda: set_placement(spec))
258async def placement_clear() -> PlacementResponse:
259 """Clear manual placement; returns to the auto planner and rebuilds the fleet."""
260 from lilbee.app.placement import set_placement
262 return await _placement_response_off_loop(lambda: set_placement(None))
265async def _placement_response_off_loop(action: Callable[[], PlacementView]) -> PlacementResponse:
266 """Run a placement action and serialize it off the event loop.
268 Placement actions and the Intel util notice both shell out to GPU probes,
269 so neither may run on the loop.
270 """
271 return await asyncio.to_thread(lambda: _placement_response(action()))
274def _placement_response(view: PlacementView) -> PlacementResponse:
275 """Serialize a placement view with the host-level Intel util notice attached."""
276 resp = PlacementResponse.from_view(view)
277 resp.notice = _intel_notice_text(view.gpus)
278 return resp
281def _intel_notice_text(devices: Sequence[GpuInfo]) -> str | None:
282 """Formatted Intel util fix for the JSON surfaces, or None when util reads fine."""
283 from lilbee.cli.tui import messages as msg
284 from lilbee.providers.fleet.gpu_stats import probe_intel_util_hint
286 hint = probe_intel_util_hint(devices)
287 return msg.intel_util_hint_text(hint) if hint else None
290async def gpus() -> GpusResponse:
291 """Detected GPUs with free/total VRAM, plus the host-level Intel util notice."""
292 from lilbee.app.placement import get_placement
294 def _body() -> GpusResponse:
295 view = get_placement()
296 return GpusResponse(
297 gpus=PlacementResponse.from_view(view).gpus,
298 notice=_intel_notice_text(view.gpus),
299 )
301 return await asyncio.to_thread(_body)
304__all__ = [
305 "TASK_ENDPOINT_PATH",
306 "ModelCatalogSection",
307 "ModelsResponse",
308 "SseStream",
309 "add_files_stream",
310 "add_uploads_stream",
311 "agent_config",
312 "agent_config_index",
313 "ask",
314 "ask_stream",
315 "chat",
316 "chat_stream",
317 "classify_load_error",
318 "crawl_stream",
319 "delete_documents",
320 "enforce_pull_arch_compat",
321 "format_task_mismatch",
322 "get_config",
323 "get_config_defaults",
324 "get_source_content",
325 "gpu_stats_stream",
326 "gpus",
327 "health",
328 "import_stream",
329 "list_documents",
330 "list_external_models",
331 "list_models",
332 "models_catalog",
333 "models_delete",
334 "models_installed",
335 "models_pull",
336 "models_show",
337 "placement",
338 "placement_clear",
339 "placement_preview",
340 "placement_set",
341 "search",
342 "set_chat_model",
343 "set_embedding_model",
344 "set_reranker_model",
345 "set_vision_model",
346 "sse_done",
347 "sse_error",
348 "sse_event",
349 "status",
350 "sync_stream",
351 "update_config",
352 "validate_add_paths",
353 "validate_upload_names",
354 "warm_stream",
355 "wiki_build_stream",
356 "wiki_generate_stream",
357 "wiki_synthesize_stream",
358]