참고소스 수정본

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/__init__.py

@@ -0,0 +1,34 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from .base import Embedder
+from .bedrock import BedrockEmbeddings
+from .cohere import CohereEmbeddings
+from .mistral import MistralAIEmbeddings
+from .ollama import OllamaEmbeddings
+from .openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
+from .sentence_transformers import SentenceTransformerEmbeddings
+from .vertexai import VertexAIEmbeddings
+
+__all__ = [
+    "Embedder",
+    "BedrockEmbeddings",
+    "SentenceTransformerEmbeddings",
+    "OllamaEmbeddings",
+    "OpenAIEmbeddings",
+    "AzureOpenAIEmbeddings",
+    "VertexAIEmbeddings",
+    "MistralAIEmbeddings",
+    "CohereEmbeddings",
+]

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/base.py

@@ -0,0 +1,61 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from neo4j_graphrag.utils.rate_limit import (
+    DEFAULT_RATE_LIMIT_HANDLER,
+    RateLimitHandler,
+)
+
+
+class Embedder(ABC):
+    """
+    Interface for embedding models.
+    An embedder passed into a retriever must implement this interface.
+
+    Args:
+        rate_limit_handler (Optional[RateLimitHandler]): Handler for rate limiting. Defaults to retry with exponential backoff.
+    """
+
+    def __init__(self, rate_limit_handler: Optional[RateLimitHandler] = None):
+        if rate_limit_handler is not None:
+            self._rate_limit_handler = rate_limit_handler
+        else:
+            self._rate_limit_handler = DEFAULT_RATE_LIMIT_HANDLER
+
+    @abstractmethod
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text.
+
+        Args:
+            text (str): Text to convert to vector embedding
+
+        Returns:
+            list[float]: A vector embedding.
+        """
+
+    async def async_embed_query(self, text: str) -> list[float]:
+        """Asynchronously embed query text.
+
+        Args:
+            text (str): Text to convert to vector embedding
+
+        Returns:
+            list[float]: A vector embedding.
+        """
+        return self.embed_query(text)

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/bedrock.py

@@ -0,0 +1,135 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from typing import Any, Optional
+
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+from neo4j_graphrag.utils.rate_limit import (
+    RateLimitHandler,
+    async_rate_limit_handler,
+    rate_limit_handler,
+)
+
+try:
+    import boto3
+except ImportError:
+    boto3 = None
+
+DEFAULT_MODEL_ID = os.getenv("BEDROCK_EMBED_MODEL_ID", "amazon.titan-embed-text-v2:0")
+try:
+    DEFAULT_DIMENSIONS = int(os.getenv("BEDROCK_EMBED_DIMENSIONS", "1024"))
+except ValueError:
+    DEFAULT_DIMENSIONS = 1024
+
+
+class BedrockEmbeddings(Embedder):
+    """Embedder that uses Amazon Bedrock's embedding models via the boto3 SDK.
+
+    Supports Amazon Titan Embed models available through Bedrock.
+
+    Args:
+        model_id: Bedrock model ID. Defaults to the ``BEDROCK_EMBED_MODEL_ID``
+            environment variable, or "amazon.titan-embed-text-v2:0" if not set.
+        dimensions: Output embedding dimensionality. Defaults to the
+            ``BEDROCK_EMBED_DIMENSIONS`` environment variable, or 1024 if not set.
+        normalize: Whether to normalize the embedding vector. Defaults to True.
+        region_name: AWS region. Defaults to boto3 session default.
+        rate_limit_handler: Optional rate limit handler.
+        **kwargs: Arguments passed to ``boto3.client("bedrock-runtime", ...)``.
+
+    Example:
+
+    .. code-block:: python
+
+        from neo4j_graphrag.embeddings import BedrockEmbeddings
+
+        embedder = BedrockEmbeddings(
+            model_id="amazon.titan-embed-text-v2:0",
+            dimensions=1024,
+            region_name="us-east-1",
+        )
+        vector = embedder.embed_query("my question")
+    """
+
+    def __init__(
+        self,
+        model_id: str = DEFAULT_MODEL_ID,
+        dimensions: int = DEFAULT_DIMENSIONS,
+        normalize: bool = True,
+        region_name: Optional[str] = None,
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        **kwargs: Any,
+    ) -> None:
+        if boto3 is None:
+            raise ImportError(
+                "Could not import boto3 python client. "
+                'Please install it with `pip install "neo4j-graphrag[bedrock]"`.'
+            )
+        super().__init__(rate_limit_handler)
+        self.model_id = model_id
+        self.dimensions = dimensions
+        self.normalize = normalize
+        client_kwargs: dict[str, Any] = {**kwargs}
+        if region_name:
+            client_kwargs["region_name"] = region_name
+        self.client = boto3.client("bedrock-runtime", **client_kwargs)
+
+    def _invoke_embedding(self, text: str) -> list[float]:
+        """Invoke the Bedrock embedding model and return the embedding vector."""
+        body = json.dumps(
+            {
+                "inputText": text,
+                "dimensions": self.dimensions,
+                "normalize": self.normalize,
+            }
+        )
+        response = self.client.invoke_model(
+            body=body,
+            modelId=self.model_id,
+            accept="application/json",
+            contentType="application/json",
+        )
+        response_body_stream = response.get("body")
+        if response_body_stream is None:
+            raise ValueError("No body in Bedrock API response")
+        response_body = json.loads(response_body_stream.read())
+        embedding = response_body.get("embedding")
+        if not embedding:
+            raise ValueError("No embedding returned from Bedrock API")
+        return list(embedding)
+
+    @rate_limit_handler
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        try:
+            return self._invoke_embedding(text)
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with Bedrock: {e}"
+            ) from e
+
+    @async_rate_limit_handler
+    async def async_embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        try:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(None, self._invoke_embedding, text)
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with Bedrock: {e}"
+            ) from e

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/cohere.py

@@ -0,0 +1,57 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+from neo4j_graphrag.utils.rate_limit import RateLimitHandler, rate_limit_handler
+
+try:
+    import cohere
+except ImportError:
+    cohere = None  # type: ignore
+
+
+class CohereEmbeddings(Embedder):
+    def __init__(
+        self,
+        model: str = "",
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        **kwargs: Any,
+    ) -> None:
+        if cohere is None:
+            raise ImportError(
+                """Could not import cohere python client.
+                Please install it with `pip install "neo4j-graphrag[cohere]"`."""
+            )
+        super().__init__(rate_limit_handler)
+        self.model = model
+        self.client = cohere.Client(**kwargs)
+
+    @rate_limit_handler
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        try:
+            response = self.client.embed(
+                texts=[text],
+                model=self.model,
+                **kwargs,
+            )
+            return response.embeddings[0]  # type: ignore
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with Cohere: {e}"
+            ) from e

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/mistral.py

@@ -0,0 +1,84 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+import os
+from typing import Any, Optional
+
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+from neo4j_graphrag.utils.rate_limit import RateLimitHandler, rate_limit_handler
+
+try:
+    from mistralai import Mistral
+except ImportError:
+    Mistral = None  # type: ignore
+
+
+class MistralAIEmbeddings(Embedder):
+    """
+    Mistral AI embeddings class.
+    This class uses the Mistral AI Python client to generate vector embeddings for text data.
+
+    Args:
+        model (str): The name of the Mistral AI text embedding model to use. Defaults to "mistral-embed".
+    """
+
+    def __init__(
+        self,
+        model: str = "mistral-embed",
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        **kwargs: Any,
+    ) -> None:
+        if Mistral is None:
+            raise ImportError(
+                """Could not import mistralai.
+                Please install it with `pip install "neo4j-graphrag[mistralai]"`."""
+            )
+        super().__init__(rate_limit_handler)
+        api_key = kwargs.pop("api_key", None)
+        if api_key is None:
+            api_key = os.getenv("MISTRAL_API_KEY", "")
+        self.model = model
+        self.mistral_client = Mistral(api_key=api_key, **kwargs)
+
+    @rate_limit_handler
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        """
+        Generate embeddings for a given query using a Mistral AI text embedding model.
+
+        Args:
+            text (str): The text to generate an embedding for.
+            **kwargs (Any): Additional keyword arguments to pass to the Mistral AI client.
+        """
+        try:
+            embeddings_batch_response = self.mistral_client.embeddings.create(
+                model=self.model, inputs=[text], **kwargs
+            )
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with MistralAI: {e}"
+            ) from e
+
+        if embeddings_batch_response is None or not embeddings_batch_response.data:
+            raise EmbeddingsGenerationError("Failed to retrieve embeddings.")
+
+        embedding = embeddings_batch_response.data[0].embedding
+
+        if not isinstance(embedding, list):
+            raise EmbeddingsGenerationError("Embedding is not a list of floats.")
+
+        return embedding

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/ollama.py

@@ -0,0 +1,106 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+from neo4j_graphrag.utils.rate_limit import (
+    RateLimitHandler,
+    rate_limit_handler,
+    async_rate_limit_handler,
+)
+
+
+class OllamaEmbeddings(Embedder):
+    """
+    Ollama embeddings class.
+    This class uses the ollama Python client to generate vector embeddings for text data.
+
+    Args:
+        model (str): The name of the Ollama text embedding model to use.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            import ollama
+        except ImportError:
+            raise ImportError(
+                """Could not import ollama python client.
+                Please install it with `pip install "neo4j_graphrag[ollama]"`."""
+            )
+        super().__init__(rate_limit_handler)
+        self.model = model
+        self.client = ollama.Client(**kwargs)
+        self.async_client = ollama.AsyncClient(**kwargs)
+
+    @rate_limit_handler
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        """
+        Generate embeddings for a given query using an Ollama text embedding model.
+
+        Args:
+            text (str): The text to generate an embedding for.
+            **kwargs (Any): Additional keyword arguments to pass to the Ollama client.
+        """
+        embeddings_response = self.client.embed(
+            model=self.model,
+            input=text,
+            **kwargs,
+        )
+
+        if embeddings_response is None or not embeddings_response.embeddings:
+            raise EmbeddingsGenerationError("Failed to retrieve embeddings.")
+
+        embeddings = embeddings_response.embeddings
+        # client always returns a sequence of sequences
+        embedding = embeddings[0]
+        if not isinstance(embedding, list):
+            raise EmbeddingsGenerationError("Embedding is not a list of floats.")
+
+        return embedding
+
+    @async_rate_limit_handler
+    async def async_embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        """
+        Asynchronously generate embeddings for a given query using an Ollama text embedding model.
+
+        Args:
+            text (str): The text to generate an embedding for.
+            **kwargs (Any): Additional keyword arguments to pass to the Ollama client.
+        """
+        embeddings_response = await self.async_client.embed(
+            model=self.model,
+            input=text,
+            **kwargs,
+        )
+
+        if embeddings_response is None or not embeddings_response.embeddings:
+            raise EmbeddingsGenerationError("Failed to retrieve embeddings.")
+
+        embeddings = embeddings_response.embeddings
+        # client always returns a sequence of sequences
+        embedding = embeddings[0]
+        if not isinstance(embedding, list):
+            raise EmbeddingsGenerationError("Embedding is not a list of floats.")
+
+        return embedding

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/openai.py

