From 5aefedf3e3466515fd383398ee6a5977fa100c3d Mon Sep 17 00:00:00 2001 From: antejavor Date: Wed, 22 Jul 2026 08:25:10 +0200 Subject: [PATCH] unstructured2graph: derive entity workspace instead of hardcoding "base" connect_chunks_to_entities was called with the literal "base", which is just LightRAG's default-workspace fallback duplicated as a magic string. Add an entity_workspace param (default None) that auto-derives the resolved value from lightrag_wrapper.get_lightrag().chunk_entity_relation_graph.workspace, falling back to "base" with a logged warning if that ever fails. Fixes #235 --- .../src/unstructured2graph/loaders.py | 13 +++- unstructured2graph/tests/test_loaders.py | 60 ++++++++++++++++++- 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/unstructured2graph/src/unstructured2graph/loaders.py b/unstructured2graph/src/unstructured2graph/loaders.py index 58fa5c63..2dae5395 100644 --- a/unstructured2graph/src/unstructured2graph/loaders.py +++ b/unstructured2graph/src/unstructured2graph/loaders.py @@ -114,6 +114,7 @@ async def from_unstructured( lightrag_wrapper: MemgraphLightRAGWrapper, only_chunks: bool = False, link_chunks: bool = False, + entity_workspace: str | None = None, partition_kwargs: dict[str, Any] | None = None, ): """ @@ -124,6 +125,9 @@ async def from_unstructured( lightrag_wrapper: MemgraphLightRAGWrapper instance (requires lightrag-memgraph) only_chunks: If True, only create chunk nodes without LightRAG processing link_chunks: If True, link chunks in order with NEXT relationship + entity_workspace: Node label LightRAG entities were written under. If None + (default), auto-derived from lightrag_wrapper's resolved LightRAG + workspace, falling back to "base" if that fails. partition_kwargs: Additional keyword arguments to pass to unstructured's partition function (e.g., strategy, languages, pdf_infer_table_structure, ocr_languages, headers, ssl_verify, etc.) @@ -133,6 +137,13 @@ async def from_unstructured( # 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. + resolved_entity_workspace = entity_workspace + if not only_chunks and resolved_entity_workspace is None: + try: + resolved_entity_workspace = lightrag_wrapper.get_lightrag().chunk_entity_relation_graph.workspace + except Exception as e: + logger.warning(f"Could not auto-derive LightRAG entity workspace, falling back to 'base': {e}") + resolved_entity_workspace = "base" chunked_documents = make_chunks(sources, partition_kwargs=partition_kwargs) total_chunks = sum(len(document.chunks) for document in chunked_documents) start_time = time.time() @@ -161,7 +172,7 @@ async def from_unstructured( if not only_chunks: await lightrag_wrapper.ainsert(input=chunk.text, file_paths=[chunk.hash]) if not only_chunks: - connect_chunks_to_entities(memgraph, "Chunk", "base") + connect_chunks_to_entities(memgraph, "Chunk", resolved_entity_workspace) processed_chunks += len(document.chunks) elapsed_time = time.time() - start_time diff --git a/unstructured2graph/tests/test_loaders.py b/unstructured2graph/tests/test_loaders.py index 9bbcc7d3..7561cb7b 100644 --- a/unstructured2graph/tests/test_loaders.py +++ b/unstructured2graph/tests/test_loaders.py @@ -1,14 +1,26 @@ """Simple test for unstructured2graph loaders.""" import os +from unittest.mock import AsyncMock, MagicMock, patch 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__)) +def _fake_document(): + return ChunkedDocument(chunks=[Chunk(text="a", hash="h1")], source="fake.txt") + + +def _lightrag_wrapper_with_workspace(workspace): + wrapper = MagicMock() + wrapper.ainsert = AsyncMock() + wrapper.get_lightrag.return_value.chunk_entity_relation_graph.workspace = workspace + return wrapper + + def test_parse_source_with_simple_text(tmp_path): """Test that parse_source can handle a simple text file.""" # Create a simple text file @@ -74,6 +86,52 @@ def test_partition_kwargs_passed_through(tmp_path): assert isinstance(chunks, list) +@pytest.mark.asyncio +async def test_entity_workspace_explicit_override_wins(): + memgraph = MagicMock() + lightrag_wrapper = _lightrag_wrapper_with_workspace("auto-derived") + + with ( + patch("unstructured2graph.loaders.make_chunks", return_value=[_fake_document()]), + patch("unstructured2graph.loaders.connect_chunks_to_entities") as mock_connect, + ): + await from_unstructured( + ["fake.txt"], memgraph, lightrag_wrapper, only_chunks=False, entity_workspace="explicit" + ) + + mock_connect.assert_called_once_with(memgraph, "Chunk", "explicit") + + +@pytest.mark.asyncio +async def test_entity_workspace_auto_derived_from_lightrag_wrapper(): + memgraph = MagicMock() + lightrag_wrapper = _lightrag_wrapper_with_workspace("tenant-42") + + with ( + patch("unstructured2graph.loaders.make_chunks", return_value=[_fake_document()]), + patch("unstructured2graph.loaders.connect_chunks_to_entities") as mock_connect, + ): + await from_unstructured(["fake.txt"], memgraph, lightrag_wrapper, only_chunks=False) + + mock_connect.assert_called_once_with(memgraph, "Chunk", "tenant-42") + + +@pytest.mark.asyncio +async def test_entity_workspace_falls_back_to_base_when_auto_derive_fails(): + memgraph = MagicMock() + lightrag_wrapper = MagicMock() + lightrag_wrapper.ainsert = AsyncMock() + lightrag_wrapper.get_lightrag.side_effect = RuntimeError("not initialized") + + with ( + patch("unstructured2graph.loaders.make_chunks", return_value=[_fake_document()]), + patch("unstructured2graph.loaders.connect_chunks_to_entities") as mock_connect, + ): + await from_unstructured(["fake.txt"], memgraph, lightrag_wrapper, only_chunks=False) + + mock_connect.assert_called_once_with(memgraph, "Chunk", "base") + + @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")