phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +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 5e709732a7
commit 200e1232e2
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
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
from vertexai.preview.language_models import ChatModel, InputOutputTextPair
@ -637,12 +635,10 @@ def completion(
)
## 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["model"] = model
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
from vertexai.language_models import TextGenerationModel
@ -663,7 +659,7 @@ def completion(
input=prompt, api_key=None, original_response=completion_response
)
## 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["model"] = model
response = model_response