nucliadb 6.2.0.post2675__py3-none-any.whl → 6.2.1__py3-none-any.whl

nucliadb/common/nidx.py

@@ -21,7 +21,7 @@
 import os
 from typing import Optional
 
-from nidx_protos.nidx_pb2_grpc import NidxApiStub, NidxSearcherStub
+from nidx_protos.nidx_pb2_grpc import NidxApiStub, NidxIndexerStub, NidxSearcherStub
 
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.settings import settings
@@ -37,12 +37,10 @@ from nucliadb_utils.settings import FileBackendConfig, indexing_settings, storag
 from nucliadb_utils.storages.settings import settings as extended_storage_settings
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
-NIDX_ENABLED = bool(os.environ.get("NIDX_ENABLED"))
-
 
 class NidxUtility:
-    api_client = None
-    searcher_client = None
+    api_client: NidxApiStub
+    searcher_client: NidxSearcherStub
 
     async def initialize(self):
         raise NotImplementedError()
@@ -98,6 +96,9 @@ class NidxBindingUtility(NidxUtility):
 
         self.config = {
             "METADATA__DATABASE_URL": ingest_settings.driver_pg_url,
+            "SEARCHER__METADATA_REFRESH_INTERVAL": str(
+                indexing_settings.index_searcher_refresh_interval
+            ),
             **_storage_config("INDEXER", None),
             **_storage_config("STORAGE", "nidx"),
         }
@@ -123,16 +124,8 @@ class NidxBindingUtility(NidxUtility):
         self.binding.wait_for_sync()
 
 
-class NidxServiceUtility(NidxUtility):
-    """Implements Nidx utility connecting to the network service"""
-
+class NidxNatsIndexer:
     def __init__(self):
