From 47e25f950545fab31a955453f1c68991fecd2b2b Mon Sep 17 00:00:00 2001 From: antejavor Date: Wed, 22 Jul 2026 08:29:33 +0200 Subject: [PATCH] unstructured2graph: make chunk ingestion idempotent and self-provisioning from_unstructured had two load-bearing TODOs: Chunk nodes were inserted via CREATE (duplicating on every re-run) and a non-unique index on Chunk.hash was left for the caller to create themselves, so nothing actually prevented hash collisions. - create_nodes_from_list gains an optional merge_key: when given, nodes are upserted via MERGE ... ON CREATE SET instead of duplicated via CREATE. - Add create_unique_constraint, and have from_unstructured call it idempotently on Chunk.hash instead of requiring the caller to call create_index first. Ingestion is now safe to retry/rerun over the same sources without manual dedup or a mandatory setup step. Fixes #233 --- .../src/unstructured2graph/__init__.py | 2 + .../src/unstructured2graph/loaders.py | 6 +- .../src/unstructured2graph/memgraph.py | 46 ++++++++++++-- unstructured2graph/tests/test_loaders.py | 19 +++++- unstructured2graph/tests/test_memgraph.py | 60 +++++++++++++++++++ 5 files changed, 123 insertions(+), 10 deletions(-) create mode 100644 unstructured2graph/tests/test_memgraph.py diff --git a/unstructured2graph/src/unstructured2graph/__init__.py b/unstructured2graph/src/unstructured2graph/__init__.py index db78e5f5..cd48a6f1 100644 --- a/unstructured2graph/src/unstructured2graph/__init__.py +++ b/unstructured2graph/src/unstructured2graph/__init__.py @@ -18,6 +18,7 @@ create_index, create_label_index, create_nodes_from_list, + create_unique_constraint, create_vector_search_index, link_nodes_in_order, ) @@ -31,6 +32,7 @@ "create_index", "create_label_index", "create_nodes_from_list", + "create_unique_constraint", "create_vector_search_index", "from_unstructured", "link_nodes_in_order", diff --git a/unstructured2graph/src/unstructured2graph/loaders.py b/unstructured2graph/src/unstructured2graph/loaders.py index 58fa5c63..16b00e68 100644 --- a/unstructured2graph/src/unstructured2graph/loaders.py +++ b/unstructured2graph/src/unstructured2graph/loaders.py @@ -16,6 +16,7 @@ from .memgraph import ( connect_chunks_to_entities, create_nodes_from_list, + create_unique_constraint, link_nodes_in_order, ) @@ -129,10 +130,9 @@ async def from_unstructured( ocr_languages, headers, ssl_verify, etc.) """ - # TODO(gitbuda): Create all required indexes. - # TODO(gitbuda): Make the calls idempotent. # TODO(gitbuda): Implement batching on the Cypher side as well under memgraph.compute_embeddings # NOTE: LightRAG uses { source_id: "chunk-ID..." } to reference its chunks. + create_unique_constraint(memgraph, "Chunk", "hash") chunked_documents = make_chunks(sources, partition_kwargs=partition_kwargs) total_chunks = sum(len(document.chunks) for document in chunked_documents) start_time = time.time() @@ -147,7 +147,7 @@ async def from_unstructured( for chunk in document.chunks: logger.debug(f"Chunk: {chunk.hash} - {chunk.text}") memgraph_node_props.append({"hash": chunk.hash, "text": chunk.text}) - create_nodes_from_list(memgraph, memgraph_node_props, "Chunk", 100) + create_nodes_from_list(memgraph, memgraph_node_props, "Chunk", 100, merge_key="hash") if link_chunks: hash_pairs = [ diff --git a/unstructured2graph/src/unstructured2graph/memgraph.py b/unstructured2graph/src/unstructured2graph/memgraph.py index 65954c04..2ba00b35 100644 --- a/unstructured2graph/src/unstructured2graph/memgraph.py +++ b/unstructured2graph/src/unstructured2graph/memgraph.py @@ -6,9 +6,21 @@ logger = logging.getLogger(__name__) -def create_nodes_from_list(memgraph: Memgraph, nodes: list[dict], node_label: str, batch_size: int) -> None: +def create_nodes_from_list( + memgraph: Memgraph, + nodes: list[dict], + node_label: str, + batch_size: int, + merge_key: str | None = None, +) -> None: """ Import data from the given list of dictionaries to Memgraph by batching. + + Args: + merge_key: If given, nodes are upserted via MERGE keyed on this + property (a no-op for nodes that already exist), making re-runs + over the same data safe. If None (default), nodes are inserted + via CREATE, so re-running over the same data duplicates them. """ if not nodes: logger.warning(f"No nodes provided to create_nodes_from_list for label {node_label}") @@ -17,11 +29,20 @@ def create_nodes_from_list(memgraph: Memgraph, nodes: list[dict], node_label: st num_nodes = len(nodes) max_retries = 3 retry_delay = 3 - properties_string = ", ".join([f"{key}: data.{key}" for key in nodes[0]]) - insert_query = f""" - UNWIND $batch AS data - CREATE (n:{node_label} {{{properties_string}}}) - """ + if merge_key: + set_keys = [key for key in nodes[0] if key != merge_key] + set_string = ", ".join(f"n.{key} = data.{key}" for key in set_keys) + on_create_clause = f" ON CREATE SET {set_string}" if set_string else "" + insert_query = f""" + UNWIND $batch AS data + MERGE (n:{node_label} {{{merge_key}: data.{merge_key}}}){on_create_clause} + """ + else: + properties_string = ", ".join([f"{key}: data.{key}" for key in nodes[0]]) + insert_query = f""" + UNWIND $batch AS data + CREATE (n:{node_label} {{{properties_string}}}) + """ for offset in range(0, num_nodes, batch_size): batch_nodes = nodes[offset : offset + batch_size] for attempt in range(max_retries): @@ -74,6 +95,19 @@ def create_index(memgraph: Memgraph, label: str, property: str): logger.warning(f"Error creating index: {e}") +def create_unique_constraint(memgraph: Memgraph, label: str, property: str): + """ + Idempotently ensure a uniqueness constraint on :label(property). Unlike + CREATE INDEX, this actually rejects duplicate values instead of merely + speeding up lookups, and is safe to call on every run. + """ + try: + memgraph.query(f"CREATE CONSTRAINT ON (n:{label}) ASSERT n.{property} IS UNIQUE;") + logger.info(f"Ensured uniqueness constraint on :{label}({property})") + except Exception as e: + logger.warning(f"Error creating uniqueness constraint on :{label}({property}): {e}") + + def create_label_index(memgraph: Memgraph, label: str): """ Create a label index for efficient node lookups by label. diff --git a/unstructured2graph/tests/test_loaders.py b/unstructured2graph/tests/test_loaders.py index 9bbcc7d3..e8be2b30 100644 --- a/unstructured2graph/tests/test_loaders.py +++ b/unstructured2graph/tests/test_loaders.py @@ -1,10 +1,11 @@ """Simple test for unstructured2graph loaders.""" import os +from unittest.mock import MagicMock import pytest -from unstructured2graph import Chunk, ChunkedDocument, make_chunks, parse_source +from unstructured2graph import Chunk, ChunkedDocument, from_unstructured, make_chunks, parse_source SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__)) @@ -74,6 +75,22 @@ def test_partition_kwargs_passed_through(tmp_path): assert isinstance(chunks, list) +@pytest.mark.asyncio +async def test_from_unstructured_creates_unique_constraint_and_upserts_chunks(tmp_path): + """from_unstructured must self-provision its Chunk.hash constraint and + upsert (not duplicate-insert) Chunk nodes, so re-runs are safe.""" + test_file = tmp_path / "test.txt" + test_file.write_text("Some content for idempotent ingestion.") + memgraph = MagicMock() + lightrag_wrapper = MagicMock() + + await from_unstructured([str(test_file)], memgraph, lightrag_wrapper, only_chunks=True) + + queries = [call.args[0] for call in memgraph.query.call_args_list] + assert any("CONSTRAINT" in q and "Chunk" in q and "hash" in q for q in queries) + assert any("MERGE (n:Chunk {hash: data.hash})" in q for q in queries) + + @pytest.mark.skip(reason="Requires sample-data files and network access - run locally with full deps") def test_chunking_of_different_sources(): pypdf_samples_dir = os.path.join(SCRIPT_DIR, "..", "sample-data", "pdf", "sample-files") diff --git a/unstructured2graph/tests/test_memgraph.py b/unstructured2graph/tests/test_memgraph.py new file mode 100644 index 00000000..777a24f9 --- /dev/null +++ b/unstructured2graph/tests/test_memgraph.py @@ -0,0 +1,60 @@ +"""Unit tests for unstructured2graph.memgraph Cypher-building helpers.""" + +from unittest.mock import MagicMock + +from unstructured2graph.memgraph import create_nodes_from_list, create_unique_constraint + + +def test_create_nodes_from_list_defaults_to_create(): + """Without merge_key, behavior is unchanged: a plain CREATE per node.""" + memgraph = MagicMock() + + create_nodes_from_list(memgraph, [{"hash": "h1", "text": "a"}], "Chunk", 100) + + query = memgraph.query.call_args[0][0] + assert "CREATE (n:Chunk" in query + assert "MERGE" not in query + + +def test_create_nodes_from_list_merges_on_merge_key(): + """With merge_key, re-running over the same data is a no-op (MERGE).""" + memgraph = MagicMock() + + create_nodes_from_list(memgraph, [{"hash": "h1", "text": "a"}], "Chunk", 100, merge_key="hash") + + query = memgraph.query.call_args[0][0] + assert "MERGE (n:Chunk {hash: data.hash})" in query + assert "ON CREATE SET n.text = data.text" in query + assert "CREATE (n:Chunk {" not in query + + +def test_create_nodes_from_list_merge_key_only_property(): + """merge_key as the only property still produces valid Cypher (no dangling ON CREATE SET).""" + memgraph = MagicMock() + + create_nodes_from_list(memgraph, [{"hash": "h1"}], "Chunk", 100, merge_key="hash") + + query = memgraph.query.call_args[0][0] + assert "MERGE (n:Chunk {hash: data.hash})" in query + assert "ON CREATE SET" not in query + + +def test_create_unique_constraint_issues_constraint_query(): + memgraph = MagicMock() + + create_unique_constraint(memgraph, "Chunk", "hash") + + query = memgraph.query.call_args[0][0] + assert "CONSTRAINT" in query + assert "Chunk" in query + assert "hash" in query + assert "UNIQUE" in query + + +def test_create_unique_constraint_is_idempotent_on_repeated_calls(): + """A second call (constraint already exists) must not raise.""" + memgraph = MagicMock() + memgraph.query.side_effect = [None, Exception("constraint already exists")] + + create_unique_constraint(memgraph, "Chunk", "hash") + create_unique_constraint(memgraph, "Chunk", "hash") # should log a warning, not raise