neo4j NER+RE support (#464)

Author: vangheem

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 1.0.8
+
+### Enhancements
+
+* **Update Neo4J Entity Support** to support NER + RE(Relationship extraction)
+
 ## 1.0.7
 
 ### Fixes

pyproject.toml

@@ -143,6 +143,7 @@ test = [
     "vertexai",
     "pyiceberg",
     "pyarrow",
+    "networkx"
 ]
 # Add constraints needed for CI
 ci = [

test/unit/connectors/test_neo4j.py

@@ -0,0 +1,72 @@
+import networkx as nx
+import pytest
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG, GRAPH_DB_TAG
+from unstructured_ingest.processes.connectors import neo4j
+from unstructured_ingest.processes.connectors.neo4j import (
+    CONNECTOR_TYPE,
+    Label,
+    Neo4jUploadStager,
+    Relationship,
+)
+
+
+@pytest.mark.tags(DESTINATION_TAG, CONNECTOR_TYPE, GRAPH_DB_TAG)
+def test_neo4j_stager_with_entities_no_re():
+    stager = Neo4jUploadStager()
+
+    graph = nx.MultiDiGraph()
+    doc_node = neo4j._Node(id_="root", properties={"id": "root"}, labels=[Label.DOCUMENT])
+    graph.add_node(doc_node)
+
+    element_node = neo4j._Node(
+        id_="element_id",
+        properties={"id": "element_id", "text": "This is a test"},
+        labels=[Label.UNSTRUCTURED_ELEMENT],
+    )
+    graph.add_edge(element_node, doc_node, relationship=Relationship.PART_OF_DOCUMENT)
+
+    stager._add_entity_data(
+        {"metadata": {"entities": [{"type": "PERSON", "entity": "Steve Jobs"}]}},
+        graph,
+        element_node,
+    )
+    assert len(graph.nodes) == 4
+    assert len(graph.edges) == 3
+
+
+@pytest.mark.tags(DESTINATION_TAG, CONNECTOR_TYPE, GRAPH_DB_TAG)
+def test_neo4j_stager_with_entities():
+    stager = Neo4jUploadStager()
+
+    graph = nx.MultiDiGraph()
+    doc_node = neo4j._Node(id_="root", properties={"id": "root"}, labels=[Label.DOCUMENT])
+    graph.add_node(doc_node)
+
+    element_node = neo4j._Node(
+        id_="element_id",
+        properties={"id": "element_id", "text": "This is a test"},
+        labels=[Label.UNSTRUCTURED_ELEMENT],
+    )
+    graph.add_edge(element_node, doc_node, relationship=Relationship.PART_OF_DOCUMENT)
+
+    stager._add_entity_data(
+        {
+            "metadata": {
+                "entities": {
+                    "items": [
+                        {"type": "PERSON", "entity": "Steve Jobs"},
+                        {"type": "COMPANY", "entity": "Apple"},
+                    ],
+                    "relationships": [
+                        {"from": "Steve Jobs", "to": "Apple", "relationship": "founded"},
+                        {"from": "Steve Jobs", "to": "Apple", "relationship": "worked_for"},
+                    ],
+                }
+            }
+        },
+        graph,
+        element_node,
+    )
+    assert len(graph.nodes) == 6
+    assert len(graph.edges) == 7

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.0.7"  # pragma: no cover
+__version__ = "1.0.8"  # pragma: no cover

unstructured_ingest/data_types/entities.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel, Field
+
+
+class Entity(BaseModel):
+    type: str
+    entity: str
+
+
+class EntityRelationship(BaseModel):
+    to: str
+    from_: str = Field(..., alias="from")
+    relationship: str
+
+
+class EntitiesData(BaseModel):
+    items: list[Entity] = Field(default_factory=list)
+    relationships: list[EntityRelationship] = Field(default_factory=list)

unstructured_ingest/processes/connectors/neo4j.py

@@ -9,8 +9,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Literal, Optional
 
-from pydantic import BaseModel, ConfigDict, Field, Secret, field_validator
+from pydantic import BaseModel, ConfigDict, Field, Secret, ValidationError, field_validator
 
+from unstructured_ingest.data_types.entities import EntitiesData, Entity, EntityRelationship
 from unstructured_ingest.data_types.file_data import FileData
 from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.interfaces import (
@@ -97,7 +98,6 @@ def run(  # type: ignore
         **kwargs: Any,
     ) -> Path:
         elements = get_json_data(elements_filepath)
-
         nx_graph = self._create_lexical_graph(
             elements, self._create_document_node(file_data=file_data)
         )
@@ -109,28 +109,54 @@ def run(  # type: ignore
 
         return output_filepath
 
-    def _add_entities(self, element: dict, graph: "Graph", element_node: _Node) -> None:
-        entities = element.get("metadata", {}).get("entities", [])
-        if not entities:
-            return None
-        if not isinstance(entities, list):
-            return None
-
+    def _add_entities(self, entities: list[Entity], graph: "Graph", element_node: _Node) -> None:
         for entity in entities:
-            if not isinstance(entity, dict):
-                continue
-            if "entity" not in entity or "type" not in entity:
-                continue
             entity_node = _Node(
-                labels=[Label.ENTITY], properties={"id": entity["entity"]}, id_=entity["entity"]
+                labels=[Label.ENTITY], properties={"id": entity.entity}, id_=entity.entity
             )
             graph.add_edge(
                 entity_node,
-                _Node(labels=[Label.ENTITY], properties={"id": entity["type"]}, id_=entity["type"]),
+                _Node(labels=[Label.ENTITY], properties={"id": entity.type}, id_=entity.type),
                 relationship=Relationship.ENTITY_TYPE,
             )
             graph.add_edge(element_node, entity_node, relationship=Relationship.HAS_ENTITY)
 
+    def _add_entity_relationships(
+        self, relationships: list[EntityRelationship], graph: "Graph"
+    ) -> None:
+        for relationship in relationships:
+            from_node = _Node(
+                labels=[Label.ENTITY],
+                properties={"id": relationship.from_},
+                id_=relationship.from_,
+            )
+            to_node = _Node(
+                labels=[Label.ENTITY], properties={"id": relationship.to}, id_=relationship.to
+            )
+            graph.add_edge(from_node, to_node, relationship=relationship.relationship)
+
+    def _add_entity_data(self, element: dict, graph: "Graph", element_node: _Node) -> None:
+        entities = element.get("metadata", {}).get("entities", {})
+        if not entities:
+            return None
+        try:
+            if isinstance(entities, list):
+                self._add_entities(
+                    [Entity.model_validate(e) for e in entities if isinstance(e, dict)],
+                    graph,
+                    element_node,
+                )
+            elif isinstance(entities, dict):
+                entity_data = EntitiesData.model_validate(entities)
+                self._add_entities(entity_data.items, graph, element_node)
+                self._add_entity_relationships(entity_data.relationships, graph)
+        except ValidationError:
+            logger.warning(
+                "Failed to add entities to the graph. "
+                "Please check the format of the entities in the input data."
+            )
+        return None
+
     def _create_lexical_graph(self, elements: list[dict], document_node: _Node) -> "Graph":
         import networkx as nx
 
@@ -149,7 +175,7 @@ def _create_lexical_graph(self, elements: list[dict], document_node: _Node) -> "
             previous_node = element_node
             graph.add_edge(element_node, document_node, relationship=Relationship.PART_OF_DOCUMENT)
 
-            self._add_entities(element, graph, element_node)
+            self._add_entity_data(element, graph, element_node)
 
             if self._is_chunk(element):
                 for origin_element in format_and_truncate_orig_elements(element, include_text=True):
@@ -165,7 +191,7 @@ def _create_lexical_graph(self, elements: list[dict], document_node: _Node) -> "
                         document_node,
                         relationship=Relationship.PART_OF_DOCUMENT,
                     )
-                    self._add_entities(origin_element, graph, origin_element_node)
+                    self._add_entity_data(origin_element, graph, origin_element_node)
 
         return graph
 
@@ -208,7 +234,9 @@ def from_nx(cls, nx_graph: "MultiDiGraph") -> _GraphData:
             _Edge(
                 source=u,
                 destination=v,
-                relationship=Relationship(data_dict["relationship"]),
+                relationship=Relationship(data_dict["relationship"])
+                if data_dict["relationship"] in Relationship
+                else data_dict["relationship"],
             )
             for u, v, data_dict in nx_graph.edges(data=True)
         ]
@@ -242,7 +270,7 @@ class _Edge(BaseModel):
 
     source: _Node
     destination: _Node
-    relationship: Relationship
+    relationship: Relationship | str
 
 
 class Label(Enum):
@@ -380,7 +408,7 @@ async def _merge_graph(self, graph_data: _GraphData, client: AsyncDriver) -> Non
         )
         logger.info(f"Finished merging {len(graph_data.nodes)} graph nodes.")
 
-        edges_by_relationship: defaultdict[tuple[Relationship, Label, Label], list[_Edge]] = (
+        edges_by_relationship: defaultdict[tuple[Relationship | str, Label, Label], list[_Edge]] = (
             defaultdict(list)
         )
         for edge in graph_data.edges:
@@ -463,16 +491,19 @@ def _create_nodes_query(nodes: list[_Node], label: Label) -> tuple[str, dict]:
     @staticmethod
     def _create_edges_query(
         edges: list[_Edge],
-        relationship: Relationship,
+        relationship: Relationship | str,
         source_label: Label,
         destination_label: Label,
     ) -> tuple[str, dict]:
         logger.info(f"Preparing MERGE query for {len(edges)} {relationship} relationships.")
+        relationship = (
+            relationship.value if isinstance(relationship, Relationship) else relationship
+        )
         query_string = f"""
             UNWIND $edges AS edge
             MATCH (u: `{source_label.value}` {{id: edge.source}})
             MATCH (v: `{destination_label.value}` {{id: edge.destination}})
-            MERGE (u)-[:`{relationship.value}`]->(v)
+            MERGE (u)-[:`{relationship}`]->(v)
             """
         parameters = {
             "edges": [