Updates based on PR review

deepset-ai · Sep 26, 2024 · 381ff4e · 381ff4e
1 parent ca126f8
commit 381ff4e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 33 deletions.
diff --git a/...oogle_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py b/...oogle_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py
@@ -175,25 +175,15 @@ def from_dict(cls, data: Dict[str, Any]) -> "VertexAIGeminiChatGenerator":
             Deserialized component.
         """
 
-        def _tool_config_from_dict(config_dict: Dict[str, Any]) -> Tool:
+        def _tool_config_from_dict(config_dict: Dict[str, Any]) -> ToolConfig:
             """Deserializes the ToolConfig object from a dictionary."""
-
-            allowed_function_names = config_dict["function_calling_config"].get("allowed_function_names")
-            mode = config_dict["function_calling_config"]["mode"]
-            if allowed_function_names:
-                tool_config = ToolConfig(
-                    function_calling_config=ToolConfig.FunctionCallingConfig(
-                        mode=mode,
-                        allowed_function_names=allowed_function_names,
-                    )
-                )
-            else:
-                tool_config = ToolConfig(
-                    function_calling_config=ToolConfig.FunctionCallingConfig(
-                        mode=config_dict["function_calling_config"]["mode"],
-                    )
+            function_calling_config = config_dict["function_calling_config"]
+            return ToolConfig(
+                function_calling_config=ToolConfig.FunctionCallingConfig(
+                    mode=function_calling_config["mode"],
+                    allowed_function_names=function_calling_config.get("allowed_function_names"),
                 )
-            return tool_config
+            )
 
         if (tools := data["init_parameters"].get("tools")) is not None:
             data["init_parameters"]["tools"] = [Tool.from_dict(t) for t in tools]

diff --git a/...ons/google_vertex/src/haystack_integrations/components/generators/google_vertex/gemini.py b/...ons/google_vertex/src/haystack_integrations/components/generators/google_vertex/gemini.py
@@ -184,24 +184,15 @@ def from_dict(cls, data: Dict[str, Any]) -> "VertexAIGeminiGenerator":
            Deserialized component.
         """
 
-        def _tool_config_from_dict(config_dict: Dict[str, Any]) -> Tool:
+        def _tool_config_from_dict(config_dict: Dict[str, Any]) -> ToolConfig:
             """Deserializes the ToolConfig object from a dictionary."""
-            allowed_function_names = config_dict["function_calling_config"].get("allowed_function_names")
-            mode = config_dict["function_calling_config"]["mode"]
-            if allowed_function_names:
-                tool_config = ToolConfig(
-                    function_calling_config=ToolConfig.FunctionCallingConfig(
-                        mode=mode,
-                        allowed_function_names=allowed_function_names,
-                    )
+            function_calling_config = config_dict["function_calling_config"]
+            return ToolConfig(
+                function_calling_config=ToolConfig.FunctionCallingConfig(
+                    mode=function_calling_config["mode"],
+                    allowed_function_names=function_calling_config.get("allowed_function_names"),
                 )
-            else:
-                tool_config = ToolConfig(
-                    function_calling_config=ToolConfig.FunctionCallingConfig(
-                        mode=config_dict["function_calling_config"]["mode"],
-                    )
-                )
-            return tool_config
+            )
 
         if (tools := data["init_parameters"].get("tools")) is not None:
             data["init_parameters"]["tools"] = [Tool.from_dict(t) for t in tools]