tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-16 18:16:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

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
litellm/main.py
+2 -6
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