Add VertexAI prefix to GeminiGenerator and GeminiChatGenerator compon…

…ents (#166)
deepset-ai · Jan 3, 2024 · b252d72 · b252d72
1 parent 065042e
commit b252d72
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 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
@@ -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
@@ -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: