phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(vertex_ai.py): parse out text response from response object

Browse source
This commit is contained in:
Krrish Dholakia 2023-10-24 08:52:04 -07:00
parent d244978247
commit b0842e328c
2 changed files with 5 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/vertex_ai.py
View file

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

6
litellm/tests/test_completion.py
View file

@ -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()