From a2ef089667d4033b2a6b98f793ec7810fec2e243 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 16 May 2024 13:52:46 -0700 Subject: [PATCH] fix - choices index for gemini/ provider --- litellm/llms/gemini.py | 4 ++-- litellm/tests/test_completion.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/llms/gemini.py b/litellm/llms/gemini.py index 60220fd290..a55b39aef9 100644 --- a/litellm/llms/gemini.py +++ b/litellm/llms/gemini.py @@ -260,7 +260,7 @@ def completion( message_obj = Message(content=item.content.parts[0].text) else: message_obj = Message(content=None) - choice_obj = Choices(index=idx + 1, message=message_obj) + choice_obj = Choices(index=idx, message=message_obj) choices_list.append(choice_obj) model_response["choices"] = choices_list except Exception as e: @@ -352,7 +352,7 @@ async def async_completion( message_obj = Message(content=item.content.parts[0].text) else: message_obj = Message(content=None) - choice_obj = Choices(index=idx + 1, message=message_obj) + choice_obj = Choices(index=idx, message=message_obj) choices_list.append(choice_obj) model_response["choices"] = choices_list except Exception as e: diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 120a65acbf..4e5d62b5f8 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -3246,6 +3246,7 @@ def test_completion_gemini(): response = completion(model=model_name, messages=messages) # Add any assertions,here to check the response print(response) + assert response.choices[0]["index"] == 0 except litellm.APIError as e: pass except Exception as e: