unstructured2graph: derive entity workspace instead of hardcoding "base"

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

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",