""" Batch processing logic for embedding services. """ import asyncio import inspect import os from typing import Any, cast import httpx import openai from ...config.logfire_config import safe_span, search_logger from ..credential_service import credential_service from ..llm_provider_service import create_embedding_client from ..threading_service import get_threading_service from .embedding_exceptions import ( EmbeddingAPIError, ) from .models import EmbeddingBatchResult async def create_embeddings_batch( texts: list[str], progress_callback: Any | None = None, ) -> EmbeddingBatchResult: """ Create embeddings for multiple texts with graceful failure handling and provider failover. This function attempts to use the primary embedding provider, and on failure, transparently switches to a configured fallback provider. Args: texts: List of texts to create embeddings for progress_callback: Optional callback for progress reporting Returns: EmbeddingBatchResult with successful embeddings and failure details """ if not texts: return EmbeddingBatchResult() # Dynamic offline mode handling using sentence-transformers if os.getenv("OFFLINE_MODE", "false").lower() == "true": search_logger.info("OFFLINE_MODE is enabled. Generating embeddings locally using 'all-MiniLM-L6-v2'.") try: from sentence_transformers import SentenceTransformer model = SentenceTransformer("all-MiniLM-L6-v2") # SentenceTransformer encode executes synchronously, wrap in executor to keep it async friendly loop = asyncio.get_event_loop() embeddings_np = await loop.run_in_executor( None, lambda: model.encode(texts, show_progress_bar=False) ) result = EmbeddingBatchResult() for text_item, emb in zip(texts, embeddings_np, strict=False): result.add_success(emb.tolist(), text_item) return result except Exception as e: search_logger.error(f"Failed to generate local embeddings: {e}", exc_info=True) result = EmbeddingBatchResult() final_error = EmbeddingAPIError(f"Local embedding failure: {str(e)}", original_error=e) for text_item in texts: result.add_failure(text_item, final_error) return result # Validate that all items in texts are strings validated_texts = [] for i, text in enumerate(texts): if not isinstance(text, str): search_logger.error(f"Invalid text type at index {i}: {type(text)}, value: {text}", exc_info=True) try: validated_texts.append(str(text)) except Exception as e: search_logger.error(f"Failed to convert text at index {i} to string: {e}", exc_info=True) validated_texts.append("") else: validated_texts.append(text) texts = validated_texts result = EmbeddingBatchResult() threading_service = get_threading_service() with safe_span("create_embeddings_batch", text_count=len(texts), total_chars=sum(len(t) for t in texts)) as span: try: configs = await credential_service.get_embedding_provider_configs() if not configs: raise ValueError("No valid embedding providers configured.") last_exception = None for idx, config in enumerate(configs): client: openai.AsyncOpenAI | None = None provider_name = config.get("provider", "unknown") is_last_provider = idx == len(configs) - 1 try: search_logger.info(f"Attempting embedding creation with provider: {provider_name}") client = await create_embedding_client(config) rag_settings = await credential_service.get_credentials_by_category("rag_strategy") batch_size = int(rag_settings.get("EMBEDDING_BATCH_SIZE", "100")) embedding_dimensions = int(rag_settings.get("EMBEDDING_DIMENSIONS", "768")) all_batches_succeeded_for_provider = True for i in range(0, len(texts), batch_size): batch = texts[i : i + batch_size] batch_index = i // batch_size try: # PERFORMANCE: Replaced sum(len(text.split())...) with a faster loop and .count(' ') # which avoids allocating lists for every text chunk. batch_tokens_raw = 0 for text in batch: batch_tokens_raw += text.count(" ") + 1 batch_tokens = int(batch_tokens_raw * 1.3) rate_limit_callback = None if progress_callback: async def rate_limit_callback(data: dict, res=result): processed = res.success_count + res.failure_count message = f"Rate limited: {data.get('message', 'Waiting...')}" await progress_callback(message, (processed / len(texts)) * 100) async with threading_service.rate_limited_operation( batch_tokens, rate_limit_callback ): # Re-introduced rate limiting retry_count = 0 max_retries = 3 while retry_count < max_retries: try: embedding_model = config.get("embedding_model") if provider_name == "google": # Native Google API call (using proven v1beta + header variant) async with httpx.AsyncClient(timeout=20.0) as http_client: # Use gemini-embedding-001 which is proven stable # Fallback to config model if not explicit, then to a stable default stable_model = config.get("embedding_model") if not stable_model: raise ValueError( "embedding_model is not configured in provider settings" ) api_key_to_use = ( (config.get("api_key") or os.getenv("GEMINI_API_KEY") or "") .strip() .strip('"') .strip("'") ) url = f"https://generativelanguage.googleapis.com/v1beta/models/{stable_model}:embedContent" headers = {"x-goog-api-key": api_key_to_use} for text_item in batch: payload = { "content": {"parts": [{"text": text_item}]}, "outputDimensionality": 768, } resp = await http_client.post(url, headers=headers, json=payload) if resp.status_code == 200: data = resp.json() result.add_success(data["embedding"]["values"], text_item) else: search_logger.error( f"Google native API failed: Status {resp.status_code}, Body: {resp.text}" ) raise EmbeddingAPIError( f"Google error {resp.status_code}: {resp.text}" ) else: # Standard OpenAI-compatible call if provider_name != "google": response = await client.embeddings.create( model=cast(str, embedding_model), input=batch, dimensions=embedding_dimensions, ) else: response = await client.embeddings.create( model=cast(str, embedding_model), input=batch ) for item, text_item in zip(response.data, batch, strict=False): result.add_success(item.embedding, text_item) break except openai.RateLimitError as e: error_message = str(e) if "insufficient_quota" in error_message: search_logger.error( f"Provider {provider_name} has insufficient quota.", exc_info=True ) raise retry_count += 1 if retry_count >= max_retries: search_logger.error( f"Rate limit retries exceeded for provider {provider_name}. Batch {batch_index}.", exc_info=True, ) raise wait_time = 2**retry_count search_logger.warning( f"Rate limit hit for {provider_name}. Batch {batch_index}. Waiting {wait_time}s before retry {retry_count}/{max_retries}" ) await asyncio.sleep(wait_time) except Exception as e: # Re-raise specific exceptions that should trigger provider failover if isinstance( e, openai.AuthenticationError | openai.PermissionDeniedError | openai.APIConnectionError | openai.RateLimitError, ): raise all_batches_succeeded_for_provider = False search_logger.error( f"Batch {batch_index} failed for provider {provider_name}: {e}", exc_info=True ) for text in batch: result.add_failure( text, EmbeddingAPIError(f"Batch {batch_index} failed: {e}", original_error=e), batch_index, ) if progress_callback: processed = result.success_count + result.failure_count progress = (processed / len(texts)) * 100 message = f"Processed {processed}/{len(texts)} texts" if result.has_failures: message += f" ({result.failure_count} failed)" await progress_callback(message, progress) await asyncio.sleep(0.01) if all_batches_succeeded_for_provider: span.set_attribute("provider_used", provider_name) return result except Exception as e: last_exception = e search_logger.warning( f"Provider '{provider_name}' failed with {type(e).__name__}: {e}. Trying next if available." ) if is_last_provider: search_logger.error( f"All embedding providers failed. Final source of failure '{provider_name}': {e}", exc_info=True, ) raise e finally: if client: # Safe close that handles both real AsyncOpenAI clients and MagicMocks try: # Try standard close method close_method = getattr(client, "close", None) if callable(close_method): is_coroutine = inspect.iscoroutinefunction(close_method) or inspect.isawaitable( close_method ) if is_coroutine: await close_method() else: close_method() # Fallback for older clients or mocks elif hasattr(client, "aclose"): await client.aclose() except Exception as cleanup_err: search_logger.warning(f"Error closing client: {cleanup_err}") if last_exception: raise last_exception raise ValueError("No embedding providers were attempted. Please verify API Key configurations in Settings.") except Exception as e: span.set_attribute("catastrophic_failure", True) search_logger.error(f"Catastrophic failure in batch embedding: {e}", exc_info=True) processed_count = result.success_count + result.failure_count if processed_count < len(texts): final_error = EmbeddingAPIError(f"Catastrophic failure: {str(e)}", original_error=e) for text in texts[processed_count:]: result.add_failure(text, final_error) return result