init

AlgoriAgent/tests/prompt_engine_test.py

@@ -0,0 +1,135 @@
+# -*- 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()