From 96240b7a491e680b3e0b929f7d7a202aa63dcb36 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Wed, 19 Feb 2025 16:31:30 -0800 Subject: [PATCH] fixing import path --- llama_stack/providers/remote/safety/nvidia/config.py | 2 +- llama_stack/providers/tests/safety/fixtures.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/llama_stack/providers/remote/safety/nvidia/config.py b/llama_stack/providers/remote/safety/nvidia/config.py index ef93a8d41..44e1a986f 100644 --- a/llama_stack/providers/remote/safety/nvidia/config.py +++ b/llama_stack/providers/remote/safety/nvidia/config.py @@ -9,7 +9,7 @@ from typing import Any, Dict, Optional from pydantic import BaseModel, Field, SecretStr, field_validator -from llama_models.schema_utils import json_schema_type +from llama_stack.schema_utils import json_schema_type class ShieldType(Enum): diff --git a/llama_stack/providers/tests/safety/fixtures.py b/llama_stack/providers/tests/safety/fixtures.py index 67d1ecac4..3f85473c5 100644 --- a/llama_stack/providers/tests/safety/fixtures.py +++ b/llama_stack/providers/tests/safety/fixtures.py @@ -108,7 +108,7 @@ def safety_nvidia() -> ProviderFixture: ) -SAFETY_FIXTURES = ["llama_guard", "bedrock", "remote"] +SAFETY_FIXTURES = ["llama_guard", "bedrock", "remote", "nvidia"] @pytest_asyncio.fixture(scope="session")