@@ -0,0 +1,108 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+import abc
+from typing import TYPE_CHECKING, Any, Optional
+
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+from neo4j_graphrag.utils.rate_limit import RateLimitHandler, rate_limit_handler
+
+if TYPE_CHECKING:
+    import openai
+
+
+class BaseOpenAIEmbeddings(Embedder, abc.ABC):
+    """
+    Abstract base class for OpenAI embeddings.
+    """
+
+    client: openai.OpenAI
+
+    def __init__(
+        self,
+        model: str = "text-embedding-ada-002",
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            import openai
+        except ImportError:
+            raise ImportError(
+                """Could not import openai python client.
+                Please install it with `pip install "neo4j-graphrag[openai]"`."""
+            )
+        super().__init__(rate_limit_handler)
+        self.openai = openai
+        self.model = model
+        self.client = self._initialize_client(**kwargs)
+
+    @abc.abstractmethod
+    def _initialize_client(self, **kwargs: Any) -> Any:
+        """
+        Initialize the OpenAI client.
+        Must be implemented by subclasses.
+        """
+        pass
+
+    @rate_limit_handler
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        """
+        Generate embeddings for a given query using an OpenAI text embedding model.
+
+        Args:
+            text (str): The text to generate an embedding for.
+            **kwargs (Any): Additional arguments to pass to the OpenAI embedding generation function.
+        """
+        try:
+            response = self.client.embeddings.create(
+                input=text, model=self.model, **kwargs
+            )
+            embedding: list[float] = response.data[0].embedding
+            return embedding
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with OpenAI: {e}"
+            ) from e
+
+
+class OpenAIEmbeddings(BaseOpenAIEmbeddings):
+    """
+    OpenAI embeddings class.
+    This class uses the OpenAI python client to generate embeddings for text data.
+
+    Args:
+        model (str): The name of the OpenAI embedding model to use. Defaults to "text-embedding-ada-002".
+        kwargs: All other parameters will be passed to the openai.OpenAI init.
+    """
+
+    def _initialize_client(self, **kwargs: Any) -> Any:
+        return self.openai.OpenAI(**kwargs)
+
+
+class AzureOpenAIEmbeddings(BaseOpenAIEmbeddings):
+    """
+    Azure OpenAI embeddings class.
+    This class uses the Azure OpenAI python client to generate embeddings for text data.
+
+    Args:
+        model (str): The name of the Azure OpenAI embedding model to use. Defaults to "text-embedding-ada-002".
+        kwargs: All other parameters will be passed to the openai.AzureOpenAI init.
+    """
+
+    def _initialize_client(self, **kwargs: Any) -> Any:
+        return self.openai.AzureOpenAI(**kwargs)

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/sentence_transformers.py

