diff --git a/llama_stack/providers/remote/safety/nvidia/nvidia.py b/llama_stack/providers/remote/safety/nvidia/nvidia.py index 2a6bbeaff..13bc212a1 100644 --- a/llama_stack/providers/remote/safety/nvidia/nvidia.py +++ b/llama_stack/providers/remote/safety/nvidia/nvidia.py @@ -126,10 +126,10 @@ class NeMoGuardrails: Raises: requests.HTTPError: If the POST request fails. """ - messages = [await convert_message_to_openai_dict_new(message) for message in messages] + request_messages = [await convert_message_to_openai_dict_new(message) for message in messages] request_data = { "model": self.model, - "messages": 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..a237c56b0 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!", "tool_calls": None}, ], "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!", "tool_calls": None}, ], "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!", "tool_calls": None}, ], "temperature": 1.0, "top_p": 1,