From dfb1d34e264b16b0ce992e6cab48a4059f02941f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 12:31:00 -0800 Subject: [PATCH 1/4] fix(utils.py): support returning caching streaming response for function calling streaming calls --- litellm/tests/test_streaming.py | 48 ++++++++- litellm/utils.py | 166 +++++++++++++++++--------------- 2 files changed, 134 insertions(+), 80 deletions(-) diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 66e8be4cbea..dc1e1a097a9 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -1907,6 +1907,8 @@ def test_azure_streaming_and_function_calling(): @pytest.mark.asyncio async def test_azure_astreaming_and_function_calling(): + import uuid + tools = [ { "type": "function", @@ -1927,7 +1929,20 @@ async def test_azure_astreaming_and_function_calling(): }, } ] - messages = [{"role": "user", "content": "What is the weather like in Boston?"}] + messages = [ + { + "role": "user", + "content": f"What is the weather like in Boston? {uuid.uuid4()}", + } + ] + from litellm.caching import Cache + + litellm.cache = Cache( + type="redis", + host=os.environ["REDIS_HOST"], + port=os.environ["REDIS_PORT"], + password=os.environ["REDIS_PASSWORD"], + ) try: response = await litellm.acompletion( model="azure/gpt-4-nov-release", @@ -1938,6 +1953,7 @@ async def test_azure_astreaming_and_function_calling(): api_base=os.getenv("AZURE_FRANCE_API_BASE"), api_key=os.getenv("AZURE_FRANCE_API_KEY"), api_version="2024-02-15-preview", + caching=True, ) # Add any assertions here to check the response idx = 0 @@ -1957,6 +1973,36 @@ async def test_azure_astreaming_and_function_calling(): validate_final_streaming_function_calling_chunk(chunk=chunk) idx += 1 + ## CACHING TEST + print("\n\nCACHING TESTS\n\n") + response = await litellm.acompletion( + model="azure/gpt-4-nov-release", + tools=tools, + tool_choice="auto", + messages=messages, + stream=True, + api_base=os.getenv("AZURE_FRANCE_API_BASE"), + api_key=os.getenv("AZURE_FRANCE_API_KEY"), + api_version="2024-02-15-preview", + caching=True, + ) + # Add any assertions here to check the response + idx = 0 + async for chunk in response: + print(f"chunk: {chunk}") + if idx == 0: + assert ( + chunk.choices[0].delta.tool_calls[0].function.arguments is not None + ) + assert isinstance( + chunk.choices[0].delta.tool_calls[0].function.arguments, str + ) + validate_first_streaming_function_calling_chunk(chunk=chunk) + elif idx == 1: + validate_second_streaming_function_calling_chunk(chunk=chunk) + elif chunk.choices[0].finish_reason is not None: # last chunk + validate_final_streaming_function_calling_chunk(chunk=chunk) + idx += 1 except Exception as e: pytest.fail(f"Error occurred: {e}") raise e diff --git a/litellm/utils.py b/litellm/utils.py index 7de5199b46f..ad9b65998f9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -213,6 +213,13 @@ class Function(OpenAIObject): name: str +class ChatCompletionDeltaToolCall(OpenAIObject): + id: str + function: Function + type: str + index: int + + class ChatCompletionMessageToolCall(OpenAIObject): id: str function: Function @@ -269,7 +276,15 @@ class Delta(OpenAIObject): self.content = content self.role = role self.function_call = function_call - self.tool_calls = tool_calls + if tool_calls is not None: + if isinstance(tool_calls, dict): + self.tool_calls = [] + for tool_call in tool_calls: + if tool_call.get("index", None) is None: + tool_call["index"] = 0 + self.tool_calls.append(ChatCompletionDeltaToolCall(**tool_call)) + else: + self.tool_calls = tool_calls def __contains__(self, key): # Define custom behavior for the 'in' operator @@ -5847,6 +5862,18 @@ async def convert_to_streaming_response_async(response_object: Optional[dict] = choice_list = [] for idx, choice in enumerate(response_object["choices"]): + if ( + choice["message"].get("tool_calls", None) is not None + and isinstance(choice["message"]["tool_calls"], list) + and len(choice["message"]["tool_calls"]) > 0 + and isinstance(choice["message"]["tool_calls"][0], dict) + ): + pydantic_tool_calls = [] + for index, t in enumerate(choice["message"]["tool_calls"]): + if "index" not in t: + t["index"] = index + pydantic_tool_calls.append(ChatCompletionDeltaToolCall(**t)) + choice["message"]["tool_calls"] = pydantic_tool_calls delta = Delta( content=choice["message"].get("content", None), role=choice["message"]["role"], @@ -8646,6 +8673,7 @@ class CustomStreamWrapper: "text": chunk.choices[0].delta.content, "is_finished": True, "finish_reason": chunk.choices[0].finish_reason, + "original_chunk": chunk, } completion_obj["content"] = response_obj["text"] @@ -8676,80 +8704,11 @@ class CustomStreamWrapper: model_response.choices[0].logprobs = response_obj["logprobs"] model_response.model = self.model - print_verbose( - f"model_response: {model_response}; completion_obj: {completion_obj}" - ) print_verbose( f"model_response finish reason 3: {model_response.choices[0].finish_reason}" ) + ## FUNCTION CALL PARSING if ( - len(completion_obj["content"]) > 0 - ): # cannot set content of an OpenAI Object to be an empty string - hold, model_response_str = self.check_special_tokens( - chunk=completion_obj["content"], - finish_reason=model_response.choices[0].finish_reason, - ) # filter out bos/eos tokens from openai-compatible hf endpoints - print_verbose( - f"hold - {hold}, model_response_str - {model_response_str}" - ) - if hold is False: - ## check if openai/azure chunk - original_chunk = response_obj.get("original_chunk", None) - if original_chunk: - model_response.id = original_chunk.id - if len(original_chunk.choices) > 0: - try: - delta = dict(original_chunk.choices[0].delta) - print_verbose(f"original delta: {delta}") - model_response.choices[0].delta = Delta(**delta) - except Exception as e: - model_response.choices[0].delta = Delta() - else: - return - model_response.system_fingerprint = ( - original_chunk.system_fingerprint - ) - print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}") - if self.sent_first_chunk == False: - model_response.choices[0].delta["role"] = "assistant" - self.sent_first_chunk = True - elif self.sent_first_chunk == True and hasattr( - model_response.choices[0].delta, "role" - ): - _initial_delta = model_response.choices[ - 0 - ].delta.model_dump() - _initial_delta.pop("role", None) - model_response.choices[0].delta = Delta(**_initial_delta) - print_verbose( - f"model_response.choices[0].delta: {model_response.choices[0].delta}" - ) - else: - ## else - completion_obj["content"] = model_response_str - if self.sent_first_chunk == False: - completion_obj["role"] = "assistant" - self.sent_first_chunk = True - model_response.choices[0].delta = Delta(**completion_obj) - print_verbose(f"returning model_response: {model_response}") - return model_response - else: - return - elif model_response.choices[0].finish_reason: - # flush any remaining holding chunk - if len(self.holding_chunk) > 0: - if model_response.choices[0].delta.content is None: - model_response.choices[0].delta.content = self.holding_chunk - else: - model_response.choices[0].delta.content = ( - self.holding_chunk + model_response.choices[0].delta.content - ) - self.holding_chunk = "" - model_response.choices[0].finish_reason = map_finish_reason( - model_response.choices[0].finish_reason - ) # ensure consistent output to openai - return model_response - elif ( response_obj is not None and response_obj.get("original_chunk", None) is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints @@ -8783,26 +8742,75 @@ class CustomStreamWrapper: original_chunk.choices[0].delta.tool_calls, list ): for t in original_chunk.choices[0].delta.tool_calls: - if ( - getattr( - t.function, - "arguments", - ) - is None + if hasattr(t, "functions") and hasattr( + t.functions, "arguments" ): - t.function.arguments = "" + if ( + getattr( + t.function, + "arguments", + ) + is None + ): + t.function.arguments = "" model_response.choices[0].delta = Delta(**delta) except Exception as e: traceback.print_exc() model_response.choices[0].delta = Delta() else: - return + try: + delta = dict(original_chunk.choices[0].delta) + print_verbose(f"original delta: {delta}") + model_response.choices[0].delta = Delta(**delta) + except Exception as e: + model_response.choices[0].delta = Delta() else: return model_response.system_fingerprint = original_chunk.system_fingerprint if self.sent_first_chunk == False: model_response.choices[0].delta["role"] = "assistant" self.sent_first_chunk = True + + ## RETURN ARG + if ( + response_obj.get("text", None) is not None + or response_obj.get("original_chunk", None) is not None + ): + hold = False + if response_obj.get("content", None) is not None: + hold, model_response_str = self.check_special_tokens( + chunk=completion_obj["content"], + finish_reason=model_response.choices[0].finish_reason, + ) # filter out bos/eos tokens from openai-compatible hf endpoints + print_verbose( + f"hold - {hold}, model_response_str - {model_response_str}" + ) + if hold is False: + original_chunk = response_obj.get("original_chunk", None) + if original_chunk is None: + completion_obj["content"] = model_response_str + if self.sent_first_chunk == False: + completion_obj["role"] = "assistant" + self.sent_first_chunk = True + model_response.choices[0].delta = Delta(**completion_obj) + print_verbose(f"returning model_response: {model_response}") + return model_response + else: + return + elif model_response.choices[0].finish_reason is not None: + # flush any remaining holding chunk + if len(self.holding_chunk) > 0: + if model_response.choices[0].delta.content is None: + model_response.choices[0].delta.content = self.holding_chunk + else: + model_response.choices[0].delta.content = ( + self.holding_chunk + model_response.choices[0].delta.content + ) + self.holding_chunk = "" + # get any function call arguments + model_response.choices[0].finish_reason = map_finish_reason( + model_response.choices[0].finish_reason + ) # ensure consistent output to openai return model_response else: return From 5b06627c09ea7ec2651aeb36551871e824e2fce3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 12:52:53 -0800 Subject: [PATCH 2/4] fix(utils.py): fix streaming --- litellm/tests/test_amazing_vertex_completion.py | 4 +++- litellm/tests/test_streaming.py | 1 - litellm/utils.py | 6 +++--- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 35beb75fe6f..d58d68507a0 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -138,7 +138,7 @@ def test_vertex_ai(): def test_vertex_ai_stream(): load_vertex_ai_credentials() - litellm.set_verbose = False + litellm.set_verbose = True litellm.vertex_project = "reliablekeys" import random @@ -159,6 +159,8 @@ def test_vertex_ai_stream(): "code-gecko@latest", "code-bison@001", "text-bison@001", + "gemini-1.5-pro", + "gemini-1.5-pro-vision", ]: # our account does not have access to this model continue diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index dc1e1a097a9..8c3187bd4ae 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -801,7 +801,6 @@ def test_completion_bedrock_claude_stream(): raise Exception("finish reason not set for last chunk") if complete_response.strip() == "": raise Exception("Empty response received") - print(f"completion_response: {complete_response}") except RateLimitError: pass except Exception as e: diff --git a/litellm/utils.py b/litellm/utils.py index ad9b65998f9..566c6f36df0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8705,7 +8705,7 @@ class CustomStreamWrapper: model_response.model = self.model print_verbose( - f"model_response finish reason 3: {model_response.choices[0].finish_reason}" + f"model_response finish reason 3: {model_response.choices[0].finish_reason}; response_obj={response_obj}" ) ## FUNCTION CALL PARSING if ( @@ -8773,11 +8773,11 @@ class CustomStreamWrapper: ## RETURN ARG if ( - response_obj.get("text", None) is not None + completion_obj["content"] is not None or response_obj.get("original_chunk", None) is not None ): hold = False - if response_obj.get("content", None) is not None: + if completion_obj["content"] is not None: hold, model_response_str = self.check_special_tokens( chunk=completion_obj["content"], finish_reason=model_response.choices[0].finish_reason, From 788e24bd837129914c853bdb43144316653a6677 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 14:26:58 -0800 Subject: [PATCH 3/4] fix(utils.py): fix streaming logic --- litellm/main.py | 1 + litellm/tests/test_function_calling.py | 5 +- litellm/utils.py | 112 ++++++++++++++++++------- 3 files changed, 87 insertions(+), 31 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index ed8dddf05b2..1a6c8e178ea 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3678,6 +3678,7 @@ def stream_chunk_builder( response["usage"]["total_tokens"] = ( response["usage"]["prompt_tokens"] + response["usage"]["completion_tokens"] ) + return convert_to_model_response_object( response_object=response, model_response_object=model_response, diff --git a/litellm/tests/test_function_calling.py b/litellm/tests/test_function_calling.py index 2fcbdc9460e..ffef8f6594f 100644 --- a/litellm/tests/test_function_calling.py +++ b/litellm/tests/test_function_calling.py @@ -124,11 +124,12 @@ def test_parallel_function_call(): pytest.fail(f"Error occurred: {e}") -test_parallel_function_call() +# test_parallel_function_call() def test_parallel_function_call_stream(): try: + litellm.set_verbose = True # Step 1: send the conversation and available functions to the model messages = [ { @@ -217,4 +218,4 @@ def test_parallel_function_call_stream(): pytest.fail(f"Error occurred: {e}") -test_parallel_function_call_stream() +# test_parallel_function_call_stream() diff --git a/litellm/utils.py b/litellm/utils.py index 566c6f36df0..c9618e8f382 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -276,15 +276,14 @@ class Delta(OpenAIObject): self.content = content self.role = role self.function_call = function_call - if tool_calls is not None: - if isinstance(tool_calls, dict): - self.tool_calls = [] - for tool_call in tool_calls: - if tool_call.get("index", None) is None: - tool_call["index"] = 0 - self.tool_calls.append(ChatCompletionDeltaToolCall(**tool_call)) - else: - self.tool_calls = tool_calls + if tool_calls is not None and isinstance(tool_calls, dict): + self.tool_calls = [] + for tool_call in tool_calls: + if tool_call.get("index", None) is None: + tool_call["index"] = 0 + self.tool_calls.append(ChatCompletionDeltaToolCall(**tool_call)) + else: + self.tool_calls = tool_calls def __contains__(self, key): # Define custom behavior for the 'in' operator @@ -8722,6 +8721,9 @@ class CustomStreamWrapper: ): try: delta = dict(original_chunk.choices[0].delta) + model_response.system_fingerprint = ( + original_chunk.system_fingerprint + ) ## AZURE - check if arguments is not None if ( original_chunk.choices[0].delta.function_call @@ -8762,32 +8764,64 @@ class CustomStreamWrapper: delta = dict(original_chunk.choices[0].delta) print_verbose(f"original delta: {delta}") model_response.choices[0].delta = Delta(**delta) + print_verbose( + f"new delta: {model_response.choices[0].delta}" + ) except Exception as e: model_response.choices[0].delta = Delta() else: return - model_response.system_fingerprint = original_chunk.system_fingerprint - if self.sent_first_chunk == False: - model_response.choices[0].delta["role"] = "assistant" - self.sent_first_chunk = True - + print_verbose( + 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 ( - completion_obj["content"] is not None - or response_obj.get("original_chunk", None) is not None - ): - hold = False - if completion_obj["content"] is not None: - hold, model_response_str = self.check_special_tokens( - chunk=completion_obj["content"], - finish_reason=model_response.choices[0].finish_reason, - ) # filter out bos/eos tokens from openai-compatible hf endpoints - print_verbose( - f"hold - {hold}, model_response_str - {model_response_str}" - ) + "content" in completion_obj + and isinstance(completion_obj["content"], str) + and len(completion_obj["content"]) > 0 + ): # cannot set content of an OpenAI Object to be an empty string + hold, model_response_str = self.check_special_tokens( + chunk=completion_obj["content"], + finish_reason=model_response.choices[0].finish_reason, + ) # filter out bos/eos tokens from openai-compatible hf endpoints + print_verbose( + f"hold - {hold}, model_response_str - {model_response_str}" + ) if hold is False: + ## check if openai/azure chunk original_chunk = response_obj.get("original_chunk", None) - if original_chunk is None: + if original_chunk: + model_response.id = original_chunk.id + if len(original_chunk.choices) > 0: + try: + delta = dict(original_chunk.choices[0].delta) + print_verbose(f"original delta: {delta}") + model_response.choices[0].delta = Delta(**delta) + except Exception as e: + model_response.choices[0].delta = Delta() + else: + return + model_response.system_fingerprint = ( + original_chunk.system_fingerprint + ) + print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}") + if self.sent_first_chunk == False: + model_response.choices[0].delta["role"] = "assistant" + self.sent_first_chunk = True + elif self.sent_first_chunk == True and hasattr( + model_response.choices[0].delta, "role" + ): + _initial_delta = model_response.choices[ + 0 + ].delta.model_dump() + _initial_delta.pop("role", None) + model_response.choices[0].delta = Delta(**_initial_delta) + print_verbose( + f"model_response.choices[0].delta: {model_response.choices[0].delta}" + ) + else: + ## else completion_obj["content"] = model_response_str if self.sent_first_chunk == False: completion_obj["role"] = "assistant" @@ -8812,6 +8846,14 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason ) # ensure consistent output to openai return model_response + elif ( + model_response.choices[0].delta.tool_calls is not None + or model_response.choices[0].delta.function_call is not None + ): + if self.sent_first_chunk == False: + model_response.choices[0].delta["role"] = "assistant" + self.sent_first_chunk = True + return model_response else: return except StopIteration: @@ -8860,7 +8902,14 @@ class CustomStreamWrapper: print_verbose(f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}") response: Optional[ModelResponse] = self.chunk_creator(chunk=chunk) print_verbose(f"PROCESSED CHUNK POST CHUNK CREATOR: {response}") - if response is None: + + if response is None or ( + isinstance(response, ModelResponse) + and isinstance(response.choices[0], StreamingChoices) + and response.choices[0].delta.content is None + and response.choices[0].delta.function_call is None + and response.choices[0].delta.tool_calls is None + ): continue ## LOGGING threading.Thread( @@ -8904,7 +8953,11 @@ class CustomStreamWrapper: print_verbose(f"value of async chunk: {chunk}") if chunk == "None" or chunk is None: raise Exception - elif self.custom_llm_provider == "gemini" and len(chunk.parts) == 0: + elif ( + self.custom_llm_provider == "gemini" + and hasattr(chunk, "parts") + and len(chunk.parts) == 0 + ): continue # chunk_creator() does logging/stream chunk building. We need to let it know its being called in_async_func, so we don't double add chunks. # __anext__ also calls async_success_handler, which does logging @@ -8933,6 +8986,7 @@ class CustomStreamWrapper: self.rules.post_call_rules( input=self.response_uptil_now, model=self.model ) + print_verbose(f"final returned processed chunk: {processed_chunk}") return processed_chunk raise StopAsyncIteration else: # temporary patch for non-aiohttp async calls From 2a6a72a0e76b6198fced4ce6133d004fd0191a87 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 19:32:30 -0800 Subject: [PATCH 4/4] fix(utils.py): fixing sync streaming for caching --- litellm/utils.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index c9618e8f382..dd5210e08ff 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1196,7 +1196,8 @@ class Logging: start_time=start_time, end_time=end_time, ) - except: + except Exception as e: + complete_streaming_response = None else: self.sync_streaming_chunks.append(result) @@ -8903,13 +8904,7 @@ class CustomStreamWrapper: response: Optional[ModelResponse] = self.chunk_creator(chunk=chunk) print_verbose(f"PROCESSED CHUNK POST CHUNK CREATOR: {response}") - if response is None or ( - isinstance(response, ModelResponse) - and isinstance(response.choices[0], StreamingChoices) - and response.choices[0].delta.content is None - and response.choices[0].delta.function_call is None - and response.choices[0].delta.tool_calls is None - ): + if response is None: continue ## LOGGING threading.Thread(