coding-alt/AutoGPT

1# sourcery skip: snake-case-functions2"""Tests for the MilvusMemory class."""3import random4import string5import unittest6 7from autogpt.config import Config8from autogpt.memory.milvus import MilvusMemory9 10try:11 12    class TestMilvusMemory(unittest.TestCase):13        """Tests for the MilvusMemory class."""14 15        def random_string(self, length: int) -> str:16            """Generate a random string of the given length."""17            return "".join(random.choice(string.ascii_letters) for _ in range(length))18 19        def setUp(self) -> None:20            """Set up the test environment."""21            cfg = Config()22            cfg.milvus_addr = "localhost:19530"23            self.memory = MilvusMemory(cfg)24            self.memory.clear()25 26            # Add example texts to the cache27            self.example_texts = [28                "The quick brown fox jumps over the lazy dog",29                "I love machine learning and natural language processing",30                "The cake is a lie, but the pie is always true",31                "ChatGPT is an advanced AI model for conversation",32            ]33 34            for text in self.example_texts:35                self.memory.add(text)36 37            # Add some random strings to test noise38            for _ in range(5):39                self.memory.add(self.random_string(10))40 41        def test_get_relevant(self) -> None:42            """Test getting relevant texts from the cache."""43            query = "I'm interested in artificial intelligence and NLP"44            num_relevant = 345            relevant_texts = self.memory.get_relevant(query, num_relevant)46 47            print(f"Top {k} relevant texts for the query '{query}':")48            for i, text in enumerate(relevant_texts, start=1):49                print(f"{i}. {text}")50 51            self.assertEqual(len(relevant_texts), k)52            self.assertIn(self.example_texts[1], relevant_texts)53 54except:55    print(56        "Skipping tests/integration/milvus_memory_tests.py as Milvus is not installed."57    )58