phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

bump version

Browse source
This commit is contained in:
ishaan-jaff 2023-09-15 14:14:47 -07:00
parent 9e808c0c1c
commit 5ec1fc5048
3 changed files with 27 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

15
litellm/main.py
View file

@ -659,6 +659,14 @@ def completion(
chat_model = CodeChatModel.from_pretrained(model)
chat = chat_model.start_chat()
if stream:
model_response = chat.send_message_streaming(prompt, **optional_params)
response = CustomStreamWrapper(
model_response, model, custom_llm_provider="vertexai", logging_obj=logging
)
return response
completion_response = chat.send_message(prompt, **optional_params)
## LOGGING
@ -692,6 +700,13 @@ def completion(
else:
vertex_model = CodeGenerationModel.from_pretrained(model)
if stream:
model_response = vertex_model.predict_streaming(prompt, **optional_params)
response = CustomStreamWrapper(
model_response, model, custom_llm_provider="vertexai", logging_obj=logging
)
return response
completion_response = vertex_model.predict(prompt, **optional_params)
## LOGGING