From ab58a2319fa3df2c51ec68864a5da09105bf1fef Mon Sep 17 00:00:00 2001
From: antoniomuzzolini <39350879+antoniomuzzolini@users.noreply.github.com>
Date: Wed, 5 Jun 2024 16:32:06 +0200
Subject: [PATCH] fix: Handle `TypeError: Could not create Blob` in
 `GoogleAIGeminiChatGenerator` (#772)

* bugfix 654

aligned _message_to_content to VertexAIGeminiChatGenerator

* Update test_chat_gemini.py

---------

Co-authored-by: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
---
 .../components/generators/google_ai/chat/gemini.py   |  2 --
 .../tests/generators/chat/test_chat_gemini.py        | 12 ++++++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py b/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py
index c55fcab67..d3a8299fd 100644
--- a/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py
+++ b/integrations/google_ai/src/haystack_integrations/components/generators/google_ai/chat/gemini.py
@@ -234,12 +234,10 @@ def _message_to_content(self, message: ChatMessage) -> Content:
         elif message.role == ChatRole.SYSTEM:
             part = Part()
             part.text = message.content
-            return part
         elif message.role == ChatRole.FUNCTION:
             part = Part()
             part.function_response.name = message.name
             part.function_response.response = message.content
-            return part
         elif message.role == ChatRole.USER:
             part = self._convert_part(message.content)
         else:
diff --git a/integrations/google_ai/tests/generators/chat/test_chat_gemini.py b/integrations/google_ai/tests/generators/chat/test_chat_gemini.py
index 70a8766af..730b01422 100644
--- a/integrations/google_ai/tests/generators/chat/test_chat_gemini.py
+++ b/integrations/google_ai/tests/generators/chat/test_chat_gemini.py
@@ -225,3 +225,15 @@ def get_current_weather(location: str, unit: str = "celsius"):  # noqa: ARG001
 
     res = gemini_chat.run(messages=messages)
     assert len(res["replies"]) > 0
+
+
+@pytest.mark.skipif("GOOGLE_API_KEY" not in os.environ, reason="GOOGLE_API_KEY not set")
+def test_past_conversation():
+    gemini_chat = GoogleAIGeminiChatGenerator(model="gemini-pro")
+    messages = [
+        ChatMessage.from_user(content="What is 2+2?"),
+        ChatMessage.from_system(content="It's an arithmetic operation."),
+        ChatMessage.from_user(content="Yeah, but what's the result?"),
+    ]
+    res = gemini_chat.run(messages=messages)
+    assert len(res["replies"]) > 0