nucliadb 6.9.0.post5072__py3-none-any.whl → 6.9.0.post5078__py3-none-any.whl

migrations/0039_backfill_converation_splits_metadata.py

@@ -0,0 +1,106 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+
+"""Migration #39
+
+Backfill splits metadata on conversation fields
+
+"""
+
+import logging
+from typing import cast
+
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb.common.maindb.pg import PGTransaction
+from nucliadb.ingest.fields.conversation import (
+    CONVERSATION_SPLITS_METADATA,
+    Conversation,
+)
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
+from nucliadb.migrator.context import ExecutionContext
+from nucliadb_protos import resources_pb2
+from nucliadb_protos.resources_pb2 import SplitMetadata, SplitsMetadata
+from nucliadb_utils.storages.storage import Storage
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    BATCH_SIZE = 100
+    start = ""
+    while True:
+        to_fix: list[tuple[str, str]] = []
+        async with context.kv_driver.rw_transaction() as txn:
+            txn = cast(PGTransaction, txn)
+            async with txn.connection.cursor() as cur:
+                # Retrieve a bunch of conversation fields
+                await cur.execute(
+                    """
+                    SELECT key FROM resources
+                    WHERE key ~ ('^/kbs/' || %s || '/r/[^/]*/f/c/[^/]*$')
+                    AND key > %s
+                    ORDER BY key
+                    LIMIT %s""",
+                    (kbid, start, BATCH_SIZE),
+                )
+                rows = await cur.fetchall()
+                if len(rows) == 0:
+                    return
+                for row in rows:
+                    key = row[0]
+                    start = key
+                    rid = key.split("/")[4]
+                    field_id = key.split("/")[7]
+                    to_fix.append((rid, field_id))
+
+        for rid, field_id in to_fix:
+            async with context.kv_driver.rw_transaction() as txn2:
+                splits_metadata = await build_splits_metadata(
+                    txn2, context.blob_storage, kbid, rid, field_id
+                )
+                splits_metadata_key = CONVERSATION_SPLITS_METADATA.format(
+                    kbid=kbid, uuid=rid, type="c", field=field_id
+                )
+                await txn2.set(splits_metadata_key, splits_metadata.SerializeToString())
+                await txn2.commit()
+
+
+async def build_splits_metadata(
+    txn: Transaction, storage: Storage, kbid: str, rid: str, field_id: str
+) -> SplitsMetadata:
+    splits_metadata = SplitsMetadata()
+    kb_orm = KnowledgeBoxORM(txn, storage, kbid)
+    resource_obj = await kb_orm.get(rid)
+    if resource_obj is None:
+        return splits_metadata
+    field_obj: Conversation = await resource_obj.get_field(
+        field_id, resources_pb2.FieldType.CONVERSATION, load=False
+    )
+    conv_metadata = await field_obj.get_metadata()
+    for i in range(1, conv_metadata.pages + 1):
+        page = await field_obj.get_value(page=i)
+        if page is None:
+            continue
+        for message in page.messages:
+            splits_metadata.metadata.setdefault(message.ident, SplitMetadata())
+    return splits_metadata

nucliadb/ingest/fields/conversation.py

@@ -21,7 +21,7 @@ import uuid
 from typing import Any, Optional
 
 from nucliadb.ingest.fields.base import Field
-from nucliadb_protos.resources_pb2 import CloudFile, FieldConversation
+from nucliadb_protos.resources_pb2 import CloudFile, FieldConversation, SplitMetadata, SplitsMetadata
 from nucliadb_protos.resources_pb2 import Conversation as PBConversation
 from nucliadb_utils.storages.storage import StorageField
 
@@ -30,6 +30,7 @@ MAX_CONVERSATION_MESSAGES = 50 * 1024
 PAGE_SIZE = 200
 
 CONVERSATION_PAGE_VALUE = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/{page}"
+CONVERSATION_SPLITS_METADATA = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/splits_metadata"
 CONVERSATION_METADATA = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}"
 
 
@@ -54,6 +55,7 @@ class Conversation(Field[PBConversation]):
     ):
         super(Conversation, self).__init__(id, resource, pb, value)
         self.value = {}
+        self._splits_metadata: Optional[SplitsMetadata] = None
         self.metadata = None
 
     async def set_value(self, payload: PBConversation):
@@ -72,10 +74,13 @@ class Conversation(Field[PBConversation]):
             last_page = PBConversation()
             metadata.pages += 1
 
+        self._splits_metadata = await self.get_splits_metadata()
+
         # Make sure message attachment files are on our region. This is needed
         # to support the hybrid-onprem deployment as the attachments must be stored
         # at the storage services of the client's premises.
         for message in payload.messages:
+            self._splits_metadata.metadata.setdefault(message.ident, SplitMetadata())
             new_message_files = []
             for idx, file in enumerate(message.content.attachments):
                 if self.storage.needs_move(file, self.kbid):
@@ -119,6 +124,7 @@ class Conversation(Field[PBConversation]):
 
         # Finally, set the metadata
         await self.db_set_metadata(metadata)
+        await self.set_splits_metadata(self._splits_metadata)
 
     async def get_value(self, page: Optional[int] = None) -> Optional[PBConversation]:
         # If no page was requested, force fetch of metadata
@@ -205,3 +211,29 @@ class Conversation(Field[PBConversation]):
         self.metadata = payload
         self.resource.modified = True
         self._created = False
+
+    async def get_splits_metadata(self) -> SplitsMetadata:
+        if self._splits_metadata is None:
+            field_key = CONVERSATION_SPLITS_METADATA.format(
+                kbid=self.kbid,
+                uuid=self.uuid,
+                type=self.type,
+                field=self.id,
+            )
+            payload = await self.resource.txn.get(field_key)
+            if payload is None:
+                return SplitsMetadata()
+            self._splits_metadata = SplitsMetadata()
+            self._splits_metadata.ParseFromString(payload)
+        return self._splits_metadata
+
+    async def set_splits_metadata(self, payload: SplitsMetadata) -> None:
+        key = CONVERSATION_SPLITS_METADATA.format(
+            kbid=self.kbid,
+            uuid=self.uuid,
+            type=self.type,
+            field=self.id,
+        )
+        await self.resource.txn.set(key, payload.SerializeToString())
+        self._split_metadata = payload
+        self.resource.modified = True

nucliadb/ingest/orm/brain_v2.py

@@ -217,6 +217,7 @@ class ResourceBrain:
         replace_field: bool,
         skip_paragraphs_index: Optional[bool],
         skip_texts_index: Optional[bool],
+        append_splits: Optional[set[str]] = None,
     ) -> None:
         # We need to add the extracted text to the texts section of the Resource so that
         # the paragraphs can be indexed
@@ -234,6 +235,7 @@ class ResourceBrain:
             user_field_metadata,
             replace_field=replace_field,
             skip_paragraphs=skip_paragraphs_index,
+            append_splits=append_splits,
         )
 
     @observer.wrap({"type": "apply_field_paragraphs"})
@@ -246,6 +248,7 @@ class ResourceBrain:
         user_field_metadata: Optional[UserFieldMetadata],
         replace_field: bool,
         skip_paragraphs: Optional[bool],
+        append_splits: Optional[set[str]] = None,
     ) -> None:
         if skip_paragraphs is not None:
             self.brain.skip_paragraphs = skip_paragraphs
@@ -254,7 +257,12 @@ class ResourceBrain:
         paragraph_pages = ParagraphPages(page_positions) if page_positions else None
         # Splits of the field
         for subfield, field_metadata in field_computed_metadata.split_metadata.items():
-            extracted_text_str = extracted_text.split_text[subfield] if extracted_text else None
+            if should_skip_split_indexing(subfield, replace_field, append_splits):
+                continue
+            if subfield not in extracted_text.split_text:
+                # No extracted text for this split
+                continue
+            extracted_text_str = extracted_text.split_text[subfield]
             for idx, paragraph in enumerate(field_metadata.paragraphs):
                 key = f"{self.rid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
                 denied_classifications = set(user_paragraph_classifications.denied.get(key, []))
@@ -308,7 +316,7 @@ class ResourceBrain:
                 self.brain.paragraphs[field_key].paragraphs[key].CopyFrom(p)
 
         # Main field
-        extracted_text_str = extracted_text.text if extracted_text else None
+        extracted_text_str = extracted_text.text
         for idx, paragraph in enumerate(field_computed_metadata.metadata.paragraphs):
             key = f"{self.rid}/{field_key}/{paragraph.start}-{paragraph.end}"
             denied_classifications = set(user_paragraph_classifications.denied.get(key, []))
@@ -496,9 +504,12 @@ class ResourceBrain:
         replace_field: bool = False,
         # cut to specific dimension if specified
         vector_dimension: Optional[int] = None,
+        append_splits: Optional[set[str]] = None,
     ):
         fid = ids.FieldId.from_string(f"{self.rid}/{field_id}")
         for subfield, vectors in vo.split_vectors.items():
+            if should_skip_split_indexing(subfield, replace_field, append_splits):
+                continue
             _field_id = ids.FieldId(
                 rid=fid.rid,
                 type=fid.type,
@@ -792,3 +803,10 @@ class ParagraphPages:
             if len(self._materialized) > 0:
                 return self._materialized[-1]
             return 0
+
+
+def should_skip_split_indexing(
+    split: str, replace_field: bool, append_splits: Optional[set[str]]
+) -> bool:
+    # When replacing the whole field, reindex all splits. Otherwise, we're only indexing the splits that are appended
+    return not replace_field and append_splits is not None and split not in append_splits

nucliadb/ingest/orm/index_message.py

@@ -25,6 +25,7 @@ from typing import Optional, Sequence
 from nidx_protos.noderesources_pb2 import Resource as IndexMessage
 
 from nucliadb.common import datamanagers
+from nucliadb.ingest.fields.conversation import Conversation
 from nucliadb.ingest.fields.exceptions import FieldAuthorNotFound
 from nucliadb.ingest.fields.file import File
 from nucliadb.ingest.orm.brain_v2 import ResourceBrain
@@ -32,6 +33,7 @@ from nucliadb.ingest.orm.metrics import index_message_observer as observer
 from nucliadb.ingest.orm.resource import Resource, get_file_page_positions
 from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
 from nucliadb_protos.resources_pb2 import Basic, FieldID, FieldType
+from nucliadb_protos.utils_pb2 import ExtractedText
 from nucliadb_protos.writer_pb2 import BrokerMessage
 
 
@@ -69,6 +71,7 @@ class IndexMessageBuilder:
         relations: bool = True,
         replace: bool = True,
         vectorset_configs: Optional[list[VectorSetConfig]] = None,
+        append_splits: Optional[set[str]] = None,
     ):
         field = await self.resource.get_field(fieldid.field, fieldid.field_type)
         extracted_text = await field.get_extracted_text()
@@ -120,6 +123,7 @@ class IndexMessageBuilder:
                     replace_field=replace_paragraphs,
                     skip_paragraphs_index=skip_paragraphs_index,
                     skip_texts_index=skip_texts_index,
+                    append_splits=append_splits,
                 )
         if vectors:
             assert vectorset_configs is not None
@@ -137,6 +141,7 @@ class IndexMessageBuilder:
                         vectorset=vectorset_config.vectorset_id,
                         replace_field=replace,
                         vector_dimension=dimension,
+                        append_splits=append_splits,
                     )
         if relations:
             await asyncio.to_thread(
@@ -214,6 +219,19 @@ class IndexMessageBuilder:
         for fieldid in fields_to_index:
             if fieldid in message.delete_fields:
                 continue
+
+            # For conversation fields, we only replace the full field if it is not an append messages operation.
+            # All other fields are always replaced upon modification.
+            replace_field = True
+            modified_splits = None
+            if fieldid.field_type == FieldType.CONVERSATION:
+                modified_splits = await get_bm_modified_split_ids(fieldid, message, self.resource)
+                stored_splits = await get_stored_split_ids(fieldid, self.resource)
+                is_append_messages_op = modified_splits.issubset(stored_splits) and 0 < len(
+                    modified_splits
+                ) < len(stored_splits)
+                replace_field = not is_append_messages_op
+
             await self._apply_field_index_data(
                 self.brain,
                 fieldid,
@@ -222,8 +240,9 @@ class IndexMessageBuilder:
                 paragraphs=needs_paragraphs_update(fieldid, message),
                 relations=needs_relations_update(fieldid, message),
                 vectors=needs_vectors_update(fieldid, message),
-                replace=True,
+                replace=replace_field,
                 vectorset_configs=vectorsets_configs,
+                append_splits=modified_splits,
             )
         return self.brain.brain
 
@@ -354,6 +373,37 @@ def needs_vectors_update(
     return any(field_vectors.field == field_id for field_vectors in message.field_vectors)
 
 
+async def get_bm_modified_split_ids(
+    conversation_field_id: FieldID,
+    message: BrokerMessage,
+    resource: Resource,
+) -> set[str]:
+    message_etw = next(
+        (etw for etw in message.extracted_text if etw.field == conversation_field_id), None
+    )
+    if message_etw is None:
+        return set()
+    storage = resource.storage
+    if message_etw.HasField("file"):
+        raw_payload = await storage.downloadbytescf(message_etw.file)
+        message_extracted_text = ExtractedText()
+        message_extracted_text.ParseFromString(raw_payload.read())
+        raw_payload.flush()
+    else:
+        message_extracted_text = message_etw.body
+    return set(message_extracted_text.split_text.keys())
+
+
+async def get_stored_split_ids(
+    conversation_field_id: FieldID,
+    resource: Resource,
+) -> set[str]:
+    fid = conversation_field_id
+    conv: Conversation = await resource.get_field(fid.field, fid.field_type, load=False)
+    splits_metadata = await conv.get_splits_metadata()
+    return set(splits_metadata.metadata)
+
+
 def needs_relations_update(
     field_id: FieldID,
     message: BrokerMessage,

{nucliadb-6.9.0.post5072.dist-info → nucliadb-6.9.0.post5078.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.9.0.post5072
+Version: 6.9.0.post5078
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,11 +19,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.9.0.post5072
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.9.0.post5072
-Requires-Dist: nucliadb-protos>=6.9.0.post5072
-Requires-Dist: nucliadb-models>=6.9.0.post5072
-Requires-Dist: nidx-protos>=6.9.0.post5072
+Requires-Dist: nucliadb-telemetry[all]>=6.9.0.post5078
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.9.0.post5078
+Requires-Dist: nucliadb-protos>=6.9.0.post5078
+Requires-Dist: nucliadb-models>=6.9.0.post5078
+Requires-Dist: nidx-protos>=6.9.0.post5078
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.50.0
 Requires-Dist: uvicorn[standard]

{nucliadb-6.9.0.post5072.dist-info → nucliadb-6.9.0.post5078.dist-info}/RECORD

@@ -34,6 +34,7 @@ migrations/0035_rollover_nidx_texts_4.py,sha256=W0_AUd01pjMpYMDC3yqF6HzDLgcnnPpr
 migrations/0036_backfill_catalog_slug.py,sha256=toYqxH_EfUFqoVn_cOdR5Fg8bWZU5BoFMfPBSf74LKU,2957
 migrations/0037_backfill_catalog_facets.py,sha256=IH7H4OZ4tzws6xEh7Qro0bPDHDYOoVViEUj-JwPPe1U,2791
 migrations/0038_backfill_catalog_field_labels.py,sha256=F519nYngJDb1Mtwf-OQpweDPWKPxAlqdxy5E-DyQrhA,3492
+migrations/0039_backfill_converation_splits_metadata.py,sha256=NtL9S6Kx8mbSjNJLjIsc-6vVNymD0YKlF9vEkGUEDds,3958
 migrations/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 migrations/pg/0001_bootstrap.py,sha256=3O_P17l0d0h48nebN6VQLXzM_B7S7zvDpaLR0koVgWE,1274
 migrations/pg/0002_catalog.py,sha256=Rsleecu351Ty19kYZgOpqX5G3MEAY8nMxCJrAeuS2Mw,1690
@@ -154,18 +155,18 @@ nucliadb/ingest/consumer/shard_creator.py,sha256=UKIk0yaS_jC_nGQqymn9NGJWzwZEqhI
 nucliadb/ingest/consumer/utils.py,sha256=jpX8D4lKzuPCpArQLZeX_Zczq3pfen_zAf8sPJfOEZU,2642
 nucliadb/ingest/fields/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/ingest/fields/base.py,sha256=D8NzawonF7hivDW9zvQBbV938TKA6e2OCGqV4kS96RU,24405
-nucliadb/ingest/fields/conversation.py,sha256=ShdPapTIG7sA05YlG1Dj2CsAwNwibsqRSQrsuZnr8YI,7723
+nucliadb/ingest/fields/conversation.py,sha256=Jb6nDvlGQiX60VM6NzX1-okrNexlXuUSZ-Ge_yIJYVw,9098
 nucliadb/ingest/fields/exceptions.py,sha256=sZBk21BSrXFdOdo1qUdCAyD-9YMYakSLdn4_WdIPCIQ,1217
 nucliadb/ingest/fields/file.py,sha256=1v4jLg3balUua2VmSV8hHkAwPFShTUCOzufZvIUQcQw,4740
 nucliadb/ingest/fields/generic.py,sha256=elgtqv15aJUq3zY7X_g0bli_2BpcwPArVvzhe54Y4Ig,1547
 nucliadb/ingest/fields/link.py,sha256=kN_gjRUEEj5cy8K_BwPijYg3TiWhedc24apXYlTbRJs,4172
 nucliadb/ingest/fields/text.py,sha256=2grxo8twWbpXEd_iwUMBw9q0dWorVmlPONmY5d1ThwQ,1684
 nucliadb/ingest/orm/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/ingest/orm/brain_v2.py,sha256=8MAo1N_nhoGy73TvKBuaw-NaMxIanRPCNttw6dFY4jk,33647
+nucliadb/ingest/orm/brain_v2.py,sha256=mBze0LBft9pPiEtTIlfwCDVk597RSBXsx0Pte4ciTTg,34456
 nucliadb/ingest/orm/broker_message.py,sha256=XWaiZgDOz94NPOPT-hqbRr5ZkpVimUw6PjUJNftfoVw,7514
 nucliadb/ingest/orm/entities.py,sha256=kXyeF6XOpFKhEsGLcY-GLIk21Exp0cJst4XQQ9jJoug,14791
 nucliadb/ingest/orm/exceptions.py,sha256=gsp7TtVNQPiIEh-zf_UEJClwuFU0iu-5vzj0OrKMScg,1550
-nucliadb/ingest/orm/index_message.py,sha256=DlGLuuuCsXR_rqxd6CEZMYuOx1TIiq1mR5ue114rGUk,14473
+nucliadb/ingest/orm/index_message.py,sha256=mWlpQ0-KChSVIbHewVE8sXCe-7LiPIIh0cBqr3axU8o,16554
 nucliadb/ingest/orm/knowledgebox.py,sha256=OG9dmfklYf1PgTHwQd_iFZOociLEvUSMMv1ZKeUgecE,23910
 nucliadb/ingest/orm/metrics.py,sha256=OiuggTh-n3kZHA2G73NEUdIlh8c3yFrbusI88DK-Mko,1273
 nucliadb/ingest/orm/resource.py,sha256=zQeZyZ-tCxr-DhonLobfZRkz_iEew0Y-cGfXeNNIHG0,40432
@@ -384,8 +385,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vu1BGg4VqJ_x2P1u2BxqPKlSfw5orT_a3R-Ln5oPUpU,8483
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.9.0.post5072.dist-info/METADATA,sha256=IuymQS1C87M_L3u-9Bpy-Wf76fYOScOwTD18Dr5srAQ,4158
-nucliadb-6.9.0.post5072.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.9.0.post5072.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.9.0.post5072.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.9.0.post5072.dist-info/RECORD,,
+nucliadb-6.9.0.post5078.dist-info/METADATA,sha256=x9K2fNkUdKoA0BNuiVFTg6IkhDHeiAvUZntKj6P4g94,4158
+nucliadb-6.9.0.post5078.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.9.0.post5078.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.9.0.post5078.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.9.0.post5078.dist-info/RECORD,,