powermem 0.1.0__py3-none-any.whl

powermem/integrations/embeddings/config/base.py

@@ -0,0 +1,132 @@
+import os
+from abc import ABC
+from typing import Dict, Optional, Union
+
+import httpx
+from pydantic import BaseModel, Field
+
+
+class AzureConfig(BaseModel):
+    """
+    Configuration settings for Azure.
+
+    Args:
+        api_key (str): The API key used for authenticating with the Azure service.
+        azure_deployment (str): The name of the Azure deployment.
+        azure_endpoint (str): The endpoint URL for the Azure service.
+        api_version (str): The version of the Azure API being used.
+        default_headers (Dict[str, str]): Headers to include in requests to the Azure API.
+    """
+
+    api_key: str = Field(
+        description="The API key used for authenticating with the Azure service.",
+        default=None,
+    )
+    azure_deployment: str = Field(description="The name of the Azure deployment.", default=None)
+    azure_endpoint: str = Field(description="The endpoint URL for the Azure service.", default=None)
+    api_version: str = Field(description="The version of the Azure API being used.", default=None)
+    default_headers: Optional[Dict[str, str]] = Field(
+        description="Headers to include in requests to the Azure API.", default=None
+    )
+
+class BaseEmbedderConfig(ABC):
+    """
+    Config for Embeddings.
+    """
+
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        api_key: Optional[str] = None,
+        embedding_dims: Optional[int] = None,
+        # Ollama specific
+        ollama_base_url: Optional[str] = None,
+        # Openai specific
+        openai_base_url: Optional[str] = None,
+        # Huggingface specific
+        model_kwargs: Optional[dict] = None,
+        huggingface_base_url: Optional[str] = None,
+        # AzureOpenAI specific
+        azure_kwargs: Optional[AzureConfig] = {},
+        http_client_proxies: Optional[Union[Dict, str]] = None,
+        # VertexAI specific
+        vertex_credentials_json: Optional[str] = None,
+        memory_add_embedding_type: Optional[str] = None,
+        memory_update_embedding_type: Optional[str] = None,
+        memory_search_embedding_type: Optional[str] = None,
+        # Gemini specific
+        output_dimensionality: Optional[str] = None,
+        # LM Studio specific
+        lmstudio_base_url: Optional[str] = "http://localhost:1234/v1",
+        # AWS Bedrock specific
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        aws_region: Optional[str] = None,
+    ):
+        """
+        Initializes a configuration class instance for the Embeddings.
+
+        :param model: Embedding model to use, defaults to None
+        :type model: Optional[str], optional
+        :param api_key: API key to be use, defaults to None
+        :type api_key: Optional[str], optional
+        :param embedding_dims: The number of dimensions in the embedding, defaults to None
+        :type embedding_dims: Optional[int], optional
+        :param ollama_base_url: Base URL for the Ollama API, defaults to None
+        :type ollama_base_url: Optional[str], optional
+        :param model_kwargs: key-value arguments for the huggingface embedding model, defaults a dict inside init
+        :type model_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
+        :param huggingface_base_url: Huggingface base URL to be use, defaults to None
+        :type huggingface_base_url: Optional[str], optional
+        :param openai_base_url: Openai base URL to be use, defaults to "https://api.openai.com/v1"
+        :type openai_base_url: Optional[str], optional
+        :param azure_kwargs: key-value arguments for the AzureOpenAI embedding model, defaults a dict inside init
+        :type azure_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
+        :param http_client_proxies: The proxy server settings used to create self.http_client, defaults to None
+        :type http_client_proxies: Optional[Dict | str], optional
+        :param vertex_credentials_json: The path to the Vertex AI credentials JSON file, defaults to None
+        :type vertex_credentials_json: Optional[str], optional
+        :param memory_add_embedding_type: The type of embedding to use for the add memory action, defaults to None
+        :type memory_add_embedding_type: Optional[str], optional
+        :param memory_update_embedding_type: The type of embedding to use for the update memory action, defaults to None
+        :type memory_update_embedding_type: Optional[str], optional
+        :param memory_search_embedding_type: The type of embedding to use for the search memory action, defaults to None
+        :type memory_search_embedding_type: Optional[str], optional
+        :param lmstudio_base_url: LM Studio base URL to be use, defaults to "http://localhost:1234/v1"
+        :type lmstudio_base_url: Optional[str], optional
+        """
+
+        self.model = model
+        self.api_key = api_key
+        self.openai_base_url = openai_base_url
+        self.embedding_dims = embedding_dims
+
+        # AzureOpenAI specific
+        self.http_client = httpx.Client(proxies=http_client_proxies) if http_client_proxies else None
+
+        # Ollama specific
+        self.ollama_base_url = ollama_base_url
+
+        # Huggingface specific
+        self.model_kwargs = model_kwargs or {}
+        self.huggingface_base_url = huggingface_base_url
+        # AzureOpenAI specific
+        self.azure_kwargs = AzureConfig(**azure_kwargs) or {}
+
+        # VertexAI specific
+        self.vertex_credentials_json = vertex_credentials_json
+        self.memory_add_embedding_type = memory_add_embedding_type
+        self.memory_update_embedding_type = memory_update_embedding_type
+        self.memory_search_embedding_type = memory_search_embedding_type
+
+        # Gemini specific
+        self.output_dimensionality = output_dimensionality
+
+        # LM Studio specific
+        self.lmstudio_base_url = lmstudio_base_url
+
+        # AWS Bedrock specific
+        self.aws_access_key_id = aws_access_key_id
+        self.aws_secret_access_key = aws_secret_access_key
+        self.aws_region = aws_region or os.environ.get("AWS_REGION") or "us-west-2"
+

