From 7fdc9e04ac13b478fdea0177d97622723a8f1aaa Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Wed, 15 Jan 2025 15:57:53 -0800 Subject: [PATCH] fix safety --- tests/client-sdk/safety/test_safety.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/client-sdk/safety/test_safety.py b/tests/client-sdk/safety/test_safety.py index 8eadffcfc..2d79bda5e 100644 --- a/tests/client-sdk/safety/test_safety.py +++ b/tests/client-sdk/safety/test_safety.py @@ -30,7 +30,7 @@ def data_url_from_image(file_path): @pytest.fixture(scope="session") def available_shields(llama_stack_client): - return [shield.identifier for shield in llama_stack_client.shields.list()] + return [shield.identifier for shield in llama_stack_client.shields.list().data] @pytest.fixture(scope="session") @@ -54,7 +54,7 @@ def code_scanner_shield_id(available_shields): @pytest.fixture(scope="session") def model_providers(llama_stack_client): return set( - [x.provider_id for x in llama_stack_client.providers.list()["inference"]] + [x["provider_id"] for x in llama_stack_client.providers.list().inference] )