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

fix(utils.py): fix sync streaming

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-09 18:47:11 -08:00
parent 76d8ea674e
commit 3d4c5e10a7
2 changed files with 162 additions and 158 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/tests/test_completion.py
View file

@ -733,7 +733,7 @@ def test_completion_azure_deployment_id():
print(response) print(response)
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_completion_azure_deployment_id() # test_completion_azure_deployment_id()
# Only works for local endpoint # Only works for local endpoint
# def test_completion_anthropic_openai_proxy(): # def test_completion_anthropic_openai_proxy():
@ -1169,7 +1169,7 @@ def test_mistral_anyscale_stream():
for chunk in response: for chunk in response:
# print(chunk) # print(chunk)
print(chunk["choices"][0]["delta"].get("content", ""), end="") print(chunk["choices"][0]["delta"].get("content", ""), end="")
# test_mistral_anyscale_stream() test_mistral_anyscale_stream()
# test_completion_anyscale_2() # test_completion_anyscale_2()
# def test_completion_with_fallbacks_multiple_keys(): # def test_completion_with_fallbacks_multiple_keys():
# print(f"backup key 1: {os.getenv('BACKUP_OPENAI_API_KEY_1')}") # print(f"backup key 1: {os.getenv('BACKUP_OPENAI_API_KEY_1')}")

24
litellm/utils.py
View file

@ -4113,7 +4113,7 @@ class CustomStreamWrapper:
def chunk_creator(self, chunk): def chunk_creator(self, chunk):
model_response = ModelResponse(stream=True, model=self.model) model_response = ModelResponse(stream=True, model=self.model)
try: try:
while True: # loop until a non-empty string is found
# return this for all models # return this for all models
completion_obj = {"content": ""} completion_obj = {"content": ""}
if self.custom_llm_provider and self.custom_llm_provider == "anthropic": if self.custom_llm_provider and self.custom_llm_provider == "anthropic":
@ -4237,20 +4237,18 @@ class CustomStreamWrapper:
if "error" in chunk: if "error" in chunk:
exception_type(model=self.model, custom_llm_provider=self.custom_llm_provider, original_exception=chunk["error"]) exception_type(model=self.model, custom_llm_provider=self.custom_llm_provider, original_exception=chunk["error"])
completion_obj = chunk completion_obj = chunk
elif self.custom_llm_provider == "openai":
response_obj = self.handle_openai_chat_completion_chunk(chunk)
completion_obj["content"] = response_obj["text"]
print_verbose(f"completion obj content: {completion_obj['content']}")
if response_obj["is_finished"]:
model_response.choices[0].finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider == "text-completion-openai": elif self.custom_llm_provider == "text-completion-openai":
response_obj = self.handle_openai_text_completion_chunk(chunk) response_obj = self.handle_openai_text_completion_chunk(chunk)
completion_obj["content"] = response_obj["text"] completion_obj["content"] = response_obj["text"]
print_verbose(f"completion obj content: {completion_obj['content']}") print_verbose(f"completion obj content: {completion_obj['content']}")
if response_obj["is_finished"]: if response_obj["is_finished"]:
model_response.choices[0].finish_reason = response_obj["finish_reason"] model_response.choices[0].finish_reason = response_obj["finish_reason"]
else: # openai chat/azure models else: # openai chat model
raise Exception("Unmapped Model Error") response_obj = self.handle_openai_chat_completion_chunk(chunk)
completion_obj["content"] = response_obj["text"]
print_verbose(f"completion obj content: {completion_obj['content']}")
if response_obj["is_finished"]:
model_response.choices[0].finish_reason = response_obj["finish_reason"]
model_response.model = self.model model_response.model = self.model
if len(completion_obj["content"]) > 0: # cannot set content of an OpenAI Object to be an empty string if len(completion_obj["content"]) > 0: # cannot set content of an OpenAI Object to be an empty string
@ -4284,8 +4282,14 @@ class CustomStreamWrapper:
## needs to handle the empty string case (even starting chunk can be an empty string) ## needs to handle the empty string case (even starting chunk can be an empty string)
def __next__(self): def __next__(self):
while True: # loop until a non-empty string is found
if isinstance(self.completion_stream, str):
chunk = self.completion_stream
else:
chunk = next(self.completion_stream) chunk = next(self.completion_stream)
return self.chunk_creator(chunk=chunk) response = self.chunk_creator(chunk=chunk)
if response is not None:
return response
async def __anext__(self): async def __anext__(self):
try: try: