ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add fake streaming for petals

Browse Source
This commit is contained in:
ishaan-jaff 2023-09-30 10:22:04 -07:00
parent 141c9c5bac
commit 8d1f5ba69d
2 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
litellm/main.py
View File

@ -1099,10 +1099,11 @@ def completion(
encoding=encoding,
logging_obj=logging
)
if inspect.isgenerator(model_response) or (stream == True):
# don't try to access stream object,
if stream==True: ## [BETA]
# Fake streaming for petals
resp_string = model_response["choices"][0]["message"]["content"]
response = CustomStreamWrapper(
model_response, model, custom_llm_provider="petals", logging_obj=logging
resp_string, model, custom_llm_provider="petals", logging_obj=logging
)
return response
response = model_response

12
litellm/utils.py
View File

@ -3158,6 +3158,18 @@ class CustomStreamWrapper:
completion_obj["content"] = new_chunk
self.completion_stream = self.completion_stream[chunk_size:]
time.sleep(0.05)
elif self.custom_llm_provider == "petals":
if len(self.completion_stream)==0:
if self.sent_last_chunk:
raise StopIteration
else:
model_response.choices[0].finish_reason = "stop"
self.sent_last_chunk = True
chunk_size = 30
new_chunk = self.completion_stream[:chunk_size]
completion_obj["content"] = new_chunk
self.completion_stream = self.completion_stream[chunk_size:]
time.sleep(0.05)
elif self.custom_llm_provider == "palm":
# fake streaming
if len(self.completion_stream)==0: