Python Embedding Adapters

Copyclass EmbeddingBase(ABC):
    """Initialized a base embedding class
 
    :param config: Embedding configuration option class, defaults to None
    :type config: Optional[BaseEmbedderConfig], optional
    """
 
    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
        if config is None:
            self.config = BaseEmbedderConfig()
        else:
            self.config = config
 

Copy@abstractmethod
def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]]):
    """
    Get the embedding for the given text.
 
    Args:
        text (str): The text to embed.
        memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
    Returns:
        list: The embedding vector.
    """
    pass
 
def embed_batch(self, texts, memory_action="add"):
    """Embed multiple texts. Override in subclasses for native batch support.
 
    Default implementation calls embed() sequentially for each text.
    Subclasses with native batch APIs (e.g., OpenAI) should override
    this for better performance.
 
    Args:
        texts: List of text strings to embed.
        memory_action: The action context ("add", "search", "update").
 
    Returns:
        List of embedding vectors (list of floats), one per input text.
    """
    return [self.embed(text, memory_action) for text in texts]

Copydef _normalize_vector(self, embeddings):
    """Normalize the embedding to a unit vector."""
    emb = np.array(embeddings)
    norm_emb = emb / np.linalg.norm(emb)
    return norm_emb.tolist()

Copyclass GoogleGenAIEmbedding(EmbeddingBase):
    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
        super().__init__(config)
 
        self.config.model = self.config.model or "models/gemini-embedding-001"
        self.config.embedding_dims = self.config.embedding_dims or self.config.output_dimensionality or 768
 
        api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
 
        self.client = genai.Client(api_key=api_key)

Copydef embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
    """
    Get the embedding for the given text using Google Generative AI.
    Args:
        text (str): The text to embed.
        memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
    Returns:
        list: The embedding vector.
    """
    text = text.replace("\n", " ")
 
    # Create config for embedding parameters
    config = types.EmbedContentConfig(output_dimensionality=self.config.embedding_dims)
 
    # Call the embed_content method with the correct parameters
    response = self.client.models.embed_content(model=self.config.model, contents=text, config=config)
 
    return response.embeddings[0].values
 

Copydef embed_batch(self, texts, memory_action="add"):
    if not texts:
        return []
    config = types.EmbedContentConfig(output_dimensionality=self.config.embedding_dims)
    MAX_BATCH = 100
    all_embeddings = []
    for i in range(0, len(texts), MAX_BATCH):
        chunk = [t.replace("\n", " ") for t in texts[i : i + MAX_BATCH]]
        response = self.client.models.embed_content(model=self.config.model, contents=chunk, config=config)
        all_embeddings.extend(e.values for e in response.embeddings)
    if len(all_embeddings) != len(texts):
        raise ValueError(
            f"Gemini embed_batch() returned {len(all_embeddings)} embeddings for {len(texts)} texts "
            f"using model '{self.config.model}'"
        )
    return all_embeddings

Copyclass VertexAIEmbedding(EmbeddingBase):
    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
        super().__init__(config)
 
        self.config.model = self.config.model or "gemini-embedding-001"
        self.config.embedding_dims = self.config.embedding_dims or 256
 
        self.embedding_types = {
            "add": self.config.memory_add_embedding_type or "RETRIEVAL_DOCUMENT",
            "update": self.config.memory_update_embedding_type or "RETRIEVAL_DOCUMENT",
            "search": self.config.memory_search_embedding_type or "RETRIEVAL_QUERY",
        }

Copy# Set up authentication using centralized GCP authenticator
# This supports multiple authentication methods while preserving environment variable support
try:
    GCPAuthenticator.setup_vertex_ai(
        service_account_json=getattr(self.config, 'google_service_account_json', None),
        credentials_path=self.config.vertex_credentials_json,
        project_id=getattr(self.config, 'google_project_id', None)
    )
