504 lines
16 KiB
Python
504 lines
16 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Unit test for prompt engineering strategies in format function."""
|
|
import unittest
|
|
from unittest import mock
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from agentscope.message import Msg
|
|
from agentscope.models import (
|
|
OpenAIChatWrapper,
|
|
OllamaChatWrapper,
|
|
OllamaGenerationWrapper,
|
|
GeminiChatWrapper,
|
|
ZhipuAIChatWrapper,
|
|
DashScopeChatWrapper,
|
|
DashScopeMultiModalWrapper,
|
|
LiteLLMChatWrapper,
|
|
)
|
|
|
|
|
|
class ExampleTest(unittest.TestCase):
|
|
"""
|
|
ExampleTest for a unit test.
|
|
"""
|
|
|
|
def setUp(self) -> None:
|
|
"""Init for ExampleTest."""
|
|
self.inputs = [
|
|
Msg("system", "You are a helpful assistant", role="system"),
|
|
[
|
|
Msg("user", "What is the weather today?", role="user"),
|
|
Msg("assistant", "It is sunny today", role="assistant"),
|
|
],
|
|
]
|
|
|
|
self.inputs_vision = [
|
|
Msg("system", "You are a helpful assistant", role="system"),
|
|
[
|
|
Msg(
|
|
"user",
|
|
"Describe the images",
|
|
role="user",
|
|
url="https://fakeweb/test.jpg",
|
|
),
|
|
Msg(
|
|
"user",
|
|
"And this images",
|
|
"user",
|
|
url=[
|
|
"/Users/xxx/abc.png",
|
|
"/Users/xxx/def.mp3",
|
|
],
|
|
),
|
|
],
|
|
]
|
|
|
|
self.wrong_inputs = [
|
|
Msg("system", "You are a helpful assistant", role="system"),
|
|
[
|
|
"What is the weather today?",
|
|
Msg("assistant", "It is sunny today", role="assistant"),
|
|
],
|
|
]
|
|
|
|
@patch("builtins.open", mock.mock_open(read_data=b"abcdef"))
|
|
@patch("os.path.isfile")
|
|
@patch("os.path.exists")
|
|
@patch("openai.OpenAI")
|
|
def test_openai_chat_vision_with_wrong_model(
|
|
self,
|
|
mock_client: MagicMock,
|
|
mock_exists: MagicMock,
|
|
mock_isfile: MagicMock,
|
|
) -> None:
|
|
"""Unit test for the format function in openai chat api wrapper with
|
|
vision models"""
|
|
mock_exists.side_effect = lambda url: url == "/Users/xxx/abc.png"
|
|
mock_isfile.side_effect = lambda url: url == "/Users/xxx/abc.png"
|
|
|
|
# Prepare the mock client
|
|
mock_client.return_value = "client_dummy"
|
|
|
|
model = OpenAIChatWrapper(
|
|
config_name="",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
# correct format
|
|
ground_truth = [
|
|
{
|
|
"role": "system",
|
|
"content": "You are a helpful assistant",
|
|
"name": "system",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"name": "user",
|
|
"content": "Describe the images",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"name": "user",
|
|
"content": "And this images",
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*self.inputs_vision)
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
@patch("builtins.open", mock.mock_open(read_data=b"abcdef"))
|
|
@patch("os.path.isfile")
|
|
@patch("os.path.exists")
|
|
@patch("openai.OpenAI")
|
|
def test_openai_chat_vision(
|
|
self,
|
|
mock_client: MagicMock,
|
|
mock_exists: MagicMock,
|
|
mock_isfile: MagicMock,
|
|
) -> None:
|
|
"""Unit test for the format function in openai chat api wrapper with
|
|
vision models"""
|
|
mock_exists.side_effect = lambda url: url == "/Users/xxx/abc.png"
|
|
mock_isfile.side_effect = lambda url: url == "/Users/xxx/abc.png"
|
|
|
|
# Prepare the mock client
|
|
mock_client.return_value = "client_dummy"
|
|
|
|
model = OpenAIChatWrapper(
|
|
config_name="",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
# correct format
|
|
ground_truth = [
|
|
{
|
|
"role": "system",
|
|
"content": "You are a helpful assistant",
|
|
"name": "system",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"name": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": "Describe the images",
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": "https://fakeweb/test.jpg",
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"name": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": "And this images",
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": "data:image/png;base64,YWJjZGVm",
|
|
},
|
|
},
|
|
],
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*self.inputs_vision)
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
@patch("openai.OpenAI")
|
|
def test_openai_chat(self, mock_client: MagicMock) -> None:
|
|
"""Unit test for the format function in openai chat api wrapper."""
|
|
# Prepare the mock client
|
|
mock_client.return_value = "client_dummy"
|
|
|
|
model = OpenAIChatWrapper(
|
|
config_name="",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
# correct format
|
|
ground_truth = [
|
|
{
|
|
"role": "system",
|
|
"content": "You are a helpful assistant",
|
|
"name": "system",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": "What is the weather today?",
|
|
"name": "user",
|
|
},
|
|
{
|
|
"role": "assistant",
|
|
"content": "It is sunny today",
|
|
"name": "assistant",
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*self.inputs) # type: ignore[arg-type]
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs) # type: ignore[arg-type]
|
|
|
|
def test_ollama_chat(self) -> None:
|
|
"""Unit test for the format function in ollama chat api wrapper."""
|
|
model = OllamaChatWrapper(
|
|
config_name="",
|
|
model_name="llama2",
|
|
)
|
|
|
|
# correct format
|
|
ground_truth = [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"You are a helpful assistant\n"
|
|
"\n"
|
|
"## Dialogue History\n"
|
|
"user: What is the weather today?\n"
|
|
"assistant: It is sunny today"
|
|
),
|
|
},
|
|
]
|
|
prompt = model.format(*self.inputs) # type: ignore[arg-type]
|
|
self.assertEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs) # type: ignore[arg-type]
|
|
|
|
def test_ollama_generation(self) -> None:
|
|
"""Unit test for the generation function in ollama chat api wrapper."""
|
|
model = OllamaGenerationWrapper(
|
|
config_name="",
|
|
model_name="llama2",
|
|
)
|
|
|
|
# correct format
|
|
ground_truth = (
|
|
"You are a helpful assistant\n\n## Dialogue History\nuser: "
|
|
"What is the weather today?\nassistant: It is sunny today"
|
|
)
|
|
prompt = model.format(*self.inputs) # type: ignore[arg-type]
|
|
self.assertEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs) # type: ignore[arg-type]
|
|
|
|
@patch("google.generativeai.configure")
|
|
def test_gemini_chat(self, mock_configure: MagicMock) -> None:
|
|
"""Unit test for the format function in gemini chat api wrapper."""
|
|
mock_configure.return_value = "client_dummy"
|
|
|
|
model = GeminiChatWrapper(
|
|
config_name="",
|
|
model_name="gemini-pro",
|
|
api_key="xxx",
|
|
)
|
|
|
|
# correct format
|
|
ground_truth = [
|
|
{
|
|
"role": "user",
|
|
"parts": [
|
|
"You are a helpful assistant\n\n## Dialogue History\n"
|
|
"user: What is the weather today?\nassistant: It is "
|
|
"sunny today",
|
|
],
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*self.inputs) # type: ignore[arg-type]
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs) # type: ignore[arg-type]
|
|
|
|
def test_dashscope_chat(self) -> None:
|
|
"""Unit test for the format function in dashscope chat api wrapper."""
|
|
model = DashScopeChatWrapper(
|
|
config_name="",
|
|
model_name="qwen-max",
|
|
api_key="xxx",
|
|
)
|
|
|
|
ground_truth = [
|
|
{
|
|
"content": "You are a helpful assistant",
|
|
"role": "system",
|
|
},
|
|
{
|
|
"content": (
|
|
"## Dialogue History\n"
|
|
"user: What is the weather today?\n"
|
|
"assistant: It is sunny today"
|
|
),
|
|
"role": "user",
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*self.inputs)
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs) # type: ignore[arg-type]
|
|
|
|
def test_zhipuai_chat(self) -> None:
|
|
"""Unit test for the format function in zhipu chat api wrapper."""
|
|
model = ZhipuAIChatWrapper(
|
|
config_name="",
|
|
model_name="glm-4",
|
|
api_key="xxx",
|
|
)
|
|
|
|
ground_truth = [
|
|
{
|
|
"content": "You are a helpful assistant",
|
|
"role": "system",
|
|
},
|
|
{
|
|
"content": (
|
|
"## Dialogue History\n"
|
|
"user: What is the weather today?\n"
|
|
"assistant: It is sunny today"
|
|
),
|
|
"role": "user",
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*self.inputs)
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs) # type: ignore[arg-type]
|
|
|
|
def test_litellm_chat(self) -> None:
|
|
"""Unit test for the format function in litellm chat api wrapper."""
|
|
model = LiteLLMChatWrapper(
|
|
config_name="",
|
|
model_name="gpt-3.5-turbo",
|
|
api_key="xxx",
|
|
)
|
|
|
|
ground_truth = [
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
"You are a helpful assistant\n\n"
|
|
"## Dialogue History\nuser: What is the weather today?\n"
|
|
"assistant: It is sunny today"
|
|
),
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*self.inputs)
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs) # type: ignore[arg-type]
|
|
|
|
def test_dashscope_multimodal_image(self) -> None:
|
|
"""Unit test for the format function in dashscope multimodal
|
|
conversation api wrapper for image."""
|
|
model = DashScopeMultiModalWrapper(
|
|
config_name="",
|
|
model_name="qwen-vl-plus",
|
|
api_key="xxx",
|
|
)
|
|
|
|
multimodal_input = [
|
|
Msg(
|
|
"system",
|
|
"You are a helpful assistant",
|
|
role="system",
|
|
url="url1.png",
|
|
),
|
|
[
|
|
Msg(
|
|
"user",
|
|
"What is the weather today?",
|
|
role="user",
|
|
url="url2.png",
|
|
),
|
|
Msg(
|
|
"assistant",
|
|
"It is sunny today",
|
|
role="assistant",
|
|
url="url3.png",
|
|
),
|
|
],
|
|
]
|
|
|
|
ground_truth = [
|
|
{
|
|
"role": "system",
|
|
"content": [
|
|
{"image": "url1.png"},
|
|
{"text": "You are a helpful assistant"},
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"image": "url2.png"},
|
|
{"image": "url3.png"},
|
|
{
|
|
"text": (
|
|
"## Dialogue History\n"
|
|
"user: What is the weather today?\n"
|
|
"assistant: It is sunny today"
|
|
),
|
|
},
|
|
],
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*multimodal_input)
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs)
|
|
|
|
def test_dashscope_multimodal_audio(self) -> None:
|
|
"""Unit test for the format function in dashscope multimodal
|
|
conversation api wrapper for audio."""
|
|
model = DashScopeMultiModalWrapper(
|
|
config_name="",
|
|
model_name="qwen-audio-turbo",
|
|
api_key="xxx",
|
|
)
|
|
|
|
multimodal_input = [
|
|
Msg(
|
|
"system",
|
|
"You are a helpful assistant",
|
|
role="system",
|
|
url="url1.mp3",
|
|
),
|
|
[
|
|
Msg(
|
|
"user",
|
|
"What is the weather today?",
|
|
role="user",
|
|
url="url2.mp3",
|
|
),
|
|
Msg(
|
|
"assistant",
|
|
"It is sunny today",
|
|
role="assistant",
|
|
url="url3.mp3",
|
|
),
|
|
],
|
|
]
|
|
|
|
ground_truth = [
|
|
{
|
|
"role": "system",
|
|
"content": [
|
|
{"audio": "url1.mp3"},
|
|
{"text": "You are a helpful assistant"},
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"audio": "url2.mp3"},
|
|
{"audio": "url3.mp3"},
|
|
{
|
|
"text": (
|
|
"## Dialogue History\n"
|
|
"user: What is the weather today?\n"
|
|
"assistant: It is sunny today"
|
|
),
|
|
},
|
|
],
|
|
},
|
|
]
|
|
|
|
prompt = model.format(*multimodal_input)
|
|
self.assertListEqual(prompt, ground_truth)
|
|
|
|
# wrong format
|
|
with self.assertRaises(TypeError):
|
|
model.format(*self.wrong_inputs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|