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

nucliadb/ingest/orm/utils.py

@@ -18,10 +18,11 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import urllib.parse
-from typing import Optional
+from typing import Sequence
 
+from nucliadb.ingest.processing import PushPayload
+from nucliadb_models.text import PushTextFormat, Text
 from nucliadb_protos.resources_pb2 import (
-    Basic,
     ExtractedTextWrapper,
     FieldComputedMetadataWrapper,
     FieldType,
@@ -30,35 +31,6 @@ from nucliadb_protos.resources_pb2 import (
 )
 from nucliadb_protos.writer_pb2 import BrokerMessage
 
-from nucliadb.common.maindb.driver import Transaction
-from nucliadb.ingest.processing import PushPayload
-from nucliadb.ingest.settings import settings as ingest_settings
-from nucliadb_models.text import PushTextFormat, Text
-
-KB_RESOURCE_BASIC_FS = "/kbs/{kbid}/r/{uuid}/basic"  # Only used on FS driver
-KB_RESOURCE_BASIC = "/kbs/{kbid}/r/{uuid}"
-
-
-async def set_basic(txn: Transaction, kbid: str, uuid: str, basic: Basic):
-    if ingest_settings.driver == "local":
-        await txn.set(
-            KB_RESOURCE_BASIC_FS.format(kbid=kbid, uuid=uuid),
-            basic.SerializeToString(),
-        )
-    else:
-        await txn.set(
-            KB_RESOURCE_BASIC.format(kbid=kbid, uuid=uuid),
-            basic.SerializeToString(),
-        )
-
-
-async def get_basic(txn: Transaction, kbid: str, uuid: str) -> Optional[bytes]:
-    if ingest_settings.driver == "local":
-        raw_basic = await txn.get(KB_RESOURCE_BASIC_FS.format(kbid=kbid, uuid=uuid))
-    else:
-        raw_basic = await txn.get(KB_RESOURCE_BASIC.format(kbid=kbid, uuid=uuid))
-    return raw_basic
-
 
 def set_title(writer: BrokerMessage, toprocess: PushPayload, title: str):
     title = urllib.parse.unquote(title)
@@ -82,3 +54,25 @@ def set_title(writer: BrokerMessage, toprocess: PushPayload, title: str):
 
 def compute_paragraph_key(rid: str, paragraph_key: str) -> str:
     return paragraph_key.replace("N_RID", rid)
+
+
+def choose_matryoshka_dimension(dimensions: Sequence[int]) -> int:
+    """Given a list of matryoshka embedding available dimensions, choose one to
+    set the vector dimension.
+    """
+    if len(dimensions) == 0:
+        raise ValueError("Can't choose matryoshka dimension from an empty list")
+
+    threshold = 2000
+    previous = None
+    for dimension in sorted(dimensions):
+        if dimension > threshold:
+            break
+        previous = dimension
+
+    if dimension > threshold:
+        if previous is None:
+            return dimension
+        else:
+            return previous
+    return dimension

nucliadb/ingest/partitions.py

@@ -38,9 +38,7 @@ def assign_partitions(settings: Settings):
                 try:
                     settings.replica_number = int(sts_values[-1])
                 except Exception:
-                    logger.error(
-                        f"Could not extract replica number from hostname: {hostname}"
-                    )
+                    logger.error(f"Could not extract replica number from hostname: {hostname}")
                     pass
 
         if settings.replica_number == -1:
@@ -53,9 +51,5 @@ def assign_partitions(settings: Settings):
     # update settings AND Environment Varialbe (for this process and its childs) with partition list
     settings.partitions = partitions_list
     os.environ["PARTITIONS"] = json.dumps(partitions_list)
-    logger.info(
-        f"PARTITIONS: Assigned Partitions (in settings) = {settings.partitions}"
-    )
-    logger.info(
-        f"PARTITIONS: Assigned Partitions (in environment) = {os.environ['PARTITIONS']}"
-    )
+    logger.info(f"PARTITIONS: Assigned Partitions (in settings) = {settings.partitions}")
+    logger.info(f"PARTITIONS: Assigned Partitions (in environment) = {os.environ['PARTITIONS']}")

nucliadb/ingest/processing.py

@@ -29,17 +29,22 @@ from typing import TYPE_CHECKING, Any, Optional, TypeVar
 import aiohttp
 import backoff
 import jwt
-from nucliadb_protos.resources_pb2 import CloudFile
-from nucliadb_protos.resources_pb2 import FieldFile as FieldFilePB
 from pydantic import BaseModel, Field
 
 import nucliadb_models as models
 from nucliadb_models.resource import QueueType
+from nucliadb_protos.resources_pb2 import CloudFile
+from nucliadb_protos.resources_pb2 import FieldFile as FieldFilePB
 from nucliadb_telemetry import metrics
 from nucliadb_utils.exceptions import LimitsExceededError, SendToProcessError
-from nucliadb_utils.settings import nuclia_settings, storage_settings
+from nucliadb_utils.settings import (
+    FileBackendConfig,
+    is_onprem_nucliadb,
+    nuclia_settings,
+    storage_settings,
+)
 from nucliadb_utils.storages.storage import Storage
-from nucliadb_utils.utilities import Utility, set_utility
+from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
 logger = logging.getLogger(__name__)
 
@@ -70,8 +75,8 @@ class Source(SourceValue, Enum):  # type: ignore
 
 
 class ProcessingInfo(BaseModel):
-    seqid: Optional[int]
-    account_seq: Optional[int]
+    seqid: Optional[int] = None
+    account_seq: Optional[int] = None
     queue: Optional[QueueType] = None
 
 
@@ -96,9 +101,6 @@ class PushPayload(BaseModel):
     # Diff on Text Field
     textfield: dict[str, models.Text] = {}
 
-    # Diff on a Layout Field
-    layoutfield: dict[str, models.LayoutDiff] = {}
-
     # New conversations to process
     conversationfield: dict[str, models.PushConversation] = {}
 
@@ -111,11 +113,11 @@ class PushPayload(BaseModel):
     )
 
 