powermem/integrations/embeddings/configs.py

@@ -0,0 +1,31 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class EmbedderConfig(BaseModel):
+    provider: str = Field(
+        description="Provider of the embedding model (e.g., 'ollama', 'openai')",
+        default="openai",
+    )
+    config: Optional[dict] = Field(description="Configuration for the specific embedding model", default={})
+
+    @field_validator("config")
+    def validate_config(cls, v, values):
+        provider = values.data.get("provider")
+        if provider in [
+            "openai",
+            "ollama",
+            "huggingface",
+            "azure_openai",
+            "gemini",
+            "vertexai",
+            "together",
+            "lmstudio",
+            "langchain",
+            "aws_bedrock",
+            "qwen",
+        ]:
+            return v
+        else:
+            raise ValueError(f"Unsupported embedding provider: {provider}")

powermem/integrations/embeddings/factory.py

@@ -0,0 +1,48 @@
+"""
+Embedding factory for creating embedding instances
+
+This module provides a factory for creating different embedding instances.
+"""
+
+import importlib
+from typing import Optional
+
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+from powermem.integrations.embeddings.mock import MockEmbeddings
+
+
+def load_class(class_type):
+    module_path, class_name = class_type.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+
+class EmbedderFactory:
+    provider_to_class = {
+        "openai": "powermem.integrations.embeddings.openai.OpenAIEmbedding",
+        "ollama": "powermem.integrations.embeddings.ollama.OllamaEmbedding",
+        "huggingface": "powermem.integrations.embeddings.huggingface.HuggingFaceEmbedding",
+        "azure_openai": "powermem.integrations.embeddings.azure_openai.AzureOpenAIEmbedding",
+        "gemini": "powermem.integrations.embeddings.gemini.GoogleGenAIEmbedding",
+        "vertexai": "powermem.integrations.embeddings.vertexai.VertexAIEmbedding",
+        "together": "powermem.integrations.embeddings.together.TogetherEmbedding",
+        "lmstudio": "powermem.integrations.embeddings.lmstudio.LMStudioEmbedding",
+        "langchain": "powermem.integrations.embeddings.langchain.LangchainEmbedding",
+        "aws_bedrock": "powermem.integrations.embeddings.aws_bedrock.AWSBedrockEmbedding",
+        "qwen": "powermem.integrations.embeddings.qwen.QwenEmbedding",
+    }
+
+    @classmethod
+    def create(cls, provider_name, config, vector_config: Optional[dict]):
+        # Handle mock provider directly
+        if provider_name == "mock":
+            return MockEmbeddings()
+        if provider_name == "upstash_vector" and vector_config and vector_config.enable_embeddings:
+            return MockEmbeddings()
+        class_type = cls.provider_to_class.get(provider_name)
+        if class_type:
+            embedder_instance = load_class(class_type)
+            base_config = BaseEmbedderConfig(**config)
+            return embedder_instance(base_config)
+        else:
+            raise ValueError(f"Unsupported Embedder provider: {provider_name}")

