From 5faa493d357aeca6ce12e01285e89cf4c43bb8f8 Mon Sep 17 00:00:00 2001 From: yishiyiyuan Date: Wed, 27 Mar 2024 15:13:21 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9E=20fix:=20djl=20vllm=20support=20su?= =?UTF-8?q?pport=20vllm=20response=20format=20on=20sagemaker,=20which=20on?= =?UTF-8?q?ly=20return=20one=20choice.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- litellm/llms/sagemaker.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 2604875b47..94a940ddc2 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -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"]