Coverage for src/qdrant_loader/core/embedding/embedding_service.py: 75%
167 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1import asyncio
2import logging
3import time
4from collections.abc import Sequence
5from importlib import import_module
7import requests
8import tiktoken
10from qdrant_loader.config import Settings
11from qdrant_loader.core.document import Document
12from qdrant_loader.utils.logging import LoggingConfig
14logger = LoggingConfig.get_logger(__name__)
17class EmbeddingService:
18 """Service for generating embeddings using provider-agnostic API (via core)."""
20 def __init__(self, settings: Settings):
21 """Initialize the embedding service.
23 Args:
24 settings: The application settings containing API key and endpoint.
25 """
26 self.settings = settings
27 # Build LLM settings from global config and create provider
28 llm_settings = settings.llm_settings
29 factory_mod = import_module("qdrant_loader_core.llm.factory")
30 create_provider = factory_mod.create_provider
31 self.provider = create_provider(llm_settings)
32 self.model = llm_settings.models.get(
33 "embeddings", settings.global_config.embedding.model
34 )
35 self.tokenizer = (
36 llm_settings.tokenizer or settings.global_config.embedding.tokenizer
37 )
38 self.batch_size = settings.global_config.embedding.batch_size
40 # Initialize tokenizer based on configuration
41 if self.tokenizer == "none":
42 self.encoding = None
43 else:
44 try:
45 self.encoding = tiktoken.get_encoding(self.tokenizer)
46 except Exception as e:
47 logger.warning(
48 "Failed to initialize tokenizer, falling back to simple character counting",
49 error=str(e),
50 tokenizer=self.tokenizer,
51 )
52 self.encoding = None
54 self.last_request_time = 0
55 self.min_request_interval = (
56 settings.global_config.embedding.min_request_interval
57 )
59 # Retry configuration for network resilience
60 self.max_retries = 3
61 self.base_retry_delay = 1.0 # Start with 1 second
62 self.max_retry_delay = 30.0 # Cap at 30 seconds
64 async def _apply_rate_limit(self):
65 """Apply rate limiting between API requests."""
66 current_time = time.time()
67 time_since_last_request = current_time - self.last_request_time
68 if time_since_last_request < self.min_request_interval:
69 await asyncio.sleep(self.min_request_interval - time_since_last_request)
70 self.last_request_time = time.time()
72 async def _retry_with_backoff(self, operation, operation_name: str, **kwargs):
73 """Execute an operation with exponential backoff retry logic.
75 Args:
76 operation: The async operation to retry
77 operation_name: Name of the operation for logging
78 **kwargs: Additional arguments passed to the operation
80 Returns:
81 The result of the successful operation
83 Raises:
84 The last exception if all retries fail
85 """
86 last_exception = None
88 for attempt in range(self.max_retries + 1): # +1 for initial attempt
89 try:
90 if attempt > 0:
91 # Calculate exponential backoff delay
92 delay = min(
93 self.base_retry_delay * (2 ** (attempt - 1)),
94 self.max_retry_delay,
95 )
96 logger.warning(
97 f"Retrying {operation_name} after network error",
98 attempt=attempt,
99 max_retries=self.max_retries,
100 delay_seconds=delay,
101 last_error=str(last_exception) if last_exception else None,
102 )
103 await asyncio.sleep(delay)
105 # Execute the operation
106 result = await operation(**kwargs)
108 if attempt > 0:
109 logger.info(
110 f"Successfully recovered {operation_name} after retries",
111 successful_attempt=attempt + 1,
112 total_attempts=attempt + 1,
113 )
115 return result
117 except (
118 TimeoutError,
119 requests.exceptions.Timeout,
120 requests.exceptions.ConnectionError,
121 requests.exceptions.HTTPError,
122 ConnectionError,
123 OSError,
124 ) as e:
125 last_exception = e
127 if attempt == self.max_retries:
128 logger.error(
129 f"All retry attempts failed for {operation_name}",
130 total_attempts=attempt + 1,
131 final_error=str(e),
132 error_type=type(e).__name__,
133 )
134 raise
136 logger.warning(
137 f"Network error in {operation_name}, will retry",
138 attempt=attempt + 1,
139 max_retries=self.max_retries,
140 error=str(e),
141 error_type=type(e).__name__,
142 )
144 except Exception as e:
145 # For non-network errors, don't retry
146 logger.error(
147 f"Non-retryable error in {operation_name}",
148 error=str(e),
149 error_type=type(e).__name__,
150 )
151 raise
153 # This should never be reached, but just in case
154 if last_exception:
155 raise last_exception
156 raise RuntimeError(f"Unexpected error in retry logic for {operation_name}")
158 async def get_embeddings(
159 self, texts: Sequence[str | Document]
160 ) -> list[list[float]]:
161 """Get embeddings for a list of texts."""
162 if not texts:
163 return []
165 # Extract content if texts are Document objects
166 contents = [
167 text.content if isinstance(text, Document) else text for text in texts
168 ]
170 # Filter out empty, None, or invalid content
171 valid_contents = []
172 valid_indices = []
173 for i, content in enumerate(contents):
174 if content and isinstance(content, str) and content.strip():
175 valid_contents.append(content.strip())
176 valid_indices.append(i)
177 else:
178 logger.warning(
179 f"Skipping invalid content at index {i}: {repr(content)}"
180 )
182 if not valid_contents:
183 logger.warning(
184 "No valid content found in batch, returning empty embeddings"
185 )
186 return []
188 logger.debug(
189 "Starting batch embedding process",
190 total_texts=len(contents),
191 valid_texts=len(valid_contents),
192 filtered_out=len(contents) - len(valid_contents),
193 )
195 # Validate and split content based on token limits
196 # Use configurable token limits from settings
197 MAX_TOKENS_PER_REQUEST = (
198 self.settings.global_config.embedding.max_tokens_per_request
199 )
200 MAX_TOKENS_PER_CHUNK = (
201 self.settings.global_config.embedding.max_tokens_per_chunk
202 )
204 validated_contents = []
205 truncated_count = 0
206 for content in valid_contents:
207 token_count = self.count_tokens(content)
208 if token_count > MAX_TOKENS_PER_CHUNK:
209 truncated_count += 1
210 logger.warning(
211 "Content exceeds maximum token limit, truncating",
212 content_length=len(content),
213 token_count=token_count,
214 max_tokens=MAX_TOKENS_PER_CHUNK,
215 )
216 # Truncate content to fit within token limit
217 if self.encoding is not None:
218 # Use tokenizer to truncate precisely
219 tokens = self.encoding.encode(content)
220 truncated_tokens = tokens[:MAX_TOKENS_PER_CHUNK]
221 truncated_content = self.encoding.decode(truncated_tokens)
222 validated_contents.append(truncated_content)
223 else:
224 # Fallback to character-based truncation (rough estimate)
225 # Assume ~4 characters per token on average
226 max_chars = MAX_TOKENS_PER_CHUNK * 4
227 validated_contents.append(content[:max_chars])
228 else:
229 validated_contents.append(content)
231 if truncated_count > 0:
232 logger.info(
233 f"⚠️ Truncated {truncated_count} content items due to token limits. You might want to adjust chunk size and/or max tokens settings in config.yaml"
234 )
236 # Create smart batches that respect token limits
237 embeddings = []
238 current_batch = []
239 current_batch_tokens = 0
240 batch_count = 0
242 for content in validated_contents:
243 content_tokens = self.count_tokens(content)
245 # Check if adding this content would exceed the token limit
246 if current_batch and (
247 current_batch_tokens + content_tokens > MAX_TOKENS_PER_REQUEST
248 ):
249 # Process current batch
250 batch_count += 1
251 batch_embeddings = await self._process_batch(current_batch)
252 embeddings.extend(batch_embeddings)
254 # Start new batch
255 current_batch = [content]
256 current_batch_tokens = content_tokens
257 else:
258 # Add to current batch
259 current_batch.append(content)
260 current_batch_tokens += content_tokens
262 # Process final batch if it exists
263 if current_batch:
264 batch_count += 1
265 batch_embeddings = await self._process_batch(current_batch)
266 embeddings.extend(batch_embeddings)
268 logger.info(
269 f"🔗 Generated embeddings: {len(embeddings)} items in {batch_count} batches"
270 )
272 if len(valid_indices) != len(embeddings):
273 raise ValueError(
274 "Embedding count mismatch: "
275 f"expected {len(valid_indices)} embeddings for valid contents, "
276 f"got {len(embeddings)}"
277 )
279 # Reconstruct full-length result aligned with original input indices.
280 # Invalid entries (filtered out above) get an empty list as placeholder so
281 # callers can zip safely without index shift.
282 full_result: list[list[float]] = [[] for _ in contents]
283 for idx, embedding in zip(valid_indices, embeddings, strict=False):
284 full_result[idx] = embedding
285 return full_result
287 async def _process_batch(self, batch: list[str]) -> list[list[float]]:
288 """Process a single batch of content for embeddings.
290 Args:
291 batch: List of content strings to embed
293 Returns:
294 List of embedding vectors
295 """
296 if not batch:
297 return []
299 batch_num = getattr(self, "_batch_counter", 0) + 1
300 self._batch_counter = batch_num
302 # Optimized: Only calculate tokens for debug when debug logging is enabled
303 if logging.getLogger().isEnabledFor(logging.DEBUG):
304 logger.debug(
305 "Processing embedding batch",
306 batch_num=batch_num,
307 batch_size=len(batch),
308 total_tokens=sum(self.count_tokens(content) for content in batch),
309 )
311 await self._apply_rate_limit()
313 # Use retry logic for network resilience
314 return await self._retry_with_backoff(
315 self._execute_embedding_request,
316 f"embedding batch {batch_num}",
317 batch=batch,
318 batch_num=batch_num,
319 )
321 async def _execute_embedding_request(
322 self, batch: list[str], batch_num: int
323 ) -> list[list[float]]:
324 """Execute the actual embedding request (used by retry logic).
326 Args:
327 batch: List of content strings to embed
328 batch_num: Batch number for logging
330 Returns:
331 List of embedding vectors
332 """
333 try:
334 # Use core provider for embeddings
335 embeddings_client = self.provider.embeddings()
336 batch_embeddings = await embeddings_client.embed(batch)
338 logger.debug(
339 "Completed batch processing",
340 batch_num=batch_num,
341 processed_embeddings=len(batch_embeddings),
342 )
344 return batch_embeddings
346 except Exception as e:
347 logger.debug(
348 "Embedding request failed",
349 batch_num=batch_num,
350 error=str(e),
351 error_type=type(e).__name__,
352 )
353 raise # Let the retry logic handle it
355 async def get_embedding(self, text: str) -> list[float]:
356 """Get embedding for a single text."""
357 # Validate input
358 if not text or not isinstance(text, str) or not text.strip():
359 logger.warning(f"Invalid text for embedding: {repr(text)}")
360 raise ValueError(
361 "Invalid text for embedding: text must be a non-empty string"
362 )
364 clean_text = text.strip()
366 # Use retry logic for network resilience
367 return await self._retry_with_backoff(
368 self._execute_single_embedding_request, "single embedding", text=clean_text
369 )
371 async def _execute_single_embedding_request(self, text: str) -> list[float]:
372 """Execute a single embedding request (used by retry logic).
374 Args:
375 text: The text to embed
377 Returns:
378 The embedding vector
379 """
380 try:
381 await self._apply_rate_limit()
382 embeddings_client = self.provider.embeddings()
383 vectors = await embeddings_client.embed([text])
384 return vectors[0]
385 except Exception as e:
386 logger.debug(
387 "Single embedding request failed",
388 error=str(e),
389 error_type=type(e).__name__,
390 )
391 raise # Let the retry logic handle it
393 def count_tokens(self, text: str) -> int:
394 """Count the number of tokens in a text string."""
395 if self.encoding is None:
396 # Fallback to character count if no tokenizer is available
397 return len(text)
398 return len(self.encoding.encode(text))
400 def count_tokens_batch(self, texts: list[str]) -> list[int]:
401 """Count the number of tokens in a list of text strings."""
402 return [self.count_tokens(text) for text in texts]
404 def get_embedding_dimension(self) -> int:
405 """Get the dimension of the embedding vectors."""
406 # Prefer vector size from unified settings when available
407 dimension = (
408 self.settings.llm_settings.embeddings.vector_size
409 or self.settings.global_config.embedding.vector_size
410 )
411 if not dimension:
412 logger.warning(
413 "Embedding dimension not set in config; using 1024 (deprecated default). Set global.llm.embeddings.vector_size."
414 )
415 return 1024
416 return int(dimension)