Create Document node even from text input (#413)

* Create Document node even from text input

* Update doc

* Ruff

* Save document_type

* Fix type

* Reuse file_path instead of introducing another document_path parameter

* Address comments

* Fix CI

Author: stellasia

CHANGELOG.md

@@ -15,6 +15,12 @@
 - 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.
 
 
+### Added
+
+- Document node is now always created when running SimpleKGPipeline, even if `from_pdf=False`.
+- Document metadata is exposed in SimpleKGPipeline run method.
+
+
 ## 1.9.1
 
 ### Fixed

docs/source/user_guide_kg_builder.rst

@@ -219,6 +219,15 @@ chunk overlap in the text splitter component:
     )
 
 
+Run Parameters
+--------------
+
+SimpleKGPipeline also accepts addition runtime parameters:
+
+- ``document_metadata`` (dict): each item will be saved as a property attached to the ``Document`` node.
+
+
+
 Using a Config file
 ===================
 

examples/build_graph/simple_kg_builder_from_pdf.py

@@ -54,7 +54,12 @@ async def define_and_run_pipeline(
         },
         neo4j_database=DATABASE,
     )
-    return await kg_builder.run_async(file_path=str(file_path))
+    return await kg_builder.run_async(
+        file_path=str(file_path),
+        # optional, add document metadata, each item will
+        # be saved as a property of the Document node
+        # document_metadata={"author": "J. K. Rowling"},
+    )
 
 
 async def main() -> PipelineResult:

examples/build_graph/simple_kg_builder_from_text.py

@@ -79,7 +79,15 @@ async def define_and_run_pipeline(
         from_pdf=False,
         neo4j_database=DATABASE,
     )
-    return await kg_builder.run_async(text=TEXT)
+    return await kg_builder.run_async(
+        text=TEXT,
+        # optional, specify file path for the Document node
+        # if not, a random name will be generated
+        # file_path="my_document.txt"
+        # optional, add document metadata, each item will
+        # be saved as a property of the Document node
+        # document_metadata={"author": "Frank Herbert"},
+    )
 
 
 async def main() -> PipelineResult:

src/neo4j_graphrag/experimental/components/lexical_graph.py

@@ -113,6 +113,7 @@ def create_document_node(self, document_info: DocumentInfo) -> Neo4jNode:
             properties={
                 "path": document_info.path,
                 "createdAt": datetime.datetime.now(datetime.timezone.utc).isoformat(),
+                "document_type": document_info.document_type,
                 **document_metadata,
             },
         )

src/neo4j_graphrag/experimental/components/pdf_loader.py

@@ -89,5 +89,6 @@ async def run(
             document_info=DocumentInfo(
                 path=filepath,
                 metadata=self.get_document_metadata(text, metadata),
+                document_type="pdf",
             ),
         )

src/neo4j_graphrag/experimental/components/types.py

@@ -38,6 +38,7 @@ class DocumentInfo(DataModel):
     path: str
     metadata: Optional[Dict[str, str]] = None
     uid: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    document_type: Optional[str] = None
 
     @property
     def document_id(self) -> str:

src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
+from collections import defaultdict
 from typing import (
     Any,
     ClassVar,
@@ -336,37 +337,40 @@ def _get_connections(self) -> list[ConnectionDefinition]:
         return connections
 
     def get_run_params(self, user_input: dict[str, Any]) -> dict[str, Any]:
-        run_params = {}
-        if self.lexical_graph_config:
-            run_params["extractor"] = {
-                "lexical_graph_config": self.lexical_graph_config,
-            }
-            run_params["writer"] = {
-                "lexical_graph_config": self.lexical_graph_config,
-            }
-            run_params["pruner"] = {
-                "lexical_graph_config": self.lexical_graph_config,
-            }
         text = user_input.get("text")
         file_path = user_input.get("file_path")
-        if not ((text is None) ^ (file_path is None)):
-            # exactly one of text or user_input must be set
+        if text is None and file_path is None:
+            # user must provide either text or file_path or both
             raise PipelineDefinitionError(
-                "Use either 'text' (when from_pdf=False) or 'file_path' (when from_pdf=True) argument."
+                "At least one of `text` (when from_pdf=False) or `file_path` (when from_pdf=True) argument must be provided."
             )
+        run_params: dict[str, dict[str, Any]] = defaultdict(dict)
+        if self.lexical_graph_config:
+            run_params["extractor"]["lexical_graph_config"] = self.lexical_graph_config
+            run_params["writer"]["lexical_graph_config"] = self.lexical_graph_config
+            run_params["pruner"]["lexical_graph_config"] = self.lexical_graph_config
         if self.from_pdf:
             if not file_path:
                 raise PipelineDefinitionError(
                     "Expected 'file_path' argument when 'from_pdf' is True."
                 )
-            run_params["pdf_loader"] = {"filepath": file_path}
+            run_params["pdf_loader"]["filepath"] = file_path
+            run_params["pdf_loader"]["metadata"] = user_input.get("document_metadata")
         else:
             if not text:
                 raise PipelineDefinitionError(
                     "Expected 'text' argument when 'from_pdf' is False."
                 )
-            run_params["splitter"] = {"text": text}
+            run_params["splitter"]["text"] = text
             # Add full text to schema component for automatic schema extraction
             if not self.has_user_provided_schema():
-                run_params["schema"] = {"text": text}
+                run_params["schema"]["text"] = text
+            run_params["extractor"]["document_info"] = dict(
+                path=user_input.get(
+                    "file_path",
+                )
+                or "document.txt",
+                metadata=user_input.get("document_metadata"),
+                document_type="inline_text",
+            )
         return run_params

src/neo4j_graphrag/experimental/pipeline/kg_builder.py

@@ -145,16 +145,26 @@ def __init__(
         self.runner = PipelineRunner.from_config(config)
 
     async def run_async(
-        self, file_path: Optional[str] = None, text: Optional[str] = None
+        self,
+        file_path: Optional[str] = None,
+        text: Optional[str] = None,
+        document_metadata: Optional[dict[str, Any]] = None,
     ) -> PipelineResult:
         """
         Asynchronously runs the knowledge graph building process.
 
         Args:
-            file_path (Optional[str]): The path to the PDF file to process. Required if `from_pdf` is True.
+            file_path (Optional[str]): The path to the PDF file to process. Required if `from_pdf` is True. If `from_pdf` is False, can be used to set the Document node path property.
             text (Optional[str]): The text content to process. Required if `from_pdf` is False.
+            document_metadata (Optional[dict[str, Any]]): The metadata to attach to the document.
 
         Returns:
             PipelineResult: The result of the pipeline execution.
         """
-        return await self.runner.run({"file_path": file_path, "text": text})
+        return await self.runner.run(
+            {
+                "file_path": file_path,
+                "text": text,
+                "document_metadata": document_metadata,
+            }
+        )

tests/unit/experimental/components/test_lexical_graph_builder.py

@@ -78,7 +78,11 @@ async def test_lexical_graph_builder_run_with_document() -> None:
                 TextChunk(text="text chunk 1", index=1),
             ]
         ),
-        document_info=DocumentInfo(path="test_lexical_graph", uid=doc_uid),
+        document_info=DocumentInfo(
+            path="test_lexical_graph",
+            uid=doc_uid,
+            document_type="my_type",
+        ),
     )
     assert isinstance(result, GraphResult)
     graph = result.graph
@@ -89,6 +93,7 @@ async def test_lexical_graph_builder_run_with_document() -> None:
     assert document.label == DEFAULT_DOCUMENT_NODE_LABEL
     assert document.properties["path"] == "test_lexical_graph"
     assert document.properties["createdAt"] is not None
+    assert document.properties["document_type"] == "my_type"
     chunk1 = nodes[1]
     assert chunk1.label == DEFAULT_CHUNK_NODE_LABEL
     chunk2 = nodes[2]