Enhance Google Gemini integration

- Added support for Google Gemini embeddings and LLMs in the README.
- Updated example scripts to assert the presence of the GOOGLE_API_KEY environment variable.
- Modified the default embedding model and dimension to use environment variables for better configurability.
- Improved error handling in the GeminiEmbedder class to ensure valid responses from the API.

Author: balalofernandez

README.md

@@ -50,6 +50,7 @@ the extra dependencies described below:
     - **ollama**: LLMs from Ollama
     - **openai**: LLMs from OpenAI (including AzureOpenAI)
     - **google**: LLMs from Vertex AI
+    - **google-genai**: LLMs and embeddings from Google Gemini
     - **cohere**: LLMs from Cohere
     - **anthropic**: LLMs from Anthropic
     - **mistralai**: LLMs from MistralAI

examples/customize/embeddings/google_genai_embeddings.py

@@ -1,7 +1,9 @@
+import os
+
 from neo4j_graphrag.embeddings import GeminiEmbedder
 
-# set api key here on in the GOOGLE_API_KEY env var
-api_key = None
+api_key = os.getenv("GOOGLE_API_KEY")
+assert api_key is not None, "you must set GOOGLE_API_KEY to run this experiment"
 
 embedder = GeminiEmbedder(
     model="gemini-embedding-001",

examples/customize/llms/google_genai_llm.py

@@ -1,7 +1,9 @@
+import os
+
 from neo4j_graphrag.llm import GeminiLLM
 
-# set api key here on in the GOOGLE_API_KEY env var
-api_key = None
+api_key = os.getenv("GOOGLE_API_KEY")
+assert api_key is not None, "you must set GOOGLE_API_KEY to run this experiment"
 
 llm = GeminiLLM(
     model_name="gemini-2.5-flash",

src/neo4j_graphrag/embeddings/google_genai.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 # built-in dependencies
+import os
 from typing import Any, Optional
 
 # project dependencies
@@ -33,15 +34,15 @@
     genai = None
     types = None
 
-DEFAULT_EMBEDDING_MODEL = "text-embedding-004"
-DEFAULT_EMBEDDING_DIM = 768
+DEFAULT_EMBEDDING_MODEL = os.getenv("GOOGLE_GENAI_EMB_MODEL", "gemini-embedding-001")
+DEFAULT_EMBEDDING_DIM = int(os.getenv("GOOGLE_GENAI_EMB_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".
+        model: Embedding model name. Defaults to "gemini-embedding-001".
         embedding_dim: Output dimensionality. Defaults to 768.
         rate_limit_handler: Optional rate limit handler.
         **kwargs: Arguments passed to the genai.Client.
@@ -75,7 +76,7 @@ def embed_query(self, text: str, **kwargs: Any) -> list[float]:
                 ),
                 **kwargs,
             )
-            if not result.embeddings or not result.embeddings[0].values:
+            if not result or 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:
@@ -94,7 +95,7 @@ async def async_embed_query(self, text: str, **kwargs: Any) -> list[float]:
                 ),
                 **kwargs,
             )
-            if not result.embeddings or not result.embeddings[0].values:
+            if not result or 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: