From b6c0bb9c54a426805b512ac595c665adb64052df Mon Sep 17 00:00:00 2001
From: Daniel Chalef <131175+danielchalef@users.noreply.github.com>
Date: Wed, 21 May 2025 19:31:12 -0700
Subject: [PATCH] Enhance error handling in GeminiEmbedder and GeminiClient

- Added checks to raise exceptions when no embeddings or response text are returned, improving robustness.
- Included type ignore comments for mypy compatibility in embed_content calls.
---
 graphiti_core/embedder/gemini.py          | 12 +++++++++---
 graphiti_core/llm_client/gemini_client.py |  5 ++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/graphiti_core/embedder/gemini.py b/graphiti_core/embedder/gemini.py
index a9a4e49b..2055350d 100644
--- a/graphiti_core/embedder/gemini.py
+++ b/graphiti_core/embedder/gemini.py
@@ -61,18 +61,24 @@ class GeminiEmbedder(EmbedderClient):
         # Generate embeddings
         result = await self.client.aio.models.embed_content(
             model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
-            contents=[input_data],
+            contents=[input_data],  # type: ignore[arg-type]  # mypy fails on broad union type
             config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
         )
 
+        if not result.embeddings or len(result.embeddings) == 0 or not result.embeddings[0].values:
+            raise Exception('No embeddings returned')
+
         return result.embeddings[0].values
 
     async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
         # Generate embeddings
         result = await self.client.aio.models.embed_content(
             model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
-            contents=input_data_list,
+            contents=input_data_list,  # type: ignore[arg-type]  # mypy fails on broad union type
             config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
         )
 
-        return [embedding.values for embedding in result.embeddings]
+        if not result.embeddings or len(result.embeddings) == 0:
+            raise Exception('No embeddings returned')
+
+        return [embedding.values if embedding.values else [] for embedding in result.embeddings]
diff --git a/graphiti_core/llm_client/gemini_client.py b/graphiti_core/llm_client/gemini_client.py
index b3326bf8..18691ee0 100644
--- a/graphiti_core/llm_client/gemini_client.py
+++ b/graphiti_core/llm_client/gemini_client.py
@@ -139,13 +139,16 @@ class GeminiClient(LLMClient):
             # Generate content using the simple string approach
             response = await self.client.aio.models.generate_content(
                 model=self.model or DEFAULT_MODEL,
-                contents=gemini_messages,
+                contents=gemini_messages,  # type: ignore[arg-type]  # mypy fails on broad union type
                 config=generation_config,
             )
 
             # If this was a structured output request, parse the response into the Pydantic model
             if response_model is not None:
                 try:
+                    if not response.text:
+                        raise Exception('No response text')
+
                     validated_model = response_model.model_validate(json.loads(response.text))
 
                     # Return as a dictionary for API consistency