From bbb5bd19e862c0539d916f4f5cb73b5a0edf72a2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 1 Feb 2024 21:20:46 -0800 Subject: [PATCH] test: testing fixes --- litellm/tests/test_key_generate_prisma.py | 4 +--- litellm/utils.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 9d4318fe7..d8ffcf022 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -1266,9 +1266,7 @@ async def test_user_api_key_auth(prisma_client): pytest.fail(f"This should have failed!. IT's an invalid key") except ProxyException as exc: print(exc.message) - assert ( - exc.message == "Authentication Error, No API Key passed in. api_key is None" - ) + assert exc.message == "Authentication Error, No api key passed in." # Test case: Malformed API Key (missing 'Bearer ' prefix) try: diff --git a/litellm/utils.py b/litellm/utils.py index 03dfc40f2..2d428c26e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8072,6 +8072,17 @@ class CustomStreamWrapper: 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