Merge pull request #1969 from kan-bayashi/fix/support-multiple-tools-…

…in-gemini fix: fix the issues when using tools in gemini
BerriAI · Feb 21, 2024 · 851473b · 851473b
2 parents 9f753ab + 9077591
commit 851473b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py
@@ -439,8 +439,8 @@ def completion(
                 tools=tools,
             )
 
-            if tools is not None and hasattr(
-                response.candidates[0].content.parts[0], "function_call"
+            if tools is not None and bool(
+                getattr(response.candidates[0].content.parts[0], "function_call", None)
             ):
                 function_call = response.candidates[0].content.parts[0].function_call
                 args_dict = {}

diff --git a/litellm/utils.py b/litellm/utils.py
@@ -4288,18 +4288,15 @@ def _map_and_modify_arg(supported_params: dict, provider: str, model: str):
         if tools is not None and isinstance(tools, list):
             from vertexai.preview import generative_models
 
-            gtools = []
+            gtool_func_declarations = []
             for tool in tools:
-                gtool = generative_models.FunctionDeclaration(
+                gtool_func_declaration = generative_models.FunctionDeclaration(
                     name=tool["function"]["name"],
                     description=tool["function"].get("description", ""),
                     parameters=tool["function"].get("parameters", {}),
                 )
-                gtool_func_declaration = generative_models.Tool(
-                    function_declarations=[gtool]
-                )
-                gtools.append(gtool_func_declaration)
-            optional_params["tools"] = gtools
+                gtool_func_declarations.append(gtool_func_declaration)
+            optional_params["tools"] = [generative_models.Tool(function_declarations=gtool_func_declarations)]
     elif custom_llm_provider == "sagemaker":
         ## check if unsupported param passed in
         supported_params = ["stream", "temperature", "max_tokens", "top_p", "stop", "n"]