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

fix(vertex_ai.py): map finish reason

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-14 11:42:03 -08:00
parent e7666d3908
commit fe1fe70c64
3 changed files with 12 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/__init__.py
View file

@ -67,6 +67,7 @@ max_budget: float = 0.0 # set the max budget across all providers
budget_duration: Optional[ budget_duration: Optional[
str str
] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). ] = 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 = [ _openai_completion_params = [
"functions", "functions",
"function_call", "function_call",

14
litellm/llms/vertex_ai.py
View file

@ -4,7 +4,7 @@ from enum import Enum
import requests import requests
import time import time
from typing import Callable, Optional, Union 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 litellm, uuid
import httpx import httpx
@ -575,9 +575,9 @@ def completion(
model_response["model"] = model model_response["model"] = model
## CALCULATING USAGE ## CALCULATING USAGE
if model in litellm.vertex_language_models and response_obj is not None: if model in litellm.vertex_language_models and response_obj is not None:
model_response["choices"][0].finish_reason = response_obj.candidates[ model_response["choices"][0].finish_reason = map_finish_reason(
0 response_obj.candidates[0].finish_reason.name
].finish_reason.name )
usage = Usage( usage = Usage(
prompt_tokens=response_obj.usage_metadata.prompt_token_count, prompt_tokens=response_obj.usage_metadata.prompt_token_count,
completion_tokens=response_obj.usage_metadata.candidates_token_count, completion_tokens=response_obj.usage_metadata.candidates_token_count,
@ -771,9 +771,9 @@ async def async_completion(
model_response["model"] = model model_response["model"] = model
## CALCULATING USAGE ## CALCULATING USAGE
if model in litellm.vertex_language_models and response_obj is not None: if model in litellm.vertex_language_models and response_obj is not None:
model_response["choices"][0].finish_reason = response_obj.candidates[ model_response["choices"][0].finish_reason = map_finish_reason(
0 response_obj.candidates[0].finish_reason.name
].finish_reason.name )
usage = Usage( usage = Usage(
prompt_tokens=response_obj.usage_metadata.prompt_token_count, prompt_tokens=response_obj.usage_metadata.prompt_token_count,
completion_tokens=response_obj.usage_metadata.candidates_token_count, completion_tokens=response_obj.usage_metadata.candidates_token_count,

4
litellm/tests/test_amazing_vertex_completion.py
View file

@ -123,6 +123,10 @@ def test_vertex_ai():
print(response) print(response)
assert type(response.choices[0].message.content) == str assert type(response.choices[0].message.content) == str
assert len(response.choices[0].message.content) > 1 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: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")