phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-21 03:59:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Minor unit test updates

Browse source
This commit is contained in:
Jash Gulabrai 2025-04-28 12:49:50 -04:00
parent e64961697a
commit c7ab6eeedb
2 changed files with 8 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

4
llama_stack/providers/remote/safety/nvidia/nvidia.py
View file

@ -126,10 +126,10 @@ class NeMoGuardrails:
Raises: Raises:
requests.HTTPError: If the POST request fails. 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 = { request_data = {
"model": self.model, "model": self.model,
"messages": messages, "messages": request_messages,
"temperature": self.temperature, "temperature": self.temperature,
"top_p": 1, "top_p": 1,
"frequency_penalty": 0, "frequency_penalty": 0,

13
tests/unit/providers/nvidia/test_safety.py
View file

@ -4,7 +4,6 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
import json
import os import os
import unittest import unittest
from typing import Any from typing import Any
@ -139,8 +138,8 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
data={ data={
"model": shield_id, "model": shield_id,
"messages": [ "messages": [
json.loads(messages[0].model_dump_json()), {"role": "user", "content": "Hello, how are you?"},
json.loads(messages[1].model_dump_json()), {"role": "assistant", "content": "I'm doing well, thank you for asking!", "tool_calls": None},
], ],
"temperature": 1.0, "temperature": 1.0,
"top_p": 1, "top_p": 1,
@ -193,8 +192,8 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
data={ data={
"model": shield_id, "model": shield_id,
"messages": [ "messages": [
json.loads(messages[0].model_dump_json()), {"role": "user", "content": "Hello, how are you?"},
json.loads(messages[1].model_dump_json()), {"role": "assistant", "content": "I'm doing well, thank you for asking!", "tool_calls": None},
], ],
"temperature": 1.0, "temperature": 1.0,
"top_p": 1, "top_p": 1,
@ -269,8 +268,8 @@ class TestNVIDIASafetyAdapter(unittest.TestCase):
data={ data={
"model": shield_id, "model": shield_id,
"messages": [ "messages": [
json.loads(messages[0].model_dump_json()), {"role": "user", "content": "Hello, how are you?"},
json.loads(messages[1].model_dump_json()), {"role": "assistant", "content": "I'm doing well, thank you for asking!", "tool_calls": None},
], ],
"temperature": 1.0, "temperature": 1.0,
"top_p": 1, "top_p": 1,