""" Embeddings abstraction for the memory system. Provides an interface for generating embeddings with different backends. The embedding dimension is auto-detected from the model at initialization. The database schema is automatically adjusted to match the model's dimension. Configuration via environment variables - see hindsight_api.config for all env var names. """ import logging import os import warnings from abc import ABC, abstractmethod import httpx from ..config import ( DEFAULT_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL, DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU, DEFAULT_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE, DEFAULT_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_PROVIDER, DEFAULT_LITELLM_API_BASE, ENV_EMBEDDINGS_COHERE_API_KEY, ENV_EMBEDDINGS_LITELLM_SDK_API_KEY, ENV_EMBEDDINGS_LOCAL_FORCE_CPU, ENV_EMBEDDINGS_LOCAL_MODEL, ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE, ENV_EMBEDDINGS_OPENAI_API_KEY, ENV_EMBEDDINGS_OPENAI_BASE_URL, ENV_EMBEDDINGS_OPENAI_MODEL, ENV_EMBEDDINGS_PROVIDER, ENV_EMBEDDINGS_TEI_URL, ENV_LLM_API_KEY, ) logger = logging.getLogger(__name__) class Embeddings(ABC): """ Abstract base class for embedding generation. The embedding dimension is determined by the model and detected at initialization. The database schema is automatically adjusted to match the model's dimension. """ @property @abstractmethod def provider_name(self) -> str: """Return a human-readable name for this provider (e.g., 'local', 'tei').""" pass @property @abstractmethod def dimension(self) -> int: """Return the embedding dimension produced by this model.""" pass @abstractmethod async def initialize(self) -> None: """ Initialize the embedding model asynchronously. This should be called during startup to load/connect to the model and avoid cold start latency on first encode() call. """ pass @abstractmethod def encode(self, texts: list[str]) -> list[list[float]]: """ Generate embeddings for a list of texts. Args: texts: List of text strings to encode Returns: List of embedding vectors (each is a list of floats) """ pass class LocalSTEmbeddings(Embeddings): """ Local embeddings implementation using SentenceTransformers. Call initialize() during startup to load the model and avoid cold starts. The embedding dimension is auto-detected from the model. """ def __init__(self, model_name: str | None = None, force_cpu: bool = False, trust_remote_code: bool = False): """ Initialize local SentenceTransformers embeddings. Args: model_name: Name of the SentenceTransformer model to use. Default: BAAI/bge-small-en-v1.5 force_cpu: Force CPU mode (avoids MPS/XPC issues on macOS in daemon mode). Default: False trust_remote_code: Allow loading models with custom code (security risk). Required for some models with custom architectures. Default: False (disabled for security) """ self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL self.force_cpu = force_cpu self.trust_remote_code = trust_remote_code self._model = None self._dimension: int | None = None @property def provider_name(self) -> str: return "local" @property def dimension(self) -> int: if self._dimension is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") return self._dimension async def initialize(self) -> None: """Load the embedding model.""" if self._model is not None: return try: from sentence_transformers import SentenceTransformer except ImportError: raise ImportError( "sentence-transformers is required for LocalSTEmbeddings. " "Install it with: pip install sentence-transformers" ) logger.info(f"Embeddings: initializing local provider with model {self.model_name}") # Determine device based on hardware availability. # We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors) # which can cause issues when accelerate is installed but no GPU is available. import torch # Force CPU mode if configured (used in daemon mode to avoid MPS/XPC issues on macOS) if self.force_cpu: device = "cpu" logger.info("Embeddings: forcing CPU mode") else: # Check for GPU (CUDA) or Apple Silicon (MPS) # Wrap in try-except to gracefully handle any device detection issues # (e.g., in CI environments or when PyTorch is built without GPU support) device = "cpu" # Default to CPU try: has_gpu = torch.cuda.is_available() or ( hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ) if has_gpu: device = None # Let sentence-transformers auto-detect GPU/MPS except Exception as e: logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}") # Suppress verbose transformers warnings during model loading # This suppresses the "UNEXPECTED" warnings from BertModel which are harmless # but look alarming to users (e.g., "embeddings.position_ids | UNEXPECTED") with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", message=".*was not found in model state dict.*") warnings.filterwarnings("ignore", message=".*UNEXPECTED.*") # Also suppress transformers library logging temporarily transformers_logger = logging.getLogger("transformers") original_level = transformers_logger.level transformers_logger.setLevel(logging.ERROR) try: self._model = SentenceTransformer( self.model_name, device=device, model_kwargs={"low_cpu_mem_usage": False}, trust_remote_code=self.trust_remote_code, ) finally: # Restore original logging level transformers_logger.setLevel(original_level) self._dimension = self._model.get_sentence_embedding_dimension() logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})") def encode(self, texts: list[str]) -> list[list[float]]: """ Generate embeddings for a list of texts. Args: texts: List of text strings to encode Returns: List of embedding vectors """ if self._model is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False) return [emb.tolist() for emb in embeddings] class RemoteTEIEmbeddings(Embeddings): """ Remote embeddings implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API. TEI provides a high-performance inference server for embedding models. See: https://github.com/huggingface/text-embeddings-inference The embedding dimension is auto-detected from the server at initialization. """ def __init__( self, base_url: str, timeout: float = 30.0, batch_size: int = 32, max_retries: int = 3, retry_delay: float = 0.5, ): """ Initialize remote TEI embeddings client. Args: base_url: Base URL of the TEI server (e.g., "http://localhost:8080") timeout: Request timeout in seconds (default: 30.0) batch_size: Maximum batch size for embedding requests (default: 32) max_retries: Maximum number of retries for failed requests (default: 3) retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5) """ self.base_url = base_url.rstrip("/") self.timeout = timeout self.batch_size = batch_size self.max_retries = max_retries self.retry_delay = retry_delay self._client: httpx.Client | None = None self._model_id: str | None = None self._dimension: int | None = None @property def provider_name(self) -> str: return "tei" @property def dimension(self) -> int: if self._dimension is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") return self._dimension def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response: """Make an HTTP request with automatic retries on transient errors.""" import time last_error = None delay = self.retry_delay for attempt in range(self.max_retries + 1): try: if method == "GET": response = self._client.get(url, **kwargs) else: response = self._client.post(url, **kwargs) response.raise_for_status() return response except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e: last_error = e if attempt < self.max_retries: logger.warning( f"TEI request failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s..." ) time.sleep(delay) delay *= 2 # Exponential backoff except httpx.HTTPStatusError as e: # Retry on 5xx server errors if e.response.status_code >= 500 and attempt < self.max_retries: last_error = e logger.warning( f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s..." ) time.sleep(delay) delay *= 2 else: raise raise last_error async def initialize(self) -> None: """Initialize the HTTP client and verify server connectivity.""" if self._client is not None: return logger.info(f"Embeddings: initializing TEI provider at {self.base_url}") self._client = httpx.Client(timeout=self.timeout) # Verify server is reachable and get model info try: response = self._request_with_retry("GET", f"{self.base_url}/info") info = response.json() self._model_id = info.get("model_id", "unknown") # Get dimension from server info or by doing a test embedding if "max_input_length" in info and "model_dtype" in info: # Try to get dimension from info endpoint (some TEI versions expose it) # If not available, do a test embedding pass # Do a test embedding to detect dimension test_response = self._request_with_retry( "POST", f"{self.base_url}/embed", json={"inputs": ["test"]}, ) test_embeddings = test_response.json() if test_embeddings and len(test_embeddings) > 0: self._dimension = len(test_embeddings[0]) logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id}, dim: {self._dimension})") except httpx.HTTPError as e: raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}") def encode(self, texts: list[str]) -> list[list[float]]: """ Generate embeddings using the remote TEI server. Args: texts: List of text strings to encode Returns: List of embedding vectors """ if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") if not texts: return [] all_embeddings = [] # Process in batches for i in range(0, len(texts), self.batch_size): batch = texts[i : i + self.batch_size] try: response = self._request_with_retry( "POST", f"{self.base_url}/embed", json={"inputs": batch}, ) batch_embeddings = response.json() all_embeddings.extend(batch_embeddings) except httpx.HTTPError as e: raise RuntimeError(f"TEI embedding request failed: {e}") return all_embeddings class OpenAIEmbeddings(Embeddings): """ OpenAI embeddings implementation using the OpenAI API. Supports text-embedding-3-small (1536 dims), text-embedding-3-large (3072 dims), and text-embedding-ada-002 (1536 dims, legacy). The embedding dimension is auto-detected from the model at initialization. """ # Known dimensions for OpenAI embedding models MODEL_DIMENSIONS = { "text-embedding-3-small": 1536, "text-embedding-3-large": 3072, "text-embedding-ada-002": 1536, } def __init__( self, api_key: str, model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL, base_url: str | None = None, batch_size: int = 100, max_retries: int = 3, ): """ Initialize OpenAI embeddings client. Args: api_key: OpenAI API key model: OpenAI embedding model name (default: text-embedding-3-small) base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint) batch_size: Maximum batch size for embedding requests (default: 100) max_retries: Maximum number of retries for failed requests (default: 3) """ self.api_key = api_key self.model = model self.base_url = base_url self.batch_size = batch_size self.max_retries = max_retries self._client = None self._dimension: int | None = None @property def provider_name(self) -> str: return "openai" @property def dimension(self) -> int: if self._dimension is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") return self._dimension async def initialize(self) -> None: """Initialize the OpenAI client and detect dimension.""" if self._client is not None: return try: from openai import OpenAI except ImportError: raise ImportError("openai is required for OpenAIEmbeddings. Install it with: pip install openai") base_url_msg = f" at {self.base_url}" if self.base_url else "" logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}{base_url_msg}") # Build client kwargs, only including base_url if set (for Azure or custom endpoints) client_kwargs = {"api_key": self.api_key, "max_retries": self.max_retries} if self.base_url: client_kwargs["base_url"] = self.base_url self._client = OpenAI(**client_kwargs) # Try to get dimension from known models, otherwise do a test embedding if self.model in self.MODEL_DIMENSIONS: self._dimension = self.MODEL_DIMENSIONS[self.model] else: # Do a test embedding to detect dimension response = self._client.embeddings.create( model=self.model, input=["test"], ) if response.data: self._dimension = len(response.data[0].embedding) logger.info(f"Embeddings: OpenAI provider initialized (model: {self.model}, dim: {self._dimension})") def encode(self, texts: list[str]) -> list[list[float]]: """ Generate embeddings using the OpenAI API. Args: texts: List of text strings to encode Returns: List of embedding vectors """ if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") if not texts: return [] all_embeddings = [] # Process in batches for i in range(0, len(texts), self.batch_size): batch = texts[i : i + self.batch_size] response = self._client.embeddings.create( model=self.model, input=batch, ) # Sort by index to ensure correct order batch_embeddings = sorted(response.data, key=lambda x: x.index) all_embeddings.extend([e.embedding for e in batch_embeddings]) return all_embeddings class CohereEmbeddings(Embeddings): """ Cohere embeddings implementation using the Cohere API. Supports embed-english-v3.0 (1024 dims) and embed-multilingual-v3.0 (1024 dims). The embedding dimension is auto-detected from the model at initialization. """ # Known dimensions for Cohere embedding models MODEL_DIMENSIONS = { "embed-english-v3.0": 1024, "embed-multilingual-v3.0": 1024, "embed-english-light-v3.0": 384, "embed-multilingual-light-v3.0": 384, "embed-english-v2.0": 4096, "embed-multilingual-v2.0": 768, } def __init__( self, api_key: str, model: str = DEFAULT_EMBEDDINGS_COHERE_MODEL, base_url: str | None = None, batch_size: int = 96, timeout: float = 60.0, input_type: str = "search_document", ): """ Initialize Cohere embeddings client. Args: api_key: Cohere API key model: Cohere embedding model name (default: embed-english-v3.0) base_url: Custom base URL for Cohere-compatible API (e.g., Azure-hosted endpoint) batch_size: Maximum batch size for embedding requests (default: 96, Cohere's limit) timeout: Request timeout in seconds (default: 60.0) input_type: Input type for embeddings (default: search_document). Options: search_document, search_query, classification, clustering """ self.api_key = api_key self.model = model self.base_url = base_url self.batch_size = batch_size self.timeout = timeout self.input_type = input_type self._client = None self._dimension: int | None = None @property def provider_name(self) -> str: return "cohere" @property def dimension(self) -> int: if self._dimension is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") return self._dimension async def initialize(self) -> None: """Initialize the Cohere client and detect dimension.""" if self._client is not None: return try: import cohere except ImportError: raise ImportError("cohere is required for CohereEmbeddings. Install it with: pip install cohere") base_url_msg = f" at {self.base_url}" if self.base_url else "" logger.info(f"Embeddings: initializing Cohere provider with model {self.model}{base_url_msg}") # Build client kwargs, only including base_url if set (for Azure or custom endpoints) client_kwargs = {"api_key": self.api_key, "timeout": self.timeout} if self.base_url: client_kwargs["base_url"] = self.base_url self._client = cohere.Client(**client_kwargs) # Try to get dimension from known models, otherwise do a test embedding if self.model in self.MODEL_DIMENSIONS: self._dimension = self.MODEL_DIMENSIONS[self.model] else: # Do a test embedding to detect dimension response = self._client.embed( texts=["test"], model=self.model, input_type=self.input_type, ) if response.embeddings and isinstance(response.embeddings, list): self._dimension = len(response.embeddings[0]) logger.info(f"Embeddings: Cohere provider initialized (model: {self.model}, dim: {self._dimension})") def encode(self, texts: list[str]) -> list[list[float]]: """ Generate embeddings using the Cohere API. Args: texts: List of text strings to encode Returns: List of embedding vectors """ if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") if not texts: return [] all_embeddings = [] # Process in batches for i in range(0, len(texts), self.batch_size): batch = texts[i : i + self.batch_size] response = self._client.embed( texts=batch, model=self.model, input_type=self.input_type, ) all_embeddings.extend(response.embeddings) return all_embeddings class LiteLLMEmbeddings(Embeddings): """ LiteLLM embeddings implementation using LiteLLM proxy's /embeddings endpoint. LiteLLM provides a unified interface for multiple embedding providers. The proxy exposes an OpenAI-compatible /embeddings endpoint. See: https://docs.litellm.ai/docs/embedding/supported_embedding Supported providers via LiteLLM: - OpenAI (text-embedding-3-small, text-embedding-ada-002, etc.) - Cohere (embed-english-v3.0, etc.) - prefix with cohere/ - Vertex AI (textembedding-gecko, etc.) - prefix with vertex_ai/ - HuggingFace, Mistral, Voyage AI, etc. The embedding dimension is auto-detected from the model at initialization. """ def __init__( self, api_base: str = DEFAULT_LITELLM_API_BASE, api_key: str | None = None, model: str = DEFAULT_EMBEDDINGS_LITELLM_MODEL, batch_size: int = 100, timeout: float = 60.0, ): """ Initialize LiteLLM embeddings client. Args: api_base: Base URL of the LiteLLM proxy (default: http://localhost:4000) api_key: API key for the LiteLLM proxy (optional, depends on proxy config) model: Embedding model name (default: text-embedding-3-small) Use provider prefix for non-OpenAI models (e.g., cohere/embed-english-v3.0) batch_size: Maximum batch size for embedding requests (default: 100) timeout: Request timeout in seconds (default: 60.0) """ self.api_base = api_base.rstrip("/") self.api_key = api_key self.model = model self.batch_size = batch_size self.timeout = timeout self._client: httpx.Client | None = None self._dimension: int | None = None @property def provider_name(self) -> str: return "litellm" @property def dimension(self) -> int: if self._dimension is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") return self._dimension async def initialize(self) -> None: """Initialize the HTTP client and detect embedding dimension.""" if self._client is not None: return logger.info(f"Embeddings: initializing LiteLLM provider at {self.api_base} with model {self.model}") headers = {"Content-Type": "application/json"} if self.api_key: headers["Authorization"] = f"Bearer {self.api_key}" self._client = httpx.Client(timeout=self.timeout, headers=headers) # Do a test embedding to detect dimension try: response = self._client.post( f"{self.api_base}/embeddings", json={"model": self.model, "input": ["test"]}, ) response.raise_for_status() result = response.json() if result.get("data") and len(result["data"]) > 0: self._dimension = len(result["data"][0]["embedding"]) logger.info(f"Embeddings: LiteLLM provider initialized (model: {self.model}, dim: {self._dimension})") except httpx.HTTPError as e: raise RuntimeError(f"Failed to connect to LiteLLM proxy at {self.api_base}: {e}") def encode(self, texts: list[str]) -> list[list[float]]: """ Generate embeddings using the LiteLLM proxy. Args: texts: List of text strings to encode Returns: List of embedding vectors """ if self._client is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") if not texts: return [] all_embeddings = [] # Process in batches for i in range(0, len(texts), self.batch_size): batch = texts[i : i + self.batch_size] response = self._client.post( f"{self.api_base}/embeddings", json={"model": self.model, "input": batch}, ) response.raise_for_status() result = response.json() # Sort by index to ensure correct order batch_embeddings = sorted(result["data"], key=lambda x: x["index"]) all_embeddings.extend([e["embedding"] for e in batch_embeddings]) return all_embeddings class LiteLLMSDKEmbeddings(Embeddings): """ LiteLLM SDK embeddings for direct API integration. Supports embeddings via LiteLLM SDK without requiring a proxy server. Supported providers: Cohere, OpenAI, Azure OpenAI, HuggingFace, Voyage AI, Together AI, etc. Example model names: - cohere/embed-english-v3.0 - openai/text-embedding-3-small - together_ai/togethercomputer/m2-bert-80M-8k-retrieval - voyage/voyage-2 """ def __init__( self, api_key: str, model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL, api_base: str | None = None, batch_size: int = 100, timeout: float = 60.0, ): """ Initialize LiteLLM SDK embeddings client. Args: api_key: API key for the embedding provider model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0") api_base: Custom base URL for API (optional) batch_size: Maximum batch size for embedding requests (default: 100) timeout: Request timeout in seconds (default: 60.0) """ self.api_key = api_key self.model = model self.api_base = api_base self.batch_size = batch_size self.timeout = timeout self._litellm = None # Will be set during initialization self._dimension: int | None = None @property def provider_name(self) -> str: return "litellm-sdk" @property def dimension(self) -> int: if self._dimension is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") return self._dimension async def initialize(self) -> None: """Initialize the LiteLLM SDK client and detect dimension.""" if self._litellm is not None: return try: import litellm self._litellm = litellm # Store reference except ImportError: raise ImportError("litellm is required for LiteLLMSDKEmbeddings. Install it with: pip install litellm") api_base_msg = f" at {self.api_base}" if self.api_base else "" logger.info(f"Embeddings: initializing LiteLLM SDK provider with model {self.model}{api_base_msg}") # Do a test embedding to detect dimension try: # Build kwargs for embedding call embed_kwargs = { "model": self.model, "input": ["test"], "api_key": self.api_key, "encoding_format": "float", } if self.api_base: embed_kwargs["api_base"] = self.api_base # Use async embedding method (standard in litellm) response = await self._litellm.aembedding(**embed_kwargs) # Extract dimension from response if response.data and len(response.data) > 0: self._dimension = len(response.data[0]["embedding"]) else: raise RuntimeError(f"Unable to detect embedding dimension for model {self.model}") except Exception as e: raise RuntimeError(f"Failed to initialize LiteLLM SDK embeddings: {e}") logger.info(f"Embeddings: LiteLLM SDK provider initialized (model: {self.model}, dim: {self._dimension})") def encode(self, texts: list[str]) -> list[list[float]]: """ Generate embeddings using the LiteLLM SDK. Args: texts: List of text strings to encode Returns: List of embedding vectors (one per input text) """ if self._litellm is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") if not texts: return [] all_embeddings = [] # Process in batches for i in range(0, len(texts), self.batch_size): batch = texts[i : i + self.batch_size] try: # Build kwargs for embedding call embed_kwargs = { "model": self.model, "input": batch, "api_key": self.api_key, "encoding_format": "float", } if self.api_base: embed_kwargs["api_base"] = self.api_base # Use sync embedding (litellm doesn't have async in thread-safe way) response = self._litellm.embedding(**embed_kwargs) # Extract embeddings from response # Sort by index to ensure correct order batch_embeddings = sorted(response.data, key=lambda x: x.get("index", 0)) all_embeddings.extend([e["embedding"] for e in batch_embeddings]) except Exception as e: import traceback logger.error( f"Error in LiteLLM embedding for batch starting at index {i}: {e}\n" f"Traceback: {traceback.format_exc()}" ) raise return all_embeddings def create_embeddings_from_env() -> Embeddings: """ Create an Embeddings instance based on configuration. Reads configuration via get_config() to ensure consistency across the codebase. Returns: Configured Embeddings instance """ from ..config import get_config config = get_config() provider = config.embeddings_provider.lower() if provider == "tei": url = config.embeddings_tei_url if not url: raise ValueError(f"{ENV_EMBEDDINGS_TEI_URL} is required when {ENV_EMBEDDINGS_PROVIDER} is 'tei'") return RemoteTEIEmbeddings(base_url=url) elif provider == "local": return LocalSTEmbeddings( model_name=config.embeddings_local_model, force_cpu=config.embeddings_local_force_cpu, trust_remote_code=config.embeddings_local_trust_remote_code, ) elif provider == "openai": # Use dedicated embeddings API key, or fall back to LLM API key api_key = os.environ.get(ENV_EMBEDDINGS_OPENAI_API_KEY) or os.environ.get(ENV_LLM_API_KEY) if not api_key: raise ValueError( f"{ENV_EMBEDDINGS_OPENAI_API_KEY} or {ENV_LLM_API_KEY} is required " f"when {ENV_EMBEDDINGS_PROVIDER} is 'openai'" ) model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL) base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url) elif provider == "cohere": api_key = config.embeddings_cohere_api_key if not api_key: raise ValueError(f"{ENV_EMBEDDINGS_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'") return CohereEmbeddings( api_key=api_key, model=config.embeddings_cohere_model, base_url=config.embeddings_cohere_base_url, ) elif provider == "litellm": return LiteLLMEmbeddings( api_base=config.embeddings_litellm_api_base, api_key=config.embeddings_litellm_api_key, model=config.embeddings_litellm_model, ) elif provider == "litellm-sdk": api_key = config.embeddings_litellm_sdk_api_key if not api_key: raise ValueError( f"{ENV_EMBEDDINGS_LITELLM_SDK_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'litellm-sdk'" ) return LiteLLMSDKEmbeddings( api_key=api_key, model=config.embeddings_litellm_sdk_model, api_base=config.embeddings_litellm_sdk_api_base, ) else: raise ValueError( f"Unknown embeddings provider: {provider}. " f"Supported: 'local', 'tei', 'openai', 'cohere', 'litellm', 'litellm-sdk'" )