powermem/integrations/embeddings/gemini.py

@@ -0,0 +1,39 @@
+import os
+from typing import Literal, Optional
+
+from google import genai
+from google.genai import types
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+
+class GoogleGenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "models/text-embedding-004"
+        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)
+
+    def 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

powermem/integrations/embeddings/huggingface.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Literal, Optional
+
+from openai import OpenAI
+from sentence_transformers import SentenceTransformer
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+logging.getLogger("transformers").setLevel(logging.WARNING)
+logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
+logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
+
+
+class HuggingFaceEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        if config.huggingface_base_url:
+            self.client = OpenAI(base_url=config.huggingface_base_url)
+        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()
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Hugging Face.
+
+        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.
+        """
+        if self.config.huggingface_base_url:
+            return self.client.embeddings.create(input=text, model="tei").data[0].embedding
+        else:
+            return self.model.encode(text, convert_to_numpy=True).tolist()

powermem/integrations/embeddings/langchain.py

@@ -0,0 +1,35 @@
+from typing import Literal, Optional
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+try:
+    from langchain.embeddings.base import Embeddings
+except ImportError:
+    raise ImportError("langchain is not installed. Please install it using `pip install langchain`")
+
+
+class 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
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Langchain.
+
+        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.langchain_model.embed_query(text)

powermem/integrations/embeddings/lmstudio.py

@@ -0,0 +1,29 @@
+from typing import Literal, Optional
+
+from openai import OpenAI
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+
+class LMStudioEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "nomic-ai/nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf"
+        self.config.embedding_dims = self.config.embedding_dims or 1536
+        self.config.api_key = self.config.api_key or "lm-studio"
+
+        self.client = OpenAI(base_url=self.config.lmstudio_base_url, api_key=self.config.api_key)
+
+    def 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

powermem/integrations/embeddings/mock.py

@@ -0,0 +1,11 @@
+from typing import Literal, Optional
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+
+
+class MockEmbeddings(EmbeddingBase):
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Generate a mock embedding with dimension of 10.
+        """
+        return [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]

powermem/integrations/embeddings/ollama.py

@@ -0,0 +1,53 @@
+import subprocess
+import sys
+from typing import Literal, Optional
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+try:
+    from ollama import Client
+except ImportError:
+    user_input = input("The 'ollama' library is required. Install it now? [y/N]: ")
+    if user_input.lower() == "y":
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "ollama"])
+            from ollama import Client
+        except subprocess.CalledProcessError:
+            print("Failed to install 'ollama'. Please install it manually using 'pip install ollama'.")
+            sys.exit(1)
+    else:
+        print("The required 'ollama' library is not installed.")
+        sys.exit(1)
+
+
+class OllamaEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "nomic-embed-text"
+        self.config.embedding_dims = self.config.embedding_dims or 512
+
+        self.client = Client(host=self.config.ollama_base_url)
+        self._ensure_model_exists()
+
+    def _ensure_model_exists(self):
+        """
+        Ensure the specified model exists locally. If not, pull it from Ollama.
+        """
+        local_models = self.client.list()["models"]
+        if not any(model.get("name") == self.config.model or model.get("model") == self.config.model for model in local_models):
+            self.client.pull(self.config.model)
+
+    def 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.embeddings(model=self.config.model, prompt=text)
+        return response["embedding"]

powermem/integrations/embeddings/openai.py

