From bc30a8e175d27d92a18c653f96531fc664eb468b Mon Sep 17 00:00:00 2001 From: Silvano Cerza Date: Wed, 3 Jan 2024 12:31:47 +0100 Subject: [PATCH] Add VertexAI prefix to GeminiGenerator and GeminiChatGenerator components --- .../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: