nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/ingest/orm/resource.py

@@ -23,19 +23,33 @@ import asyncio
 import logging
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
-from typing import TYPE_CHECKING, Any, AsyncIterator, Optional, Type
+from typing import TYPE_CHECKING, Any, AsyncIterator, MutableMapping, Optional, Type
 
+from nucliadb.common import datamanagers
+from nucliadb.common.datamanagers.resources import KB_RESOURCE_SLUG
+from nucliadb.common.ids import FIELD_TYPE_PB_TO_STR, FieldId
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.fields.conversation import Conversation
+from nucliadb.ingest.fields.file import File
+from nucliadb.ingest.fields.generic import VALID_GENERIC_FIELDS, Generic
+from nucliadb.ingest.fields.link import Link
+from nucliadb.ingest.fields.text import Text
+from nucliadb.ingest.orm.brain import FilePagePositions, ResourceBrain
+from nucliadb.ingest.orm.metrics import processor_observer
+from nucliadb_models import content_types
+from nucliadb_models.common import CloudLink
+from nucliadb_models.content_types import GENERIC_MIME_TYPE
+from nucliadb_protos import utils_pb2, writer_pb2
 from nucliadb_protos.resources_pb2 import AllFieldIDs as PBAllFieldIDs
-from nucliadb_protos.resources_pb2 import Basic
-from nucliadb_protos.resources_pb2 import Basic as PBBasic
-from nucliadb_protos.resources_pb2 import CloudFile
-from nucliadb_protos.resources_pb2 import Conversation as PBConversation
-from nucliadb_protos.resources_pb2 import Extra as PBExtra
 from nucliadb_protos.resources_pb2 import (
+    Basic,
+    CloudFile,
     ExtractedTextWrapper,
     ExtractedVectorsWrapper,
     FieldClassifications,
     FieldComputedMetadataWrapper,
+    FieldFile,
     FieldID,
     FieldMetadata,
     FieldQuestionAnswerWrapper,
@@ -44,41 +58,27 @@ from nucliadb_protos.resources_pb2 import (
     FileExtractedData,
     LargeComputedMetadataWrapper,
     LinkExtractedData,
+    Metadata,
+    Paragraph,
+    ParagraphAnnotation,
 )
-from nucliadb_protos.resources_pb2 import Metadata
+from nucliadb_protos.resources_pb2 import Basic as PBBasic
+from nucliadb_protos.resources_pb2 import Conversation as PBConversation
+from nucliadb_protos.resources_pb2 import Extra as PBExtra
 from nucliadb_protos.resources_pb2 import Metadata as PBMetadata
 from nucliadb_protos.resources_pb2 import Origin as PBOrigin
-from nucliadb_protos.resources_pb2 import Paragraph, ParagraphAnnotation
 from nucliadb_protos.resources_pb2 import Relations as PBRelations
-from nucliadb_protos.resources_pb2 import UserVectorsWrapper
-from nucliadb_protos.train_pb2 import EnabledMetadata
-from nucliadb_protos.train_pb2 import Position as TrainPosition
 from nucliadb_protos.train_pb2 import (
+    EnabledMetadata,
     TrainField,
     TrainMetadata,
     TrainParagraph,
     TrainResource,
     TrainSentence,
 )
+from nucliadb_protos.train_pb2 import Position as TrainPosition
 from nucliadb_protos.utils_pb2 import Relation as PBRelation
 from nucliadb_protos.writer_pb2 import BrokerMessage
-
-from nucliadb.common.maindb.driver import Transaction
-from nucliadb.ingest.fields.base import Field
-from nucliadb.ingest.fields.conversation import Conversation
-from nucliadb.ingest.fields.date import Datetime
-from nucliadb.ingest.fields.file import File
-from nucliadb.ingest.fields.generic import VALID_GENERIC_FIELDS, Generic
-from nucliadb.ingest.fields.keywordset import Keywordset
-from nucliadb.ingest.fields.layout import Layout
-from nucliadb.ingest.fields.link import Link
-from nucliadb.ingest.fields.text import Text
-from nucliadb.ingest.orm.brain import FilePagePositions, ResourceBrain
-from nucliadb.ingest.orm.metrics import processor_observer
-from nucliadb.ingest.orm.utils import get_basic, set_basic
-from nucliadb_models.common import CloudLink
-from nucliadb_models.writer import GENERIC_MIME_TYPE
-from nucliadb_protos import utils_pb2, writer_pb2
 from nucliadb_utils.storages.storage import Storage
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -86,41 +86,14 @@ if TYPE_CHECKING:  # pragma: no cover
 
 logger = logging.getLogger(__name__)
 
-KB_RESOURCE_ORIGIN = "/kbs/{kbid}/r/{uuid}/origin"
-KB_RESOURCE_EXTRA = "/kbs/{kbid}/r/{uuid}/extra"
-KB_RESOURCE_SECURITY = "/kbs/{kbid}/r/{uuid}/security"
-KB_RESOURCE_METADATA = "/kbs/{kbid}/r/{uuid}/metadata"
-KB_RESOURCE_RELATIONS = "/kbs/{kbid}/r/{uuid}/relations"
-KB_RESOURCE_FIELDS = "/kbs/{kbid}/r/{uuid}/f/"
-KB_RESOURCE_ALL_FIELDS = "/kbs/{kbid}/r/{uuid}/allfields"
-KB_RESOURCE_SLUG_BASE = "/kbs/{kbid}/s/"
-KB_RESOURCE_SLUG = f"{KB_RESOURCE_SLUG_BASE}{{slug}}"
-KB_RESOURCE_CONVERSATION = "/kbs/{kbid}/r/{uuid}/c/{page}"
-GLOBAL_FIELD = "a"
 KB_FIELDS: dict[int, Type] = {
-    FieldType.LAYOUT: Layout,
     FieldType.TEXT: Text,
     FieldType.FILE: File,
     FieldType.LINK: Link,
-    FieldType.DATETIME: Datetime,
-    FieldType.KEYWORDSET: Keywordset,
     FieldType.GENERIC: Generic,
     FieldType.CONVERSATION: Conversation,
 }
 
-KB_REVERSE: dict[str, FieldType.ValueType] = {
-    "l": FieldType.LAYOUT,
-    "t": FieldType.TEXT,
-    "f": FieldType.FILE,
-    "u": FieldType.LINK,
-    "d": FieldType.DATETIME,
-    "k": FieldType.KEYWORDSET,
-    "a": FieldType.GENERIC,
-    "c": FieldType.CONVERSATION,
-}
-
-FIELD_TYPE_TO_ID = {v: k for k, v in KB_REVERSE.items()}
-
 _executor = ThreadPoolExecutor(10)
 
 
@@ -131,6 +104,8 @@ PB_TEXT_FORMAT_TO_MIMETYPE = {
     FieldText.Format.MARKDOWN: "text/markdown",
     FieldText.Format.JSON: "application/json",
     FieldText.Format.KEEP_MARKDOWN: "text/markdown",
+    FieldText.Format.JSONL: "application/x-ndjson",
+    FieldText.Format.PLAIN_BLANKLINE_SPLIT: "text/plain+blankline",
 }
 
 BASIC_IMMUTABLE_FIELDS = ("icon",)
@@ -179,32 +154,11 @@ class Resource:
         new_key = KB_RESOURCE_SLUG.format(kbid=self.kb.kbid, slug=basic.slug)
         await self.txn.set(new_key, self.uuid.encode())
 
-    @staticmethod
-    def parse_basic(payload: bytes) -> PBBasic:
-        pb = PBBasic()
-        if payload is None:
-            return None
-
-        pb.ParseFromString(payload)
-        return pb
-
-    async def exists(self) -> bool:
-        exists = True
-        if self.basic is None:
-            payload = await get_basic(self.txn, self.kb.kbid, self.uuid)
-            if payload is not None:
-                pb = PBBasic()
-                pb.ParseFromString(payload)
-                self.basic = pb
-            else:
-                exists = False
-        return exists
-
     # Basic
     async def get_basic(self) -> Optional[PBBasic]:
         if self.basic is None:
-            payload = await get_basic(self.txn, self.kb.kbid, self.uuid)
-            self.basic = self.parse_basic(payload) if payload is not None else PBBasic()
+            basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kb.kbid, rid=self.uuid)
+            self.basic = basic if basic is not None else PBBasic()
         return self.basic
 
     def set_processing_status(self, current_basic: PBBasic, basic_in_payload: PBBasic):
@@ -249,9 +203,7 @@ class Resource:
                         fields.append(field_id)
                     positions[field_id] = i
 
-                updated = [
-                    self.basic.fieldmetadata[positions[field]] for field in fields
-                ]
+                updated = [self.basic.fieldmetadata[positions[field]] for field in fields]
 
                 del self.basic.fieldmetadata[:]
                 self.basic.fieldmetadata.extend(updated)
@@ -272,11 +224,10 @@ class Resource:
                         self.indexer.apply_field_metadata(
                             field_id,
                             field_metadata,
-                            replace_field=[],
-                            replace_splits={},
                             page_positions=page_positions,
                             extracted_text=await field_obj.get_extracted_text(),
                             basic_user_field_metadata=user_field_metadata,
+                            replace_field=True,
                         )
 
         # Some basic fields are computed off field metadata.
@@ -284,27 +235,21 @@ class Resource:
         if deleted_fields is not None and len(deleted_fields) > 0:
             remove_field_classifications(self.basic, deleted_fields=deleted_fields)
 
-        await set_basic(self.txn, self.kb.kbid, self.uuid, self.basic)
+        await datamanagers.resources.set_basic(
+            self.txn, kbid=self.kb.kbid, rid=self.uuid, basic=self.basic
+        )
         self.modified = True
 
     # Origin
     async def get_origin(self) -> Optional[PBOrigin]:
         if self.origin is None:
-            pb = PBOrigin()
-            payload = await self.txn.get(
-                KB_RESOURCE_ORIGIN.format(kbid=self.kb.kbid, uuid=self.uuid)
-            )
-            if payload is None:
-                return None
-
-            pb.ParseFromString(payload)
-            self.origin = pb
+            origin = await datamanagers.resources.get_origin(self.txn, kbid=self.kb.kbid, rid=self.uuid)
+            self.origin = origin
         return self.origin
 
     async def set_origin(self, payload: PBOrigin):
