🔀 fix: change enums usage to work on all supported python versions (#329)

* Change enum usage to work on all supported python versions

* Fix flake8 error

* Version and changelog update; neo4j string enum fix

* Enable python 3.11 and 3.12 for unit_tests workflow

* Fix python 3.12 flake8 check

* Fix nltk download

* Enable python 3.11 and 3.12 for src and dst integration tests

* Revert "Enable python 3.11 and 3.12 for src and dst integration tests"

This reverts commit fc40dbc.

* Revert "Fix python 3.12 flake8 check"

This reverts commit 709647e.

* Revert "Enable python 3.11 and 3.12 for unit_tests workflow"

This reverts commit bbad641.

* Fix neo4j tests

* Revert "Fix neo4j tests"

This reverts commit 7de281c.

---------

Co-authored-by: Filip Knefel <[email protected]>
Co-authored-by: Marek Połom <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.3.14-dev0
+
+### Fixes
+
+* **Fix Neo4j Uploader string enum error**
+
 ## 0.3.13
 
 ### Fixes

test/integration/connectors/test_neo4j.py

@@ -199,13 +199,15 @@ async def validate_uploaded_graph(upload_file: Path):
     try:
         nodes_count = len((await driver.execute_query("MATCH (n) RETURN n"))[0])
         chunk_nodes_count = len(
-            (await driver.execute_query(f"MATCH (n: {Label.CHUNK}) RETURN n"))[0]
+            (await driver.execute_query(f"MATCH (n: {Label.CHUNK.value}) RETURN n"))[0]
         )
         document_nodes_count = len(
-            (await driver.execute_query(f"MATCH (n: {Label.DOCUMENT}) RETURN n"))[0]
+            (await driver.execute_query(f"MATCH (n: {Label.DOCUMENT.value}) RETURN n"))[0]
         )
         element_nodes_count = len(
-            (await driver.execute_query(f"MATCH (n: {Label.UNSTRUCTURED_ELEMENT}) RETURN n"))[0]
+            (await driver.execute_query(f"MATCH (n: {Label.UNSTRUCTURED_ELEMENT.value}) RETURN n"))[
+                0
+            ]
         )
         with check:
             assert nodes_count == expected_nodes_count
@@ -217,12 +219,18 @@ async def validate_uploaded_graph(upload_file: Path):
             assert element_nodes_count == expected_element_count
 
         records, _, _ = await driver.execute_query(
-            f"MATCH ()-[r:{Relationship.PART_OF_DOCUMENT}]->(:{Label.DOCUMENT}) RETURN r"
+            f"""
+            MATCH ()-[r:{Relationship.PART_OF_DOCUMENT.value}]->(:{Label.DOCUMENT.value})
+            RETURN r
+            """
         )
         part_of_document_count = len(records)
 
         records, _, _ = await driver.execute_query(
-            f"MATCH (:{Label.CHUNK})-[r:{Relationship.NEXT_CHUNK}]->(:{Label.CHUNK}) RETURN r"
+            f"""
+            MATCH (:{Label.CHUNK.value})-[r:{Relationship.NEXT_CHUNK.value}]->(:{Label.CHUNK.value})
+            RETURN r
+            """
         )
         next_chunk_count = len(records)
 

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.13"  # pragma: no cover
+__version__ = "0.3.14-dev0"  # pragma: no cover

unstructured_ingest/v2/processes/connectors/neo4j.py

@@ -105,7 +105,7 @@ def run(  # type: ignore
         output_filepath.parent.mkdir(parents=True, exist_ok=True)
 
         with open(output_filepath, "w") as file:
-            json.dump(_GraphData.from_nx(nx_graph).model_dump(), file, indent=4)
+            file.write(_GraphData.from_nx(nx_graph).model_dump_json())
 
         return output_filepath
 
@@ -196,7 +196,7 @@ def from_nx(cls, nx_graph: "MultiDiGraph") -> _GraphData:
 
 
 class _Node(BaseModel):
-    model_config = ConfigDict(use_enum_values=True)
+    model_config = ConfigDict()
 
     id_: str = Field(default_factory=lambda: str(uuid.uuid4()))
     labels: list[Label] = Field(default_factory=list)
@@ -207,20 +207,20 @@ def __hash__(self):
 
 
 class _Edge(BaseModel):
-    model_config = ConfigDict(use_enum_values=True)
+    model_config = ConfigDict()
 
     source_id: str
     destination_id: str
     relationship: Relationship
 
 
-class Label(str, Enum):
+class Label(Enum):
     UNSTRUCTURED_ELEMENT = "UnstructuredElement"
     CHUNK = "Chunk"
     DOCUMENT = "Document"
 
 
-class Relationship(str, Enum):
+class Relationship(Enum):
     PART_OF_DOCUMENT = "PART_OF_DOCUMENT"
     PART_OF_CHUNK = "PART_OF_CHUNK"
     NEXT_CHUNK = "NEXT_CHUNK"
@@ -263,23 +263,23 @@ async def run_async(self, path: Path, file_data: FileData, **kwargs) -> None:  #
     async def _create_uniqueness_constraints(self, client: AsyncDriver) -> None:
         for label in Label:
             logger.info(
-                f"Adding id uniqueness constraint for nodes labeled '{label}'"
+                f"Adding id uniqueness constraint for nodes labeled '{label.value}'"
                 " if it does not already exist."
             )
-            constraint_name = f"{label.lower()}_id"
+            constraint_name = f"{label.value.lower()}_id"
             await client.execute_query(
                 f"""
                 CREATE CONSTRAINT {constraint_name} IF NOT EXISTS
-                FOR (n: {label}) REQUIRE n.id IS UNIQUE
+                FOR (n: {label.value}) REQUIRE n.id IS UNIQUE
                 """
             )
 
     async def _delete_old_data_if_exists(self, file_data: FileData, client: AsyncDriver) -> None:
         logger.info(f"Deleting old data for the record '{file_data.identifier}' (if present).")
         _, summary, _ = await client.execute_query(
             f"""
-            MATCH (n: {Label.DOCUMENT} {{id: $identifier}})
-            MATCH (n)--(m: {Label.CHUNK}|{Label.UNSTRUCTURED_ELEMENT})
+            MATCH (n: {Label.DOCUMENT.value} {{id: $identifier}})
+            MATCH (n)--(m: {Label.CHUNK.value}|{Label.UNSTRUCTURED_ELEMENT.value})
             DETACH DELETE m""",
             identifier=file_data.identifier,
         )
@@ -349,7 +349,7 @@ async def _execute_queries(
 
     @staticmethod
     def _create_nodes_query(nodes: list[_Node], labels: tuple[Label, ...]) -> tuple[str, dict]:
-        labels_string = ", ".join(labels)
+        labels_string = ", ".join([label.value for label in labels])
         logger.info(f"Preparing MERGE query for {len(nodes)} nodes labeled '{labels_string}'.")
         query_string = f"""
             UNWIND $nodes AS node
@@ -366,7 +366,7 @@ def _create_edges_query(edges: list[_Edge], relationship: Relationship) -> tuple
             UNWIND $edges AS edge
             MATCH (u {{id: edge.source}})
             MATCH (v {{id: edge.destination}})
-            MERGE (u)-[:{relationship}]->(v)
+            MERGE (u)-[:{relationship.value}]->(v)
             """
         parameters = {
             "edges": [