Skip to content

perf: bypass redundant NetworkX wrapper checks to optimize graph building loops (#798) #802

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
suhaniiz wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

perf: bypass redundant NetworkX wrapper checks to optimize graph building loops (#798) #802

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 24 additions & 11 deletions backend/app/rag/graph_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,17 +93,22 @@ def build_graph(chunks: List[Dict[str, Any]]) -> nx.Graph:
"""Build an entity co-occurrence graph from document chunks."""
graph = nx.Graph()

# Cache graph reference lookups outside the loop for speed
nodes_data = graph.nodes

for chunk in chunks:
text = chunk.get("text", "")
page = chunk.get("page")
chunk_index = chunk.get("chunk_index")
entities = extract_entities(text)

for entity in entities:
if graph.has_node(entity.id):
graph.nodes[entity.id]["mentions"] += 1
graph.nodes[entity.id]["pages"].add(page)
graph.nodes[entity.id]["chunks"].add(chunk_index)
# Direct dictionary check avoids the function call overhead of graph.has_node()
if entity.id in nodes_data:
node = nodes_data[entity.id]
node["mentions"] += 1
node["pages"].add(page)
node["chunks"].add(chunk_index)
else:
graph.add_node(
entity.id,
Expand All @@ -114,22 +119,30 @@ def build_graph(chunks: List[Dict[str, Any]]) -> nx.Graph:
chunks={chunk_index},
)

# Optimize co-occurrence edge calculations inside the tight O(N^2) loop
for left_index, left in enumerate(entities):
left_id = left.id
# Cache the adjacency dictionary lookup for the source node
adj_dict = graph[left_id]

for right in entities[left_index + 1:]:
if graph.has_edge(left.id, right.id):
graph[left.id][right.id]["weight"] += 1
graph[left.id][right.id]["pages"].add(page)
graph[left.id][right.id]["chunks"].add(chunk_index)
right_id = right.id

# Direct membership lookup inside the adjacency mapping bypasses graph.has_edge()
if right_id in adj_dict:
edge = adj_dict[right_id]
edge["weight"] += 1
edge["pages"].add(page)
edge["chunks"].add(chunk_index)
else:
graph.add_edge(
left.id,
right.id,
left_id,
right_id,
weight=1,
pages={page},
chunks={chunk_index},
)

_convert_sets_for_json(graph)
return graph


Expand Down
Loading