@@ -0,0 +1,62 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import Any, Optional
+
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+from neo4j_graphrag.utils.rate_limit import RateLimitHandler
+
+
+class SentenceTransformerEmbeddings(Embedder):
+    def __init__(
+        self,
+        model: str = "all-MiniLM-L6-v2",
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            import numpy as np
+            import sentence_transformers
+            import torch
+        except ImportError:
+            raise ImportError(
+                """Could not import sentence_transformers python package.
+                Please install it with `pip install "neo4j-graphrag[sentence-transformers]"`."""
+            )
+        super().__init__(rate_limit_handler)
+        self.torch = torch
+        self.np = np
+        self.model = sentence_transformers.SentenceTransformer(model, *args, **kwargs)
+
+    def embed_query(self, text: str) -> Any:
+        try:
+            result = self.model.encode([text])
+
+            if isinstance(result, self.torch.Tensor) or isinstance(
+                result, self.np.ndarray
+            ):
+                return result.flatten().tolist()
+            elif isinstance(result, list) and all(
+                isinstance(x, self.torch.Tensor) for x in result
+            ):
+                return [item for tensor in result for item in tensor.flatten().tolist()]
+            else:
+                raise ValueError("Unexpected return type from model encoding")
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with SentenceTransformer: {e}"
+            ) from e

