fix(main.py): return n>1 response for openai text completion

This commit is contained in:
Krrish Dholakia 2023-10-09 20:44:07 -07:00
parent 41bebaa1e3
commit db20cb84d4
4 changed files with 14 additions and 3 deletions

View file

@ -60,7 +60,9 @@ from litellm.utils import (
ModelResponse,
EmbeddingResponse,
read_config_args,
RateLimitManager
RateLimitManager,
Choices,
Message
)
####### ENVIRONMENT VARIABLES ###################
@ -509,8 +511,12 @@ def completion(
},
)
## RESPONSE OBJECT
completion_response = response["choices"][0]["text"]
model_response["choices"][0]["message"]["content"] = completion_response
choices_list = []
for idx, item in enumerate(response["choices"]):
message_obj = Message(content=item["text"])
choice_obj = Choices(finish_reason=item["finish_reason"], index=idx+1, message=message_obj)
choices_list.append(choice_obj)
model_response["choices"] = choices_list
model_response["created"] = response.get("created", time.time())
model_response["model"] = model
model_response["usage"] = response.get("usage", 0)

View file

@ -498,6 +498,11 @@ def openai_text_completion_test():
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
response_3 = litellm.completion(model="text-davinci-003",
messages=[{ "content": "Hello, how are you?","role": "user"}],
n=2)
assert len(response_3.choices) > 1
except Exception as e:
pytest.fail(f"Error occurred: {e}")