feat: Add Google Gemini support

- Introduced  and  classes for embedding and LLM functionalities using Google's Gemini API.
- Updated  to include  dependency.
- Added example scripts for using  and .
- Enhanced documentation to include new classes and their usage.
- Updated tests for the new functionalities, ensuring proper error handling and response validation.

Author: balalofernandez

docs/source/api.rst

@@ -278,6 +278,12 @@ SentenceTransformerEmbeddings
 .. autoclass:: neo4j_graphrag.embeddings.sentence_transformers.SentenceTransformerEmbeddings
     :members:
 
+GeminiEmbedder
+==============
+
+.. autoclass:: neo4j_graphrag.embeddings.google_genai.GeminiEmbedder
+    :members:
+
 OpenAIEmbeddings
 ================
 
@@ -336,6 +342,13 @@ OpenAILLM
     :undoc-members: get_messages, client_class, async_client_class
 
 
+GeminiLLM
+---------
+
+.. autoclass:: neo4j_graphrag.llm.google_genai_llm.GeminiLLM
+    :members:
+
+
 AzureOpenAILLM
 --------------
 

examples/customize/embeddings/google_genai_embeddings.py

@@ -0,0 +1,11 @@
+from neo4j_graphrag.embeddings import GeminiEmbedder
+
+# set api key here on in the GOOGLE_API_KEY env var
+api_key = None
+
+embedder = GeminiEmbedder(
+    model="gemini-embedding-001",
+    api_key=api_key,
+)
+res = embedder.embed_query("my question")
+print(res[:10])

examples/customize/llms/google_genai_llm.py

@@ -0,0 +1,11 @@
+from neo4j_graphrag.llm import GeminiLLM
+
+# set api key here on in the GOOGLE_API_KEY env var
+api_key = None
+
+llm = GeminiLLM(
+    model_name="gemini-2.5-flash",
+    api_key=api_key,
+)
+res = llm.invoke("say something")
+print(res.content)

pyproject.toml

@@ -46,6 +46,7 @@ Repository = "https://github.com/neo4j/neo4j-graphrag-python"
 weaviate = ["weaviate-client>=4.6.1,<5.0.0"]
 pinecone = ["pinecone-client>=4.1.0,<5.0.0"]
 google = ["google-cloud-aiplatform>=1.66.0,<2.0.0"]
+google-genai = ["google-genai>=1.62.0,<2.0.0"]
 cohere = ["cohere>=5.9.0,<6.0.0"]
 anthropic = ["anthropic>=0.49.0,<0.50.0"]
 ollama = ["ollama>=0.4.4,<0.5.0"]

src/neo4j_graphrag/embeddings/__init__.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 from .base import Embedder
 from .cohere import CohereEmbeddings
+from .google_genai import GeminiEmbedder
 from .mistral import MistralAIEmbeddings
 from .ollama import OllamaEmbeddings
 from .openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
@@ -29,4 +30,5 @@
     "VertexAIEmbeddings",
     "MistralAIEmbeddings",
     "CohereEmbeddings",
+    "GeminiEmbedder",
 ]

src/neo4j_graphrag/embeddings/google_genai.py

@@ -0,0 +1,103 @@
+#  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
+
+# built-in dependencies
+from typing import Any, Optional
+
+# project dependencies
+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:
+    from google import genai
+    from google.genai import types
+except ImportError:
+    genai = None
+    types = None
+
+DEFAULT_EMBEDDING_MODEL = "text-embedding-004"
+DEFAULT_EMBEDDING_DIM = 768
+
+
+class GeminiEmbedder(Embedder):
+    """Embedder that uses Google's Gemini API via the google.genai SDK.
+
+    Args:
+        model: Embedding model name. Defaults to "text-embedding-004".
+        embedding_dim: Output dimensionality. Defaults to 768.
+        rate_limit_handler: Optional rate limit handler.
+        **kwargs: Arguments passed to the genai.Client.
+    """
+
+    def __init__(
+        self,
+        model: str = DEFAULT_EMBEDDING_MODEL,
+        embedding_dim: int = DEFAULT_EMBEDDING_DIM,
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+        **kwargs: Any,
+    ) -> None:
+        if genai is None or types is None:
+            raise ImportError(
+                "Could not import google-genai python client. "
+                'Please install it with `pip install "neo4j-graphrag[google-genai]"`.'
+            )
+        super().__init__(rate_limit_handler)
+        self.model = model
+        self.embedding_dim = embedding_dim
+        self.client = genai.Client(**kwargs)
+
+    @rate_limit_handler
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        try:
+            result = self.client.models.embed_content(
+                model=self.model,
+                contents=[text],
+                config=types.EmbedContentConfig(
+                    output_dimensionality=self.embedding_dim
+                ),
+                **kwargs,
+            )
+            if not result.embeddings or not result.embeddings[0].values:
+                raise ValueError("No embeddings returned from Gemini API")
+            return list(result.embeddings[0].values)
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with Gemini: {e}"
+            ) from e
+
+    @async_rate_limit_handler
+    async def async_embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        try:
+            result = await self.client.aio.models.embed_content(
+                model=self.model,
+                contents=[text],
+                config=types.EmbedContentConfig(
+                    output_dimensionality=self.embedding_dim
+                ),
+                **kwargs,
+            )
+            if not result.embeddings or not result.embeddings[0].values:
+                raise ValueError("No embeddings returned from Gemini API")
+            return list(result.embeddings[0].values)
+        except Exception as e:
+            raise EmbeddingsGenerationError(
+                f"Failed to generate embedding with Gemini: {e}"
+            ) from e

src/neo4j_graphrag/llm/__init__.py

@@ -18,6 +18,7 @@
 from .anthropic_llm import AnthropicLLM
 from .base import LLMInterface, LLMInterfaceV2
 from .cohere_llm import CohereLLM
+from .google_genai_llm import GeminiLLM
 from .mistralai_llm import MistralAILLM
 from .ollama_llm import OllamaLLM
 from .openai_llm import AzureOpenAILLM, OpenAILLM
@@ -28,6 +29,7 @@
 __all__ = [
     "AnthropicLLM",
     "CohereLLM",
+    "GeminiLLM",
     "LLMResponse",
     "LLMInterface",
     "LLMInterfaceV2",