except Exception:
    # Fall back to original behavior for backward compatibility
    credentials_path = self.config.vertex_credentials_json
    if credentials_path:
        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
    elif not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
        raise ValueError(
            "Google application credentials JSON is not provided. Please provide a valid JSON path or set the 'GOOGLE_APPLICATION_CREDENTIALS' environment variable."
        )
 
self.model = TextEmbeddingModel.from_pretrained(self.config.model)
 

Copydef embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
    """
    Get the embedding for the given text using Vertex AI.
 
    Args:
        text (str): The text to embed.
        memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
    Returns:
        list: The embedding vector.
    """
    embedding_type = "SEMANTIC_SIMILARITY"
    if memory_action is not None:
        if memory_action not in self.embedding_types:
            raise ValueError(f"Invalid memory action: {memory_action}")
 
        embedding_type = self.embedding_types[memory_action]
 
    text_input = TextEmbeddingInput(text=text, task_type=embedding_type)
    embeddings = self.model.get_embeddings(texts=[text_input], output_dimensionality=self.config.embedding_dims)
 
    return embeddings[0].values
 
def embed_batch(self, texts, memory_action="add"):
    if not texts:
        return []
    embedding_type = "SEMANTIC_SIMILARITY"
    if memory_action is not None:
        if memory_action not in self.embedding_types:
            raise ValueError(f"Invalid memory action: {memory_action}")
        embedding_type = self.embedding_types[memory_action]
    all_embeddings = []
    for i in range(0, len(texts), 250):
        chunk = texts[i : i + 250]
        inputs = [TextEmbeddingInput(text=t, task_type=embedding_type) for t in chunk]
        results = self.model.get_embeddings(texts=inputs, output_dimensionality=self.config.embedding_dims)
        all_embeddings.extend(r.values for r in results)
    if len(all_embeddings) != len(texts):
        raise ValueError(
            f"Vertex AI embed_batch() returned {len(all_embeddings)} embeddings for {len(texts)} texts"
            f" using model '{self.config.model}'"
        )
    return all_embeddings

Copy@staticmethod
def _normalize_model_name(name: str) -> str:
    return name if ":" in name else f"{name}:latest"
 
def _ensure_model_exists(self):
    """
    Ensure the specified model exists locally. If not, pull it from Ollama.
    """
    local_models = self.client.list()["models"]
    target = self._normalize_model_name(self.config.model)
    if not any(
        self._normalize_model_name(model.get("name", "")) == target
        or self._normalize_model_name(model.get("model", "")) == target
        for model in local_models
    ):
        self.client.pull(self.config.model)
 

Copydef embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
    """
    Get the embedding for the given text using Ollama.
 
    Args:
        text (str): The text to embed.
        memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
    Returns:
        list: The embedding vector.
    """
    response = self.client.embed(model=self.config.model, input=text)
    embeddings = response.get("embeddings") or []
    if not embeddings:
        raise ValueError(f"Ollama embed() returned no embeddings for model '{self.config.model}'")
    return embeddings[0]
 
def embed_batch(self, texts, memory_action="add"):
    """Embed multiple texts in a single Ollama API call."""
    if not texts:
        return []
    response = self.client.embed(model=self.config.model, input=texts)
    embeddings = response.get("embeddings") or []
    if len(embeddings) != len(texts):
        raise ValueError(f"Ollama embed() returned {len(embeddings)} embeddings for {len(texts)} texts using model '{self.config.model}'")
    return embeddings

Copydef embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
    """
    Convert the text to embeddings using FastEmbed running in the Onnx runtime
    Args:
        text (str): The text to embed.
        memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
    Returns:
        list: The embedding vector.
    """
    text = text.replace("\n", " ")
    embeddings = list(self.dense_model.embed(text))
    return embeddings[0].tolist()

Copydef embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
    """
    Get the embedding for the given text using LM Studio.
    Args:
        text (str): The text to embed.
        memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
    Returns:
        list: The embedding vector.
    """
    text = text.replace("\n", " ")
    return self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding
 

