diff --git a/litellm/tests/test_function_calling.py b/litellm/tests/test_function_calling.py index 9c6e7ca79e..f86041e3b8 100644 --- a/litellm/tests/test_function_calling.py +++ b/litellm/tests/test_function_calling.py @@ -12,14 +12,12 @@ import pytest import litellm from litellm import embedding, completion, completion_cost, Timeout from litellm import RateLimitError +import pytest litellm.num_retries = 3 litellm.cache = None -litellm.set_verbose=True +litellm.set_verbose=False import json - -# Example dummy function hard coded to return the same weather -# In production, this could be your backend API or an external API def get_current_weather(location, unit="fahrenheit"): """Get the current weather in a given location""" if "tokyo" in location.lower(): @@ -31,67 +29,79 @@ def get_current_weather(location, unit="fahrenheit"): else: return json.dumps({"location": location, "temperature": "unknown"}) -def run_conversation(): - # Step 1: send the conversation and available functions to the model - messages = [{"role": "user", "content": "What's the weather like in San Francisco, Tokyo, and Paris?"}] - tools = [ - { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", +# Example dummy function hard coded to return the same weather +# In production, this could be your backend API or an external API +def test_parallel_function_call(): + try: + # Step 1: send the conversation and available functions to the model + messages = [{"role": "user", "content": "What's the weather like in San Francisco, Tokyo, and Paris?"}] + tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, }, - "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + "required": ["location"], }, - "required": ["location"], }, - }, - } - ] - response = litellm.completion( - model="gpt-3.5-turbo-1106", - messages=messages, - tools=tools, - tool_choice="auto", # auto is default, but we'll be explicit - ) - print("Response\n", response) - response_message = response.choices[0].message - tool_calls = response_message.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"), - ) - messages.append( - { - "tool_call_id": tool_call.id, - "role": "tool", - "name": function_name, - "content": function_response, - } - ) # extend conversation with function response - second_response = litellm.completion( + } + ] + response = litellm.completion( model="gpt-3.5-turbo-1106", messages=messages, - ) # get a new response from the model where it can see the function response - return second_response -print(run_conversation()) \ No newline at end of file + tools=tools, + tool_choice="auto", # auto is default, but we'll be explicit + ) + print("Response\n", response) + response_message = response.choices[0].message + tool_calls = response_message.tool_calls + + print("length of tool calls", len(tool_calls)) + print("Expecting there to be 3 tool calls") + assert len(tool_calls) > 1 # this has to call the function for SF, Tokyo and parise + + # 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"), + ) + messages.append( + { + "tool_call_id": tool_call.id, + "role": "tool", + "name": function_name, + "content": function_response, + } + ) # extend conversation with function response + second_response = litellm.completion( + model="gpt-3.5-turbo-1106", + messages=messages, + ) # get a new response from the model where it can see the function response + print("second response\n", second_response) + return second_response + except Exception as e: + pytest.fail(f"Error occurred: {e}") + +test_parallel_function_call() \ No newline at end of file