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

test: add vertex claude to streaming valid json str test

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-07 10:18:17 -07:00
parent 89273722ba
commit 75b2fd2e7f
1 changed files with 15 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

16
litellm/tests/test_streaming.py
View file

@ -3697,9 +3697,20 @@ def test_unit_test_custom_stream_wrapper_function_call():
"gpt-3.5-turbo", "gpt-3.5-turbo",
"claude-3-5-sonnet-20240620", "claude-3-5-sonnet-20240620",
"anthropic.claude-3-sonnet-20240229-v1:0", "anthropic.claude-3-sonnet-20240229-v1:0",
"vertex_ai/claude-3-5-sonnet@20240620",
], ],
) )
def test_streaming_tool_calls_valid_json_str(model): def test_streaming_tool_calls_valid_json_str(model):
if "vertex_ai" in model:
from litellm.tests.test_amazing_vertex_completion import (
load_vertex_ai_credentials,
)
load_vertex_ai_credentials()
vertex_location = "us-east5"
else:
vertex_location = None
litellm.set_verbose = False
messages = [ messages = [
{"role": "user", "content": "Hit the snooze button."}, {"role": "user", "content": "Hit the snooze button."},
] ]
@ -3718,8 +3729,11 @@ def test_streaming_tool_calls_valid_json_str(model):
} }
] ]
stream = litellm.completion(model, messages, tools=tools, stream=True) stream = litellm.completion(
model, messages, tools=tools, stream=True, vertex_location=vertex_location
)
chunks = [*stream] chunks = [*stream]
print(f"chunks: {chunks}")
tool_call_id_arg_map = {} tool_call_id_arg_map = {}
curr_tool_call_id = None curr_tool_call_id = None
curr_tool_call_str = "" curr_tool_call_str = ""