-        if indexing_settings.index_nidx_subject is None:
-            raise ValueError("INDEX_NIDX_SUBJECT needed for nidx utility")
-
-        if not settings.nidx_api_address or not settings.nidx_searcher_address:
-            raise ValueError("NIDX_API_ADDRESS and NIDX_SEARCHER_ADDRESS are required")
-
         self.nats_connection_manager = NatsConnectionManager(
             service_name="NidxIndexer",
             nats_servers=indexing_settings.index_jetstream_servers,
@@ -142,10 +135,6 @@ class NidxServiceUtility(NidxUtility):
 
     async def initialize(self):
         await self.nats_connection_manager.initialize()
-        self.api_client = NidxApiStub(get_traced_grpc_channel(settings.nidx_api_address, "nidx_api"))
-        self.searcher_client = NidxSearcherStub(
-            get_traced_grpc_channel(settings.nidx_searcher_address, "nidx_searcher")
-        )
 
     async def finalize(self):
         await self.nats_connection_manager.finalize()
@@ -158,18 +147,68 @@ class NidxServiceUtility(NidxUtility):
         return res.seq
 
 
-async def start_nidx_utility() -> Optional[NidxUtility]:
-    if not NIDX_ENABLED:
-        return None
+class NidxGrpcIndexer:
+    def __init__(self, address):
+        self.address = address
 
-    nidx = get_nidx()
+    async def initialize(self):
+        self.client = NidxIndexerStub(get_traced_grpc_channel(self.address, "nidx_indexer"))
+
+    async def finalize(self):
+        pass
+
+    async def index(self, writer: IndexMessage) -> int:
+        await self.client.Index(writer)
+        return 0
+
+
+class NidxServiceUtility(NidxUtility):
+    """Implements Nidx utility connecting to the network service"""
+
+    def __init__(self):
+        if not settings.nidx_api_address or not settings.nidx_searcher_address:
+            raise ValueError("NIDX_API_ADDRESS and NIDX_SEARCHER_ADDRESS are required")
+
+        if indexing_settings.index_nidx_subject:
+            self.indexer = NidxNatsIndexer()
+        elif settings.nidx_indexer_address is not None:
+            self.indexer = NidxGrpcIndexer(settings.nidx_indexer_address)
+        else:
+            raise ValueError("NIDX_INDEXER_ADDRESS or INDEX_NIDX_SUBJECT are required")
+
+    async def initialize(self):
+        await self.indexer.initialize()
+        self.api_client = NidxApiStub(get_traced_grpc_channel(settings.nidx_api_address, "nidx_api"))
+        self.searcher_client = NidxSearcherStub(
+            get_traced_grpc_channel(settings.nidx_searcher_address, "nidx_searcher")
+        )
+
+    async def finalize(self):
+        await self.indexer.finalize()
+
+    async def index(self, writer: IndexMessage) -> int:
+        return await self.indexer.index(writer)
+
+
+async def start_nidx_utility() -> Optional[NidxUtility]:
+    nidx = get_utility(Utility.NIDX)
     if nidx:
         return nidx
 
     nidx_utility: NidxUtility
     if settings.standalone_mode:
-        nidx_utility = NidxBindingUtility()
+        if (
+            settings.nidx_api_address is not None
+            and settings.nidx_searcher_address is not None
+            and settings.nidx_indexer_address is not None
+        ):
+            # Standalone with nidx service (via grpc). This is used in clustered standalone mode
+            nidx_utility = NidxServiceUtility()
+        else:
+            # Normal standalone mode with binding
+            nidx_utility = NidxBindingUtility()
     else:
+        # Component deploy with nidx service via grpc & nats (cloud)
         nidx_utility = NidxServiceUtility()
 
     await nidx_utility.initialize()
@@ -178,30 +217,33 @@ async def start_nidx_utility() -> Optional[NidxUtility]:
 
 
 async def stop_nidx_utility():
-    nidx_utility = get_nidx()
+    nidx_utility = get_utility(Utility.NIDX)
     if nidx_utility:
         clean_utility(Utility.NIDX)
         await nidx_utility.finalize()
 
 
-def get_nidx() -> Optional[NidxUtility]:
-    return get_utility(Utility.NIDX)
+def get_nidx() -> NidxUtility:
+    nidx = get_utility(Utility.NIDX)
+    if nidx is None:
+        raise Exception("nidx not initialized")
+    return nidx
 
 
-def get_nidx_api_client() -> Optional["NidxApiStub"]:
+def get_nidx_api_client() -> "NidxApiStub":
     nidx = get_nidx()
-    if nidx:
+    if nidx.api_client:
         return nidx.api_client
     else:
-        return None
+        raise Exception("nidx not initialized")
 
 
-def get_nidx_searcher_client() -> Optional["NidxSearcherStub"]:
+def get_nidx_searcher_client() -> "NidxSearcherStub":
     nidx = get_nidx()
-    if nidx:
+    if nidx.searcher_client:
         return nidx.searcher_client
     else:
-        return None
+        raise Exception("nidx not initialized")
 
 
 # TODO: Remove the index node abstraction
@@ -252,9 +294,6 @@ class FakeNode(AbstractIndexNode):
         return "nidx"
 
 
-def get_nidx_fake_node() -> Optional[FakeNode]:
+def get_nidx_fake_node() -> FakeNode:
     nidx = get_nidx()
-    if nidx:
-        return FakeNode(nidx.api_client, nidx.searcher_client)
-    else:
-        return None
+    return FakeNode(nidx.api_client, nidx.searcher_client)

nucliadb/export_import/models.py

@@ -17,7 +17,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-from datetime import datetime
+import datetime
 from enum import Enum
 from typing import Any
 
@@ -57,8 +57,8 @@ class Metadata(BaseModel):
     task: TaskMetadata = TaskMetadata(status=Status.SCHEDULED)
     total: int = 0
     processed: int = 0
-    created: datetime = datetime.utcnow()
-    modified: datetime = datetime.utcnow()
+    created: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)
+    modified: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)
 
 
 class ExportMetadata(Metadata):

nucliadb/health.py

@@ -40,13 +40,6 @@ def nats_manager_healthy() -> bool:
     return nats_manager.healthy()
 
 
-def nodes_health_check() -> bool:
-    from nucliadb.common.cluster import manager
-    from nucliadb.ingest.settings import DriverConfig, settings
-
-    return len(manager.INDEX_NODES) > 0 or settings.driver == DriverConfig.LOCAL
-
-
 def pubsub_check() -> bool:
     driver: Optional[PubSubDriver] = get_utility(Utility.PUBSUB)
     if driver is None:

nucliadb/ingest/app.py

@@ -22,10 +22,6 @@ import importlib.metadata
 from typing import Awaitable, Callable
 
 from nucliadb import health
-from nucliadb.common.cluster.discovery.utils import (
-    setup_cluster_discovery,
-    teardown_cluster_discovery,
-)
 from nucliadb.common.cluster.settings import settings as cluster_settings
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.context import ApplicationContext
@@ -89,13 +85,9 @@ async def initialize() -> list[Callable[[], Awaitable[None]]]:
         )
         finalizers.append(stop_nats_manager)
 
-        await setup_cluster_discovery()
-        finalizers.append(teardown_cluster_discovery)
-
     health.register_health_checks(
         [
             health.nats_manager_healthy,
-            health.nodes_health_check,
             health.pubsub_check,
         ]
     )

nucliadb/ingest/consumer/auditing.py

@@ -113,7 +113,7 @@ class IndexAuditHandler:
 
         for shard_obj in shard_groups:
             # TODO: Uses node for auditing, don't want to suddenly change metrics
-            node, shard_id = choose_node(shard_obj, use_nidx=False)
+            node, shard_id = choose_node(shard_obj)
             shard: nodereader_pb2.Shard = await node.reader.GetShard(
                 nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))  # type: ignore
             )

nucliadb/ingest/consumer/shard_creator.py

@@ -103,7 +103,7 @@ class ShardCreatorHandler:
         async with locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=kbid)):
             # remember, a lock will do at least 1+ reads and 1 write.
             # with heavy writes, this adds some simple k/v pressure
-            node, shard_id = choose_node(current_shard, use_nidx=True)
+            node, shard_id = choose_node(current_shard)
             shard: nodereader_pb2.Shard = await node.reader.GetShard(
                 nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))  # type: ignore
             )

nucliadb/ingest/fields/base.py

@@ -21,12 +21,13 @@ from __future__ import annotations
 
 import enum
 from datetime import datetime
-from typing import Any, Generic, Optional, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, Optional, Type, TypeVar
 
 from google.protobuf.message import DecodeError, Message
 
 from nucliadb.common import datamanagers
 from nucliadb.ingest.fields.exceptions import InvalidFieldClass, InvalidPBClass
+from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
 from nucliadb_protos.resources_pb2 import (
     CloudFile,
     ExtractedTextWrapper,
@@ -41,12 +42,19 @@ from nucliadb_protos.resources_pb2 import (
     QuestionAnswers,
 )
 from nucliadb_protos.utils_pb2 import ExtractedText, VectorObject
-from nucliadb_protos.writer_pb2 import Error
+from nucliadb_protos.writer_pb2 import Error, FieldStatus
+from nucliadb_utils.storages.exceptions import CouldNotCopyNotFound
 from nucliadb_utils.storages.storage import Storage, StorageField
 
+if TYPE_CHECKING:  # pragma: no cover
+    from nucliadb.ingest.orm.resource import Resource
+
+
 SUBFIELDFIELDS = ("c",)
 
 
+# NOTE extracted vectors key is no longer a static key, it is stored in each
+# vectorset
 class FieldTypes(str, enum.Enum):
     FIELD_TEXT = "extracted_text"
     FIELD_VECTORS = "extracted_vectors"
@@ -73,7 +81,7 @@ class Field(Generic[PbType]):
     def __init__(
         self,
         id: str,
-        resource: Any,
+        resource: Resource,
         pb: Optional[Any] = None,
         value: Optional[Any] = None,
     ):
@@ -88,7 +96,7 @@ class Field(Generic[PbType]):
         self.question_answers = None
 
         self.id: str = id
-        self.resource: Any = resource
+        self.resource = resource
 
         if value is not None:
             newpb = self.pbklass()
@@ -119,11 +127,20 @@ class Field(Generic[PbType]):
     def get_storage_field(self, field_type: FieldTypes) -> StorageField:
         return self.storage.file_extracted(self.kbid, self.uuid, self.type, self.id, field_type.value)
 
-    def _get_extracted_vectors_storage_field(self, vectorset: Optional[str] = None) -> StorageField:
-        if vectorset:
+    def _get_extracted_vectors_storage_field(
+        self,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+    ) -> StorageField:
+        if storage_key_kind == VectorSetConfig.StorageKeyKind.LEGACY:
+            key = FieldTypes.FIELD_VECTORS.value
+        elif storage_key_kind == VectorSetConfig.StorageKeyKind.VECTORSET_PREFIX:
             key = FieldTypes.FIELD_VECTORSET.value.format(vectorset=vectorset)
         else:
-            key = FieldTypes.FIELD_VECTORS.value
+            raise ValueError(
+                f"Can't do anything with UNSET or unknown vectorset storage key kind: {storage_key_kind}"
+            )
+
         return self.storage.file_extracted(self.kbid, self.uuid, self.type, self.id, key)
 
     async def db_get_value(self) -> Optional[PbType]:
@@ -163,7 +180,8 @@ class Field(Generic[PbType]):
             field_id=self.id,
         )
         await self.delete_extracted_text()
-        await self.delete_vectors()
+        async for vectorset_id, vs in datamanagers.vectorsets.iter(self.resource.txn, kbid=self.kbid):
+            await self.delete_vectors(vectorset_id, vs.storage_key_kind)
         await self.delete_metadata()
         await self.delete_question_answers()
 
@@ -181,9 +199,13 @@ class Field(Generic[PbType]):
         except KeyError:
             pass
 
-    async def delete_vectors(self, vectorset: Optional[str] = None) -> None:
+    async def delete_vectors(
+        self,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+    ) -> None:
         # Try delete vectors
-        sf = self._get_extracted_vectors_storage_field(vectorset)
+        sf = self._get_extracted_vectors_storage_field(vectorset, storage_key_kind)
         try:
             await self.storage.delete_upload(sf.key, sf.bucket)
         except KeyError:
@@ -215,6 +237,25 @@ class Field(Generic[PbType]):
             error=error,
         )
 
+    async def get_status(self) -> Optional[FieldStatus]:
+        return await datamanagers.fields.get_status(
+            self.resource.txn,
+            kbid=self.kbid,
+            rid=self.uuid,
+            field_type=self.type,
+            field_id=self.id,
+        )
+
+    async def set_status(self, status: FieldStatus) -> None:
+        await datamanagers.fields.set_status(
+            self.resource.txn,
+            kbid=self.kbid,
+            rid=self.uuid,
+            field_type=self.type,
+            field_id=self.id,
+            status=status,
+        )
+
     async def get_question_answers(self, force=False) -> Optional[FieldQuestionAnswers]:
         if self.question_answers is None or force:
             sf = self.get_storage_field(FieldTypes.QUESTION_ANSWERS)
@@ -309,12 +350,17 @@ class Field(Generic[PbType]):
                 self.extracted_text = payload
         return self.extracted_text
 
-    async def set_vectors(self, payload: ExtractedVectorsWrapper) -> Optional[VectorObject]:
-        vectorset = payload.vectorset_id or None
+    async def set_vectors(
+        self,
+        payload: ExtractedVectorsWrapper,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+    ) -> Optional[VectorObject]:
         if self.type in SUBFIELDFIELDS:
             try:
                 actual_payload: Optional[VectorObject] = await self.get_vectors(
                     vectorset=vectorset,
+                    storage_key_kind=storage_key_kind,
                     force=True,
                 )
             except KeyError:
@@ -322,12 +368,29 @@ class Field(Generic[PbType]):
         else:
             actual_payload = None
 
-        sf = self._get_extracted_vectors_storage_field(vectorset)
+        sf = self._get_extracted_vectors_storage_field(vectorset, storage_key_kind)
         vo: Optional[VectorObject] = None
         if actual_payload is None:
-            # Its first extracted text
+            # Its first extracted vectors
             if payload.HasField("file"):
-                await self.storage.normalize_binary(payload.file, sf)
+                # When we receive vectors in a cloud file, it points to our
+                # storage but paths are different, we may want to move it. This
+                # can happen, for example, with LEGACY KBs where processing
+                # sends us the extracted vectors prefixed by vectorset but, to
+                # maintain bw/c, we move those to the original not prefixed
+                # path.
+                try:
+                    await self.storage.normalize_binary(payload.file, sf)
+                except CouldNotCopyNotFound:
+                    # A failure here could mean the payload has already been
+                    # moved and we're retrying due to a redelivery or another
+                    # retry mechanism
+                    already_moved = await sf.exists()
+                    if already_moved:
+                        # We assume is the correct one and do nothing else
+                        pass
+                    else:
+                        raise
                 vo = await self.storage.download_pb(sf, VectorObject)
             else:
                 await self.storage.upload_pb(sf, payload.vectors)
@@ -354,14 +417,13 @@ class Field(Generic[PbType]):
         return vo
 
     async def get_vectors(
-        self, vectorset: Optional[str] = None, force: bool = False
+        self,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+        force: bool = False,
     ) -> Optional[VectorObject]:
-        # compat with vectorsets coming from protobuffers where no value is
-        # empty string instead of None. This shouldn't be handled here but we
-        # have to make sure it gets the correct vectorset
-        vectorset = vectorset or None
         if self.extracted_vectors.get(vectorset, None) is None or force:
-            sf = self._get_extracted_vectors_storage_field(vectorset)
+            sf = self._get_extracted_vectors_storage_field(vectorset, storage_key_kind)
             payload = await self.storage.download_pb(sf, VectorObject)
             if payload is not None:
                 self.extracted_vectors[vectorset] = payload

nucliadb/ingest/orm/brain.py

@@ -100,6 +100,8 @@ class ResourceBrain:
         page_positions: Optional[FilePagePositions],
         extracted_text: Optional[ExtractedText],
         basic_user_field_metadata: Optional[UserFieldMetadata] = None,
+        *,
+        replace_field: bool = False,
     ):
         # To check for duplicate paragraphs
         unique_paragraphs: set[str] = set()
@@ -224,6 +226,11 @@ class ResourceBrain:
 
             self.brain.paragraphs[field_key].paragraphs[key].CopyFrom(p)
 
+        if replace_field:
+            field_type, field_name = field_key.split("/")
+            full_field_id = ids.FieldId(rid=self.rid, type=field_type, key=field_name).full()
+            self.brain.paragraphs_to_delete.append(full_field_id)
+
         for relations in metadata.metadata.relations:
             for relation in relations.relations:
                 self.brain.relations.append(relation)
@@ -239,9 +246,10 @@ class ResourceBrain:
         field_id: str,
         vo: utils_pb2.VectorObject,
         *,
-        vectorset: Optional[str] = None,
+        vectorset: str,
         replace_field: bool = False,
-        matryoshka_vector_dimension: Optional[int] = None,
+        # cut to specific dimension if specified
+        vector_dimension: Optional[int] = None,
     ):
         fid = ids.FieldId.from_string(f"{self.rid}/{field_id}")
         for subfield, vectors in vo.split_vectors.items():
@@ -270,7 +278,7 @@ class ResourceBrain:
                     sentence_key,
                     vector,
                     vectorset=vectorset,
-                    matryoshka_vector_dimension=matryoshka_vector_dimension,
+                    vector_dimension=vector_dimension,
                 )
 
         _field_id = ids.FieldId(
@@ -296,13 +304,12 @@ class ResourceBrain:
                 sentence_key,
                 vector,
                 vectorset=vectorset,
-                matryoshka_vector_dimension=matryoshka_vector_dimension,
+                vector_dimension=vector_dimension,
             )
 
         if replace_field:
             full_field_id = ids.FieldId(rid=self.rid, type=fid.type, key=fid.key).full()
-            self.brain.sentences_to_delete.append(full_field_id)
-            self.brain.paragraphs_to_delete.append(full_field_id)
+            self.brain.vector_prefixes_to_delete[vectorset].items.append(full_field_id)
 
     def _apply_field_vector(
         self,
@@ -311,22 +318,15 @@ class ResourceBrain:
         sentence_key: ids.VectorId,
         vector: utils_pb2.Vector,
         *,
-        vectorset: Optional[str],
-        matryoshka_vector_dimension: Optional[int] = None,
+        vectorset: str,
+        # cut vectors if a specific dimension is specified
+        vector_dimension: Optional[int] = None,
     ):
         paragraph_pb = self.brain.paragraphs[field_id].paragraphs[paragraph_key.full()]
-        if vectorset:
-            sentence_pb = paragraph_pb.vectorsets_sentences[vectorset].sentences[sentence_key.full()]
-        else:
-            sentence_pb = paragraph_pb.sentences[sentence_key.full()]
+        sentence_pb = paragraph_pb.vectorsets_sentences[vectorset].sentences[sentence_key.full()]
 
         sentence_pb.ClearField("vector")  # clear first to prevent duplicates
