unstructured-ingest 1.0.6.dev0__py3-none-any.whl → 1.0.8__py3-none-any.whl

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.0.6-dev0"  # 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 enum import Enum
 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 @@ class Neo4jUploadStager(UploadStager):
         **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 @@ class Neo4jUploadStager(UploadStager):
 
         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 @@ class Neo4jUploadStager(UploadStager):
             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 @@ class Neo4jUploadStager(UploadStager):
                         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 @@ class _GraphData(BaseModel):
             _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 @@ class Neo4jUploader(Uploader):
         )
         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 @@ class Neo4jUploader(Uploader):
     @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": [

{unstructured_ingest-1.0.6.dev0.dist-info → unstructured_ingest-1.0.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unstructured_ingest
-Version: 1.0.6.dev0
+Version: 1.0.8
 Summary: Local ETL data pipeline to get data RAG ready
 Author-email: Unstructured Technologies <devops@unstructuredai.io>
 License-Expression: Apache-2.0

{unstructured_ingest-1.0.6.dev0.dist-info → unstructured_ingest-1.0.8.dist-info}/RECORD

@@ -1,5 +1,5 @@
 unstructured_ingest/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
-unstructured_ingest/__version__.py,sha256=f9THPBYTYgzAGHqOPz3K_-VIVaQAVKMJgCpBy7w0j_k,47
+unstructured_ingest/__version__.py,sha256=Ca4MzLfEjKrGXx21Kyt3Ve65pu59qVvEIU_io-qxQ9o,42
 unstructured_ingest/error.py,sha256=qDncnJgbf5ils956RcO2CGlAKYDT5OaEM9Clv1JVTNc,1448
 unstructured_ingest/errors_v2.py,sha256=9RuRCi7lbDxCguDz07y5RiHoQiFIOWwOD7xqzJ2B3Yw,436
 unstructured_ingest/logger.py,sha256=7e_7UeK6hVOd5BQ6i9NzRUAPCS_DF839Y8TjUDywraY,1428
@@ -19,6 +19,7 @@ unstructured_ingest/cli/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5N
 unstructured_ingest/cli/utils/click.py,sha256=1_eJgrwS2DFBl1jZPLsj1vgVgR7agFBIEBe4A_n7mH4,7827
 unstructured_ingest/cli/utils/model_conversion.py,sha256=hMjAfOVvO1RXTDsw26mmersdncvddkb_rP9JTEgVVCw,7649
 unstructured_ingest/data_types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+unstructured_ingest/data_types/entities.py,sha256=ECc6EkZ5_ZUvK7uaALYOynfFmofIrHYIJZfb67hUIxA,371
 unstructured_ingest/data_types/file_data.py,sha256=J0RQa7YXhhxiLVzhPbF5Hl2nzSpxLFK9vrP6RTBWlSg,3833
 unstructured_ingest/embed/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unstructured_ingest/embed/azure_openai.py,sha256=_-I-nwd-wdCiKkSdYBL4UKrTZ2UPWsM_0T69fcObs_I,1707
@@ -77,7 +78,7 @@ unstructured_ingest/processes/connectors/kdbai.py,sha256=XhxYpKSAoFPBsDQWwNuLX03
 unstructured_ingest/processes/connectors/local.py,sha256=LluTLKv4g7FbJb4A6vuSxI9VhzKZuuQUpDS-cVNAQ2g,7426
 unstructured_ingest/processes/connectors/milvus.py,sha256=Jr9cul7By03tGAPFnFBoqncnNWwbhKd-qbmkuqnin8U,8908
 unstructured_ingest/processes/connectors/mongodb.py,sha256=1g_5bfbS6lah3nsOXqLAanR3zNYJ47_Njw_uV-uj3_U,14324
-unstructured_ingest/processes/connectors/neo4j.py,sha256=eAM2XWSLA5caKJmbcd7ctn2TapreIJEXRoHoxT1OZwA,18718
+unstructured_ingest/processes/connectors/neo4j.py,sha256=ztxvI9KY8RF5kYUuMGSzzN5mz7Fu_4Ai9P7dqCpJLc0,20267
 unstructured_ingest/processes/connectors/onedrive.py,sha256=VBkKlbJgR7uKlKTnjNybAw6ZawLKflDPpy2uVvgWYWw,19296
 unstructured_ingest/processes/connectors/outlook.py,sha256=zHM5frO7CqQG0-KcTyX49aZeSlsvVrl8kh_lR_ESgQw,9275
 unstructured_ingest/processes/connectors/pinecone.py,sha256=BdO1PS_Y6FOeL-7uPl-Eh6ij1wHOwMkopOzKQGQ9Ac0,13979
@@ -230,8 +231,8 @@ unstructured_ingest/utils/ndjson.py,sha256=nz8VUOPEgAFdhaDOpuveknvCU4x82fVwqE01q
 unstructured_ingest/utils/pydantic_models.py,sha256=BT_j15e4rX40wQbt8LUXbqfPhA3rJn1PHTI_G_A_EHY,1720
 unstructured_ingest/utils/string_and_date_utils.py,sha256=oXOI6rxXq-8ncbk7EoJK0WCcTXWj75EzKl8pfQMID3U,2522
 unstructured_ingest/utils/table.py,sha256=WZechczgVFvlodUWFcsnCGvBNh1xRm6hr0VbJTPxKAc,3669
-unstructured_ingest-1.0.6.dev0.dist-info/METADATA,sha256=2A8V0IxykBKTDpalgDOmkRw98MoN6whUPI0DwHkBuBc,8724
-unstructured_ingest-1.0.6.dev0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-unstructured_ingest-1.0.6.dev0.dist-info/entry_points.txt,sha256=gUAAFnjFPnBgThJSEbw0N5ZjxtaKlT1s9e05_arQrNw,70
-unstructured_ingest-1.0.6.dev0.dist-info/licenses/LICENSE.md,sha256=SxkKP_62uIAKb9mb1eH7FH4Kn2aYT09fgjKpJt5PyTk,11360
-unstructured_ingest-1.0.6.dev0.dist-info/RECORD,,
+unstructured_ingest-1.0.8.dist-info/METADATA,sha256=IzN7b_dpadQBKpp59jO7VfWzgQfJrF8ykGLo7epNMeY,8719
+unstructured_ingest-1.0.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+unstructured_ingest-1.0.8.dist-info/entry_points.txt,sha256=gUAAFnjFPnBgThJSEbw0N5ZjxtaKlT1s9e05_arQrNw,70
+unstructured_ingest-1.0.8.dist-info/licenses/LICENSE.md,sha256=SxkKP_62uIAKb9mb1eH7FH4Kn2aYT09fgjKpJt5PyTk,11360
+unstructured_ingest-1.0.8.dist-info/RECORD,,