diff --git a/unstructured2graph/src/unstructured2graph/__init__.py b/unstructured2graph/src/unstructured2graph/__init__.py index db47a27..78ea87f 100644 --- a/unstructured2graph/src/unstructured2graph/__init__.py +++ b/unstructured2graph/src/unstructured2graph/__init__.py @@ -18,6 +18,7 @@ create_label_index, create_nodes_from_list, create_property_index, + create_unique_constraint, create_vector_search_index, link_nodes_in_order, ) @@ -31,6 +32,7 @@ "create_label_index", "create_nodes_from_list", "create_property_index", + "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 d11b49a..3db41a9 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, ) @@ -136,10 +137,9 @@ async def from_unstructured( if not only_chunks and lightrag_wrapper is None: raise ValueError("lightrag_wrapper is required when only_chunks=False") - # 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") resolved_entity_workspace = entity_workspace if not only_chunks and resolved_entity_workspace is None: try: @@ -161,7 +161,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 c257ab5..ce21c9c 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_property_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 effaca1..4061e32 100644 --- a/unstructured2graph/tests/test_loaders.py +++ b/unstructured2graph/tests/test_loaders.py @@ -86,6 +86,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.asyncio async def test_entity_workspace_explicit_override_wins(): memgraph = MagicMock() diff --git a/unstructured2graph/tests/test_memgraph.py b/unstructured2graph/tests/test_memgraph.py new file mode 100644 index 0000000..777a24f --- /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