phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-28 04:04:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(vertex_ai.py): support safety settings for async streaming calls

Browse source
This commit is contained in:
Krrish Dholakia 2024-04-27 11:43:00 -07:00
parent 48d590c428
commit 33b5829e31
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/vertex_ai.py
View file

@ -529,7 +529,7 @@ def completion(
"instances": instances, "instances": instances,
"vertex_location": vertex_location, "vertex_location": vertex_location,
"vertex_project": vertex_project, "vertex_project": vertex_project,
"safety_settings":safety_settings, "safety_settings": safety_settings,
**optional_params, **optional_params,
} }
if optional_params.get("stream", False) is True: if optional_params.get("stream", False) is True:
@ -1025,6 +1025,7 @@ async def async_streaming(
instances=None, instances=None,
vertex_project=None, vertex_project=None,
vertex_location=None, vertex_location=None,
safety_settings=None,
**optional_params, **optional_params,
): ):
""" """
@ -1051,6 +1052,7 @@ async def async_streaming(
response = await llm_model._generate_content_streaming_async( response = await llm_model._generate_content_streaming_async(
contents=content, contents=content,
generation_config=optional_params, generation_config=optional_params,
safety_settings=safety_settings,
tools=tools, tools=tools,
) )