참고/neo4j-graphrag-python-main/src/neo4j_graphrag/embeddings/vertexai.py

@@ -0,0 +1,77 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+from neo4j_graphrag.utils.rate_limit import RateLimitHandler, rate_limit_handler
+
+try:
+    from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
+except (ImportError, AttributeError):
+    TextEmbeddingModel = TextEmbeddingInput = None  # type: ignore[misc, assignment]
+
+
+if TYPE_CHECKING:
+    from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
+
+
+class VertexAIEmbeddings(Embedder):
+    """
+    Vertex AI embeddings class.
+    This class uses the Vertex AI Python client to generate vector embeddings for text data.
+
+    Args:
+        model (str): The name of the Vertex AI text embedding model to use. Defaults to "text-embedding-004".
+    """
+
+    def __init__(
+        self,
+        model: str = "text-embedding-004",
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+    ) -> None:
+        if TextEmbeddingModel is None:
+            raise ImportError(
+                """Could not import Vertex AI Python client.
+                Please install it with `pip install "neo4j-graphrag[google]"`."""
+            )
+        super().__init__(rate_limit_handler)
+        self.model = TextEmbeddingModel.from_pretrained(model)
+
+    @rate_limit_handler
+    def embed_query(
+        self, text: str, task_type: str = "RETRIEVAL_QUERY", **kwargs: Any
+    ) -> list[float]:
+        """
+        Generate embeddings for a given query using a Vertex AI text embedding model.
+
+        Args:
+            text (str): The text to generate an embedding for.
+            task_type (str): The type of the text embedding task. Defaults to "RETRIEVAL_QUERY". See https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#tasktype for a full list.
+            **kwargs (Any): Additional keyword arguments to pass to the Vertex AI client's get_embeddings method.
+        """
+        try:
+            # type annotation needed for mypy
+            inputs: list[str | TextEmbeddingInput] = [
+                TextEmbeddingInput(text, task_type)
+            ]
+            embeddings = self.model.get_embeddings(inputs, **kwargs)
+            return list(embeddings[0].values)
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with VertexAI: {e}"
+            ) from e