allow setting the engine

Author: bjchambers

libs/community/langchain_community/graph_vectorstores/visualize.py

@@ -36,13 +36,15 @@ def _split_prefix(s: str, max_chars: int = 50) -> str:
 
 def render_graphviz(
     documents: Iterable[Document],
+    engine: Optional[str] = None,
     node_color: Optional[str] = None,
     node_colors: Optional[Dict[str, Optional[str]]] = {},
 ) -> "graphviz.Digraph":
     """Render a collection of GraphVectorStore documents to GraphViz format.
 
     Args:
         documents: The documents to render.
+        engine: GraphViz layout engine to use. `None` uses the default.
         node_color: General node color. Defaults to `white`.
         node_colors: Dictionary specifying colors of specific nodes. Useful for
             emphasizing nodes that were selected by MMR, or differ from other
@@ -78,7 +80,7 @@ def render_graphviz(
 
     tags = set()
 
-    graph = graphviz.Digraph()
+    graph = graphviz.Digraph(engine=engine)
     graph.attr(rankdir="LR")
     graph.attr("node", style="filled")
     for document in documents:

libs/community/tests/unit_tests/graph_vectorstores/test_visualize.py

@@ -49,6 +49,8 @@ def test_visualize_simple_graph():
         "}\n"
     )
 
+    assert render_graphviz([doc1, doc2], engine="fdp").engine == "fdp"
+
     assert render_graphviz([doc1, doc2], node_colors={"a": "gold"}).source == (
         "digraph {\n"
         "\trankdir=LR\n"