nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

nucliadb/ingest/orm/utils.py

@@ -18,7 +18,10 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import urllib.parse
+from typing import Sequence
 
+from nucliadb.ingest.processing import PushPayload
+from nucliadb_models.text import PushTextFormat, Text
 from nucliadb_protos.resources_pb2 import (
     ExtractedTextWrapper,
     FieldComputedMetadataWrapper,
@@ -28,9 +31,6 @@ from nucliadb_protos.resources_pb2 import (
 )
 from nucliadb_protos.writer_pb2 import BrokerMessage
 
-from nucliadb.ingest.processing import PushPayload
-from nucliadb_models.text import PushTextFormat, Text
-
 
 def set_title(writer: BrokerMessage, toprocess: PushPayload, title: str):
     title = urllib.parse.unquote(title)
@@ -56,7 +56,7 @@ def compute_paragraph_key(rid: str, paragraph_key: str) -> str:
     return paragraph_key.replace("N_RID", rid)
 
 
-def choose_matryoshka_dimension(dimensions: list[int]) -> int:
+def choose_matryoshka_dimension(dimensions: Sequence[int]) -> int:
     """Given a list of matryoshka embedding available dimensions, choose one to
     set the vector 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 FileBackendConfig, 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__)
 
@@ -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] = {}
 
@@ -112,6 +114,10 @@ class PushPayload(BaseModel):
 
 
 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:
@@ -129,22 +135,41 @@ async def start_processing_engine():
     set_utility(Utility.PROCESSING, processing_engine)
 
 
-def to_processing_driver_type(file_backend_driver: FileBackendConfig) -> int:
+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 file_backend_driver == FileBackendConfig.GCS:
-        return 0
-    elif file_backend_driver == FileBackendConfig.S3:
-        return 1
-    elif file_backend_driver in (FileBackendConfig.LOCAL, FileBackendConfig.PG):
-        return 2
-    else:
+    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 2
+        return ProcessingDriverType.LOCAL
 
 
 class ProcessingEngine:
@@ -162,37 +187,25 @@ class ProcessingEngine:
         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
-        self.driver = to_processing_driver_type(driver)
+        self.driver: ProcessingDriverType = to_processing_driver_type(driver)
         self._exit_stack = AsyncExitStack()
 
     async def initialize(self):
@@ -215,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,
@@ -239,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,
@@ -248,6 +261,7 @@ class ProcessingEngine:
             "content_type": file.file.content_type,
             "password": file.password,
             "language": file.language,
+            "extract_strategy": file.extract_strategy,
         }
         return jwt.encode(payload, self.nuclia_jwt_key, algorithm="HS256")
 
@@ -265,6 +279,8 @@ class ProcessingEngine:
         headers["X-LANGUAGE"] = file.language
         headers["X-FILENAME"] = base64.b64encode(file.file.filename.encode()).decode()  # type: ignore
         headers["X-MD5"] = file.file.md5
+        if file.extract_strategy is not None:
+            headers["X-EXTRACT-STRATEGY"] = file.extract_strategy
         headers["CONTENT_TYPE"] = file.file.content_type
         headers["CONTENT-LENGTH"] = str(len(file.file.payload))  # type: ignore
         headers["X-STF-NUAKEY"] = f"Bearer {self.nuclia_service_account}"
@@ -304,6 +320,7 @@ class ProcessingEngine:
             "content_type": file_field.file.content_type,
             "language": file_field.language,
             "password": file_field.password,
+            "extract_strategy": file_field.extract_strategy,
         }
         return jwt.encode(payload, self.nuclia_jwt_key, algorithm="HS256")
 
@@ -314,9 +331,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
@@ -325,19 +340,17 @@ 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:
                 headers["CONTENT-LENGTH"] = str(file.file.size)
+            if file.extract_strategy != "":
+                headers["X-EXTRACT-STRATEGY"] = file.extract_strategy
             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:
@@ -371,9 +384,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:
@@ -393,9 +404,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"}
@@ -403,15 +412,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()
@@ -441,9 +448,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.
@@ -473,7 +478,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
 
@@ -495,9 +500,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])
@@ -509,16 +512,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])