Coverage for src/lilbee/retrieval/embedder.py: 100%
110 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Thin wrapper around LLM provider embeddings API."""
3import logging
4import threading
6import numpy as np
8from lilbee.core.config import Config, cfg
9from lilbee.core.vectors import Vector
10from lilbee.data.extract.chunk import CHARS_PER_TOKEN
11from lilbee.modelhub.install_state import InstallState, install_state
12from lilbee.modelhub.registry import ModelRegistry
13from lilbee.providers.base import LLMProvider
14from lilbee.providers.model_ref import ProviderModelRef, parse_model_ref
15from lilbee.retrieval.embedding_profiles import EmbeddingProfile, resolve_embedding_profile
16from lilbee.runtime.progress import DetailedProgressCallback, EmbedEvent, EventType, noop_callback
18log = logging.getLogger(__name__)
21def _name_base(ref: ProviderModelRef) -> str:
22 return ref.name.split(":")[0].lower().replace(" ", "-")
25def _remote_sees_model(ref: ProviderModelRef, provider: LLMProvider) -> bool:
26 try:
27 available = provider.list_models()
28 except Exception:
29 log.debug("provider list_models failed during availability check", exc_info=True)
30 return False
31 base = _name_base(ref)
32 return any(base in m.lower().replace(" ", "-") for m in available)
35def is_model_installed(model: str) -> bool:
36 """True if *model* loads from the registry or from a loose GGUF file.
38 Provider-free by design: paint paths need an installed check that cannot
39 build the services container (a cold build eager-starts the worker pool).
40 """
41 return install_state(model, ModelRegistry(cfg.models_dir)) is not InstallState.MISSING
44def is_model_available(model: str, provider: LLMProvider) -> bool:
45 """Return True if *model* resolves via *provider* or the native registry.
47 Remote-prefixed refs (``ollama/`` and API providers) skip the native
48 probe since they resolve through the SDK backend at call time.
49 """
50 if not model:
51 return False
52 ref = parse_model_ref(model)
53 if _remote_sees_model(ref, provider):
54 return True
55 if ref.is_remote:
56 return False
57 return is_model_installed(model)
60class Embedder:
61 """Embedding wrapper: truncates, batches, validates vectors, and counts truncations."""
63 def __init__(self, config: Config, provider: LLMProvider) -> None:
64 self._config = config
65 self._provider = provider
66 self._truncated_total = 0
67 self._truncated_lock = threading.Lock()
69 @property
70 def embed_char_budget(self) -> int:
71 """Effective char limit, never below the chunker's max chunk size.
73 ``max_embed_chars`` guards the embed model's context; clamping it up to
74 ``chunk_size * CHARS_PER_TOKEN`` stops a finished full-budget chunk from
75 silently losing its tail to a limit set below what the chunker emits.
76 """
77 return max(self._config.max_embed_chars, self._config.chunk_size * CHARS_PER_TOKEN)
79 @property
80 def batch_char_budget(self) -> int:
81 """Per-request char cap: a full packed batch of maximum-size chunks.
83 The target sequences-per-request is ``embed_batch_sequences`` (tunable to
84 keep a multi-GPU fleet's batching slots full); the engine still re-splits
85 to its physical batch, so it is an upper bound, not a guarantee.
86 """
87 return self._config.embed_batch_sequences * self._config.chunk_size * CHARS_PER_TOKEN
89 @property
90 def truncated_total(self) -> int:
91 """Cumulative count of chunks truncated since process start (thread-safe read)."""
92 with self._truncated_lock:
93 return self._truncated_total
95 def truncate(self, text: str, reserved: int = 0) -> str:
96 """Truncate text to the embed char budget, counting any truncation.
98 *reserved* is charged against the budget so an instruction prefix
99 prepended afterwards still leaves the sent text inside the guard.
100 """
101 budget = max(1, self.embed_char_budget - reserved)
102 if len(text) <= budget:
103 return text
104 log.debug("Truncating chunk from %d to %d chars for embedding", len(text), budget)
105 with self._truncated_lock:
106 self._truncated_total += 1
107 return text[:budget]
109 def validate_vector(self, vector: Vector) -> None:
110 """Validate embedding vector dimension and values."""
111 if len(vector) != self._config.embedding_dim:
112 raise ValueError(
113 f"Embedding dimension mismatch: expected {self._config.embedding_dim}, "
114 f"got {len(vector)}"
115 )
116 arr = np.asarray(vector, dtype=np.float64)
117 bad = np.where(~np.isfinite(arr))[0]
118 if bad.size:
119 i = int(bad[0])
120 raise ValueError(f"Embedding contains invalid value at index {i}: {vector[i]}")
122 def validate_model(self) -> bool:
123 """Availability gate for the startup and ingest paths: warns when it fails.
125 Same probe as :meth:`embedding_available`, but this is the entry-point
126 check whose whole job is to not fail silently -- a missing model here
127 means every chunk of the run ahead will fail to embed, and the operator
128 needs to hear that once, up front, rather than as a per-file error much
129 later. Callers that can genuinely degrade (search falling back to
130 keyword) ask :meth:`embedding_available` and stay quiet.
131 """
132 if self.embedding_available():
133 return True
134 if not self._config.embedding_model:
135 # Unconfigured is a state, not a failure: nothing was chosen, so
136 # there is nothing to warn about. One INFO line for the curious.
137 log.info("No embedding model configured; embedding is off until one is set.")
138 return False
139 log.warning(
140 "Embedding model %r is not available; embedding will fail. "
141 "Pull it or set a different embedding_model.",
142 self._config.embedding_model,
143 )
144 return False
146 def embedding_available(self) -> bool:
147 """Return True if the embedding model can be resolved.
149 Checks the provider model list and the native registry path
150 resolution. Returns True if either finds the model.
151 """
152 return is_model_available(self._config.embedding_model, self._provider)
154 def _profile(self) -> EmbeddingProfile:
155 """Instruction profile for the configured embedder (symmetric if unrecognized)."""
156 return resolve_embedding_profile(self._config.embedding_model)
158 def embed(self, text: str) -> Vector:
159 """Embed a single document string, return vector."""
160 return self._embed_with([text], self._profile().doc_prefix)[0]
162 def embed_query(self, text: str) -> Vector:
163 """Embed a single query string, applying the model's query instruction if any."""
164 return self._embed_with([text], self._profile().query_instruction)[0]
166 def embed_batch(
167 self,
168 texts: list[str],
169 *,
170 source: str = "",
171 on_progress: DetailedProgressCallback = noop_callback,
172 ) -> list[Vector]:
173 """Embed document texts with adaptive batching, return list of vectors."""
174 return self._embed_with(
175 texts, self._profile().doc_prefix, source=source, on_progress=on_progress
176 )
178 def embed_query_batch(
179 self,
180 texts: list[str],
181 *,
182 source: str = "",
183 on_progress: DetailedProgressCallback = noop_callback,
184 ) -> list[Vector]:
185 """Embed query texts (query instruction applied), adaptive batching."""
186 return self._embed_with(
187 texts, self._profile().query_instruction, source=source, on_progress=on_progress
188 )
190 def _embed_with(
191 self,
192 texts: list[str],
193 prefix: str,
194 *,
195 source: str = "",
196 on_progress: DetailedProgressCallback = noop_callback,
197 ) -> list[Vector]:
198 """Adaptive-batched embed of *texts*, each prefixed (query/document instruction).
200 Fires ``embed`` progress events per batch when *on_progress* is provided.
201 """
202 if not texts:
203 return []
204 total_chunks = len(texts)
205 max_batch_chars = self.batch_char_budget
206 vectors: list[Vector] = []
207 batch: list[str] = []
208 batch_chars = 0
209 for text in texts:
210 # The prefix is charged against the budget: it is part of what the
211 # embed model receives, so counting only the text would ship
212 # budget+len(prefix) chars and lose the tail the clamp protects.
213 truncated = prefix + self.truncate(text, reserved=len(prefix))
214 chunk_len = len(truncated)
215 if batch and batch_chars + chunk_len > max_batch_chars:
216 vectors.extend(self._provider.embed(batch))
217 on_progress(
218 EventType.EMBED,
219 EmbedEvent(file=source, chunk=len(vectors), total_chunks=total_chunks),
220 )
221 batch = []
222 batch_chars = 0
223 batch.append(truncated)
224 batch_chars += chunk_len
225 if batch:
226 vectors.extend(self._provider.embed(batch))
227 on_progress(
228 EventType.EMBED,
229 EmbedEvent(file=source, chunk=len(vectors), total_chunks=total_chunks),
230 )
231 for vec in vectors:
232 self.validate_vector(vec)
233 return vectors