From 2d085f61bb829498d1958f498608ddc3750e4e9e Mon Sep 17 00:00:00 2001 From: Pavel Jakovlev Date: Sun, 27 Jul 2025 13:57:57 +0300 Subject: [PATCH 1/2] fix: BFS max_depth parameter now properly controls traversal depth - Fixed hardcoded {1,3} path lengths in both edge_bfs_search and node_bfs_search - Replaced with {1,} to use the bfs_max_depth parameter correctly - Added tests to verify the fix works - Resolves issue where bfs_max_depth=1 would still traverse 3 hops Closes #772 --- graphiti_core/search/search_utils.py | 4 +- tests/utils/search/test_bfs_depth_fix.py | 94 ++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 2 deletions(-) create mode 100644 tests/utils/search/test_bfs_depth_fix.py diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 5d6828f7..bc80bc28 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -295,7 +295,7 @@ async def edge_bfs_search( query = ( """ UNWIND $bfs_origin_node_uuids AS origin_uuid - MATCH path = (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity) + MATCH path = (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,$depth}(n:Entity) UNWIND relationships(path) AS rel MATCH (n:Entity)-[r:RELATES_TO]-(m:Entity) WHERE r.uuid = rel.uuid @@ -446,7 +446,7 @@ async def node_bfs_search( query = ( """ UNWIND $bfs_origin_node_uuids AS origin_uuid - MATCH (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity) + MATCH (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,$depth}(n:Entity) WHERE n.group_id = origin.group_id AND origin.group_id IN $group_ids """ diff --git a/tests/utils/search/test_bfs_depth_fix.py b/tests/utils/search/test_bfs_depth_fix.py new file mode 100644 index 00000000..6655c08f --- /dev/null +++ b/tests/utils/search/test_bfs_depth_fix.py @@ -0,0 +1,94 @@ +"""Test for BFS max_depth parameter fix.""" + +from unittest.mock import AsyncMock + +import pytest + +from graphiti_core.search.search_filters import SearchFilters +from graphiti_core.search.search_utils import edge_bfs_search, node_bfs_search + + +@pytest.mark.asyncio +async def test_edge_bfs_search_uses_depth_parameter(): + """Test that edge_bfs_search uses the bfs_max_depth parameter in the query.""" + # Mock driver + mock_driver = AsyncMock() + mock_driver.execute_query.return_value = ([], None, None) + + # Mock search filter + search_filter = SearchFilters() + + # Call edge_bfs_search with depth=2 + await edge_bfs_search( + driver=mock_driver, + bfs_origin_node_uuids=['test-uuid'], + bfs_max_depth=2, + search_filter=search_filter, + group_ids=['test-group'], + limit=10, + ) + + # Verify the query was called + assert mock_driver.execute_query.called + call_args = mock_driver.execute_query.call_args + + # Check that depth parameter is passed + assert 'depth' in call_args.kwargs + assert call_args.kwargs['depth'] == 2 + + # Check that the query contains the variable depth pattern + query = call_args.args[0] + assert '{1,$depth}' in query, f"Query should contain '{{1,$depth}}' but got: {query}" + + +@pytest.mark.asyncio +async def test_node_bfs_search_uses_depth_parameter(): + """Test that node_bfs_search uses the bfs_max_depth parameter in the query.""" + # Mock driver + mock_driver = AsyncMock() + mock_driver.execute_query.return_value = ([], None, None) + + # Mock search filter + search_filter = SearchFilters() + + # Call node_bfs_search with depth=1 + await node_bfs_search( + driver=mock_driver, + bfs_origin_node_uuids=['test-uuid'], + search_filter=search_filter, + bfs_max_depth=1, + group_ids=['test-group'], + limit=10, + ) + + # Verify the query was called + assert mock_driver.execute_query.called + call_args = mock_driver.execute_query.call_args + + # Check that depth parameter is passed + assert 'depth' in call_args.kwargs + assert call_args.kwargs['depth'] == 1 + + # Check that the query contains the variable depth pattern + query = call_args.args[0] + assert '{1,$depth}' in query, f"Query should contain '{{1,$depth}}' but got: {query}" + + +@pytest.mark.asyncio +async def test_different_depth_values(): + """Test that different bfs_max_depth values are correctly passed.""" + mock_driver = AsyncMock() + mock_driver.execute_query.return_value = ([], None, None) + search_filter = SearchFilters() + + # Test depth=5 + await edge_bfs_search( + driver=mock_driver, + bfs_origin_node_uuids=['test-uuid'], + bfs_max_depth=5, + search_filter=search_filter, + group_ids=['test-group'], + ) + + call_args = mock_driver.execute_query.call_args + assert call_args.kwargs['depth'] == 5 From f99ec00a0c0506a3d17bbddfdf0ff109a5442f50 Mon Sep 17 00:00:00 2001 From: Pavel Jakovlev Date: Tue, 29 Jul 2025 19:18:26 +0300 Subject: [PATCH 2/2] Moving bfs depth unit tests to search utils tests file. --- tests/utils/search/search_utils_test.py | 81 +++++++++++++++++++- tests/utils/search/test_bfs_depth_fix.py | 94 ------------------------ 2 files changed, 80 insertions(+), 95 deletions(-) delete mode 100644 tests/utils/search/test_bfs_depth_fix.py diff --git a/tests/utils/search/search_utils_test.py b/tests/utils/search/search_utils_test.py index 6b97daab..47f91c6c 100644 --- a/tests/utils/search/search_utils_test.py +++ b/tests/utils/search/search_utils_test.py @@ -4,7 +4,7 @@ import pytest from graphiti_core.nodes import EntityNode from graphiti_core.search.search_filters import SearchFilters -from graphiti_core.search.search_utils import hybrid_node_search +from graphiti_core.search.search_utils import edge_bfs_search, hybrid_node_search, node_bfs_search @pytest.mark.asyncio @@ -161,3 +161,82 @@ async def test_hybrid_node_search_with_limit_and_duplicates(): mock_similarity_search.assert_called_with( mock_driver, [0.1, 0.2, 0.3], SearchFilters(), ['1'], 4 ) + + +@pytest.mark.asyncio +async def test_edge_bfs_search_uses_depth_parameter(): + """Test that edge_bfs_search uses the bfs_max_depth parameter in the query.""" + # Mock driver + mock_driver = AsyncMock() + mock_driver.execute_query.return_value = ([], None, None) + + # Mock search filter + search_filter = SearchFilters() + + # Call edge_bfs_search with depth=2 + await edge_bfs_search( + driver=mock_driver, + bfs_origin_node_uuids=['test-uuid'], + bfs_max_depth=2, + search_filter=search_filter, + group_ids=['test-group'], + limit=10, + ) + + # Verify the query was called + assert mock_driver.execute_query.called + call_args = mock_driver.execute_query.call_args + + # Check that the query contains the variable depth pattern + query = call_args.args[0] + assert '*1..2' in query, f"Query should contain '*1..2' but got: {query}" + + +@pytest.mark.asyncio +async def test_node_bfs_search_uses_depth_parameter(): + """Test that node_bfs_search uses the bfs_max_depth parameter in the query.""" + # Mock driver + mock_driver = AsyncMock() + mock_driver.execute_query.return_value = ([], None, None) + + # Mock search filter + search_filter = SearchFilters() + + # Call node_bfs_search with depth=1 + await node_bfs_search( + driver=mock_driver, + bfs_origin_node_uuids=['test-uuid'], + search_filter=search_filter, + bfs_max_depth=1, + group_ids=['test-group'], + limit=10, + ) + + # Verify the query was called + assert mock_driver.execute_query.called + call_args = mock_driver.execute_query.call_args + + # Check that the query contains the variable depth pattern + query = call_args.args[0] + assert '*1..1' in query, f"Query should contain '*1..1' but got: {query}" + + +@pytest.mark.asyncio +async def test_different_depth_values(): + """Test that different bfs_max_depth values are correctly passed.""" + mock_driver = AsyncMock() + mock_driver.execute_query.return_value = ([], None, None) + search_filter = SearchFilters() + + # Test depth=5 + await edge_bfs_search( + driver=mock_driver, + bfs_origin_node_uuids=['test-uuid'], + bfs_max_depth=5, + search_filter=search_filter, + group_ids=['test-group'], + ) + + call_args = mock_driver.execute_query.call_args + query = call_args.args[0] + assert '*1..5' in query, f"Query should contain '*1..5' but got: {query}" diff --git a/tests/utils/search/test_bfs_depth_fix.py b/tests/utils/search/test_bfs_depth_fix.py deleted file mode 100644 index 6655c08f..00000000 --- a/tests/utils/search/test_bfs_depth_fix.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Test for BFS max_depth parameter fix.""" - -from unittest.mock import AsyncMock - -import pytest - -from graphiti_core.search.search_filters import SearchFilters -from graphiti_core.search.search_utils import edge_bfs_search, node_bfs_search - - -@pytest.mark.asyncio -async def test_edge_bfs_search_uses_depth_parameter(): - """Test that edge_bfs_search uses the bfs_max_depth parameter in the query.""" - # Mock driver - mock_driver = AsyncMock() - mock_driver.execute_query.return_value = ([], None, None) - - # Mock search filter - search_filter = SearchFilters() - - # Call edge_bfs_search with depth=2 - await edge_bfs_search( - driver=mock_driver, - bfs_origin_node_uuids=['test-uuid'], - bfs_max_depth=2, - search_filter=search_filter, - group_ids=['test-group'], - limit=10, - ) - - # Verify the query was called - assert mock_driver.execute_query.called - call_args = mock_driver.execute_query.call_args - - # Check that depth parameter is passed - assert 'depth' in call_args.kwargs - assert call_args.kwargs['depth'] == 2 - - # Check that the query contains the variable depth pattern - query = call_args.args[0] - assert '{1,$depth}' in query, f"Query should contain '{{1,$depth}}' but got: {query}" - - -@pytest.mark.asyncio -async def test_node_bfs_search_uses_depth_parameter(): - """Test that node_bfs_search uses the bfs_max_depth parameter in the query.""" - # Mock driver - mock_driver = AsyncMock() - mock_driver.execute_query.return_value = ([], None, None) - - # Mock search filter - search_filter = SearchFilters() - - # Call node_bfs_search with depth=1 - await node_bfs_search( - driver=mock_driver, - bfs_origin_node_uuids=['test-uuid'], - search_filter=search_filter, - bfs_max_depth=1, - group_ids=['test-group'], - limit=10, - ) - - # Verify the query was called - assert mock_driver.execute_query.called - call_args = mock_driver.execute_query.call_args - - # Check that depth parameter is passed - assert 'depth' in call_args.kwargs - assert call_args.kwargs['depth'] == 1 - - # Check that the query contains the variable depth pattern - query = call_args.args[0] - assert '{1,$depth}' in query, f"Query should contain '{{1,$depth}}' but got: {query}" - - -@pytest.mark.asyncio -async def test_different_depth_values(): - """Test that different bfs_max_depth values are correctly passed.""" - mock_driver = AsyncMock() - mock_driver.execute_query.return_value = ([], None, None) - search_filter = SearchFilters() - - # Test depth=5 - await edge_bfs_search( - driver=mock_driver, - bfs_origin_node_uuids=['test-uuid'], - bfs_max_depth=5, - search_filter=search_filter, - group_ids=['test-group'], - ) - - call_args = mock_driver.execute_query.call_args - assert call_args.kwargs['depth'] == 5