phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(utils.py): fix parallel tool calling when streaming

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-29 10:56:21 -08:00
parent 9024a47dc2
commit b6bc75e27a
3 changed files with 74 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

89
litellm/tests/test_function_calling.py
View file

@ -13,7 +13,7 @@ import litellm
from litellm import embedding, completion, completion_cost, Timeout
from litellm import RateLimitError
import pytest
litellm.num_retries = 3
litellm.num_retries = 0
litellm.cache = None
# litellm.set_verbose=True
import json
@ -97,6 +97,7 @@ def test_parallel_function_call():
"content": function_response,
}
) # extend conversation with function response
print(f"messages: {messages}")
second_response = litellm.completion(
model="gpt-3.5-turbo-1106",
messages=messages,
@ -108,7 +109,7 @@ def test_parallel_function_call():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_parallel_function_call()
test_parallel_function_call()
@ -143,51 +144,53 @@ def test_parallel_function_call_stream():
tools=tools,
stream=True,
tool_choice="auto", # auto is default, but we'll be explicit
complete_response = True
)
print("Response\n", response)
for chunk in response:
print(chunk)
# response_message = response.choices[0].message
# tool_calls = response_message.tool_calls
# for chunk in response:
# print(chunk)
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
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,
# temperature=0.2,
# seed=22
# ) # get a new response from the model where it can see the function response
# print("second response\n", second_response)
# return second_response
# 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
print(f"messages: {messages}")
second_response = litellm.completion(
model="gpt-3.5-turbo-1106",
messages=messages,
temperature=0.2,
seed=22
) # 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}")