scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge pull request #3111 from BerriAI/litellm_streaming_fix

Browse source 
fix streaming special character flushing logic
This commit is contained in:
Krish Dholakia 2024-04-17 19:08:08 -07:00 • committed by GitHub
commit 9c7179e66f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 65 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

79
litellm/tests/test_function_calling.py
View file

@ -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}")

13
litellm/tests/test_streaming.py
View file

@ -220,6 +220,19 @@ tools_schema = [
# test_completion_cohere_stream()
def test_completion_azure_stream_special_char():
litellm.set_verbose = True
messages = [
{"role": "user", "content": "Respond with the '<' sign and nothing else."}
]
response = completion(model="azure/chatgpt-v-2", messages=messages, stream=True)
response_str = ""
for part in response:
response_str += part.choices[0].delta.content or ""
assert len(response_str) > 0
def test_completion_cohere_stream_bad_key():
try:
litellm.cache = None

12
litellm/utils.py
View file

@ -8856,7 +8856,16 @@ class CustomStreamWrapper:
raise e
def check_special_tokens(self, chunk: str, finish_reason: Optional[str]):
"""
Output parse <s> / </s> special tokens for sagemaker + hf streaming.
"""
hold = False
# if (
# self.custom_llm_provider != "huggingface"
# and self.custom_llm_provider != "sagemaker"
# ):
# return hold, chunk
if finish_reason:
for token in self.special_tokens:
if token in chunk:
@ -8872,6 +8881,7 @@ class CustomStreamWrapper:
for token in self.special_tokens:
if len(curr_chunk) < len(token) and curr_chunk in token:
hold = True
self.holding_chunk = curr_chunk
elif len(curr_chunk) >= len(token):
if token in curr_chunk:
self.holding_chunk = curr_chunk.replace(token, "")
@ -9953,6 +9963,7 @@ class CustomStreamWrapper:
f"model_response.choices[0].delta: {model_response.choices[0].delta}; completion_obj: {completion_obj}"
)
print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}")
## RETURN ARG
if (
"content" in completion_obj
@ -10025,7 +10036,6 @@ class CustomStreamWrapper:
elif self.received_finish_reason is not None:
if self.sent_last_chunk == True:
raise StopIteration
# flush any remaining holding chunk
if len(self.holding_chunk) > 0:
if model_response.choices[0].delta.content is None: