"""Tests for hybrid search functionality. Hybrid search combines keyword search over document metadata with semantic (vector) search via Ollama embeddings, then fuses results using Reciprocal Rank Fusion (RRF) per RAGraph research. """ from unittest.mock import MagicMock, Mock, patch import pytest from cyberbrain26.data_plane.sqlite_store import SQLiteDatabase from cyberbrain26.knowledge_plane.memory_service import MemoryService def test_rrf_fuse_basic(): """Test basic RRF fusion functionality.""" results_a = [ {"id": "doc1", "score": 10}, {"id": "doc2", "score": 8}, {"id": "doc3", "score": 5}, ] results_b = [ {"id": "doc2", "score": 12}, {"id": "doc4", "score": 9}, ] fused = MemoryService._rrf_fuse(results_a, results_b, k=60) # All docs should be present assert len(fused) == 4 fused_ids = {r["id"] for r in fused} assert fused_ids == {"doc1", "doc2", "doc3", "doc4"} # doc2 has highest rank in both (1 and 1) -> highest RRF score doc2_score = next(r["rrf_score"] for r in fused if r["id"] == "doc2") doc1_score = next(r["rrf_score"] for r in fused if r["id"] == "doc1") assert doc2_score > doc1_score # doc2 should score higher than doc1 def test_rrf_fuse_no_overlap(): """Test RRF fusion when there is no overlap between results.""" results_a = [ {"id": "doc1", "score": 10}, {"id": "doc2", "score": 8}, ] results_b = [ {"id": "doc3", "score": 9}, {"id": "doc4", "score": 7}, ] fused = MemoryService._rrf_fuse(results_a, results_b, k=60) assert len(fused) == 4 fused_ids = {r["id"] for r in fused} assert fused_ids == {"doc1", "doc2", "doc3", "doc4"} def _make_service_with_connection(rows): """Build a MemoryService whose database.connection() returns mock rows.""" from contextlib import contextmanager mock_db = MagicMock(spec=SQLiteDatabase) mock_connection = MagicMock() mock_result = MagicMock() mock_result.fetchall.return_value = rows mock_connection.execute.return_value = mock_result @contextmanager def _mock_connection(): yield mock_connection mock_db.connection = _mock_connection service = MemoryService( database=mock_db, audit_log=Mock(spec="AuditLog"), ) return service def test_keyword_search_basic(): """Test basic keyword search functionality.""" rows = [ { "document_id": "doc1", "source_uri": "file:///test1.txt", "media_type": "text/plain", "metadata_json": '{"topic": "AI", "tags": ["machine learning", "neural networks"], "summary": "AI research paper about neural networks"}', }, { "document_id": "doc2", "source_uri": "file:///test2.txt", "media_type": "text/plain", "metadata_json": '{"topic": "Security", "tags": ["cryptography", "encryption"], "summary": "Security paper on encryption methods"}', }, ] service = _make_service_with_connection(rows) results = service.keyword_search("neural", limit=5) # Should find doc1 because "neural" appears in tags assert len(results) == 1 assert results[0]["id"] == "doc1" def test_keyword_search_multiple_matches(): """Test keyword search with multiple matches.""" rows = [ { "document_id": "doc1", "source_uri": "file:///test1.txt", "media_type": "text/plain", "metadata_json": '{"topic": "AI research on neural networks", "tags": ["deep learning"], "summary": "Neural networks for deep learning"}', }, { "document_id": "doc2", "source_uri": "file:///test2.txt", "media_type": "text/plain", "metadata_json": '{"topic": "Machine Learning", "tags": ["neural networks", "AI"], "summary": "Machine learning techniques"}', }, { "document_id": "doc3", "source_uri": "file:///test3.txt", "media_type": "text/plain", "metadata_json": '{"topic": "Data Science", "tags": ["statistics"], "summary": "Statistical analysis"}', }, ] service = _make_service_with_connection(rows) results = service.keyword_search("neural", limit=10) # Should find 2 documents because "neural" appears in metadata of doc1 and doc2 assert len(results) == 2 # doc1 should have highest score (topic match) assert results[0]["id"] == "doc1" # doc2 should have second highest score (tag match) assert results[1]["id"] == "doc2" def test_keyword_search_no_matches(): """Test keyword search with no matching documents.""" rows = [ { "document_id": "doc1", "source_uri": "file:///test1.txt", "media_type": "text/plain", "metadata_json": '{"topic": "Security", "tags": ["cryptography"], "summary": "Security paper"}', }, ] service = _make_service_with_connection(rows) results = service.keyword_search("neural", limit=10) assert len(results) == 0 def test_hybrid_search_integration(): """Test integration of keyword and semantic search with RRF fusion.""" service = MemoryService( database=MagicMock(spec=SQLiteDatabase), audit_log=Mock(spec="AuditLog"), ) # Mock keyword_search and semantic_search methods on the service instance service.keyword_search = Mock(return_value=[ {"id": "doc1", "score": 12, "text": "text1"}, {"id": "doc2", "score": 8, "text": "text2"}, ]) service.semantic_search = Mock(return_value=[ {"id": "doc2", "score": 10, "text": "text2"}, {"id": "doc3", "score": 9, "text": "text3"}, ]) results = service.hybrid_search("test", limit=10) # Should have all 3 documents assert len(results) == 3 result_ids = {r["id"] for r in results} assert result_ids == {"doc1", "doc2", "doc3"} # doc2 has highest ranks in both (1 and 1) -> should have highest RRF score doc2_score = next(r["rrf_score"] for r in results if r["id"] == "doc2") doc1_score = next(r["rrf_score"] for r in results if r["id"] == "doc1") doc3_score = next(r["rrf_score"] for r in results if r["id"] == "doc3") assert doc2_score > doc1_score assert doc2_score > doc3_score # Verify mocks were called with correct args service.keyword_search.assert_called_once_with("test", limit=20) service.semantic_search.assert_called_once_with("test", limit=20, topic=None, tag=None) def test_hybrid_search_with_topic_filter(): """Test hybrid search with topic filter.""" service = MemoryService( database=MagicMock(spec=SQLiteDatabase), audit_log=Mock(spec="AuditLog"), ) service.keyword_search = Mock(return_value=[ {"id": "doc1", "score": 12, "text": "text1"}, {"id": "doc2", "score": 8, "text": "text2"}, ]) service.semantic_search = Mock(return_value=[ {"id": "doc2", "score": 10, "text": "text2"}, {"id": "doc3", "score": 9, "text": "text3"}, ]) # With topic filter, semantic search should receive it results = service.hybrid_search("test", topic="ai", limit=10) # Both mock methods should have been called service.keyword_search.assert_called_once_with("test", limit=20) service.semantic_search.assert_called_once_with("test", limit=20, topic="ai", tag=None) if __name__ == "__main__": pytest.main([__file__, "-v"])