Spaces:
Sleeping
Sleeping
File size: 2,021 Bytes
d6152d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
# sourcery skip: snake-case-functions
"""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()
# Add example texts to the cache
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)
# Add some random strings to test noise
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."
)
|