Replace pygraphviz with neo4j-viz for graph visualization #306
Conversation
- Replace pygraphviz imports with neo4j-viz in pipeline.py - Update draw() method to save visualizations as HTML files - Implement get_neo4j_viz_graph() method for interactive visualizations - Update example in visualization.py to use HTML output - Update tests to work with neo4j-viz implementation - Add color coding for better visual distinction
| else: | ||
| # For other formats, we'll use the render method and save the image | ||
| G.render() | ||
| # Note: neo4j-viz doesn't support direct saving to image formats |
There was a problem hiding this comment.
Is the comment saying that it can't save to PNG? Because code in the 'else' looks very similar to the one in the former 'if' block to me.
| @@ -24,7 +24,9 @@ | |||
| try: | |||
| from mistralai import Mistral | |||
| except ImportError: | |||
| Mistral = None # type: ignore | |||
| # Define placeholder type for type checking | |||
| class Mistral: # type: ignore | |||
There was a problem hiding this comment.
Is this better compared to the previous approach? (out of curiosity)
| def get_pygraphviz_graph(self, hide_unused_outputs: bool = True) -> pgv.AGraph: |
There was a problem hiding this comment.
I don't remember why we haven't made this method private, but since we are removing it we should mention it in the changelog.
| with open(path, "w") as f: | ||
| f.write(G.render()._repr_html_()) | ||
|
|
||
| def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph: |
There was a problem hiding this comment.
And maybe make this method private?
| "Follow installation instruction in pygraphviz documentation " | ||
| "to get it up and running on your system." | ||
| "Could not import neo4j-viz. " | ||
| "Install it with 'pip install neo4j-viz'." |
There was a problem hiding this comment.
I think using extra groups instead reduces a bit dependency issues:
| "Install it with 'pip install neo4j-viz'." | |
| 'Install it with 'pip install "neo4j-graphrag[experimental]".' |
| node_ids[n] = node_counter | ||
| nodes.append( | ||
| Node( | ||
| id=node_counter, |
There was a problem hiding this comment.
According to the doc, it looks like id can be a string? That could simplify the id and node matching part.
| G.remove_node(n) | ||
| return G | ||
| used_nodes = set() | ||
| for rel in relationships: |
There was a problem hiding this comment.
I think there is a logic issue here: we do not want to filter out isolated nodes (this should not happen I think?), but node of type "output" with no outgoing relationship.
| assert len(g.nodes()) == 3 | ||
| g = pipe.get_pygraphviz_graph(hide_unused_outputs=False) | ||
| # - 2 outputs 'a.result' and 'b.result' (neo4j-viz implementation includes both) | ||
| assert len(g.nodes) == 4 |
There was a problem hiding this comment.
Related to my comment above about logic, the test should not need to be updated.
Neo4j has its own viz library, let's use that!
Type of Change
How Has This Been Tested?
Checklist
The following requirements should have been met (depending on the changes in the branch):