nucliadb 6.9.2.post5275__py3-none-any.whl → 6.9.2.post5276__py3-none-any.whl

nucliadb/common/datamanagers/atomic.py

@@ -88,6 +88,7 @@ class resources:
     get_resource_uuid_from_slug = ro_txn_wrap(resources_dm.get_resource_uuid_from_slug)
     resource_exists = ro_txn_wrap(resources_dm.resource_exists)
     slug_exists = ro_txn_wrap(resources_dm.slug_exists)
+    get_all_field_ids = ro_txn_wrap(resources_dm.get_all_field_ids)
 
 
 class labelset:

nucliadb/search/search/chat/prompt.py

@@ -26,6 +26,7 @@ from typing import Deque, Dict, List, Optional, Sequence, Tuple, Union, cast
 import yaml
 from pydantic import BaseModel
 
+from nucliadb.common import datamanagers
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB, FieldId, ParagraphId
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.common.models_utils import from_proto
@@ -589,18 +590,7 @@ async def field_extension_prompt_context(
         if resource_uuid not in ordered_resources:
             ordered_resources.append(resource_uuid)
 
-    # Fetch the extracted texts of the specified fields for each resource
-    extend_fields = strategy.fields
-    extend_field_ids = []
-    for resource_uuid in ordered_resources:
-        for field_id in extend_fields:
-            try:
-                fid = FieldId.from_string(f"{resource_uuid}/{field_id.strip('/')}")
-                extend_field_ids.append(fid)
-            except ValueError:  # pragma: no cover
-                # Invalid field id, skiping
-                continue
-
+    extend_field_ids = await get_matching_field_ids(kbid, ordered_resources, strategy)
     tasks = [hydrate_field_text(kbid, fid) for fid in extend_field_ids]
     field_extracted_texts = await run_concurrently(tasks)
 
@@ -630,6 +620,43 @@ async def field_extension_prompt_context(
             context[paragraph.id] = _clean_paragraph_text(paragraph)
 
 
+async def get_matching_field_ids(
+    kbid: str, ordered_resources: list[str], strategy: FieldExtensionStrategy
+) -> list[FieldId]:
+    extend_field_ids: list[FieldId] = []
+    # Fetch the extracted texts of the specified fields for each resource
+    for resource_uuid in ordered_resources:
+        for field_id in strategy.fields:
+            try:
+                fid = FieldId.from_string(f"{resource_uuid}/{field_id.strip('/')}")
+                extend_field_ids.append(fid)
+            except ValueError:  # pragma: no cover
+                # Invalid field id, skiping
+                continue
+    if len(strategy.data_augmentation_field_prefixes) > 0:
+        for resource_uuid in ordered_resources:
+            all_field_ids = await datamanagers.atomic.resources.get_all_field_ids(
+                kbid=kbid, rid=resource_uuid, for_update=False
+            )
+            if all_field_ids is None:
+                continue
+            for fieldid in all_field_ids.fields:
+                # Generated fields are always text fields starting with "da-"
+                if any(
+                    (
+                        fieldid.field_type == resources_pb2.FieldType.TEXT
+                        and fieldid.field.startswith(f"da-{prefix}-")
+                    )
+                    for prefix in strategy.data_augmentation_field_prefixes
+                ):
+                    extend_field_ids.append(
+                        FieldId.from_pb(
+                            rid=resource_uuid, field_type=fieldid.field_type, key=fieldid.field
+                        )
+                    )
+    return extend_field_ids
+
+
 async def get_orm_field(kbid: str, field_id: FieldId) -> Optional[Field]:
     resource = await cache.get_resource(kbid, field_id.rid)
     if resource is None:  # pragma: no cover

{nucliadb-6.9.2.post5275.dist-info → nucliadb-6.9.2.post5276.dist-info}/METADATA

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

{nucliadb-6.9.2.post5275.dist-info → nucliadb-6.9.2.post5276.dist-info}/RECORD

@@ -95,7 +95,7 @@ nucliadb/common/cluster/standalone/utils.py,sha256=af3r-x_GF7A6dwIAhZLR-r-SZQEVx
 nucliadb/common/context/__init__.py,sha256=IKAHuiCjbOEsqfLozWwJ6mRFzFncsZMyxNC5E_XZ5EM,6016
 nucliadb/common/context/fastapi.py,sha256=mH_8n5t7quNSPivNM2JS5EQf2sTVJsdzXW6LaY7EHAA,1629
 nucliadb/common/datamanagers/__init__.py,sha256=xKc6ZMqKUs20R90jJT4xkQ8TFMNwQnhhuWnBBqVnKdM,2084
-nucliadb/common/datamanagers/atomic.py,sha256=-3T2--A0_qFWmy573t-DS29EHFcMTyjGvcvky6J9tAM,3341
+nucliadb/common/datamanagers/atomic.py,sha256=tZEOUqrbUjNKEgxNn8IpukK0BuBQMtf-rKUq0KuT_YY,3409
 nucliadb/common/datamanagers/cluster.py,sha256=iU0b7AESm1Yi8Wp3pIKgqixZGNMjeBrxSpvEKsaZKgY,1831
 nucliadb/common/datamanagers/entities.py,sha256=gI-0mbMlqrr9FiyhexEh6czhgYcMxE2s9m4o866EK9o,5340
 nucliadb/common/datamanagers/exceptions.py,sha256=Atz_PP_GGq4jgJaWcAkcRbHBoBaGcC9yJvFteylKtTE,883
@@ -271,7 +271,7 @@ nucliadb/search/search/chat/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn
 nucliadb/search/search/chat/ask.py,sha256=siMKgXEuAOUPa9oxi72idEOhN9BXZ-NVvWY3rqMMM3I,41610
 nucliadb/search/search/chat/exceptions.py,sha256=Siy4GXW2L7oPhIR86H3WHBhE9lkV4A4YaAszuGGUf54,1356
 nucliadb/search/search/chat/images.py,sha256=PA8VWxT5_HUGfW1ULhKTK46UBsVyINtWWqEM1ulzX1E,3095
-nucliadb/search/search/chat/prompt.py,sha256=dqk6969zEymfGuDl0S1uAGM_Ii5kgdcCgNkuVrM0_8c,54406
+nucliadb/search/search/chat/prompt.py,sha256=Ti0dO8RDY9YeRPowwiWbojjd-IgpAKkHEat79jGFidk,55671
 nucliadb/search/search/chat/query.py,sha256=Ds_0vYfCdbfTZerSM0AykkfKYYJj88j_M2XA87FfjJQ,17101
 nucliadb/search/search/hydrator/__init__.py,sha256=3Pc-rcax4TI174qcrllnReE728DoJTaA8tpvBUFf98g,7005
 nucliadb/search/search/hydrator/fields.py,sha256=LhKw-aNU5eJqfZADtq3iB7AGXm0l_QabAAoSHJTk8Is,5962
@@ -385,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.2.post5275.dist-info/METADATA,sha256=Rw1v7kl_32369UkLp0l5JKfZ8ipe5j0wd1hNBDFTxiw,4118
-nucliadb-6.9.2.post5275.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.9.2.post5275.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.9.2.post5275.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.9.2.post5275.dist-info/RECORD,,
+nucliadb-6.9.2.post5276.dist-info/METADATA,sha256=92ehiGl-d963I9_9rxJo1lyNxd7om9lcGXMtunuwhc8,4118
+nucliadb-6.9.2.post5276.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.9.2.post5276.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.9.2.post5276.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.9.2.post5276.dist-info/RECORD,,