Copydef embed_batch(self, texts, memory_action="add"):
    if not texts:
        return []
    cleaned = [t.replace("\n", " ") for t in texts]
    response = self.client.embeddings.create(input=cleaned, model=self.config.model)
    sorted_data = sorted(response.data, key=lambda x: x.index)
    embeddings = [item.embedding for item in sorted_data]
    if len(embeddings) != len(texts):
        raise ValueError(
            f"LM Studio embed_batch() returned {len(embeddings)} embeddings for {len(texts)} texts"
            f" using model '{self.config.model}'"
        )
    return embeddings

Copyclass HuggingFaceEmbedding(EmbeddingBase):
    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
        super().__init__(config)
 
        if self.config.huggingface_base_url:
            api_key = self.config.api_key or os.getenv("HUGGINGFACE_API_KEY") or "hf"
            self.client = OpenAI(base_url=self.config.huggingface_base_url, api_key=api_key)
            self.config.model = self.config.model or "tei"
        else:
            self.config.model = self.config.model or "multi-qa-MiniLM-L6-cos-v1"
 
            self.model = SentenceTransformer(self.config.model, **self.config.model_kwargs)
 
            self.config.embedding_dims = self.config.embedding_dims or self.model.get_sentence_embedding_dimension()
 

Copydef embed_batch(self, texts, memory_action="add"):
    if not texts:
        return []
    if self.config.huggingface_base_url:
        response = self.client.embeddings.create(input=texts, model=self.config.model, **self.config.model_kwargs)
        sorted_data = sorted(response.data, key=lambda x: x.index)
        embeddings = [item.embedding for item in sorted_data]
        if len(embeddings) != len(texts):
            raise ValueError(
                f"HuggingFace embed_batch() returned {len(embeddings)} embeddings for {len(texts)} texts"
                f" using model '{self.config.model}'"
            )
        return embeddings
    else:
        result = self.model.encode(texts, convert_to_numpy=True).tolist()
        if len(result) != len(texts):
            raise ValueError(
                f"HuggingFace embed_batch() returned {len(result)} embeddings for {len(texts)} texts"
                f" using model '{self.config.model}'"
            )
        return result

Copyclass TogetherEmbedding(EmbeddingBase):
    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
        super().__init__(config)
 
        self.config.model = self.config.model or "intfloat/multilingual-e5-large-instruct"
        api_key = self.config.api_key or os.getenv("TOGETHER_API_KEY")
        self.config.embedding_dims = self.config.embedding_dims or 1024
        self.client = Together(api_key=api_key)
 

Copydef embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
    """
    Get the embedding for the given text using OpenAI.
 
    Args:
        text (str): The text to embed.
        memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
    Returns:
        list: The embedding vector.
    """
 
    return self.client.embeddings.create(model=self.config.model, input=text).data[0].embedding
 
def embed_batch(self, texts, memory_action="add"):
    if not texts:
        return []
    response = self.client.embeddings.create(model=self.config.model, input=texts)
    sorted_data = sorted(response.data, key=lambda x: x.index)
    embeddings = [item.embedding for item in sorted_data]
    if len(embeddings) != len(texts):
        raise ValueError(
            f"Together embed_batch() returned {len(embeddings)} embeddings for {len(texts)} texts"
            f" using model '{self.config.model}'"
        )
    return embeddings

Copyclass LangchainEmbedding(EmbeddingBase):
    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
        super().__init__(config)
 
        if self.config.model is None:
            raise ValueError("`model` parameter is required")
 
        if not isinstance(self.config.model, Embeddings):
            raise ValueError("`model` must be an instance of Embeddings")
 
        self.langchain_model = self.config.model
 

Copytry:
    from langchain.embeddings.base import Embeddings
except ImportError:
    raise ImportError("langchain is not installed. Please install it using `pip install langchain`")