File size: 979 Bytes
4c8597f |
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 |
import unittest
import numpy as np
from memory_manager import MemoryManager
def dummy_embed(text):
arr = np.array([ord(c) for c in text], dtype=np.float32)
if arr.size == 0:
return np.zeros(10)
arr = arr / (np.linalg.norm(arr) + 1e-9)
if arr.size 10:
arr = np.pad(arr, (0, 10 - arr.size))
else:
arr = arr[:10]
return arr
class MemoryManagerTest(unittest.TestCase):
def setUp(self):
self.mem = MemoryManager(short_term_size=10, embedding_function=dummy_embed)
def test_add_and_retrieve_recent(self):
self.mem.add_memory("Test memory 1")
self.mem.add_memory("Another memory")
recent = self.mem.get_recent(2)
self.assertEqual(len(recent), 2)
def test_retrieve_similar(self):
self.mem.add_memory("Hello world")
result = self.mem.retrieve_similar("Hello", top_k=1)
self.assertTrue(isinstance(result, list))
if __name__ == "__main__":
unittest.main()
|