init

AlgoriAgent/tests/model_test.py

@@ -0,0 +1,116 @@
+# -*- coding: utf-8 -*-
+"""
+Unit tests for model wrapper classes and functions
+"""
+
+from typing import Any, Union, List, Sequence
+import unittest
+from unittest.mock import patch, MagicMock
+
+from agentscope.message import Msg
+from agentscope.models import (
+    ModelResponse,
+    ModelWrapperBase,
+    OpenAIChatWrapper,
+    PostAPIModelWrapperBase,
+    _get_model_wrapper,
+    read_model_configs,
+    load_model_by_config_name,
+    clear_model_configs,
+)
+
+
+class TestModelWrapperSimple(ModelWrapperBase):
+    """A simple model wrapper class for test usage"""
+
+    def __call__(self, *args: Any, **kwargs: Any) -> ModelResponse:
+        return ModelResponse(text=self.config_name)
+
+    def format(
+        self,
+        *args: Union[Msg, Sequence[Msg]],
+    ) -> Union[List[dict], str]:
+        return ""
+
+
+class BasicModelTest(unittest.TestCase):
+    """Test cases for basic model wrappers"""
+
+    def test_model_registry(self) -> None:
+        """Test the automatic registration mechanism of model wrapper."""
+        # get model wrapper class by class name
+        self.assertEqual(
+            _get_model_wrapper(model_type="TestModelWrapperSimple"),
+            TestModelWrapperSimple,
+        )
+        # get model wrapper class by model type
+        self.assertEqual(
+            _get_model_wrapper(model_type="openai_chat"),
+            OpenAIChatWrapper,
+        )
+        # return PostAPIModelWrapperBase if model_type is not supported
+        self.assertEqual(
+            _get_model_wrapper(model_type="unknown_model_wrapper"),
+            PostAPIModelWrapperBase,
+        )
+
+    @patch("loguru.logger.warning")
+    def test_load_model_configs(self, mock_logging: MagicMock) -> None:
+        """Test to load model configs"""
+        configs = [
+            {
+                "model_type": "openai_chat",
+                "config_name": "gpt-4",
+                "model_name": "gpt-4",
+                "api_key": "xxx",
+                "organization": "xxx",
+                "generate_args": {"temperature": 0.5},
+            },
+            {
+                "model_type": "post_api",
+                "config_name": "my_post_api",
+                "api_url": "https://xxx",
+                "headers": {},
+                "json_args": {},
+            },
+        ]
+        # load a list of configs
+        read_model_configs(configs=configs, clear_existing=True)
+        model = load_model_by_config_name("gpt-4")
+        self.assertEqual(model.config_name, "gpt-4")
+        model = load_model_by_config_name("my_post_api")
+        self.assertEqual(model.config_name, "my_post_api")
+        self.assertRaises(
+            ValueError,
+            load_model_by_config_name,
+            "non_existent_id",
+        )
+
+        # load a single config
+        read_model_configs(configs=configs[0], clear_existing=True)
+        model = load_model_by_config_name("gpt-4")
+        self.assertEqual(model.config_name, "gpt-4")
+        self.assertRaises(ValueError, load_model_by_config_name, "my_post_api")
+
+        # load model with the same id
+        read_model_configs(configs=configs[0], clear_existing=False)
+        mock_logging.assert_called_once_with(
+            "config_name [gpt-4] already exists.",
+        )
+
+        read_model_configs(
+            configs={
+                "model_type": "TestModelWrapperSimple",
+                "config_name": "test_model_wrapper",
+                "args": {},
+            },
+        )
+        test_model = load_model_by_config_name("test_model_wrapper")
+        response = test_model()
+        self.assertEqual(response.text, "test_model_wrapper")
+        clear_model_configs()
+        self.assertRaises(
+            ValueError,
+            load_model_by_config_name,
+            "test_model_wrapper",
+        )