diff --git a/llama_stack/providers/remote/safety/nvidia/nvidia.py b/llama_stack/providers/remote/safety/nvidia/nvidia.py index 1ff4a6ad9..13bc212a1 100644 --- a/llama_stack/providers/remote/safety/nvidia/nvidia.py +++ b/llama_stack/providers/remote/safety/nvidia/nvidia.py @@ -12,8 +12,8 @@ import requests from llama_stack.apis.inference import Message from llama_stack.apis.safety import RunShieldResponse, Safety, SafetyViolation, ViolationLevel from llama_stack.apis.shields import Shield -from llama_stack.distribution.library_client import convert_pydantic_to_json_value from llama_stack.providers.datatypes import ShieldsProtocolPrivate +from llama_stack.providers.utils.inference.openai_compat import convert_message_to_openai_dict_new from .config import NVIDIASafetyConfig @@ -28,7 +28,6 @@ class NVIDIASafetyAdapter(Safety, ShieldsProtocolPrivate): Args: config (NVIDIASafetyConfig): The configuration containing the guardrails service URL and config ID. """ - print(f"Initializing NVIDIASafetyAdapter({config.guardrails_service_url})...") self.config = config async def initialize(self) -> None: @@ -127,9 +126,10 @@ class NeMoGuardrails: Raises: requests.HTTPError: If the POST request fails. """ + request_messages = [await convert_message_to_openai_dict_new(message) for message in messages] request_data = { "model": self.model, - "messages": convert_pydantic_to_json_value(messages), + "messages": request_messages, "temperature": self.temperature, "top_p": 1, "frequency_penalty": 0, diff --git a/tests/unit/providers/nvidia/test_safety.py b/tests/unit/providers/nvidia/test_safety.py index e7e1cb3dc..8c74f178b 100644 --- a/tests/unit/providers/nvidia/test_safety.py +++ b/tests/unit/providers/nvidia/test_safety.py @@ -4,7 +4,6 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import json import os import unittest from typing import Any @@ -139,8 +138,8 @@ class TestNVIDIASafetyAdapter(unittest.TestCase): data={ "model": shield_id, "messages": [ - json.loads(messages[0].model_dump_json()), - json.loads(messages[1].model_dump_json()), + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, ], "temperature": 1.0, "top_p": 1, @@ -193,8 +192,8 @@ class TestNVIDIASafetyAdapter(unittest.TestCase): data={ "model": shield_id, "messages": [ - json.loads(messages[0].model_dump_json()), - json.loads(messages[1].model_dump_json()), + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, ], "temperature": 1.0, "top_p": 1, @@ -269,8 +268,8 @@ class TestNVIDIASafetyAdapter(unittest.TestCase): data={ "model": shield_id, "messages": [ - json.loads(messages[0].model_dump_json()), - json.loads(messages[1].model_dump_json()), + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, ], "temperature": 1.0, "top_p": 1,