mirror of
https://github.com/NVIDIA-NeMo/DataDesigner
synced 2026-05-24 09:48:29 +00:00
f8c201e085
* update script * update headers * refactor a bit and add test script * update headers * update for edge case * update headers * add step to get file creation date * use git history to get copyright year * generation type is printed with inference parameters * fix unit test
36 lines
1.7 KiB
Python
36 lines
1.7 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from data_designer.engine.models.utils import prompt_to_messages, str_to_message
|
|
|
|
|
|
def test_str_to_message():
|
|
assert str_to_message("hello") == {"content": "hello", "role": "user"}
|
|
assert str_to_message("hello", role="system") == {"content": "hello", "role": "system"}
|
|
assert str_to_message([{"type": "text", "text": "hello"}]) == {
|
|
"content": [{"type": "text", "text": "hello"}],
|
|
"role": "user",
|
|
}
|
|
assert str_to_message([{"type": "text", "text": "hello"}], role="system") == {
|
|
"content": [{"type": "text", "text": "hello"}],
|
|
"role": "system",
|
|
}
|
|
|
|
|
|
def test_prompt_to_messages():
|
|
stub_system_prompt = "some system prompt"
|
|
mult_modal_context = {"type": "image_url", "image_url": {"url": "http://example.com/image.png"}}
|
|
assert prompt_to_messages(user_prompt="hello") == [{"content": "hello", "role": "user"}]
|
|
assert prompt_to_messages(user_prompt="hello", system_prompt=stub_system_prompt) == [
|
|
{"content": stub_system_prompt, "role": "system"},
|
|
{"content": "hello", "role": "user"},
|
|
]
|
|
assert prompt_to_messages(user_prompt="hello", multi_modal_context=[mult_modal_context]) == [
|
|
{"content": [{"type": "text", "text": "hello"}, mult_modal_context], "role": "user"}
|
|
]
|
|
assert prompt_to_messages(
|
|
user_prompt="hello", system_prompt=stub_system_prompt, multi_modal_context=[mult_modal_context]
|
|
) == [
|
|
{"content": stub_system_prompt, "role": "system"},
|
|
{"content": [{"type": "text", "text": "hello"}, mult_modal_context], "role": "user"},
|
|
]
|