diff --git a/unstructured2graph/src/unstructured2graph/loaders.py b/unstructured2graph/src/unstructured2graph/loaders.py index eed9a65..d11b49a 100644 --- a/unstructured2graph/src/unstructured2graph/loaders.py +++ b/unstructured2graph/src/unstructured2graph/loaders.py @@ -114,6 +114,7 @@ async def from_unstructured( lightrag_wrapper: MemgraphLightRAGWrapper | None = None, only_chunks: bool = False, link_chunks: bool = False, + entity_workspace: str | None = None, partition_kwargs: dict[str, Any] | None = None, ): """ @@ -125,6 +126,9 @@ async def from_unstructured( Required unless only_chunks=True, since it's only used for entity extraction. 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.) @@ -136,6 +140,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() @@ -163,7 +174,7 @@ async def from_unstructured( if not only_chunks: for chunk in document.chunks: await lightrag_wrapper.ainsert(input=chunk.text, file_paths=[chunk.hash]) - 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 53543e1..effaca1 100644 --- a/unstructured2graph/tests/test_loaders.py +++ b/unstructured2graph/tests/test_loaders.py @@ -10,6 +10,17 @@ 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 @@ -75,13 +86,58 @@ 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.asyncio async def test_connect_chunks_to_entities_called_once_per_document(): """connect_chunks_to_entities is a full graph scan; it must run once per document, not once per chunk.""" memgraph = MagicMock() - lightrag_wrapper = MagicMock() - lightrag_wrapper.ainsert = AsyncMock() + lightrag_wrapper = _lightrag_wrapper_with_workspace("base") fake_document = ChunkedDocument( chunks=[Chunk(text="a", hash="h1"), Chunk(text="b", hash="h2"), Chunk(text="c", hash="h3")], source="fake.txt",