diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 894114559d..235c13c59c 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -952,7 +952,7 @@ def completion( original_response=json.dumps(response_body), additional_args={"complete_input_dict": data}, ) - print_verbose(f"raw model_response: {response}") + print_verbose(f"raw model_response: {response_body}") ## RESPONSE OBJECT outputText = "default" if provider == "ai21": @@ -1065,6 +1065,7 @@ def completion( outputText = response_body.get("results")[0].get("outputText") response_metadata = response.get("ResponseMetadata", {}) + if response_metadata.get("HTTPStatusCode", 500) >= 400: raise BedrockError( message=outputText, @@ -1100,11 +1101,13 @@ def completion( prompt_tokens = response_metadata.get( "x-amzn-bedrock-input-token-count", len(encoding.encode(prompt)) ) + _text_response = model_response["choices"][0]["message"].get("content", "") completion_tokens = response_metadata.get( "x-amzn-bedrock-output-token-count", len( encoding.encode( - model_response["choices"][0]["message"].get("content", "") + _text_response, + disallowed_special=(), ) ), ) diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index 0fd2cfabb5..2754ac6561 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -529,6 +529,7 @@ def test_chat_bedrock_stream(): @pytest.mark.asyncio async def test_async_chat_bedrock_stream(): try: + litellm.set_verbose = True customHandler = CompletionCustomHandler() litellm.callbacks = [customHandler] response = await litellm.acompletion( diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index d0d8a720a7..271a53dd49 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -518,7 +518,7 @@ async def test_acompletion_gemini_stream(): litellm.set_verbose = True print("Streaming gemini response") messages = [ - {"role": "system", "content": "You are a helpful assistant."}, + # {"role": "system", "content": "You are a helpful assistant."}, { "role": "user", "content": "What do you know?", @@ -1297,6 +1297,7 @@ def test_completion_watsonx_stream(): except Exception as e: pytest.fail(f"Error occurred: {e}") + # test_completion_sagemaker_stream() diff --git a/litellm/utils.py b/litellm/utils.py index be14b13776..6b8aaaa5d8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -410,10 +410,11 @@ class Message(OpenAIObject): except: # if using pydantic v1 return self.dict() - + + class Delta(OpenAIObject): - tool_calls: Optional[List[ChatCompletionDeltaToolCall]] - + tool_calls: Optional[List[ChatCompletionDeltaToolCall]] = None + def __init__( self, content=None, @@ -475,7 +476,7 @@ class Choices(OpenAIObject): ) # set finish_reason for all responses self.index = index if message is None: - self.message = Message(content=None) + self.message = Message() else: if isinstance(message, Message): self.message = message