-class PushResponse(BaseModel):
-    seqid: Optional[int] = None
-
-
 async def start_processing_engine():
+    processing_engine = get_utility(Utility.PROCESSING)
+    if processing_engine is not None:
+        return
+
     if nuclia_settings.dummy_processing:
         processing_engine = DummyProcessingEngine()
     else:
@@ -133,6 +135,43 @@ async def start_processing_engine():
     set_utility(Utility.PROCESSING, processing_engine)
 
 
+async def stop_processing_engine():
+    utility = get_utility(Utility.PROCESSING)
+    if utility is not None:
+        await utility.finalize()
+        clean_utility(Utility.PROCESSING)
+
+
+class ProcessingDriverType(Enum):
+    # XXX IMPORTANT XXX: Make sure the values are in sync with
+    # the ones defined in nuclia/learning/processing repository
+    GCS = 0
+    S3 = 1
+    LOCAL = 2
+
+
+def to_processing_driver_type(file_backend_driver: FileBackendConfig) -> ProcessingDriverType:
+    """
+    Outputs a nuclia-internal backend driver identifier that is used by processing
+    to store the blobs of processed metadata in the right bucket folder.
+    """
+    if is_onprem_nucliadb():
+        # On-prem installations are always regarded as local storage from the processing perspective,
+        # as Nuclia processing engine will not have direct access to the storage.
+        return ProcessingDriverType.LOCAL
+
+    try:
+        return {
+            FileBackendConfig.GCS: ProcessingDriverType.GCS,
+            FileBackendConfig.S3: ProcessingDriverType.S3,
+        }[file_backend_driver]
+    except KeyError:
+        logger.error(
+            f"Not a valid file backend driver to processing, fallback to local: {file_backend_driver}"
+        )
+        return ProcessingDriverType.LOCAL
+
+
 class ProcessingEngine:
     def __init__(
         self,
@@ -143,52 +182,30 @@ class ProcessingEngine:
         onprem: Optional[bool] = False,
         nuclia_jwt_key: Optional[str] = None,
         days_to_keep: int = 3,
-        driver: str = "gcs",
+        driver: FileBackendConfig = FileBackendConfig.GCS,
     ):
         self.nuclia_service_account = nuclia_service_account
         self.nuclia_zone = nuclia_zone
         if nuclia_public_url is not None:
-            self.nuclia_public_url: Optional[str] = nuclia_public_url.format(
-                zone=nuclia_zone
-            )
+            self.nuclia_public_url: Optional[str] = nuclia_public_url.format(zone=nuclia_zone)
         else:
             self.nuclia_public_url = None
 
         self.onprem = onprem
         if self.onprem:
-            self.nuclia_upload_url = (
-                f"{self.nuclia_public_url}/api/v1/processing/upload"
-            )
+            self.nuclia_upload_url = f"{self.nuclia_public_url}/api/v1/processing/upload"
         else:
-            self.nuclia_upload_url = (
-                f"{nuclia_processing_cluster_url}/api/v1/processing/upload"
-            )
-        self.nuclia_internal_push = (
-            f"{nuclia_processing_cluster_url}/api/v1/internal/processing/push"
-        )
+            self.nuclia_upload_url = f"{nuclia_processing_cluster_url}/api/v1/processing/upload"
+        self.nuclia_internal_push = f"{nuclia_processing_cluster_url}/api/v1/internal/processing/push"
         self.nuclia_internal_delete = (
             f"{nuclia_processing_cluster_url}/api/v1/internal/processing/requests"
         )
-        self.nuclia_external_push_v2 = (
-            f"{self.nuclia_public_url}/api/v1/processing/push"
-        )
-        self.nuclia_external_delete = (
-            f"{self.nuclia_public_url}/api/v1/processing/requests"
-        )
+        self.nuclia_external_push_v2 = f"{self.nuclia_public_url}/api/v1/processing/push"
+        self.nuclia_external_delete = f"{self.nuclia_public_url}/api/v1/processing/requests"
 
         self.nuclia_jwt_key = nuclia_jwt_key
         self.days_to_keep = days_to_keep
-        if driver == "gcs":
-            self.driver = 0
-        elif driver == "s3":
-            self.driver = 1
-        elif driver in ("local", "pg"):
-            self.driver = 2
-        else:
-            logger.error(
-                f"Not a valid driver to processing, fallback to local: {driver}"
-            )
-            self.driver = 2
+        self.driver: ProcessingDriverType = to_processing_driver_type(driver)
         self._exit_stack = AsyncExitStack()
 
     async def initialize(self):
@@ -211,7 +228,7 @@ class ProcessingEngine:
             "iat": now,
             "md5": cf.md5,
             "source": 1,  # To indicate that this files comes internally
-            "driver": self.driver,
+            "driver": self.driver.value,
             "jti": uuid.uuid4().hex,
             "bucket_name": cf.bucket_name,
             "filename": cf.filename,
@@ -235,7 +252,7 @@ class ProcessingEngine:
             "iat": now,
             "md5": file.file.md5,
             "source": 1,  # To indicate that this files comes internally
-            "driver": self.driver,
+            "driver": self.driver.value,
             "jti": uuid.uuid4().hex,
             "bucket_name": file.file.bucket_name,
             "filename": file.file.filename,
@@ -310,9 +327,7 @@ class ProcessingEngine:
         max_tries=MAX_TRIES,
     )
     @processing_observer.wrap({"type": "file_field_upload_internal"})
-    async def convert_internal_filefield_to_str(
-        self, file: FieldFilePB, storage: Storage
-    ) -> str:
+    async def convert_internal_filefield_to_str(self, file: FieldFilePB, storage: Storage) -> str:
         """It's already an internal file that needs to be uploaded"""
         if self.onprem is False:
             # Upload the file to processing upload
@@ -321,9 +336,7 @@ class ProcessingEngine:
             headers = {}
             headers["X-PASSWORD"] = file.password
             headers["X-LANGUAGE"] = file.language
-            headers["X-FILENAME"] = base64.b64encode(
-                file.file.filename.encode()
-            ).decode()
+            headers["X-FILENAME"] = base64.b64encode(file.file.filename.encode()).decode()
             headers["X-MD5"] = file.file.md5
             headers["CONTENT-TYPE"] = file.file.content_type
             if file.file.size:
@@ -331,9 +344,7 @@ class ProcessingEngine:
             headers["X-STF-NUAKEY"] = f"Bearer {self.nuclia_service_account}"
 
             iterator = storage.downloadbytescf_iterator(file.file)
-            async with self.session.post(
-                self.nuclia_upload_url, data=iterator, headers=headers
-            ) as resp:
+            async with self.session.post(self.nuclia_upload_url, data=iterator, headers=headers) as resp:
                 if resp.status == 200:
                     jwttoken = await resp.text()
                 elif resp.status == 402:
@@ -367,9 +378,7 @@ class ProcessingEngine:
             headers["X-STF-NUAKEY"] = f"Bearer {self.nuclia_service_account}"
 
             iterator = storage.downloadbytescf_iterator(cf)
-            async with self.session.post(
-                self.nuclia_upload_url, data=iterator, headers=headers
-            ) as resp:
+            async with self.session.post(self.nuclia_upload_url, data=iterator, headers=headers) as resp:
                 if resp.status == 200:
                     jwttoken = await resp.text()
                 elif resp.status == 402:
@@ -389,9 +398,7 @@ class ProcessingEngine:
         jitter=backoff.random_jitter,
         max_tries=MAX_TRIES,
     )
-    async def send_to_process(
-        self, item: PushPayload, partition: int
-    ) -> ProcessingInfo:
+    async def send_to_process(self, item: PushPayload, partition: int) -> ProcessingInfo:
         op_type = "process_external" if self.onprem else "process_internal"
         with processing_observer({"type": op_type}):
             headers = {"CONTENT-TYPE": "application/json"}
@@ -399,15 +406,13 @@ class ProcessingEngine:
                 # Upload the payload
                 item.partition = partition
                 resp = await self.session.post(
-                    url=self.nuclia_internal_push, data=item.json(), headers=headers
+                    url=self.nuclia_internal_push, data=item.model_dump_json(), headers=headers
                 )
             else:
-                headers.update(
-                    {"X-STF-NUAKEY": f"Bearer {self.nuclia_service_account}"}
-                )
+                headers.update({"X-STF-NUAKEY": f"Bearer {self.nuclia_service_account}"})
                 # Upload the payload
                 resp = await self.session.post(
-                    url=self.nuclia_external_push_v2, data=item.json(), headers=headers
+                    url=self.nuclia_external_push_v2, data=item.model_dump_json(), headers=headers
                 )
             if resp.status == 200:
                 data = await resp.json()
@@ -437,9 +442,7 @@ class ProcessingEngine:
             queue=QueueType(queue_type) if queue_type is not None else None,
         )
 
-    async def delete_from_processing(
-        self, *, kbid: str, resource_id: Optional[str] = None
-    ) -> None:
+    async def delete_from_processing(self, *, kbid: str, resource_id: Optional[str] = None) -> None:
         """
         Delete a resource from processing. This prevents inflight resources from being processed
         and wasting resources.
@@ -469,7 +472,7 @@ class ProcessingEngine:
 
 class DummyProcessingEngine(ProcessingEngine):
     def __init__(self):
-        self.calls: list[list[Any]] = []  # type: ignore
+        self.calls: list[list[Any]] = []
         self.values = defaultdict(list)
         self.onprem = True
 
@@ -491,9 +494,7 @@ class DummyProcessingEngine(ProcessingEngine):
         self.values["convert_external_filefield_to_str"].append(file_field)
         return f"convert_external_filefield_to_str,{index}"
 
-    async def convert_internal_filefield_to_str(
-        self, file: FieldFilePB, storage: Storage
-    ) -> str:
+    async def convert_internal_filefield_to_str(self, file: FieldFilePB, storage: Storage) -> str:
         self.calls.append([file, storage])
         index = len(self.values["convert_internal_filefield_to_str"])
         self.values["convert_internal_filefield_to_str"].append([file, storage])
@@ -505,16 +506,10 @@ class DummyProcessingEngine(ProcessingEngine):
         self.values["convert_internal_cf_to_str"].append([cf, storage])
         return f"convert_internal_cf_to_str,{index}"
 
-    async def send_to_process(
-        self, item: PushPayload, partition: int
-    ) -> ProcessingInfo:
+    async def send_to_process(self, item: PushPayload, partition: int) -> ProcessingInfo:
         self.calls.append([item, partition])
         self.values["send_to_process"].append([item, partition])
-        return ProcessingInfo(
-            seqid=len(self.calls), account_seq=0, queue=QueueType.SHARED
-        )
+        return ProcessingInfo(seqid=len(self.calls), account_seq=0, queue=QueueType.SHARED)
 
-    async def delete_from_processing(
-        self, *, kbid: str, resource_id: Optional[str] = None
-    ) -> None:
+    async def delete_from_processing(self, *, kbid: str, resource_id: Optional[str] = None) -> None:
         self.calls.append([kbid, resource_id])