diff --git a/litellm/__init__.py b/litellm/__init__.py index 3d62e17a3..00d7488c9 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -67,6 +67,7 @@ max_budget: float = 0.0 # set the max budget across all providers budget_duration: Optional[ str ] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). +_openai_finish_reasons = ["stop", "length", "function_call", "content_filter", "null"] _openai_completion_params = [ "functions", "function_call", diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py index d42bd003f..816ded394 100644 --- a/litellm/llms/vertex_ai.py +++ b/litellm/llms/vertex_ai.py @@ -4,7 +4,7 @@ from enum import Enum import requests import time from typing import Callable, Optional, Union -from litellm.utils import ModelResponse, Usage, CustomStreamWrapper +from litellm.utils import ModelResponse, Usage, CustomStreamWrapper, map_finish_reason import litellm, uuid import httpx @@ -575,9 +575,9 @@ def completion( model_response["model"] = model ## CALCULATING USAGE if model in litellm.vertex_language_models and response_obj is not None: - model_response["choices"][0].finish_reason = response_obj.candidates[ - 0 - ].finish_reason.name + model_response["choices"][0].finish_reason = map_finish_reason( + response_obj.candidates[0].finish_reason.name + ) usage = Usage( prompt_tokens=response_obj.usage_metadata.prompt_token_count, completion_tokens=response_obj.usage_metadata.candidates_token_count, @@ -771,9 +771,9 @@ async def async_completion( model_response["model"] = model ## CALCULATING USAGE if model in litellm.vertex_language_models and response_obj is not None: - model_response["choices"][0].finish_reason = response_obj.candidates[ - 0 - ].finish_reason.name + model_response["choices"][0].finish_reason = map_finish_reason( + response_obj.candidates[0].finish_reason.name + ) usage = Usage( prompt_tokens=response_obj.usage_metadata.prompt_token_count, completion_tokens=response_obj.usage_metadata.candidates_token_count, diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 0188d2358..5e48d7941 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -123,6 +123,10 @@ def test_vertex_ai(): print(response) assert type(response.choices[0].message.content) == str assert len(response.choices[0].message.content) > 1 + print( + f"response.choices[0].finish_reason: {response.choices[0].finish_reason}" + ) + assert response.choices[0].finish_reason in litellm._openai_finish_reasons except Exception as e: pytest.fail(f"Error occurred: {e}")