test(test_amazing_vertex_completion.py): fix vertex anthropic streaming text

This commit is contained in:
Krrish Dholakia 2024-04-15 14:37:22 -07:00
parent 1ec7118e1f
commit 40d9167cae

View file

@ -145,7 +145,7 @@ def test_vertex_ai_anthropic():
# reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd."
# )
def test_vertex_ai_anthropic_streaming():
load_vertex_ai_credentials()
# load_vertex_ai_credentials()
# litellm.set_verbose = True
@ -153,6 +153,8 @@ def test_vertex_ai_anthropic_streaming():
vertex_ai_project = "adroit-crow-413218"
vertex_ai_location = "asia-southeast1"
json_obj = get_vertex_ai_creds_json()
vertex_credentials = json.dumps(json_obj)
response = completion(
model="vertex_ai/" + model,
@ -205,12 +207,14 @@ async def test_vertex_ai_anthropic_async():
# )
@pytest.mark.asyncio
async def test_vertex_ai_anthropic_async_streaming():
load_vertex_ai_credentials()
# load_vertex_ai_credentials()
litellm.set_verbose = True
model = "claude-3-sonnet@20240229"
vertex_ai_project = "adroit-crow-413218"
vertex_ai_location = "asia-southeast1"
json_obj = get_vertex_ai_creds_json()
vertex_credentials = json.dumps(json_obj)
response = await acompletion(
model="vertex_ai/" + model,
@ -218,6 +222,8 @@ async def test_vertex_ai_anthropic_async_streaming():
temperature=0.7,
vertex_ai_project=vertex_ai_project,
vertex_ai_location=vertex_ai_location,
vertex_credentials=vertex_credentials,
stream=True,
)
async for chunk in response: