Skip to content

Commit 08248bf

Browse files
committed
Fix type errors in pipeline.py by adding required parameters to Node and Relationship
1 parent 9a92182 commit 08248bf

File tree

1 file changed

+17
-2
lines changed

1 file changed

+17
-2
lines changed

src/neo4j_graphrag/experimental/pipeline/pipeline.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from neo4j_graphrag.utils.logging import prettify
2424

2525
try:
26-
from neo4j_viz import Node, Relationship, VisualizationGraph as NeoVizGraph
26+
from neo4j_viz import Node, Relationship, VisualizationGraph as NeoVizGraph, CaptionAlignment
2727

2828
HAS_NEO4J_VIZ = True
2929
except ImportError:
@@ -244,6 +244,11 @@ def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph:
244244
caption=f"{node.component.__class__.__name__}: {n}({comp_inputs})",
245245
size=20, # Component nodes are larger
246246
color="#4C8BF5", # Blue for component nodes
247+
caption_align=CaptionAlignment.CENTER,
248+
caption_size=12,
249+
pinned=False,
250+
x=0,
251+
y=0
247252
)
248253
)
249254
node_counter += 1
@@ -258,12 +263,20 @@ def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph:
258263
caption=o,
259264
size=10, # Output nodes are smaller
260265
color="#34A853", # Green for output nodes
266+
caption_align=CaptionAlignment.CENTER,
267+
caption_size=10,
268+
pinned=False,
269+
x=0,
270+
y=0
261271
)
262272
)
263273
# Connect component to its output
264274
relationships.append(
265275
Relationship(
266-
source=node_ids[n], target=node_ids[param_node_name], caption=""
276+
source=node_ids[n], target=node_ids[param_node_name], caption="",
277+
caption_align=CaptionAlignment.CENTER,
278+
caption_size=10,
279+
color="#000000"
267280
)
268281
)
269282
node_counter += 1
@@ -285,6 +298,8 @@ def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph:
285298
target=node_ids[component_name],
286299
caption=param,
287300
color="#EA4335", # Red for parameter connections
301+
caption_align=CaptionAlignment.CENTER,
302+
caption_size=10
288303
)
289304
)
290305

0 commit comments

Comments
 (0)