-
-        # cut vectors if a specific dimension is specified
-        if matryoshka_vector_dimension is not None:
-            sentence_pb.vector.extend(vector.vector[:matryoshka_vector_dimension])
-        else:
-            sentence_pb.vector.extend(vector.vector)
+        sentence_pb.vector.extend(vector.vector[:vector_dimension])
 
         # we only care about start/stop position of the paragraph for a given sentence here
         # the key has the sentence position
@@ -490,22 +490,29 @@ class ResourceBrain:
     ):
         if metadata.mime_type != "":
             labels["mt"].add(metadata.mime_type)
+
+        base_classification_relation = Relation(
+            relation=Relation.ABOUT,
+            source=relation_node_document,
+            to=RelationNode(
+                ntype=RelationNode.NodeType.LABEL,
+            ),
+        )
         for classification in metadata.classifications:
             label = f"{classification.labelset}/{classification.label}"
             if label not in user_canceled_labels:
                 labels["l"].add(label)
-                relation_node_label = RelationNode(
-                    value=label,
-                    ntype=RelationNode.NodeType.LABEL,
-                )
-                self.brain.relations.append(
-                    Relation(
-                        relation=Relation.ABOUT,
-                        source=relation_node_document,
-                        to=relation_node_label,
-                    )
-                )
+                relation = Relation()
+                relation.CopyFrom(base_classification_relation)
+                relation.to.value = label
+                self.brain.relations.append(relation)
+
         # Data Augmentation + Processor entities
+        base_entity_relation = Relation(
+            relation=Relation.ENTITY,
+            source=relation_node_document,
+            to=RelationNode(ntype=RelationNode.NodeType.ENTITY),
+        )
         use_legacy_entities = True
         for data_augmentation_task_id, entities in metadata.entities.items():
             # If we recieved the entities from the processor here, we don't want to use the legacy entities
@@ -521,38 +528,30 @@ class ResourceBrain:
                 labels["e"].add(
                     f"{entity_label}/{entity_text}"
                 )  # Add data_augmentation_task_id as a prefix?
-                relation_node_entity = RelationNode(
-                    value=entity_text,
-                    ntype=RelationNode.NodeType.ENTITY,
-                    subtype=entity_label,
-                )
-                rel = Relation(
-                    relation=Relation.ENTITY,
-                    source=relation_node_document,
-                    to=relation_node_entity,
-                )
-                self.brain.relations.append(rel)
+                relation = Relation()
+                relation.CopyFrom(base_entity_relation)
+                relation.to.value = entity_text
+                relation.to.subtype = entity_label
+                self.brain.relations.append(relation)
 
         # Legacy processor entities
         # TODO: Remove once processor doesn't use this anymore and remove the positions and ner fields from the message
+        def _parse_entity(klass_entity: str) -> tuple[str, str]:
+            try:
+                klass, entity = klass_entity.split("/", 1)
+                return klass, entity
+            except ValueError:
+                raise AttributeError(f"Entity should be with type {klass_entity}")
+
         if use_legacy_entities:
-            for klass_entity, _ in metadata.positions.items():
+            for klass_entity in metadata.positions.keys():
                 labels["e"].add(klass_entity)
-                entity_array = klass_entity.split("/")
-                if len(entity_array) == 1:
-                    raise AttributeError(f"Entity should be with type {klass_entity}")
-                elif len(entity_array) > 1:
-                    klass = entity_array[0]
-                    entity = "/".join(entity_array[1:])
-                relation_node_entity = RelationNode(
-                    value=entity, ntype=RelationNode.NodeType.ENTITY, subtype=klass
-                )
-                rel = Relation(
-                    relation=Relation.ENTITY,
-                    source=relation_node_document,
-                    to=relation_node_entity,
-                )
-                self.brain.relations.append(rel)
+                klass, entity = _parse_entity(klass_entity)
+                relation = Relation()
+                relation.CopyFrom(base_entity_relation)
+                relation.to.value = entity
+                relation.to.subtype = klass
+                self.brain.relations.append(relation)
 
     def apply_field_labels(
         self,