117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""zhipuai test"""
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
import agentscope
|
|
from agentscope.models import load_model_by_config_name
|
|
|
|
|
|
class TestZhipuAIChatWrapper(unittest.TestCase):
|
|
"""Test ZhipuAI Chat Wrapper"""
|
|
|
|
def setUp(self) -> None:
|
|
self.api_key = "test_api_key.secret_key"
|
|
self.messages = [
|
|
{"role": "user", "content": "Hello, ZhipuAI!"},
|
|
{"role": "assistant", "content": "How can I assist you?"},
|
|
]
|
|
|
|
@patch("agentscope.models.zhipu_model.zhipuai")
|
|
def test_chat(self, mock_zhipuai: MagicMock) -> None:
|
|
"""
|
|
Test chat"""
|
|
mock_response = MagicMock()
|
|
mock_response.model_dump.return_value = {
|
|
"choices": [
|
|
{"message": {"content": "Hello, this is a mocked response!"}},
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 5,
|
|
"total_tokens": 105,
|
|
},
|
|
}
|
|
mock_response.choices[
|
|
0
|
|
].message.content = "Hello, this is a mocked response!"
|
|
mock_zhipuai_client = MagicMock()
|
|
mock_zhipuai.ZhipuAI.return_value = mock_zhipuai_client
|
|
|
|
mock_zhipuai_client.chat.completions.create.return_value = (
|
|
mock_response
|
|
)
|
|
|
|
agentscope.init(
|
|
model_configs={
|
|
"config_name": "test_config",
|
|
"model_type": "zhipuai_chat",
|
|
"model_name": "glm-4",
|
|
"api_key": self.api_key,
|
|
},
|
|
)
|
|
|
|
model = load_model_by_config_name("test_config")
|
|
|
|
response = model(messages=self.messages)
|
|
|
|
self.assertEqual(response.text, "Hello, this is a mocked response!")
|
|
|
|
mock_zhipuai_client.chat.completions.create.assert_called_once()
|
|
|
|
|
|
class TestZhipuAIEmbeddingWrapper(unittest.TestCase):
|
|
"""Test ZhipuAI Embedding Wrapper"""
|
|
|
|
def setUp(self) -> None:
|
|
self.api_key = "test_api_key"
|
|
self.model_name = "embedding-2"
|
|
self.text_to_embed = "This is a test sentence for embedding."
|
|
|
|
@patch("agentscope.models.zhipu_model.zhipuai")
|
|
def test_embedding(self, mock_zhipuai: MagicMock) -> None:
|
|
"""Test embedding API"""
|
|
mock_embedding_response = MagicMock()
|
|
mock_embedding_response.model_dump.return_value = {
|
|
"data": [
|
|
{"embedding": [0.1, 0.2, 0.3]},
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 2,
|
|
"total_tokens": 12,
|
|
},
|
|
}
|
|
|
|
mock_zhipuai_client = MagicMock()
|
|
mock_zhipuai.ZhipuAI.return_value = mock_zhipuai_client
|
|
mock_zhipuai_client.embeddings.create.return_value = (
|
|
mock_embedding_response
|
|
)
|
|
|
|
agentscope.init(
|
|
model_configs={
|
|
"config_name": "test_embedding",
|
|
"model_type": "zhipuai_embedding",
|
|
"model_name": self.model_name,
|
|
"api_key": self.api_key,
|
|
},
|
|
)
|
|
|
|
model = load_model_by_config_name("test_embedding")
|
|
|
|
response = model(self.text_to_embed)
|
|
|
|
expected_embedding = [[0.1, 0.2, 0.3]]
|
|
self.assertEqual(response.embedding, expected_embedding)
|
|
|
|
mock_zhipuai_client.embeddings.create.assert_called_once_with(
|
|
input=self.text_to_embed,
|
|
model=self.model_name,
|
|
**{},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|