From b0842e328c02e3fffb218de59c4dfaa8c0268042 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 24 Oct 2023 08:52:04 -0700 Subject: [PATCH] fix(vertex_ai.py): parse out text response from response object --- litellm/llms/vertex_ai.py | 2 +- litellm/tests/test_completion.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py index d73905261..211489fdb 100644 --- a/litellm/llms/vertex_ai.py +++ b/litellm/llms/vertex_ai.py @@ -112,7 +112,7 @@ def completion( model_response = chat.send_message_streaming(prompt, **optional_params) return model_response - completion_response = chat.send_message(prompt, **optional_params) + completion_response = chat.send_message(prompt, **optional_params).text elif mode == "text": ## LOGGING logging_obj.pre_call(input=prompt, api_key=None) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 9d2cf01be..90f2b10f2 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1026,12 +1026,14 @@ def test_completion_sagemaker(): # test_completion_custom_api_base() # def test_vertex_ai(): -# test_models = litellm.vertex_chat_models + litellm.vertex_code_chat_models + litellm.vertex_text_models + litellm.vertex_code_text_models +# # test_models = litellm.vertex_chat_models + litellm.vertex_code_chat_models + litellm.vertex_text_models + litellm.vertex_code_text_models +# test_models = ["chat-bison"] # for model in test_models: # try: # print("making request", model) -# response = completion(model=model, messages=[{"role": "user", "content": "write code for saying hi"}]) +# response = completion(model="chat-bison", messages=[{'role': 'user', 'content': 'hi'}]) # print(response) +# assert type(response.choices[0].message.content) == str # except Exception as e: # pytest.fail(f"Error occurred: {e}") # test_vertex_ai()