init

AlgoriAgent/tests/ollama_test.py

@@ -0,0 +1,173 @@
+# -*- coding: utf-8 -*-
+"""Unit test for Ollama model APIs."""
+import unittest
+from unittest.mock import patch, MagicMock
+import agentscope
+from agentscope.models import load_model_by_config_name
+from agentscope._runtime import _Runtime
+from agentscope.file_manager import _FileManager
+from agentscope.utils.monitor import MonitorFactory
+
+
+def flush() -> None:
+    """
+    ** Only for unittest usage. Don't use this function in your code. **
+    Clear the runtime dir and destroy all singletons.
+    """
+    _Runtime._flush()  # pylint: disable=W0212
+    _FileManager._flush()  # pylint: disable=W0212
+    MonitorFactory.flush()
+
+
+class OllamaModelWrapperTest(unittest.TestCase):
+    """Unit test for Ollama model APIs."""
+
+    def setUp(self) -> None:
+        """Init for OllamaModelWrapperTest."""
+        self.dummy_response = {
+            "model": "llama2",
+            "created_at": "2024-03-12T04:16:48.911377Z",
+            "message": {
+                "role": "assistant",
+                "content": (
+                    "Hello! It's nice to meet you. Is there something I can "
+                    "help you with or would you like to chat?",
+                ),
+            },
+            "done": True,
+            "total_duration": 20892900042,
+            "load_duration": 20019679292,
+            "prompt_eval_count": 22,
+            "prompt_eval_duration": 149094000,
+            "eval_count": 26,
+            "eval_duration": 721982000,
+        }
+
+        self.dummy_embedding = {
+            "embedding": [1.0, 2.0, 3.0],
+        }
+
+        self.dummy_generate = {
+            "model": "llama2",
+            "created_at": "2024-03-12T03:42:19.621919Z",
+            "response": "\n1 + 1 = 2",
+            "done": True,
+            "context": [
+                518,
+                25580,
+                29962,
+                3532,
+                14816,
+                29903,
+                29958,
+                5299,
+                829,
+                14816,
+                29903,
+                6778,
+                13,
+                13,
+                29896,
+                29974,
+                29896,
+                29922,
+                518,
+                29914,
+                25580,
+                29962,
+                13,
+                13,
+                29896,
+                718,
+                29871,
+                29896,
+                353,
+                29871,
+                29906,
+            ],
+            "total_duration": 6146120041,
+            "load_duration": 6677375,
+            "prompt_eval_count": 9,
+            "prompt_eval_duration": 5913554000,
+            "eval_count": 9,
+            "eval_duration": 223689000,
+        }
+        flush()
+
+    @patch("ollama.chat")
+    def test_ollama_chat(self, mock_chat: MagicMock) -> None:
+        """Unit test for ollama chat API."""
+        # prepare the mock
+        mock_chat.return_value = self.dummy_response
+
+        # run test
+        agentscope.init(
+            model_configs={
+                "config_name": "my_ollama_chat",
+                "model_type": "ollama_chat",
+                "model_name": "llama2",
+                "options": {
+                    "temperature": 0.5,
+                },
+                "keep_alive": "5m",
+            },
+        )
+
+        model = load_model_by_config_name("my_ollama_chat")
+        response = model(messages=[{"role": "user", "content": "Hi!"}])
+
+        self.assertEqual(response.raw, self.dummy_response)
+
+    @patch("ollama.embeddings")
+    def test_ollama_embedding(self, mock_embeddings: MagicMock) -> None:
+        """Unit test for ollama embeddings API."""
+        # prepare the mock
+        mock_embeddings.return_value = self.dummy_embedding
+
+        # run test
+        agentscope.init(
+            model_configs={
+                "config_name": "my_ollama_embedding",
+                "model_type": "ollama_embedding",
+                "model_name": "llama2",
+                "options": {
+                    "temperature": 0.5,
+                },
+                "keep_alive": "5m",
+            },
+        )
+
+        model = load_model_by_config_name("my_ollama_embedding")
+        response = model(prompt="Hi!")
+
+        self.assertEqual(response.raw, self.dummy_embedding)
+
+    @patch("ollama.generate")
+    def test_ollama_generate(self, mock_generate: MagicMock) -> None:
+        """Unit test for ollama generate API."""
+        # prepare the mock
+        mock_generate.return_value = self.dummy_generate
+
+        # run test
+        agentscope.init(
+            model_configs={
+                "config_name": "my_ollama_generate",
+                "model_type": "ollama_generate",
+                "model_name": "llama2",
+                "options": None,
+                "keep_alive": "5m",
+            },
+        )
+
+        model = load_model_by_config_name("my_ollama_generate")
+        response = model(prompt="1+1=")
+
+        self.assertEqual(response.raw, self.dummy_generate)
+
+    def tearDown(self) -> None:
+        """Clean up after each test."""
+        flush()
+
+
+if __name__ == "__main__":
+    unittest.main()