From b252d72755d294400009a0f842c79637f80f7727 Mon Sep 17 00:00:00 2001
From: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
Date: Wed, 3 Jan 2024 12:36:38 +0100
Subject: [PATCH] Add VertexAI prefix to GeminiGenerator and
 GeminiChatGenerator components (#166)

---
 .../src/google_vertex_haystack/generators/chat/gemini.py      | 4 ++--
 .../src/google_vertex_haystack/generators/gemini.py           | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/chat/gemini.py b/integrations/google_vertex/src/google_vertex_haystack/generators/chat/gemini.py
index 765706301..5a6137765 100644
--- a/integrations/google_vertex/src/google_vertex_haystack/generators/chat/gemini.py
+++ b/integrations/google_vertex/src/google_vertex_haystack/generators/chat/gemini.py
@@ -21,7 +21,7 @@
 
 
 @component
-class GeminiChatGenerator:
+class VertexAIGeminiChatGenerator:
     def __init__(
         self,
         *,
@@ -100,7 +100,7 @@ def to_dict(self) -> Dict[str, Any]:
         return data
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "GeminiChatGenerator":
+    def from_dict(cls, data: Dict[str, Any]) -> "VertexAIGeminiChatGenerator":
         if (tools := data["init_parameters"].get("tools")) is not None:
             data["init_parameters"]["tools"] = [Tool.from_dict(t) for t in tools]
         if (generation_config := data["init_parameters"].get("generation_config")) is not None:
diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/gemini.py b/integrations/google_vertex/src/google_vertex_haystack/generators/gemini.py
index 995a52614..698b07b01 100644
--- a/integrations/google_vertex/src/google_vertex_haystack/generators/gemini.py
+++ b/integrations/google_vertex/src/google_vertex_haystack/generators/gemini.py
@@ -21,7 +21,7 @@
 
 
 @component
-class GeminiGenerator:
+class VertexAIGeminiGenerator:
     def __init__(
         self,
         *,
@@ -111,7 +111,7 @@ def to_dict(self) -> Dict[str, Any]:
         return data
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "GeminiGenerator":
+    def from_dict(cls, data: Dict[str, Any]) -> "VertexAIGeminiGenerator":
         if (tools := data["init_parameters"].get("tools")) is not None:
             data["init_parameters"]["tools"] = [Tool.from_dict(t) for t in tools]
         if (generation_config := data["init_parameters"].get("generation_config")) is not None: