init
This commit is contained in:
CakeCN
503
AlgoriAgent/tests/format_test.py
Normal file
503
AlgoriAgent/tests/format_test.py
Normal file
@@ -0,0 +1,503 @@
|
||||
# -*- 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()
|
||||
Reference in New Issue
Block a user