| |
| """Tests for the MilvusMemory class.""" |
| import random |
| import string |
| import unittest |
|
|
| from autogpt.config import Config |
| from autogpt.memory.milvus import MilvusMemory |
|
|
| try: |
|
|
| class TestMilvusMemory(unittest.TestCase): |
| """Tests for the MilvusMemory class.""" |
|
|
| def random_string(self, length: int) -> str: |
| """Generate a random string of the given length.""" |
| return "".join(random.choice(string.ascii_letters) for _ in range(length)) |
|
|
| def setUp(self) -> None: |
| """Set up the test environment.""" |
| cfg = Config() |
| cfg.milvus_addr = "localhost:19530" |
| self.memory = MilvusMemory(cfg) |
| self.memory.clear() |
|
|
| |
| self.example_texts = [ |
| "The quick brown fox jumps over the lazy dog", |
| "I love machine learning and natural language processing", |
| "The cake is a lie, but the pie is always true", |
| "ChatGPT is an advanced AI model for conversation", |
| ] |
|
|
| for text in self.example_texts: |
| self.memory.add(text) |
|
|
| |
| for _ in range(5): |
| self.memory.add(self.random_string(10)) |
|
|
| def test_get_relevant(self) -> None: |
| """Test getting relevant texts from the cache.""" |
| query = "I'm interested in artificial intelligence and NLP" |
| num_relevant = 3 |
| relevant_texts = self.memory.get_relevant(query, num_relevant) |
|
|
| print(f"Top {k} relevant texts for the query '{query}':") |
| for i, text in enumerate(relevant_texts, start=1): |
| print(f"{i}. {text}") |
|
|
| self.assertEqual(len(relevant_texts), k) |
| self.assertIn(self.example_texts[1], relevant_texts) |
|
|
| except: |
| print( |
| "Skipping tests/integration/milvus_memory_tests.py as Milvus is not installed." |
| ) |
|
|
