From 2d085f61bb829498d1958f498608ddc3750e4e9e Mon Sep 17 00:00:00 2001
From: Pavel Jakovlev <pavel.jakovlev@nortal.com>
Date: Sun, 27 Jul 2025 13:57:57 +0300
Subject: [PATCH] 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