a209ef1ae2
* fix: parse query params from base_url in OpenAI embeddings client The OpenAI-compatible LLM provider already parses query parameters (e.g. ?api-version=xxx for Azure OpenAI) from the base_url and passes them as default_query to the OpenAI client. However, the OpenAI embeddings provider did not do this, causing Azure OpenAI embeddings to fail with 404 errors at runtime. This applies the same URL parsing logic from the LLM provider to the embeddings provider, enabling Azure OpenAI embeddings to work correctly. * ci: add workflow to build fork Docker image * ci: add slim image build (no local models) * ci: remove fork build workflow per review request --------- Co-authored-by: Antoine Khater <ak@ptgroup.eu>
946 lines
34 KiB
Python
946 lines
34 KiB
Python
"""
|
|
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
|
|
from urllib.parse import parse_qs, urlparse, urlunparse
|
|
|
|
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)
|
|
# Parse query parameters from base_url (e.g. ?api-version=xxx for Azure OpenAI)
|
|
# and pass them as default_query so they're included in every request.
|
|
client_kwargs = {"api_key": self.api_key, "max_retries": self.max_retries}
|
|
if self.base_url:
|
|
parsed = urlparse(self.base_url)
|
|
if parsed.query:
|
|
clean_url = urlunparse(parsed._replace(query=""))
|
|
client_kwargs["base_url"] = clean_url
|
|
default_query = {k: v[0] for k, v in parse_qs(parsed.query).items()}
|
|
client_kwargs["default_query"] = default_query
|
|
self.base_url = clean_url
|
|
else:
|
|
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'"
|
|
)
|