@@ -0,0 +1,49 @@
+import os
+import warnings
+from typing import Literal, Optional
+
+from openai import OpenAI
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+
+class OpenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "text-embedding-3-small"
+        self.config.embedding_dims = self.config.embedding_dims or 1536
+
+        api_key = self.config.api_key or os.getenv("OPENAI_API_KEY")
+        base_url = (
+            self.config.openai_base_url
+            or os.getenv("OPENAI_API_BASE")
+            or os.getenv("OPENAI_BASE_URL")
+            or "https://api.openai.com/v1"
+        )
+        if os.environ.get("OPENAI_API_BASE"):
+            warnings.warn(
+                "The environment variable 'OPENAI_API_BASE' is deprecated and will be removed in the 0.1.80. "
+                "Please use 'OPENAI_BASE_URL' instead.",
+                DeprecationWarning,
+            )
+
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
+
+    def 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.
+        """
+        text = text.replace("\n", " ")
+        return (
+            self.client.embeddings.create(input=[text], model=self.config.model, dimensions=self.config.embedding_dims)
+            .data[0]
+            .embedding
+        )

powermem/integrations/embeddings/qwen.py

@@ -0,0 +1,102 @@
+import os
+from typing import Literal, Optional
+
+try:
+    from dashscope import TextEmbedding
+    from dashscope.api_entities.dashscope_response import DashScopeAPIResponse
+except ImportError:
+    TextEmbedding = None
+    DashScopeAPIResponse = None
+import dashscope
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+
+class QwenEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        # Set default model and dimensions
+        self.config.model = self.config.model or "text-embedding-v4"
+        self.config.embedding_dims = self.config.embedding_dims or 1536
+
+        # Check if dashscope is available
+        if TextEmbedding is None:
+            raise ImportError(
+                "DashScope SDK is not installed. Please install it with: pip install dashscope"
+            )
+
+        # Set API key
+        api_key = self.config.api_key or os.getenv("DASHSCOPE_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "API key is required. Set DASHSCOPE_API_KEY environment variable or pass api_key in config.")
+
+        # Set API key for DashScope SDK
+        dashscope.api_key = api_key
+
+        # Set base URL (if needed)
+        base_url = os.getenv("DASHSCOPE_BASE_URL") or "https://dashscope.aliyuncs.com/api/v1"
+        if base_url:
+            os.environ["DASHSCOPE_BASE_URL"] = base_url
+
+    def embed(self, text: str, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Qwen.
+
+        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.
+        """
+        # Clean text
+        text = text.replace("\n", " ").strip()
+
+        # Determine embedding type based on memory action
+        # Default values for DashScope text-embedding-v4:
+        # - "document" for add/update (RETRIEVAL_DOCUMENT equivalent)
+        # - "query" for search (RETRIEVAL_QUERY equivalent)
+        if memory_action == "add":
+            embedding_type = self.config.memory_add_embedding_type or "document"
+        elif memory_action == "search":
+            embedding_type = self.config.memory_search_embedding_type or "query"
+        elif memory_action == "update":
+            embedding_type = self.config.memory_update_embedding_type or "document"
+        else:
+            # Default to "document" if memory_action is None or unknown
+            embedding_type = "document"
+
+        try:
+            # Prepare parameters
+            params = {
+                "model": self.config.model,
+                "input": text,
+            }
+
+            # Add dimension parameter if specified
+            if hasattr(self.config, 'embedding_dims') and self.config.embedding_dims:
+                params["dimension"] = self.config.embedding_dims
+
+            # Add embedding type (always set, either from config or default)
+            params["text_type"] = embedding_type
+
+            # Call the API
+            response = TextEmbedding.call(**params)
+
+            if response.status_code != 200:
+                raise Exception(f"API request failed with status {response.status_code}: {response.message}")
+
+            # Extract embedding vector
+            # response.output is a dict, not an object with attributes
+            if isinstance(response.output, dict) and 'embeddings' in response.output:
+                embedding = response.output['embeddings'][0]['embedding']
+            else:
+                # Fallback for different response structures
+                embedding = response.output.get('embeddings', [{}])[0].get('embedding', [])
+
+            return embedding
+
+        except Exception as e:
+            raise Exception(f"Failed to generate embedding: {e}")

powermem/integrations/embeddings/together.py

@@ -0,0 +1,31 @@
+import os
+from typing import Literal, Optional
+
+from together import Together
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+
+class TogetherEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "togethercomputer/m2-bert-80M-8k-retrieval"
+        api_key = self.config.api_key or os.getenv("TOGETHER_API_KEY")
+        # TODO: check if this is correct
+        self.config.embedding_dims = self.config.embedding_dims or 768
+        self.client = Together(api_key=api_key)
+
+    def 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

powermem/integrations/embeddings/vertexai.py

@@ -0,0 +1,54 @@
+import os
+from typing import Literal, Optional
+
+from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+
+class VertexAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "text-embedding-004"
+        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",
+        }
+
+        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)
+
+    def 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