-        await self.txn.set(
-            KB_RESOURCE_ORIGIN.format(kbid=self.kb.kbid, uuid=self.uuid),
-            payload.SerializeToString(),
+        await datamanagers.resources.set_origin(
+            self.txn, kbid=self.kb.kbid, rid=self.uuid, origin=payload
         )
         self.modified = True
         self.origin = payload
@@ -312,42 +257,27 @@ class Resource:
     # Extra
     async def get_extra(self) -> Optional[PBExtra]:
         if self.extra is None:
-            pb = PBExtra()
-            payload = await self.txn.get(
-                KB_RESOURCE_EXTRA.format(kbid=self.kb.kbid, uuid=self.uuid)
-            )
-            if payload is None:
-                return None
-            pb.ParseFromString(payload)
-            self.extra = pb
+            extra = await datamanagers.resources.get_extra(self.txn, kbid=self.kb.kbid, rid=self.uuid)
+            self.extra = extra
         return self.extra
 
     async def set_extra(self, payload: PBExtra):
-        key = KB_RESOURCE_EXTRA.format(kbid=self.kb.kbid, uuid=self.uuid)
-        await self.txn.set(
-            key,
-            payload.SerializeToString(),
-        )
+        await datamanagers.resources.set_extra(self.txn, kbid=self.kb.kbid, rid=self.uuid, extra=payload)
         self.modified = True
         self.extra = payload
 
     # Security
     async def get_security(self) -> Optional[utils_pb2.Security]:
         if self.security is None:
-            pb = utils_pb2.Security()
-            key = KB_RESOURCE_SECURITY.format(kbid=self.kb.kbid, uuid=self.uuid)
-            payload = await self.txn.get(key)
-            if payload is None:
-                return None
-            pb.ParseFromString(payload)
-            self.security = pb
+            security = await datamanagers.resources.get_security(
+                self.txn, kbid=self.kb.kbid, rid=self.uuid
+            )
+            self.security = security
         return self.security
 
     async def set_security(self, payload: utils_pb2.Security) -> None:
-        key = KB_RESOURCE_SECURITY.format(kbid=self.kb.kbid, uuid=self.uuid)
-        await self.txn.set(
-            key,
-            payload.SerializeToString(),
+        await datamanagers.resources.set_security(
+            self.txn, kbid=self.kb.kbid, rid=self.uuid, security=payload
         )
         self.modified = True
         self.security = payload
@@ -355,29 +285,24 @@ class Resource:
     # Relations
     async def get_relations(self) -> Optional[PBRelations]:
         if self.relations is None:
-            pb = PBRelations()
-            payload = await self.txn.get(
-                KB_RESOURCE_RELATIONS.format(kbid=self.kb.kbid, uuid=self.uuid)
+            relations = await datamanagers.resources.get_relations(
+                self.txn, kbid=self.kb.kbid, rid=self.uuid
             )
-            if payload is None:
-                return None
-            pb.ParseFromString(payload)
-            self.relations = pb
+            self.relations = relations
         return self.relations
 
     async def set_relations(self, payload: list[PBRelation]):
         relations = PBRelations()
         for relation in payload:
             relations.relations.append(relation)
-        await self.txn.set(
-            KB_RESOURCE_RELATIONS.format(kbid=self.kb.kbid, uuid=self.uuid),
-            relations.SerializeToString(),
+        await datamanagers.resources.set_relations(
+            self.txn, kbid=self.kb.kbid, rid=self.uuid, relations=relations
         )
         self.modified = True
         self.relations = relations
 
     @processor_observer.wrap({"type": "generate_index_message"})
-    async def generate_index_message(self) -> ResourceBrain:
+    async def generate_index_message(self, reindex: bool = False) -> ResourceBrain:
         brain = ResourceBrain(rid=self.uuid)
         origin = await self.get_origin()
         basic = await self.get_basic()
@@ -387,7 +312,7 @@ class Resource:
         await self.compute_global_tags(brain)
         fields = await self.get_fields(force=True)
         for (type_id, field_id), field in fields.items():
-            fieldid = FieldID(field_type=type_id, field=field_id)  # type: ignore
+            fieldid = FieldID(field_type=type_id, field=field_id)
             await self.compute_global_text_field(fieldid, brain)
 
             field_metadata = await field.get_field_metadata()
@@ -403,251 +328,66 @@ class Resource:
                         (
                             fm
                             for fm in basic.fieldmetadata
-                            if fm.field.field == field_id
-                            and fm.field.field_type == type_id
+                            if fm.field.field == field_id and fm.field.field_type == type_id
                         ),
                         None,
                     )
                 brain.apply_field_metadata(
                     field_key,
                     field_metadata,
-                    replace_field=[],
-                    replace_splits={},
                     page_positions=page_positions,
                     extracted_text=await field.get_extracted_text(),
                     basic_user_field_metadata=user_field_metadata,
+                    replace_field=reindex,
                 )
 
             if self.disable_vectors is False:
+                # XXX: while we don't remove the "default" vectorset concept, we
+                # need to do use None as the default one
                 vo = await field.get_vectors()
                 if vo is not None:
-                    brain.apply_field_vectors(field_key, vo, False, [])
+                    async with datamanagers.with_ro_transaction() as ro_txn:
+                        dimension = await datamanagers.kb.get_matryoshka_vector_dimension(
+                            ro_txn, kbid=self.kb.kbid
+                        )
+                    brain.apply_field_vectors(
+                        field_key,
+                        vo,
+                        matryoshka_vector_dimension=dimension,
+                        replace_field=reindex,
+                    )
 
-                vu = await field.get_user_vectors()
-                if vu is not None:
-                    vectors_to_delete = {}  # type: ignore
-                    brain.apply_user_vectors(field_key, vu, vectors_to_delete)  # type: ignore
+                vectorset_configs = []
+                async with datamanagers.with_ro_transaction() as ro_txn:
+                    async for vectorset_id, vectorset_config in datamanagers.vectorsets.iter(
+                        ro_txn, kbid=self.kb.kbid
+                    ):
+                        vectorset_configs.append(vectorset_config)
+                for vectorset_config in vectorset_configs:
+                    vo = await field.get_vectors(vectorset=vectorset_config.vectorset_id)
+                    if vo is not None:
+                        dimension = vectorset_config.vectorset_index_config.vector_dimension
+                        brain.apply_field_vectors(
+                            field_key,
+                            vo,
+                            vectorset=vectorset_config.vectorset_id,
+                            matryoshka_vector_dimension=dimension,
+                            replace_field=reindex,
+                        )
         return brain
 
-    async def generate_field_vectors(
-        self,
-        bm: BrokerMessage,
-        type_id: FieldType.ValueType,
-        field_id: str,
-        field: Field,
-    ):
-        vo = await field.get_vectors()
-        if vo is None:
-            return
-        evw = ExtractedVectorsWrapper()
-        evw.field.field = field_id
-        evw.field.field_type = type_id  # type: ignore
-        evw.vectors.CopyFrom(vo)
-        bm.field_vectors.append(evw)
-
-    async def generate_user_vectors(
-        self,
-        bm: BrokerMessage,
-        type_id: FieldType.ValueType,
-        field_id: str,
-        field: Field,
-    ):
-        uv = await field.get_user_vectors()
-        if uv is None:
-            return
-        uvw = UserVectorsWrapper()
-        uvw.field.field = field_id
-        uvw.field.field_type = type_id  # type: ignore
-        uvw.vectors.CopyFrom(uv)
-        bm.user_vectors.append(uvw)
-
-    async def generate_field_large_computed_metadata(
-        self,
-        bm: BrokerMessage,
-        type_id: FieldType.ValueType,
-        field_id: str,
-        field: Field,
-    ):
-        lcm = await field.get_large_field_metadata()
-        if lcm is None:
-            return
-        lcmw = LargeComputedMetadataWrapper()
-        lcmw.field.field = field_id
-        lcmw.field.field_type = type_id  # type: ignore
-        lcmw.real.CopyFrom(lcm)
-        bm.field_large_metadata.append(lcmw)
-
-    async def generate_field_computed_metadata(
-        self,
-        bm: BrokerMessage,
-        type_id: FieldType.ValueType,
-        field_id: str,
-        field: Field,
-    ):
-        fcmw = FieldComputedMetadataWrapper()
-        fcmw.field.field = field_id
-        fcmw.field.field_type = type_id  # type: ignore
-
-        field_metadata = await field.get_field_metadata()
-        if field_metadata is not None:
-            fcmw.metadata.CopyFrom(field_metadata)
-            fcmw.field.field = field_id
-            fcmw.field.field_type = type_id  # type: ignore
-            bm.field_metadata.append(fcmw)
-            # Make sure cloud files are removed for exporting
-
-    async def generate_extracted_text(
-        self,
-        bm: BrokerMessage,
-        type_id: FieldType.ValueType,
-        field_id: str,
-        field: Field,
-    ):
-        etw = ExtractedTextWrapper()
-        etw.field.field = field_id
-        etw.field.field_type = type_id  # type: ignore
-        extracted_text = await field.get_extracted_text()
-        if extracted_text is not None:
-            etw.body.CopyFrom(extracted_text)
-            bm.extracted_text.append(etw)
-
-    async def generate_field(
-        self,
-        bm: BrokerMessage,
-        type_id: FieldType.ValueType,
-        field_id: str,
-        field: Field,
-    ):
-        # Used for exporting a field
-        if type_id == FieldType.TEXT:
-            value = await field.get_value()
-            bm.texts[field_id].CopyFrom(value)
-        elif type_id == FieldType.LINK:
-            value = await field.get_value()
-            bm.links[field_id].CopyFrom(value)
-        elif type_id == FieldType.FILE:
-            value = await field.get_value()
-            bm.files[field_id].CopyFrom(value)
-        elif type_id == FieldType.CONVERSATION:
-            value = await self.get_full_conversation(field)  # type: ignore
-            bm.conversations[field_id].CopyFrom(value)
-        elif type_id == FieldType.KEYWORDSET:
-            value = await field.get_value()
-            bm.keywordsets[field_id].CopyFrom(value)
-        elif type_id == FieldType.DATETIME:
-            value = await field.get_value()
-            bm.datetimes[field_id].CopyFrom(value)
-        elif type_id == FieldType.LAYOUT:
-            value = await field.get_value()
-            bm.layouts[field_id].CopyFrom(value)
-
-    async def get_full_conversation(
-        self,
-        conversation_field: Conversation,
-    ) -> Optional[PBConversation]:
-        """
-        Messages of a conversations may be stored across several pages.
-        This method fetches them all and returns a single complete conversation.
-        """
-        full_conv = PBConversation()
-        n_page = 1
-        while True:
-            page = await conversation_field.get_value(page=n_page)
-            if page is None:
-                break
-            full_conv.messages.extend(page.messages)
-            n_page += 1
-        return full_conv
-
-    async def generate_broker_message(self) -> BrokerMessage:
-        # full means downloading all the pointers
-        # minuts the ones to external files that are not PB
-        # Go for all fields and recreate brain
-        bm = BrokerMessage()
-        bm.kbid = self.kb.kbid
-        bm.uuid = self.uuid
-        basic = await self.get_basic()
-        if basic is not None:
-            bm.basic.CopyFrom(basic)
-        bm.slug = bm.basic.slug
-        origin = await self.get_origin()
-        if origin is not None:
-            bm.origin.CopyFrom(origin)
-        relations = await self.get_relations()
-        if relations is not None:
-            for relation in relations.relations:
-                bm.relations.append(relation)
-
-        fields = await self.get_fields(force=True)
-        for (type_id, field_id), field in fields.items():
-            # Value
-            await self.generate_field(bm, type_id, field_id, field)
-
-            # Extracted text
-            await self.generate_extracted_text(bm, type_id, field_id, field)
-
-            # Field Computed Metadata
-            await self.generate_field_computed_metadata(bm, type_id, field_id, field)
-
-            if type_id == FieldType.FILE and isinstance(field, File):
-                field_extracted_data = await field.get_file_extracted_data()
-                if field_extracted_data is not None:
-                    bm.file_extracted_data.append(field_extracted_data)
-
-            elif type_id == FieldType.LINK and isinstance(field, Link):
-                link_extracted_data = await field.get_link_extracted_data()
-                if link_extracted_data is not None:
-                    bm.link_extracted_data.append(link_extracted_data)
-
-            # Field vectors
-            await self.generate_field_vectors(bm, type_id, field_id, field)
-
-            # User vectors
-            await self.generate_user_vectors(bm, type_id, field_id, field)
-
-            # Large metadata
-            await self.generate_field_large_computed_metadata(
-                bm, type_id, field_id, field
-            )
-
-        return bm
-
     # Fields
-    async def get_fields(
-        self, force: bool = False
-    ) -> dict[tuple[FieldType.ValueType, str], Field]:
+    async def get_fields(self, force: bool = False) -> dict[tuple[FieldType.ValueType, str], Field]:
         # Get all fields
         for type, field in await self.get_fields_ids(force=force):
             if (type, field) not in self.fields:
                 self.fields[(type, field)] = await self.get_field(field, type)
         return self.fields
 
-    async def _deprecated_scan_fields_ids(
-        self,
-    ) -> AsyncIterator[tuple[FieldType.ValueType, str]]:
-        logger.warning("Scanning fields ids. This is not optimal.")
-        prefix = KB_RESOURCE_FIELDS.format(kbid=self.kb.kbid, uuid=self.uuid)
-        allfields = set()
-        async for key in self.txn.keys(prefix, count=-1):
-            # The [6:8] `slicing purpose is to match exactly the two
-            # splitted parts corresponding to type and field, and nothing else!
-            type, field = key.split("/")[6:8]
-            type_id = KB_REVERSE.get(type)
-            if type_id is None:
-                raise AttributeError("Invalid field type")
-            result = (type_id, field)
-            if result not in allfields:
-                # fields can have errors that are stored in a subkey:
-                # - field key       -> kbs/kbid/r/ruuid/f/myfield
-                # - field error key -> kbs/kbid/r/ruuid/f/myfield/errors
-                # and that would return duplicates here.
-                yield result
-            allfields.add(result)
-
     async def _inner_get_fields_ids(self) -> list[tuple[FieldType.ValueType, str]]:
         # Use a set to make sure we don't have duplicate field ids
         result = set()
-        all_fields = await self.get_all_field_ids()
+        all_fields = await self.get_all_field_ids(for_update=False)
         if all_fields is not None:
             for f in all_fields.fields:
                 result.add((f.field_type, f.field))
@@ -664,9 +404,7 @@ class Resource:
                     result.add((FieldType.GENERIC, generic))
         return list(result)
 
-    async def get_fields_ids(
-        self, force: bool = False
-    ) -> list[tuple[FieldType.ValueType, str]]:
+    async def get_fields_ids(self, force: bool = False) -> list[tuple[FieldType.ValueType, str]]:
         """
         Get all ids of the fields of the resource and cache them.
         """
@@ -710,32 +448,26 @@ class Resource:
             if field in self.all_fields_keys:
                 self.all_fields_keys.remove(field)
 
-        field_key = self.generate_field_id(FieldID(field_type=type, field=key))  # type: ignore
-        vo = await field_obj.get_vectors()
-        if vo is not None:
-            self.indexer.delete_vectors(field_key=field_key, vo=vo)
+        field_key = self.generate_field_id(FieldID(field_type=type, field=key))
 
         metadata = await field_obj.get_field_metadata()
         if metadata is not None:
-            self.indexer.delete_metadata(field_key=field_key, metadata=metadata)
+            self.indexer.delete_field(field_key=field_key)
 
         await field_obj.delete()
 
     def has_field(self, type: FieldType.ValueType, field: str) -> bool:
         return (type, field) in self.fields
 
-    async def get_all_field_ids(self) -> Optional[PBAllFieldIDs]:
-        key = KB_RESOURCE_ALL_FIELDS.format(kbid=self.kb.kbid, uuid=self.uuid)
-        payload = await self.txn.get(key)
-        if payload is None:
-            return None
-        all_fields = PBAllFieldIDs()
-        all_fields.ParseFromString(payload)
-        return all_fields
+    async def get_all_field_ids(self, *, for_update: bool) -> Optional[PBAllFieldIDs]:
+        return await datamanagers.resources.get_all_field_ids(
+            self.txn, kbid=self.kb.kbid, rid=self.uuid, for_update=for_update
+        )
 
     async def set_all_field_ids(self, all_fields: PBAllFieldIDs):
-        key = KB_RESOURCE_ALL_FIELDS.format(kbid=self.kb.kbid, uuid=self.uuid)
-        await self.txn.set(key, all_fields.SerializeToString())
+        return await datamanagers.resources.set_all_field_ids(
+            self.txn, kbid=self.kb.kbid, rid=self.uuid, allfields=all_fields
+        )
 
     async def update_all_field_ids(
         self,
@@ -745,7 +477,7 @@ class Resource:
         errors: Optional[list[writer_pb2.Error]] = None,
     ):
         needs_update = False
-        all_fields = await self.get_all_field_ids()
+        all_fields = await self.get_all_field_ids(for_update=True)
         if all_fields is None:
             needs_update = True
             all_fields = PBAllFieldIDs()
@@ -772,26 +504,12 @@ class Resource:
     @processor_observer.wrap({"type": "apply_fields"})
     async def apply_fields(self, message: BrokerMessage):
         message_updated_fields = []
-        for field, layout in message.layouts.items():
-            fid = FieldID(field_type=FieldType.LAYOUT, field=field)
-            await self.set_field(fid.field_type, fid.field, layout)
-            message_updated_fields.append(fid)
 
         for field, text in message.texts.items():
             fid = FieldID(field_type=FieldType.TEXT, field=field)
             await self.set_field(fid.field_type, fid.field, text)
             message_updated_fields.append(fid)
 
-        for field, keywordset in message.keywordsets.items():
-            fid = FieldID(field_type=FieldType.KEYWORDSET, field=field)
-            await self.set_field(fid.field_type, fid.field, keywordset)
-            message_updated_fields.append(fid)
-
-        for field, datetimeobj in message.datetimes.items():
-            fid = FieldID(field_type=FieldType.DATETIME, field=field)
-            await self.set_field(fid.field_type, fid.field, datetimeobj)
-            message_updated_fields.append(fid)
-
         for field, link in message.links.items():
             fid = FieldID(field_type=FieldType.LINK, field=field)
             await self.set_field(fid.field_type, fid.field, link)
@@ -810,13 +528,11 @@ class Resource:
         for fieldid in message.delete_fields:
             await self.delete_field(fieldid.field_type, fieldid.field)
 
-        if (
-            len(message_updated_fields)
-            or len(message.delete_fields)
-            or len(message.errors)
-        ):
+        if len(message_updated_fields) or len(message.delete_fields) or len(message.errors):
             await self.update_all_field_ids(
-                updated=message_updated_fields, deleted=message.delete_fields, errors=message.errors  # type: ignore
+                updated=message_updated_fields,
+                deleted=message.delete_fields,  # type: ignore
+                errors=message.errors,  # type: ignore
             )
 
     @processor_observer.wrap({"type": "apply_extracted"})
@@ -852,13 +568,15 @@ class Resource:
 
         for link_extracted_data in message.link_extracted_data:
             await self._apply_link_extracted_data(link_extracted_data)
-            await self.maybe_update_title_metadata(link_extracted_data)
+            await self.maybe_update_resource_title_from_link(link_extracted_data)
             extracted_languages.append(link_extracted_data.language)
 
         for file_extracted_data in message.file_extracted_data:
             await self._apply_file_extracted_data(file_extracted_data)
             extracted_languages.append(file_extracted_data.language)
 
+        await self.maybe_update_resource_title_from_file_extracted_data(message)
+
         # Metadata should go first
         for field_metadata in message.field_metadata:
             await self._apply_field_computed_metadata(field_metadata)
@@ -869,10 +587,9 @@ class Resource:
         # Upload to binary storage
         # Vector indexing
         if self.disable_vectors is False:
+            await self.get_fields(force=True)
             for field_vectors in message.field_vectors:
                 await self._apply_extracted_vectors(field_vectors)
-            for user_vectors in message.user_vectors:
-                await self._apply_user_vectors(user_vectors)
 
         # Only uploading to binary storage
         for field_large_metadata in message.field_large_metadata:
@@ -896,9 +613,7 @@ class Resource:
             extracted_text.field,
         )
 
-    async def _apply_question_answers(
-        self, question_answers: FieldQuestionAnswerWrapper
-    ):
+    async def _apply_question_answers(self, question_answers: FieldQuestionAnswerWrapper):
         field = question_answers.field
         field_obj = await self.get_field(field.field, field.field_type, load=False)
         await field_obj.set_question_answers(question_answers)
@@ -918,19 +633,27 @@ class Resource:
 
         maybe_update_basic_summary(self.basic, link_extracted_data.description)
 
-    async def maybe_update_title_metadata(self, link_extracted_data: LinkExtractedData):
+    async def maybe_update_resource_title_from_link(self, link_extracted_data: LinkExtractedData):
+        """
+        When parsing link extracted data, we want to replace the resource title for the first link
+        that gets processed and has a title, and only if the current title is a URL, which we take
+        as a hint that the title was not set by the user.
+        """
         assert self.basic is not None
         if not link_extracted_data.title:
             return
         if not (self.basic.title.startswith("http") or self.basic.title == ""):
             return
-
         title = link_extracted_data.title
-        self.basic.title = title
+        await self.update_resource_title(title)
+
+    async def update_resource_title(self, computed_title: str) -> None:
+        assert self.basic is not None
+        self.basic.title = computed_title
         # Extracted text
         field = await self.get_field("title", FieldType.GENERIC, load=False)
         etw = ExtractedTextWrapper()
-        etw.body.text = title
+        etw.body.text = computed_title
         await field.set_extracted_text(etw)
 
         # Field computed metadata
@@ -942,11 +665,8 @@ class Resource:
         fcm = await field.get_field_metadata(force=True)
         if fcm is not None:
             fcmw.metadata.CopyFrom(fcm)
-
         fcmw.metadata.metadata.ClearField("paragraphs")
-        paragraph = Paragraph(
-            start=0, end=len(title), kind=Paragraph.TypeParagraph.TITLE
-        )
+        paragraph = Paragraph(start=0, end=len(computed_title), kind=Paragraph.TypeParagraph.TITLE)
         fcmw.metadata.metadata.paragraphs.append(paragraph)
 
         await field.set_field_metadata(fcmw)
@@ -963,9 +683,54 @@ class Resource:
         maybe_update_basic_icon(self.basic, file_extracted_data.icon)
         maybe_update_basic_thumbnail(self.basic, file_extracted_data.file_thumbnail)
 
-    async def _apply_field_computed_metadata(
-        self, field_metadata: FieldComputedMetadataWrapper
-    ):
+    async def _should_update_resource_title_from_file_metadata(self) -> bool:
+        """
+        We only want to update resource title from file metadata if the title is empty,
+        equal to the resource uuid or equal to any of the file filenames in the resource.
+        """
+        basic = await self.get_basic()
+        if basic is None:
+            return True
+        current_title = basic.title
+        if current_title == "":
+            # If the title is empty, we should update it
+            return True
+        if current_title == self.uuid:
+            # If the title is the same as the resource uuid, we should update it
+            return True
+        fields = await self.get_fields(force=True)
+        filenames = set()
+        for (field_type, _), field_obj in fields.items():
+            if field_type == FieldType.FILE:
+                field_value: Optional[FieldFile] = await field_obj.get_value()
+                if field_value is not None:
+                    if field_value.file.filename not in ("", None):
+                        filenames.add(field_value.file.filename)
+        if current_title in filenames:
+            # If the title is equal to any of the file filenames, we should update it
+            return True
+        return False
+
+    async def maybe_update_resource_title_from_file_extracted_data(self, message: BrokerMessage):
+        """
+        Update the resource title with the first file that has a title extracted.
+        """
+        if not await self._should_update_resource_title_from_file_metadata():
+            return
+        for fed in message.file_extracted_data:
+            if fed.title == "":
+                # Skip if the extracted title is empty
+                continue
+            fid = FieldId.from_pb(rid=self.uuid, field_type=FieldType.FILE, key=fed.field)
+            logger.info(
+                "Updating resource title from file extracted data",
+                extra={"kbid": self.kb.kbid, "field": fid.full(), "new_title": fed.title},
+            )
+            await self.update_resource_title(fed.title)
+            # Break after the first file with a title is found
+            break
+
+    async def _apply_field_computed_metadata(self, field_metadata: FieldComputedMetadataWrapper):
         assert self.basic is not None
         maybe_update_basic_summary(self.basic, field_metadata.metadata.metadata.summary)
 
@@ -974,17 +739,11 @@ class Resource:
             field_metadata.field.field_type,
             load=False,
         )
-        (
-            metadata,
-            replace_field,
-            replace_splits,
-        ) = await field_obj.set_field_metadata(field_metadata)
+        metadata = await field_obj.set_field_metadata(field_metadata)
         field_key = self.generate_field_id(field_metadata.field)
 
         page_positions: Optional[FilePagePositions] = None
-        if field_metadata.field.field_type == FieldType.FILE and isinstance(
-            field_obj, File
-        ):
+        if field_metadata.field.field_type == FieldType.FILE and isinstance(field_obj, File):
             page_positions = await get_file_page_positions(field_obj)
 
         user_field_metadata = next(
@@ -1002,29 +761,24 @@ class Resource:
             self.indexer.apply_field_metadata,
             field_key,
             metadata,
-            replace_field=replace_field,
-            replace_splits=replace_splits,
             page_positions=page_positions,
             extracted_text=extracted_text,
             basic_user_field_metadata=user_field_metadata,
+            replace_field=True,
         )
         loop = asyncio.get_running_loop()
         await loop.run_in_executor(_executor, apply_field_metadata)
 
-        maybe_update_basic_thumbnail(
-            self.basic, field_metadata.metadata.metadata.thumbnail
-        )
+        maybe_update_basic_thumbnail(self.basic, field_metadata.metadata.metadata.thumbnail)
 
         add_field_classifications(self.basic, field_metadata)
 
     async def _apply_extracted_vectors(self, field_vectors: ExtractedVectorsWrapper):
-        if not self.has_field(
-            field_vectors.field.field_type, field_vectors.field.field
-        ):
+        # Store vectors in the resource
+
+        if not self.has_field(field_vectors.field.field_type, field_vectors.field.field):
             # skipping because field does not exist
-            logger.warning(
-                f'Field "{field_vectors.field.field}" does not exist, skipping vectors'
-            )
+            logger.warning(f'Field "{field_vectors.field.field}" does not exist, skipping vectors')
             return
 
         field_obj = await self.get_field(
@@ -1032,49 +786,44 @@ class Resource:
             field_vectors.field.field_type,
             load=False,
         )
-        (
-            vo,
-            replace_field_sentences,
-            replace_splits_sentences,
-        ) = await field_obj.set_vectors(field_vectors)
+        vo = await field_obj.set_vectors(field_vectors)
+
+        # Prepare vectors to be indexed
+
         field_key = self.generate_field_id(field_vectors.field)
         if vo is not None:
-            apply_field_vectors = partial(
+            vectorset_id = field_vectors.vectorset_id or None
+            if vectorset_id is None:
+                dimension = await datamanagers.kb.get_matryoshka_vector_dimension(
+                    self.txn, kbid=self.kb.kbid
+                )
+            else:
+                config = await datamanagers.vectorsets.get(
+                    self.txn, kbid=self.kb.kbid, vectorset_id=vectorset_id
+                )
+                if config is None:
+                    logger.warning(
+                        f"Trying to apply a resource on vectorset '{vectorset_id}' that doesn't exist."
+                    )
+                    return
+                dimension = config.vectorset_index_config.vector_dimension
+                if not dimension:
+                    raise ValueError(f"Vector dimension not set for vectorset '{vectorset_id}'")
+
+            apply_field_vectors_partial = partial(
                 self.indexer.apply_field_vectors,
                 field_key,
                 vo,
-                replace_field_sentences,
-                replace_splits_sentences,
+                vectorset=vectorset_id,
+                replace_field=True,
+                matryoshka_vector_dimension=dimension,
             )
             loop = asyncio.get_running_loop()
-            await loop.run_in_executor(_executor, apply_field_vectors)
+            await loop.run_in_executor(_executor, apply_field_vectors_partial)
         else:
             raise AttributeError("VO not found on set")
 
-    async def _apply_user_vectors(self, user_vectors: UserVectorsWrapper):
-        field_obj = await self.get_field(
-            user_vectors.field.field,
-            user_vectors.field.field_type,
-            load=False,
-        )
-        uv, vectors_to_delete = await field_obj.set_user_vectors(user_vectors)
-        field_key = self.generate_field_id(user_vectors.field)
-        if uv is not None:
-            # We need to make sure that the vectors replaced are not on the new vectors
-            # So we extend the vectors to delete with the one replaced by the update
-            for vectorset, vectors in vectors_to_delete.items():
-                for vector in vectors.vectors:
-                    if vector not in user_vectors.vectors_to_delete[vectorset].vectors:
-                        user_vectors.vectors_to_delete[vectorset].vectors.append(vector)
-            self.indexer.apply_user_vectors(
-                field_key, uv, user_vectors.vectors_to_delete
-            )
-        else:
-            raise AttributeError("User Vectors not found on set")
-
-    async def _apply_field_large_metadata(
-        self, field_large_metadata: LargeComputedMetadataWrapper
-    ):
+    async def _apply_field_large_metadata(self, field_large_metadata: LargeComputedMetadataWrapper):
         field_obj = await self.get_field(
             field_large_metadata.field.field,
             field_large_metadata.field.field_type,
@@ -1083,7 +832,7 @@ class Resource:
         await field_obj.set_large_field_metadata(field_large_metadata)
 
     def generate_field_id(self, field: FieldID) -> str:
-        return f"{FIELD_TYPE_TO_ID[field.field_type]}/{field.field}"
+        return f"{FIELD_TYPE_PB_TO_STR[field.field_type]}/{field.field}"
 
     async def compute_security(self, brain: ResourceBrain):
         security = await self.get_security()
@@ -1102,7 +851,7 @@ class Resource:
         brain.set_resource_metadata(basic=basic, origin=origin)
         for type, field in await self.get_fields_ids(force=True):
             fieldobj = await self.get_field(field, type, load=False)
-            fieldid = FieldID(field_type=type, field=field)  # type: ignore
+            fieldid = FieldID(field_type=type, field=field)
             fieldkey = self.generate_field_id(fieldid)
             extracted_metadata = await fieldobj.get_field_metadata()
             valid_user_field_metadata = None
@@ -1113,16 +862,16 @@ class Resource:
                 ):
                     valid_user_field_metadata = user_field_metadata
                     break
+
+            generated_by = await fieldobj.generated_by()
             brain.apply_field_labels(
                 fieldkey,
                 extracted_metadata,
                 self.uuid,
+                generated_by,
                 basic.usermetadata,
                 valid_user_field_metadata,
             )
-            if type == FieldType.KEYWORDSET:
-                field_data = await fieldobj.db_get_value()
-                brain.process_keywordset_fields(fieldkey, field_data)
 
     @processor_observer.wrap({"type": "compute_global_text"})
     async def compute_global_text(self):
@@ -1159,12 +908,10 @@ class Resource:
                 for fieldmetadata in self.basic.fieldmetadata:
                     field_id = self.generate_field_id(fieldmetadata.field)
                     for annotationparagraph in fieldmetadata.paragraphs:
-                        userdefinedparagraphclass[
-                            annotationparagraph.key
-                        ] = annotationparagraph
+                        userdefinedparagraphclass[annotationparagraph.key] = annotationparagraph
 
         for (type_id, field_id), field in fields.items():
-            fieldid = FieldID(field_type=type_id, field=field_id)  # type: ignore
+            fieldid = FieldID(field_type=type_id, field=field_id)
             field_key = self.generate_field_id(fieldid)
             fm = await field.get_field_metadata()
             extracted_text = None
@@ -1179,9 +926,7 @@ class Resource:
             if fm is None:
                 continue
 
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [
-                (None, fm.metadata)
-            ]
+            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
             for subfield_metadata, splitted_metadata in fm.split_metadata.items():
                 field_metadatas.append((subfield_metadata, splitted_metadata))
 
@@ -1192,7 +937,7 @@ class Resource:
 
                 entities: dict[str, str] = {}
                 if enabled_metadata.entities:
-                    entities.update(field_metadata.ner)
+                    _update_entities_dict(entities, field_metadata)
 
                 precomputed_vectors = {}
                 if vo is not None:
@@ -1203,9 +948,7 @@ class Resource:
                         vectors = vo.vectors
                         base_vector_key = f"{self.uuid}/{field_key}"
                     for index, vector in enumerate(vectors.vectors):
-                        vector_key = (
-                            f"{base_vector_key}/{index}/{vector.start}-{vector.end}"
-                        )
+                        vector_key = f"{base_vector_key}/{index}/{vector.start}-{vector.end}"
                         precomputed_vectors[vector_key] = vector.vector
 
                 if extracted_text is not None:
@@ -1216,11 +959,11 @@ class Resource:
 
                 for paragraph in field_metadata.paragraphs:
                     if subfield is not None:
-                        paragraph_key = f"{self.uuid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
-                    else:
                         paragraph_key = (
-                            f"{self.uuid}/{field_key}/{paragraph.start}-{paragraph.end}"
+                            f"{self.uuid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
                         )
+                    else:
+                        paragraph_key = f"{self.uuid}/{field_key}/{paragraph.start}-{paragraph.end}"
 
                     if enabled_metadata.labels:
                         metadata.labels.ClearField("field")
@@ -1234,7 +977,9 @@ class Resource:
                         if subfield is not None:
                             sentence_key = f"{self.uuid}/{field_key}/{subfield}/{index}/{sentence.start}-{sentence.end}"
                         else:
-                            sentence_key = f"{self.uuid}/{field_key}/{index}/{sentence.start}-{sentence.end}"
+                            sentence_key = (
+                                f"{self.uuid}/{field_key}/{index}/{sentence.start}-{sentence.end}"
+                            )
 
                         if vo is not None:
                             metadata.ClearField("vector")
@@ -1273,12 +1018,10 @@ class Resource:
                 for fieldmetadata in self.basic.fieldmetadata:
                     field_id = self.generate_field_id(fieldmetadata.field)
                     for annotationparagraph in fieldmetadata.paragraphs:
-                        userdefinedparagraphclass[
-                            annotationparagraph.key
-                        ] = annotationparagraph
+                        userdefinedparagraphclass[annotationparagraph.key] = annotationparagraph
 
         for (type_id, field_id), field in fields.items():
-            fieldid = FieldID(field_type=type_id, field=field_id)  # type: ignore
+            fieldid = FieldID(field_type=type_id, field=field_id)
             field_key = self.generate_field_id(fieldid)
             fm = await field.get_field_metadata()
             extracted_text = None
@@ -1289,9 +1032,7 @@ class Resource:
             if fm is None:
                 continue
 
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [
-                (None, fm.metadata)
-            ]
+            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
             for subfield_metadata, splitted_metadata in fm.split_metadata.items():
                 field_metadatas.append((subfield_metadata, splitted_metadata))
 
@@ -1302,7 +1043,7 @@ class Resource:
 
                 entities: dict[str, str] = {}
                 if enabled_metadata.entities:
-                    entities.update(field_metadata.ner)
+                    _update_entities_dict(entities, field_metadata)
 
                 if extracted_text is not None:
                     if subfield is not None:
@@ -1312,11 +1053,11 @@ class Resource:
 
                 for paragraph in field_metadata.paragraphs:
                     if subfield is not None:
-                        paragraph_key = f"{self.uuid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
-                    else:
                         paragraph_key = (
-                            f"{self.uuid}/{field_key}/{paragraph.start}-{paragraph.end}"
+                            f"{self.uuid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
                         )
+                    else:
+                        paragraph_key = f"{self.uuid}/{field_key}/{paragraph.start}-{paragraph.end}"
 
                     if enabled_metadata.labels:
                         metadata.labels.ClearField("paragraph")
@@ -1344,9 +1085,7 @@ class Resource:
 
                         yield pb_paragraph
 
-    async def iterate_fields(
-        self, enabled_metadata: EnabledMetadata
-    ) -> AsyncIterator[TrainField]:
+    async def iterate_fields(self, enabled_metadata: EnabledMetadata) -> AsyncIterator[TrainField]:
         fields = await self.get_fields(force=True)
         metadata = TrainMetadata()
         if enabled_metadata.labels:
@@ -1356,7 +1095,7 @@ class Resource:
                 metadata.labels.resource.extend(self.basic.usermetadata.classifications)
 
         for (type_id, field_id), field in fields.items():
-            fieldid = FieldID(field_type=type_id, field=field_id)  # type: ignore
+            fieldid = FieldID(field_type=type_id, field=field_id)
             fm = await field.get_field_metadata()
             extracted_text = None
 
@@ -1366,9 +1105,7 @@ class Resource:
             if fm is None:
                 continue
 
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [
-                (None, fm.metadata)
-            ]
+            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
             for subfield_metadata, splitted_metadata in fm.split_metadata.items():
                 field_metadatas.append((subfield_metadata, splitted_metadata))
 
@@ -1385,7 +1122,7 @@ class Resource:
 
                 if enabled_metadata.entities:
                     metadata.ClearField("entities")
-                    metadata.entities.update(splitted_metadata.ner)
+                    _update_entities_dict(metadata.entities, splitted_metadata)
 
                 pb_field = TrainField()
                 pb_field.uuid = self.uuid
@@ -1393,9 +1130,7 @@ class Resource:
                 pb_field.metadata.CopyFrom(metadata)
                 yield pb_field
 
-    async def generate_train_resource(
-        self, enabled_metadata: EnabledMetadata
-    ) -> TrainResource:
+    async def generate_train_resource(self, enabled_metadata: EnabledMetadata) -> TrainResource:
         fields = await self.get_fields(force=True)
         metadata = TrainMetadata()
         if enabled_metadata.labels:
@@ -1422,9 +1157,7 @@ class Resource:
             if fm is None:
                 continue
 
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [
-                (None, fm.metadata)
-            ]
+            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
             for subfield_metadata, splitted_metadata in fm.split_metadata.items():
                 field_metadatas.append((subfield_metadata, splitted_metadata))
 
@@ -1433,7 +1166,7 @@ class Resource:
                     metadata.labels.field.extend(splitted_metadata.classifications)
 
                 if enabled_metadata.entities:
-                    metadata.entities.update(splitted_metadata.ner)
+                    _update_entities_dict(metadata.entities, splitted_metadata)
 
         pb_resource = TrainResource()
         pb_resource.uuid = self.uuid
@@ -1462,33 +1195,35 @@ def remove_field_classifications(basic: PBBasic, deleted_fields: list[FieldID]):
     Clean classifications of fields that have been deleted
     """
     field_classifications = [
-        fc
-        for fc in basic.computedmetadata.field_classifications
-        if fc.field not in deleted_fields
+        fc for fc in basic.computedmetadata.field_classifications if fc.field not in deleted_fields
     ]
     basic.computedmetadata.ClearField("field_classifications")
     basic.computedmetadata.field_classifications.extend(field_classifications)
 
 
-def add_field_classifications(
-    basic: PBBasic, fcmw: FieldComputedMetadataWrapper
-) -> bool:
+def add_field_classifications(basic: PBBasic, fcmw: FieldComputedMetadataWrapper) -> bool:
     """
     Returns whether some new field classifications were added
     """
-    if len(fcmw.metadata.metadata.classifications) == 0:
+    if len(fcmw.metadata.metadata.classifications) == 0 and all(
+        len(split.classifications) == 0 for split in fcmw.metadata.split_metadata.values()
+    ):
         return False
+
     remove_field_classifications(basic, [fcmw.field])
     fcfs = FieldClassifications()
     fcfs.field.CopyFrom(fcmw.field)
     fcfs.classifications.extend(fcmw.metadata.metadata.classifications)
+
+    for split_id, split in fcmw.metadata.split_metadata.items():
+        if split_id not in fcmw.metadata.deleted_splits:
+            fcfs.classifications.extend(split.classifications)
+
     basic.computedmetadata.field_classifications.append(fcfs)
     return True
 
 
-def add_entities_to_metadata(
-    entities: dict[str, str], local_text: str, metadata: TrainMetadata
-) -> None:
+def add_entities_to_metadata(entities: dict[str, str], local_text: str, metadata: TrainMetadata) -> None:
     for entity_key, entity_value in entities.items():
         if entity_key not in local_text:
             # Add the entity only if found in text
@@ -1502,9 +1237,7 @@ def add_entities_to_metadata(
         for _ in range(local_text.count(entity_key)):
             start = local_text.index(entity_key, last_occurrence_end)
             end = start + len(entity_key)
-            metadata.entity_positions[poskey].positions.append(
-                TrainPosition(start=start, end=end)
-            )
+            metadata.entity_positions[poskey].positions.append(TrainPosition(start=start, end=end))
             last_occurrence_end = end
 
 
@@ -1519,15 +1252,22 @@ def maybe_update_basic_icon(basic: PBBasic, mimetype: Optional[str]) -> bool:
     if basic.icon not in (None, "", "application/octet-stream", GENERIC_MIME_TYPE):
         # Icon already set or detected
         return False
+
     if not mimetype:
         return False
+
+    if not content_types.valid(mimetype):
+        logger.warning(
+            "Invalid mimetype. Skipping icon update.",
+            extra={"mimetype": mimetype, "rid": basic.uuid, "slug": basic.slug},
+        )
+        return False
+
     basic.icon = mimetype
     return True
 
 
-def maybe_update_basic_thumbnail(
-    basic: PBBasic, thumbnail: Optional[CloudFile]
-) -> bool:
+def maybe_update_basic_thumbnail(basic: PBBasic, thumbnail: Optional[CloudFile]) -> bool:
     if basic.thumbnail or thumbnail is None:
         return False
     basic.thumbnail = CloudLink.format_reader_download_uri(thumbnail.uri)
@@ -1569,3 +1309,23 @@ def extract_field_metadata_languages(
     for _, splitted_metadata in field_metadata.metadata.split_metadata.items():
         languages.add(splitted_metadata.language)
     return list(languages)
+
+
+def _update_entities_dict(target_entites_dict: MutableMapping[str, str], field_metadata: FieldMetadata):
+    """
+    Update the entities dict with the entities from the field metadata.
+    Method created to ease the transition from legacy ner field to new entities field.
+    """
+    # Data Augmentation + Processor entities
+    # This will overwrite entities detected from more than one data augmentation task
+    # TODO: Change TrainMetadata proto to accept multiple entities with the same text
+    entity_map = {
+        entity.text: entity.label
+        for data_augmentation_task_id, entities_wrapper in field_metadata.entities.items()
+        for entity in entities_wrapper.entities
+    }
+    target_entites_dict.update(entity_map)
+
+    # Legacy processor entities
+    # TODO: Remove once processor doesn't use this anymore and remove the positions and ner fields from the message
+    target_entites_dict.update(field_metadata.ner)