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: 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 ## RESPONSE OBJECT
completion_response = json.loads(response) completion_response = json.loads(response)
try: 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 = "" completion_output = ""
if "generation" in completion_response_choices: if "generation" in completion_response_choices:
completion_output += completion_response_choices["generation"] completion_output += completion_response_choices["generation"]
@ -580,7 +583,10 @@ async def async_completion(
## RESPONSE OBJECT ## RESPONSE OBJECT
completion_response = json.loads(response) completion_response = json.loads(response)
try: 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 = "" completion_output = ""
if "generation" in completion_response_choices: if "generation" in completion_response_choices:
completion_output += completion_response_choices["generation"] completion_output += completion_response_choices["generation"]