From fabc398acd1908db4e16b08e11ad5047d644c454 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 17 Apr 2024 18:28:59 -0700
Subject: [PATCH] test(test_function_calling.py): fix test to check if tool
 call in response

---
 litellm/tests/test_function_calling.py | 79 +++++++++++++-------------
 1 file changed, 41 insertions(+), 38 deletions(-)

diff --git a/litellm/tests/test_function_calling.py b/litellm/tests/test_function_calling.py
index 2a815edc8e..f76a082f60 100644
--- a/litellm/tests/test_function_calling.py
+++ b/litellm/tests/test_function_calling.py
@@ -266,47 +266,50 @@ def test_groq_parallel_function_call():
         )
         print("Response\n", response)
         response_message = response.choices[0].message
-        tool_calls = response_message.tool_calls
+        if hasattr(response_message, "tool_calls"):
+            tool_calls = response_message.tool_calls
 
-        assert isinstance(response.choices[0].message.tool_calls[0].function.name, str)
-        assert isinstance(
-            response.choices[0].message.tool_calls[0].function.arguments, str
-        )
+            assert isinstance(
+                response.choices[0].message.tool_calls[0].function.name, str
+            )
+            assert isinstance(
+                response.choices[0].message.tool_calls[0].function.arguments, str
+            )
 
-        print("length of tool calls", len(tool_calls))
+            print("length of tool calls", len(tool_calls))
 
-        # Step 2: check if the model wanted to call a function
-        if tool_calls:
-            # Step 3: call the function
-            # Note: the JSON response may not always be valid; be sure to handle errors
-            available_functions = {
-                "get_current_weather": get_current_weather,
-            }  # only one function in this example, but you can have multiple
-            messages.append(
-                response_message
-            )  # extend conversation with assistant's reply
-            print("Response message\n", response_message)
-            # Step 4: send the info for each function call and function response to the model
-            for tool_call in tool_calls:
-                function_name = tool_call.function.name
-                function_to_call = available_functions[function_name]
-                function_args = json.loads(tool_call.function.arguments)
-                function_response = function_to_call(
-                    location=function_args.get("location"),
-                    unit=function_args.get("unit"),
-                )
+            # Step 2: check if the model wanted to call a function
+            if tool_calls:
+                # Step 3: call the function
+                # Note: the JSON response may not always be valid; be sure to handle errors
+                available_functions = {
+                    "get_current_weather": get_current_weather,
+                }  # only one function in this example, but you can have multiple
                 messages.append(
-                    {
-                        "tool_call_id": tool_call.id,
-                        "role": "tool",
-                        "name": function_name,
-                        "content": function_response,
-                    }
-                )  # extend conversation with function response
-            print(f"messages: {messages}")
-            second_response = litellm.completion(
-                model="groq/llama2-70b-4096", messages=messages
-            )  # get a new response from the model where it can see the function response
-            print("second response\n", second_response)
+                    response_message
+                )  # extend conversation with assistant's reply
+                print("Response message\n", response_message)
+                # Step 4: send the info for each function call and function response to the model
+                for tool_call in tool_calls:
+                    function_name = tool_call.function.name
+                    function_to_call = available_functions[function_name]
+                    function_args = json.loads(tool_call.function.arguments)
+                    function_response = function_to_call(
+                        location=function_args.get("location"),
+                        unit=function_args.get("unit"),
+                    )
+                    messages.append(
+                        {
+                            "tool_call_id": tool_call.id,
+                            "role": "tool",
+                            "name": function_name,
+                            "content": function_response,
+                        }
+                    )  # extend conversation with function response
+                print(f"messages: {messages}")
+                second_response = litellm.completion(
+                    model="groq/llama2-70b-4096", messages=messages
+                )  # get a new response from the model where it can see the function response
+                print("second response\n", second_response)
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")