diff --git a/tests/test_memory_query_router.py b/tests/test_memory_query_router.py new file mode 100644 index 000000000..4a1b42d47 --- /dev/null +++ b/tests/test_memory_query_router.py @@ -0,0 +1,97 @@ +""" +Tests for hybrid memory query router + +Issue: #663 +""" + +import unittest +from tools.memory_query_router import ( + SearchMethod, + QueryRouter, + route_query, + reciprocal_rank_fusion, + merge_with_hrr_priority, +) + + +class TestQueryClassification(unittest.TestCase): + + def setUp(self): + self.router = QueryRouter() + + def test_contradiction_routes_hrr(self): + c = self.router.classify("What contradicts this statement?") + self.assertEqual(c.method, SearchMethod.HRR) + self.assertGreater(c.confidence, 0.9) + + def test_compositional_routes_hrr(self): + c = self.router.classify("How does Python relate to machine learning?") + self.assertEqual(c.method, SearchMethod.HRR) + + c = self.router.classify("What is associated with quantum computing?") + self.assertEqual(c.method, SearchMethod.HRR) + + def test_exact_keywords_routes_fts5(self): + c = self.router.classify('Find documents containing "FastAPI tutorial"') + self.assertEqual(c.method, SearchMethod.FTS5) + + def test_short_query_routes_fts5(self): + c = self.router.classify("Python syntax") + self.assertEqual(c.method, SearchMethod.FTS5) + + def test_temporal_routes_fts5(self): + c = self.router.classify("Recent changes to the config") + self.assertEqual(c.method, SearchMethod.FTS5) + + def test_semantic_routes_vector(self): + c = self.router.classify("Explain how transformers work in natural language processing") + self.assertEqual(c.method, SearchMethod.VECTOR) + + +class TestReciprocalRankFusion(unittest.TestCase): + + def test_basic_fusion(self): + results = { + "hrr": [("a", 0.9), ("b", 0.8)], + "vector": [("b", 0.85), ("c", 0.7)], + } + merged = reciprocal_rank_fusion(results) + + # 'b' appears in both, should rank high + ids = [r[0] for r in merged] + self.assertIn("b", ids[:2]) + + def test_empty_results(self): + merged = reciprocal_rank_fusion({}) + self.assertEqual(len(merged), 0) + + +class TestHRRPriority(unittest.TestCase): + + def test_compositional_hrr_first(self): + hrr = [("a", 0.9), ("b", 0.8)] + vector = [("c", 0.85), ("d", 0.7)] + fts5 = [("e", 0.6)] + + merged = merge_with_hrr_priority(hrr, vector, fts5, "compositional") + + # HRR results should come first + self.assertEqual(merged[0][0], "a") + self.assertEqual(merged[1][0], "b") + + +class TestHybridDecision(unittest.TestCase): + + def test_low_confidence_uses_hybrid(self): + from tools.memory_query_router import should_use_hybrid + # Ambiguous query + self.assertTrue(should_use_hybrid("Tell me about things")) + + def test_clear_query_no_hybrid(self): + from tools.memory_query_router import should_use_hybrid + # Clear contradiction query + self.assertFalse(should_use_hybrid("What contradicts X?")) + + +if __name__ == "__main__": + unittest.main()