phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

vertex ai fix

Browse source
This commit is contained in:
ishaan-jaff 2023-08-29 09:48:52 -07:00
parent 1d782a1153
commit fd77baef83
1 changed files with 2 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/main.py
View file

@ -612,8 +612,6 @@ def completion(
} }
response = model_response response = model_response
elif model in litellm.vertex_chat_models: elif model in litellm.vertex_chat_models:
# import vertexai/if it fails then pip install vertexai# import cohere/if it fails then pip install cohere
install_and_import("vertexai")
import vertexai import vertexai
from vertexai.preview.language_models import ChatModel, InputOutputTextPair from vertexai.preview.language_models import ChatModel, InputOutputTextPair
@ -637,12 +635,10 @@ def completion(
) )
## RESPONSE OBJECT ## RESPONSE OBJECT
model_response["choices"][0]["message"]["content"] = completion_response model_response["choices"][0]["message"]["content"] = str(completion_response)
model_response["created"] = time.time() model_response["created"] = time.time()
model_response["model"] = model model_response["model"] = model
elif model in litellm.vertex_text_models: elif model in litellm.vertex_text_models:
# import vertexai/if it fails then pip install vertexai# import cohere/if it fails then pip install cohere
install_and_import("vertexai")
import vertexai import vertexai
from vertexai.language_models import TextGenerationModel from vertexai.language_models import TextGenerationModel
@ -663,7 +659,7 @@ def completion(
input=prompt, api_key=None, original_response=completion_response input=prompt, api_key=None, original_response=completion_response
) )
## RESPONSE OBJECT ## RESPONSE OBJECT
model_response["choices"][0]["message"]["content"] = completion_response model_response["choices"][0]["message"]["content"] = str(completion_response)
model_response["created"] = time.time() model_response["created"] = time.time()
model_response["model"] = model model_response["model"] = model
response = model_response response = model_response