fix: add max_items limit to fetch_all_edges to prevent unbounded memory growth (#513)

fetch_all_nodes already had a max_items guard (default 2000) but
fetch_all_edges had no such safeguard, allowing unbounded memory growth
on graphs with large numbers of edges.

Add _MAX_EDGES = 5000 constant and mirror the same loop-guard pattern
from fetch_all_nodes: cap the result list, emit a warning log, and
break pagination once the limit is reached.
This commit is contained in:
harshitanand 2026-04-10 13:51:52 +05:30
parent 1536a79334
commit 0fa41a906b
No known key found for this signature in database
GPG Key ID: 702340683D146B97
1 changed files with 7 additions and 1 deletions

View File

@ -19,6 +19,7 @@ logger = get_logger('mirofish.zep_paging')
_DEFAULT_PAGE_SIZE = 100
_MAX_NODES = 2000
_MAX_EDGES = 5000
_DEFAULT_MAX_RETRIES = 3
_DEFAULT_RETRY_DELAY = 2.0 # seconds, doubles each retry
@ -106,10 +107,11 @@ def fetch_all_edges(
client: Zep,
graph_id: str,
page_size: int = _DEFAULT_PAGE_SIZE,
max_items: int = _MAX_EDGES,
max_retries: int = _DEFAULT_MAX_RETRIES,
retry_delay: float = _DEFAULT_RETRY_DELAY,
) -> list[Any]:
"""分页获取图谱所有边,返回完整列表。每页请求自带重试。"""
"""分页获取图谱所有边,最多返回 max_items 条(默认 5000。每页请求自带重试。"""
all_edges: list[Any] = []
cursor: str | None = None
page_num = 0
@ -132,6 +134,10 @@ def fetch_all_edges(
break
all_edges.extend(batch)
if len(all_edges) >= max_items:
all_edges = all_edges[:max_items]
logger.warning(f"Edge count reached limit ({max_items}), stopping pagination for graph {graph_id}")
break
if len(batch) < page_size:
break