Function to visualize a GraphSchema object or dict

CHANGELOG.md

@@ -6,11 +6,13 @@
 
 - Added a `ToolsRetriever` retriever that uses an LLM to decide on what tools to use to find the relevant data.
 - Added `convert_to_tool` method to the `Retriever` interface to convert a Retriever to a Tool so it can be used within the ToolsRetriever. This is useful when you might want to have both a VectorRetriever and a Text2CypherRetreiver as a fallback.
+- Added `schema_visualization` function to visualize a graph schema using neo4j-viz.
 
 ### Fixed
 
 - Fixed an edge case where the LLM can output a property with type 'map', which was causing errors during import as it is not a valid property type in Neo4j.
 
+
 ## 1.9.1
 
 ### Fixed
@@ -30,7 +32,6 @@
 
 - Added automatic rate limiting with retry logic and exponential backoff for all LLM providers using tenacity. The `RateLimitHandler` interface allows for custom rate limiting strategies, including the ability to disable rate limiting entirely.
 
-
 ## 1.8.0
 
 ### Added

docs/source/api.rst

@@ -83,6 +83,12 @@ SchemaFromTextExtractor
 .. autoclass:: neo4j_graphrag.experimental.components.schema.SchemaFromTextExtractor
     :members: run
 
+schema_visualization
+--------------------
+
+.. autofunction:: neo4j_graphrag.experimental.utils.schema.schema_visualization
+
+
 EntityRelationExtractor
 =======================
 

docs/source/user_guide_kg_builder.rst

@@ -852,6 +852,28 @@ You can also save and reload the extracted schema:
     restored_schema = GraphSchema.from_file("my_schema.json")  # or my_schema.yaml
 
 
+Schema Visualization
+--------------------
+
+It is possible to visualize a validated schema or a schema dict using the `schema_visualization` function. This function
+returns a VisualizationGraph object (from the neo4j-viz package) that can visualized like this:
+
+.. code:: python
+
+        from neo4j_graphrag.experimental.utils.schema import schema_visualization
+
+        VG = schema_visualization(schema)
+        html = VG.render()
+
+        # in Jupyter:
+        display(html)
+
+        # to save the generated HTML
+        with open("my_schema.html", "w") as f:
+            f.write(html.data)
+
+
+
 Entity and Relation Extractor
 =============================
 

pyproject.toml

@@ -38,7 +38,7 @@ pyyaml = "^6.0.2"
 types-pyyaml = "^6.0.12.20240917"
 # optional deps
 langchain-text-splitters = {version = "^0.3.0", optional = true }
-neo4j-viz = {version = "^0.2.2", optional = true }
+neo4j-viz = {version = "^0.4.2", optional = true }
 weaviate-client = {version = "^4.6.1", optional = true }
 pinecone-client = {version = "^4.1.0", optional = true }
 google-cloud-aiplatform = {version = "^1.66.0", optional = true }
@@ -74,7 +74,6 @@ sphinx = { version = "^7.2.6", python = "^3.9" }
 langchain-openai = {version = "^0.2.2", optional = true }
 langchain-huggingface = {version = "^0.1.0", optional = true }
 enum-tools = {extras = ["sphinx"], version = "^0.12.0"}
-neo4j-viz = "^0.2.2"
 
 [tool.poetry.extras]
 weaviate = ["weaviate-client"]

src/neo4j_graphrag/experimental/utils/__init__.py

@@ -0,0 +1,14 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.