Coverage for src/lilbee/app/services.py: 100%
173 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Typed service container: single point of access for all singletons.
3All runtime dependencies (provider, store, embedder, reranker, concepts,
4clusterer, searcher, worker pool) are created lazily on first call to
5``get_services()`` and cached for the process lifetime. Tests call
6``reset_services()`` between runs.
8``build_services(config)`` is the construction seam: it builds a full container
9against an arbitrary Config without touching the process-global singleton. The
10library API (:class:`lilbee.Lilbee`) builds one per instance and installs it for
11the duration of each call via :func:`services_scope`, so ingest code reaching for
12``get_services()`` resolves the caller's container. The override is a ContextVar,
13so ``reset_services`` / ``set_services`` / ``peek_services`` (which operate only
14on the global singleton) never see it.
15"""
17from __future__ import annotations
19import asyncio
20import atexit
21import logging
22import signal
23import sys
24import threading
25from collections.abc import Callable
26from contextlib import contextmanager
27from contextvars import ContextVar
28from dataclasses import dataclass, field
29from typing import TYPE_CHECKING
31if TYPE_CHECKING:
32 from collections.abc import Iterator
34 from lilbee.catalog.hf_client import HfClient
35 from lilbee.core.config import Config
36 from lilbee.data.store import Store
37 from lilbee.modelhub.model_manager import ModelManager
38 from lilbee.modelhub.model_manager.discovery import KnownModelCache
39 from lilbee.modelhub.registry import ModelRegistry
40 from lilbee.providers.base import LLMProvider
41 from lilbee.providers.roles import WorkerRole
42 from lilbee.retrieval.clustering import Clusterer
43 from lilbee.retrieval.concepts import ConceptGraph
44 from lilbee.retrieval.embedder import Embedder
45 from lilbee.retrieval.query import Searcher
46 from lilbee.retrieval.reranker import Reranker
47 from lilbee.runtime.ingest_lock import IngestLockRegistry
48 from lilbee.sessions import SessionStore
51log = logging.getLogger(__name__)
53_SIGNAL_EXIT_BASE = 128
55_HARD_EXIT_THREAD_NAME = "hard-exit-teardown"
58def _default_session_store() -> SessionStore:
59 """Build the file-backed session store, importing it lazily.
61 ``lilbee.sessions`` pulls in the config/catalog import chain, so importing it
62 at this module's top would form a cycle during CLI config load.
63 """
64 from lilbee.sessions import SessionStore
66 return SessionStore()
69@dataclass
70class CrawlerSyncState:
71 """Process-wide sync coordination state (lock + last-run timestamp)."""
73 lock: threading.Lock = field(default_factory=threading.Lock)
74 last_run: float = 0.0
77@dataclass(frozen=True)
78class Services:
79 """Holds all runtime service instances.
81 Inference lifecycle (cancel, per-role reload, spawn notifications) is owned
82 by the provider, which manages the llama-server fleet. Services exposes thin
83 pass-throughs so callers (Ctrl+C, the chat-stream cancel action, the settings
84 and model-bar pickers, the TUI task bar) need not reach into the provider's
85 API. ``cancel_inference()`` is the canonical cancel entry point.
86 """
88 provider: LLMProvider
89 store: Store
90 embedder: Embedder
91 reranker: Reranker
92 concepts: ConceptGraph
93 clusterer: Clusterer
94 searcher: Searcher
95 registry: ModelRegistry
96 hf_client: HfClient
97 ingest_lock_registry: IngestLockRegistry
98 model_manager: ModelManager
99 crawler_semaphore: asyncio.Semaphore | None
100 crawler_sync_state: CrawlerSyncState
101 known_models: KnownModelCache
102 session_store: SessionStore = field(default_factory=_default_session_store)
104 def cancel_inference(self) -> None:
105 """Interrupt any in-flight generation. Idempotent.
107 The fleet engine severs its live chat streams (llama-server stops
108 generating when the connection drops); providers with nothing in
109 flight treat this as a no-op.
110 """
111 self.provider.cancel_inference()
113 def reload_role(self, role_name: WorkerRole, *, wait: bool = False) -> None:
114 """Respawn only *role_name*'s model server so it picks up changed cfg.
116 Other roles' servers and any in-flight stream they own are untouched. Use
117 when one role-bound model setting changed (e.g. embedding_model). The
118 respawn runs off the caller's thread, so this returns immediately, unless
119 ``wait=True`` (the caller is already off the event loop and wants to block
120 until the new model has loaded).
121 """
122 self.provider.reload_role(role_name, wait=wait)
124 def add_pool_listener(
125 self,
126 *,
127 on_spawning: Callable[[WorkerRole], None] | None = None,
128 on_spawned: Callable[[WorkerRole], None] | None = None,
129 ) -> None:
130 """Subscribe to server spawn lifecycle events.
132 Forwards to :meth:`LLMProvider.add_spawn_listener`. The TUI uses this to
133 surface "Starting <role>..." / "<role> ready" notifications when a role's
134 server (re)spawns (cold start after a non-eager boot, or a reload).
135 """
136 self.provider.add_spawn_listener(on_spawning=on_spawning, on_spawned=on_spawned)
139class _ServicesState:
140 """The cached process-global singleton plus the per-task scoped override.
142 ``singleton`` is set on first ``get_services()`` call. Concurrency
143 contract: creation is serialized by ``_singleton_create_lock`` (several
144 worker threads can first-touch services at once, and a duplicate build
145 would collide in xberg's process-global backend registry), and the
146 Services dataclass is logically immutable post-construction, so concurrent
147 reads are safe without a lock. Tests that need a custom container call
148 ``set_services(make_mock_services(...))``; ``peek_services()`` is the
149 read-only inspector for cleanup fixtures.
151 ``override`` shadows the singleton for the entering task: set by
152 :func:`services_scope` (the library API's per-call binding), read by
153 :func:`get_services`, and invisible to ``reset_services`` /
154 ``set_services`` / ``peek_services``, which only touch the singleton.
155 """
157 def __init__(self) -> None:
158 self.singleton: Services | None = None
159 self.override: ContextVar[Services | None] = ContextVar(
160 "lilbee_services_override", default=None
161 )
162 # Whether this process is an interactive session (the TUI). Recorded by
163 # the interactive entry point before anything builds the container, and
164 # read once at build so the provider it creates holds its fleet resident
165 # for the session. Build-time intent only; the state that matters after
166 # that lives on the provider itself.
167 self.interactive: bool = False
170_state = _ServicesState()
173def build_services(
174 config: Config,
175 *,
176 provider: LLMProvider | None = None,
177 registry: ModelRegistry | None = None,
178 interactive: bool = False,
179) -> Services:
180 """Build a full Services container bound to *config*, without caching it.
182 ``get_services()`` calls this with the process-global cfg; the library API
183 calls it per instance. Service modules are imported inside the function to
184 keep CLI startup fast (they transitively pull in lancedb / xberg). Pass
185 *provider* to reuse a caller-supplied one; otherwise it is built from
186 *config* via the provider factory. Pass *registry* to reuse one already built
187 (get_services builds it for embedding-dim reconciliation). Embedding-dim
188 reconciliation is a global-cfg concern owned by :func:`get_services`, not
189 done here.
191 Side effect: binds *provider* into xberg's process-global OCR/embedding/
192 tokenizer backends. The registry is a single global slot, so every container
193 (singleton or per-instance library) binds its own, not only get_services.
194 """
195 from lilbee.catalog.hf_client import HfClient
196 from lilbee.data.store import Store
197 from lilbee.modelhub.model_manager import ModelManager
198 from lilbee.modelhub.model_manager.discovery import KnownModelCache
199 from lilbee.modelhub.registry import ModelRegistry
200 from lilbee.providers.factory import create_provider
201 from lilbee.retrieval.clustering import Clusterer
202 from lilbee.retrieval.concepts import ConceptGraph
203 from lilbee.retrieval.embedder import Embedder
204 from lilbee.retrieval.query import Searcher
205 from lilbee.retrieval.reranker import Reranker
206 from lilbee.runtime.ingest_lock import IngestLockRegistry
208 provider = provider or create_provider(config, hold_warm=interactive)
209 from lilbee.data.extract.backends import sync_xberg_backends
211 sync_xberg_backends(provider)
212 registry = registry or ModelRegistry(config.models_dir)
213 store = Store(config)
214 embedder = Embedder(config, provider)
215 reranker = Reranker(config)
216 concepts = ConceptGraph(config, store)
217 clusterer = Clusterer(config, store)
218 searcher = Searcher(config, provider, store, embedder, reranker, concepts)
219 hf_client = HfClient()
220 ingest_lock_registry = IngestLockRegistry()
221 model_manager = ModelManager(config.models_dir)
222 crawler_semaphore = (
223 asyncio.Semaphore(config.crawl_max_concurrent) if config.crawl_max_concurrent > 0 else None
224 )
225 crawler_sync_state = CrawlerSyncState()
226 known_models = KnownModelCache()
227 return Services(
228 provider=provider,
229 store=store,
230 embedder=embedder,
231 reranker=reranker,
232 concepts=concepts,
233 clusterer=clusterer,
234 searcher=searcher,
235 registry=registry,
236 hf_client=hf_client,
237 ingest_lock_registry=ingest_lock_registry,
238 model_manager=model_manager,
239 crawler_semaphore=crawler_semaphore,
240 crawler_sync_state=crawler_sync_state,
241 known_models=known_models,
242 )
245# Serializes first-touch singleton creation: several worker threads (e.g.
246# concurrent downloads) can call get_services() before the singleton exists,
247# and a duplicate build re-registers xberg's process-global backends mid-flight,
248# raising 'already registered' in the losing thread.
249_singleton_create_lock = threading.Lock()
252def get_services() -> Services:
253 """Return the active container: a scoped override if set, else the cached singleton.
255 Creates the singleton on first call (against the process-global cfg). A
256 config-file embedding_model with no embedding_dim would otherwise build the
257 store at the stale 768 default, so the width is pinned to the embedder before
258 the store is built.
259 """
260 override = _state.override.get()
261 if override is not None:
262 return override
263 if _state.singleton is not None:
264 return _state.singleton
266 with _singleton_create_lock:
267 if _state.singleton is not None:
268 return _state.singleton
270 from lilbee.app.settings import reconcile_embedding_dim
271 from lilbee.core.config import cfg
272 from lilbee.modelhub.registry import ModelRegistry
274 registry = ModelRegistry(cfg.models_dir)
275 # Pin the store width to the embedder before Store(); pass the registry so
276 # resolution doesn't re-enter this half-built get_services.
277 reconcile_embedding_dim(registry)
278 _state.singleton = build_services(cfg, registry=registry, interactive=_state.interactive)
279 # Eager start is the default: pay the spawn cost per role server at TUI mount
280 # so the first user action lands on a warm fleet. Roles whose model is unset
281 # are skipped, so a setup with only chat + embed never spawns rerank or
282 # vision. Set ``cfg.worker_pool_eager_start = false`` for headless scripts
283 # where mount time matters more than first-call latency.
284 if cfg.worker_pool_eager_start:
285 from contextlib import suppress
287 with suppress(Exception):
288 _state.singleton.provider.warm_up_pool()
289 return _state.singleton
292@contextmanager
293def services_scope(services: Services) -> Iterator[None]:
294 """Bind *services* as the container ``get_services()`` returns for this block.
296 Isolated to the entering task via a ContextVar (it propagates into
297 ``to_ingest_thread`` workers), and never affects the global singleton, so
298 ``reset_services`` is unnecessary and unused around a scoped call.
299 """
300 token = _state.override.set(services)
301 try:
302 yield
303 finally:
304 _state.override.reset(token)
307def mark_interactive_session() -> None:
308 """Record that this process is an interactive session before services build.
310 The TUI owns the process for its whole lifetime, so the fleet it builds keeps
311 its weights resident rather than idle-unloading under a user who is still in
312 the app; closing lilbee releases it. Called by the interactive entry point
313 before anything touches ``get_services``, so the provider is constructed with
314 that intent; a one-shot CLI or the MCP server never calls it.
315 """
316 _state.interactive = True
319def set_services(services: Services | None) -> None:
320 """Replace the cached Services singleton (for testing)."""
321 _state.singleton = services
324def peek_services() -> Services | None:
325 """Return the cached Services container, or None if not yet initialized.
327 Public read-only accessor for test cleanup helpers that need to
328 inspect the singleton without forcing initialization.
329 """
330 return _state.singleton
333# Serializes the singleton swap in reset_services: a signal's teardown thread
334# and an exiting caller's can race, and both tearing down the same container
335# would double-close the store.
336_reset_swap_lock = threading.Lock()
339def reset_services() -> None:
340 """Shut down and discard all cached instances.
342 Swap the module reference to ``None`` *before* tearing the old instances
343 down, so a new caller never observes a half-closed container. The swap is
344 locked so concurrent callers (a signal's teardown thread plus an exiting
345 one) tear the container down exactly once. On the shared HTTP daemon every
346 entry point that would call this mid-flight is refused, so it only ever
347 runs single-client (CLI, TUI, stdio MCP).
348 """
349 with _reset_swap_lock:
350 old = _state.singleton
351 _state.singleton = None
352 if old is not None:
353 old.provider.shutdown()
354 old.store.close()
357def reset_store() -> None:
358 """Close and rebuild only the Store and its dependents; keep providers loaded.
360 Used after a data-dir wipe (``/reset``) where the LanceDB handle is invalid
361 but the running provider/embedder/reranker are still good. Avoids the
362 multi-second reload cost of ``reset_services()``.
363 """
364 svc = _state.singleton
365 if svc is None:
366 return
367 from dataclasses import replace
369 from lilbee.core.config import cfg
370 from lilbee.data.store import Store
371 from lilbee.retrieval.clustering import Clusterer
372 from lilbee.retrieval.concepts import ConceptGraph
373 from lilbee.retrieval.query import Searcher
375 # Build the replacement, swap the reference, then close the old store last so
376 # a new caller never observes a closed handle mid-swap.
377 old_store = svc.store
378 store = Store(cfg)
379 concepts = ConceptGraph(cfg, store)
380 clusterer = Clusterer(cfg, store)
381 searcher = Searcher(cfg, svc.provider, store, svc.embedder, svc.reranker, concepts)
382 _state.singleton = replace(
383 svc,
384 store=store,
385 concepts=concepts,
386 clusterer=clusterer,
387 searcher=searcher,
388 )
389 old_store.close()
392class _EngineLifecycle:
393 """Owns the hard-exit hooks that stop the engine fleet."""
395 def __init__(self) -> None:
396 self._installed = False
398 @staticmethod
399 def _hard_exit_signals() -> tuple[signal.Signals, ...]: # pragma: no cover - platform split
400 """Signals whose default disposition kills us without running atexit."""
401 if sys.platform == "win32": # Windows has no SIGHUP
402 return (signal.SIGTERM,)
403 return (signal.SIGTERM, signal.SIGHUP)
405 def install(self) -> None:
406 """Route hard-exit signals through teardown. Idempotent; no-op off the main thread."""
407 if self._installed:
408 return
409 try:
410 for sig in self._hard_exit_signals():
411 signal.signal(sig, self._on_hard_exit)
412 except ValueError:
413 return
414 self._installed = True
416 def reset(self) -> None:
417 """Forget that handlers were installed."""
418 self._installed = False
420 def _on_hard_exit(self, signum: int, frame: object) -> None:
421 """Stop the fleet on its own thread, then exit with the signal status.
423 Signal handlers all run on the main thread, and a second signal can
424 interrupt this one mid-teardown: the kernel pairs SIGCONT with SIGHUP
425 for an orphaned process group, and Textual's SIGCONT handler raises
426 once the event loop is gone, which aborted the reap half-done and
427 orphaned a loaded fleet. A dedicated non-daemon thread cannot be
428 interrupted by signals, and the interpreter waits for it even as the
429 SystemExit unwinds the main thread.
430 """
431 del frame
432 threading.Thread(
433 target=_teardown_for_signal, args=(signum,), name=_HARD_EXIT_THREAD_NAME
434 ).start()
435 raise SystemExit(_SIGNAL_EXIT_BASE + signum)
438def wait_for_hard_exit_teardown() -> None:
439 """Block until any teardown thread (signal-driven or exit-driven) finishes.
441 Lets ``serve`` hold its OS locks through the fleet stop, so a successor
442 cannot acquire them while this server's models still occupy memory.
443 """
444 for thread in threading.enumerate():
445 if thread.name == _HARD_EXIT_THREAD_NAME:
446 thread.join()
449def reset_services_on_exit() -> None:
450 """Tear the container down on a thread no signal reaches, and wait for it.
452 Engine release waits on the fleet build lock before it releases anything, so
453 a Ctrl-C on the main thread skips the release, and atexit cannot retry: the
454 singleton is already cleared. The teardown thread is non-daemon and takes no
455 signals, so an interrupt breaks only the join here.
456 """
457 if peek_services() is None:
458 return
459 threading.Thread(target=reset_services, name=_HARD_EXIT_THREAD_NAME).start()
460 wait_for_hard_exit_teardown()
463def _teardown_for_signal(signum: int) -> None:
464 """Log the fatal signal, then stop services; runs off the signal handler's thread."""
465 log.info(
466 "Received signal %s; stopping the engine fleet before exit", signal.Signals(signum).name
467 )
468 reset_services()
471_lifecycle = _EngineLifecycle()
474def install_engine_lifecycle_hooks() -> None:
475 """Make a terminal close or ``kill`` stop the engine fleet instead of orphaning it."""
476 _lifecycle.install()
479atexit.register(reset_services_on_exit)