Minor unit test updates

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

View file

@ -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,