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