phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

🐞 fix: djl vllm support

Browse source 
support vllm response format on sagemaker, which only return one choice.
This commit is contained in:
yishiyiyuan 2024-03-27 15:13:21 +08:00
parent d7601a4844
commit 5faa493d35
1 changed files with 8 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

10
litellm/llms/sagemaker.py
View file

@ -365,7 +365,10 @@ def completion(
## RESPONSE OBJECT
completion_response = json.loads(response)
try:
completion_response_choices = completion_response[0]
if isinstance(completion_response, list):
completion_response_choices = completion_response[0]
else:
completion_response_choices = completion_response
completion_output = ""
if "generation" in completion_response_choices:
completion_output += completion_response_choices["generation"]
@ -580,7 +583,10 @@ async def async_completion(
## RESPONSE OBJECT
completion_response = json.loads(response)
try:
completion_response_choices = completion_response[0]
if isinstance(completion_response, list):
completion_response_choices = completion_response[0]
else:
completion_response_choices = completion_response
completion_output = ""
if "generation" in completion_response_choices:
completion_output += completion_response_choices["generation"]