File size: 2,541 Bytes
c3a3710 | 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 59 60 61 62 63 64 65 66 67 | import pytest
import pytest_asyncio
import os
from mnemocore.core.config import get_config, reset_config
from mnemocore.core.engine import HAIMEngine
@pytest.fixture
def test_engine(tmp_path):
from mnemocore.core.hnsw_index import HNSWIndexManager
HNSWIndexManager._instance = None
reset_config()
data_dir = tmp_path / "data"
data_dir.mkdir()
os.environ["HAIM_DATA_DIR"] = str(data_dir)
os.environ["HAIM_ENCODING_MODE"] = "binary"
os.environ["HAIM_DIMENSIONALITY"] = "1024"
os.environ["HAIM_CONTEXT_ENABLED"] = "True"
# Unrelated binary vectors have similarity ~0.5.
os.environ["HAIM_CONTEXT_SHIFT_THRESHOLD"] = "0.6"
os.environ["HAIM_CONTEXT_ROLLING_WINDOW_SIZE"] = "3"
os.environ["HAIM_TIERS_HOT_LTP_THRESHOLD_MIN"] = "0.01"
reset_config()
engine = HAIMEngine()
yield engine
del os.environ["HAIM_DATA_DIR"]
del os.environ["HAIM_ENCODING_MODE"]
del os.environ["HAIM_DIMENSIONALITY"]
del os.environ["HAIM_CONTEXT_ENABLED"]
del os.environ["HAIM_CONTEXT_SHIFT_THRESHOLD"]
del os.environ["HAIM_CONTEXT_ROLLING_WINDOW_SIZE"]
del os.environ["HAIM_TIERS_HOT_LTP_THRESHOLD_MIN"]
reset_config()
@pytest.mark.asyncio
async def test_topic_tracker_and_context(test_engine):
await test_engine.initialize()
# Store some nodes for retrieval
await test_engine.store("I love learning about Machine Learning and AI algorithms.")
await test_engine.store("Neural networks and deep learning models are very powerful.")
await test_engine.store("Pineapples belong on pizza. Yes, it is true.")
# 1. Querying related to AI multiple times to build context
await test_engine.query("Tell me about Machine Learning.")
await test_engine.query("How do Neural Networks work?")
# Check context nodes for AI context
ctx_nodes1 = await test_engine.get_context_nodes(top_k=2)
assert len(ctx_nodes1) > 0, "Should retrieve context nodes"
# 2. Shift topic drastically
await test_engine.query("What about pizza toppings? Do pineapples belong?")
# The history should have reset due to shift drop
assert len(test_engine.topic_tracker.history) == 1
# The retrieved context nodes should now be heavily biased toward Pizza
ctx_nodes2 = await test_engine.get_context_nodes(top_k=1)
assert ctx_nodes2[0][0] != ctx_nodes1[0][0], "Context should have decisively shifted"
await test_engine.close()
|