136 lines
4.1 KiB
Python
136 lines
4.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Unit test for prompt engine."""
|
|
import unittest
|
|
from typing import Any
|
|
|
|
from agentscope.models import read_model_configs, ModelResponse
|
|
from agentscope.models import load_model_by_config_name
|
|
from agentscope.models import OpenAIWrapperBase
|
|
from agentscope.prompt import PromptEngine
|
|
|
|
|
|
class PromptEngineTest(unittest.TestCase):
|
|
"""Unit test for prompt engine."""
|
|
|
|
def setUp(self) -> None:
|
|
"""Init for PromptEngineTest."""
|
|
self.name = "white"
|
|
self.sys_prompt = (
|
|
"You're a player in a chess game, and you are playing {name}."
|
|
)
|
|
self.dialog_history = [
|
|
{"name": "white player", "content": "Move to E4."},
|
|
{"name": "black player", "content": "Okay, I moved to F4."},
|
|
{"name": "white player", "content": "Move to F5."},
|
|
]
|
|
self.hint = "Now decide your next move."
|
|
self.prefix = "{name} player: "
|
|
|
|
read_model_configs(
|
|
[
|
|
{
|
|
"model_type": "post_api",
|
|
"config_name": "open-source",
|
|
"api_url": "http://xxx",
|
|
"headers": {"Autherization": "Bearer {API_TOKEN}"},
|
|
"parameters": {
|
|
"temperature": 0.5,
|
|
},
|
|
},
|
|
{
|
|
"model_type": "openai_chat",
|
|
"config_name": "gpt-4",
|
|
"model_name": "gpt-4",
|
|
"api_key": "xxx",
|
|
"organization": "xxx",
|
|
},
|
|
],
|
|
clear_existing=True,
|
|
)
|
|
|
|
def test_list_prompt(self) -> None:
|
|
"""Test for list prompt."""
|
|
|
|
class TestModelWrapperBase(OpenAIWrapperBase):
|
|
"""Test model wrapper."""
|
|
|
|
def __init__(self) -> None:
|
|
self.max_length = 1000
|
|
|
|
def __call__(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> ModelResponse:
|
|
return ModelResponse(text="")
|
|
|
|
def _register_default_metrics(self) -> None:
|
|
pass
|
|
|
|
model = TestModelWrapperBase()
|
|
engine = PromptEngine(model)
|
|
|
|
prompt = engine.join(
|
|
self.sys_prompt,
|
|
self.dialog_history,
|
|
self.hint,
|
|
format_map={"name": self.name},
|
|
)
|
|
|
|
self.assertEqual(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": "You're a player in a chess game, and you are "
|
|
"playing white.",
|
|
},
|
|
{
|
|
"name": "white player",
|
|
"role": "assistant",
|
|
"content": "Move to E4.",
|
|
},
|
|
{
|
|
"name": "black player",
|
|
"role": "assistant",
|
|
"content": "Okay, I moved to F4.",
|
|
},
|
|
{
|
|
"name": "white player",
|
|
"role": "assistant",
|
|
"content": "Move to F5.",
|
|
},
|
|
{
|
|
"role": "assistant",
|
|
"content": "Now decide your next move.",
|
|
},
|
|
],
|
|
prompt,
|
|
)
|
|
|
|
def test_str_prompt(self) -> None:
|
|
"""Test for string prompt."""
|
|
model = load_model_by_config_name("open-source")
|
|
engine = PromptEngine(model)
|
|
|
|
prompt = engine.join(
|
|
self.sys_prompt,
|
|
self.dialog_history,
|
|
self.hint,
|
|
self.prefix,
|
|
format_map={"name": self.name},
|
|
)
|
|
|
|
self.assertEqual(
|
|
"""You're a player in a chess game, and you are playing white.
|
|
white player: Move to E4.
|
|
black player: Okay, I moved to F4.
|
|
white player: Move to F5.
|
|
Now decide your next move.
|
|
white player: """,
|
|
prompt,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|