arkindex-base-worker 0.4.0a1__py3-none-any.whl → 0.4.0b1__py3-none-any.whl

{arkindex_base_worker-0.4.0a1.dist-info → arkindex_base_worker-0.4.0b1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.4.0a1
+Version: 0.4.0b1
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>
@@ -41,17 +41,17 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: peewee ~=3.17
-Requires-Dist: Pillow ==10.3.0
+Requires-Dist: Pillow ==10.4.0
 Requires-Dist: python-gnupg ==0.5.2
-Requires-Dist: shapely ==2.0.3
+Requires-Dist: shapely ==2.0.5
 Requires-Dist: teklia-toolbox ==0.1.5
 Requires-Dist: zstandard ==0.22.0
 Provides-Extra: docs
-Requires-Dist: black ==24.4.0 ; extra == 'docs'
-Requires-Dist: mkdocs-material ==9.5.17 ; extra == 'docs'
-Requires-Dist: mkdocstrings-python ==1.9.2 ; extra == 'docs'
+Requires-Dist: black ==24.4.2 ; extra == 'docs'
+Requires-Dist: mkdocs-material ==9.5.31 ; extra == 'docs'
+Requires-Dist: mkdocstrings-python ==1.10.7 ; extra == 'docs'
 Provides-Extra: tests
-Requires-Dist: pytest ==8.1.1 ; extra == 'tests'
+Requires-Dist: pytest ==8.3.2 ; extra == 'tests'
 Requires-Dist: pytest-mock ==3.14.0 ; extra == 'tests'
 Requires-Dist: pytest-responses ==0.5.1 ; extra == 'tests'
 

{arkindex_base_worker-0.4.0a1.dist-info → arkindex_base_worker-0.4.0b1.dist-info}/RECORD

@@ -3,40 +3,40 @@ arkindex_worker/cache.py,sha256=FTlB0coXofn5zTNRTcVIvh709mcw4a1bPGqkwWjKs3w,1124
 arkindex_worker/image.py,sha256=5ymIGaTm2D7Sp2YYQkbuheuGnx5VJo0_AzYAEIvNGhs,14267
 arkindex_worker/models.py,sha256=bPQzGZNs5a6z6DEcygsa8T33VOqPlMUbwKzHqlKzwbw,9923
 arkindex_worker/utils.py,sha256=KXWIACda7D3IpdToaAplLoAgnCK8bKWw7aWUyq-IWUA,7211
-arkindex_worker/worker/__init__.py,sha256=3sJ_EPB7yG-kPfgunbm2B7B7DzoeOi5ZNpQwC_3QuZ0,19429
-arkindex_worker/worker/base.py,sha256=c9u37W1BNHt5RoQV2ZrYUYv6tBs-CjiSgUAAg7p7GA0,18876
-arkindex_worker/worker/classification.py,sha256=JVz-6YEeuavOy7zGfQi4nE_wpj9hwMUZDXTem-hXQY8,10328
-arkindex_worker/worker/corpus.py,sha256=ZHAAYE4PRPXqqaZm71wjrsxYETFqU6TAz-3VYgIXzac,1794
-arkindex_worker/worker/dataset.py,sha256=roX2IMMNA-icteTtRADiFSZiZSRPClqS62ZPJm9s2JI,2923
-arkindex_worker/worker/element.py,sha256=AWK3YJSHWy3j4ajntJloi_2X4zxsgXZ6c6dzphgq3OI,33848
-arkindex_worker/worker/entity.py,sha256=suhycfikC9oTPEWmX48_cnvFEw-Wu5zBA8n_00K4KUk,14714
+arkindex_worker/worker/__init__.py,sha256=belqRtbs0raTdFJoQJoGBoDJkUOrEE3wyXv90f85bTs,19760
+arkindex_worker/worker/base.py,sha256=JStHpwSP3bis9LLvV2C2n6GTWtLUVIDA9JPgPJEt17o,18717
+arkindex_worker/worker/classification.py,sha256=4YAY4weF6kMSMsoYiz6oia3SN21PzRR1bAdhMJCGBbw,10361
+arkindex_worker/worker/corpus.py,sha256=s9bCxOszJMwRq1WWAmKjWq888mjDfbaJ18Wo7h-rNOw,1827
+arkindex_worker/worker/dataset.py,sha256=UXElhhARca9m7Himp-yxD5dAqWbdxDKWOUJUGgeCZXI,2934
+arkindex_worker/worker/element.py,sha256=kMaJNXEfZbFBK4YYc3XLqyGvPyNvJs7mJG2T_a1c7D0,34294
+arkindex_worker/worker/entity.py,sha256=BbQp56kxTPmOQI482TUFZ8KOXISj7KtQAyHRT0CmedM,14744
 arkindex_worker/worker/image.py,sha256=t_Az6IGnj0EZyvcA4XxfPikOUjn_pztgsyxTkFZhaXU,621
-arkindex_worker/worker/metadata.py,sha256=Bouuc_JaXogKykVXOTKDVP3tX--OUQeHoazxIGrGrJI,6702
-arkindex_worker/worker/task.py,sha256=cz3wJNPgogZv1lm_3lm7WScitQtYQtL6H6I7Xokq208,1475
-arkindex_worker/worker/training.py,sha256=hkwCBjVE4bByXzHUmCZF73Bl5JxARdXWjYgFE6ydAT0,10749
-arkindex_worker/worker/transcription.py,sha256=6R7ofcGnNqX4rjT0kRKIE-G9FHq2TJ1tfztNM5sTqYE,20464
-arkindex_worker/worker/version.py,sha256=cs2pdlDxpKRO2Oldvcu54w-D_DQhf1cdeEt4tKX_QYs,1927
+arkindex_worker/worker/metadata.py,sha256=PFO0oJc8N91HIpj4yHLscwGW5UFRXtuyQYfEXW27-WQ,6724
+arkindex_worker/worker/task.py,sha256=1O9zrWXxe3na3TOcoHX5Pxn1875v7EU08BSsCPnb62g,1519
+arkindex_worker/worker/training.py,sha256=qnBFEk11JOWWPLTbjF-lZ9iFBdTPpQzZAzQ9a03J1j4,10874
+arkindex_worker/worker/transcription.py,sha256=9TC3E6zu_CnQKWsaTAzI83TrSfMuzh3KSMOCLdbEG18,20497
+arkindex_worker/worker/version.py,sha256=JIT7OI3Mo7RPkNrjOB9hfqrsG-FYygz_zi4l8PbkuAo,1960
 hooks/pre_gen_project.py,sha256=xQJERv3vv9VzIqcBHI281eeWLWREXUF4mMw7PvJHHXM,269
 tests/__init__.py,sha256=6aeTMHf4q_dKY4jIZWg1KT70VKaLvVlzCxh-Uu_cWiQ,241
 tests/conftest.py,sha256=-ZQTV4rg7TgW84-5Ioqndqv8byNILfDOpyUt8wecEiI,21967
-tests/test_base_worker.py,sha256=qG45O3nPbASXN5a5RadXU1BAXn3EIaTK6Hvjj3s4Ozs,24292
+tests/test_base_worker.py,sha256=LdFV0LFdNU2IOyEKlX59MB1kuyxHCuhy4Tm7eE_iPiU,24281
 tests/test_cache.py,sha256=ii0gyr0DrG7ChEs7pmT8hMdSguAOAcCze4bRMiFQxuk,10640
 tests/test_dataset_worker.py,sha256=d9HG36qnO5HXu9vQ0UTBvdTSRR21FVq1FNoXM-vZbPk,22105
 tests/test_element.py,sha256=2G9M15TLxQRmvrWM9Kw2ucnElh4kSv_oF_5FYwwAxTY,13181
 tests/test_image.py,sha256=Fs9vKYgQ7mEFylbzI4YIO_JyOLeAcs-WxUXpzewxCd8,16188
-tests/test_merge.py,sha256=Q4zCbtZbe0wBfqE56gvAD06c6pDuhqnjKaioFqIgAQw,8331
+tests/test_merge.py,sha256=FMdpsm_ncHNmIvOrJ1vcwlyn8o9-SPcpFTcbAsXwK-w,8320
 tests/test_utils.py,sha256=vpeHMeL7bJQonv5ZEbJmlJikqVKn5VWlVEbvmYFzDYA,1650
 tests/test_elements_worker/__init__.py,sha256=Fh4nkbbyJSMv_VtjQxnWrOqTnxXaaWI8S9WU0VrzCHs,179
 tests/test_elements_worker/test_classifications.py,sha256=DYRKhPpplFp144GCXKyFG1hz4Ra9vk5FiAN6dhfMP6k,25511
 tests/test_elements_worker/test_cli.py,sha256=a23i1pUDbXi23MUtbWwGEcLLrmc_YlrbDgOG3h66wLM,2620
 tests/test_elements_worker/test_corpus.py,sha256=c_LUHvkJIYgk_wXF06VQPNOoWfiZ06XpjOXrJ7MRiBc,4479
 tests/test_elements_worker/test_dataset.py,sha256=lSXqubhg1EEq2Y2goE8Y2RYaqIpM9Iejq6fGNW2BczU,11411
-tests/test_elements_worker/test_elements.py,sha256=2_kdeo99biCH3Uez6HB8ltS_iIizZ7ir5uOkFjIXfjM,84812
+tests/test_elements_worker/test_elements.py,sha256=HH8jUU4xHp5gXcrGJLQlo4kLFh7oYfMxO3QQEYo2itg,84885
 tests/test_elements_worker/test_entities.py,sha256=jirb_IKAMqMhwxeDgjO-rsr1fTP9GdXwuyhncUjCJFM,33494
 tests/test_elements_worker/test_image.py,sha256=_E3UGdDOwTo1MW5KMS81PrdeSPBPWinWYoQPNy2F9Ro,2077
 tests/test_elements_worker/test_metadata.py,sha256=-cZhlVAh4o2uRnHz8fPf_thfavRnJrtJYN_p4BmHISU,17566
 tests/test_elements_worker/test_task.py,sha256=7Sr3fbjdgWUXJUhJEiC9CwnbhQIQX3rCInmHMIrmA38,5573
-tests/test_elements_worker/test_training.py,sha256=wVYWdMdeSA6T2XyhH5AJJNGemYq3LOViiZvj0dblACA,9468
+tests/test_elements_worker/test_training.py,sha256=Qxi9EzGr_uKcn2Fh5aE6jNrq1K8QKLiOiSew4upASPs,8721
 tests/test_elements_worker/test_transcriptions.py,sha256=7HDkIW8IDK7pKAfpSdAPB7YOyKyeBJTn2_alvVK46SA,72411
 tests/test_elements_worker/test_worker.py,sha256=AwdP8uSXNQ_SJavXxJV2s3_J3OiCafShVjMV1dgt4xo,17162
 worker-demo/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -44,8 +44,8 @@ worker-demo/tests/conftest.py,sha256=XzNMNeg6pmABUAH8jN6eZTlZSFGLYjS3-DTXjiRN6Yc
 worker-demo/tests/test_worker.py,sha256=3DLd4NRK4bfyatG5P_PK4k9P9tJHx9XQq5_ryFEEFVg,304
 worker-demo/worker_demo/__init__.py,sha256=2BPomV8ZMNf3YXJgloatKeHQCE6QOkwmsHGkO6MkQuM,125
 worker-demo/worker_demo/worker.py,sha256=Rt-DjWa5iBP08k58NDZMfeyPuFbtNcbX6nc5jFX7GNo,440
-arkindex_base_worker-0.4.0a1.dist-info/LICENSE,sha256=NVshRi1efwVezMfW7xXYLrdDr2Li1AfwfGOd5WuH1kQ,1063
-arkindex_base_worker-0.4.0a1.dist-info/METADATA,sha256=PBTlbhWTCvvkkcGqQew6yvJIdncf9mKZ71yI_QSX2iM,3269
-arkindex_base_worker-0.4.0a1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-arkindex_base_worker-0.4.0a1.dist-info/top_level.txt,sha256=58NuslgxQC2vT4DiqZEgO4JqJRrYa2yeNI9QvkbfGQU,40
-arkindex_base_worker-0.4.0a1.dist-info/RECORD,,
+arkindex_base_worker-0.4.0b1.dist-info/LICENSE,sha256=NVshRi1efwVezMfW7xXYLrdDr2Li1AfwfGOd5WuH1kQ,1063
+arkindex_base_worker-0.4.0b1.dist-info/METADATA,sha256=02rPRlcFlghY1Trb-_trpdCCMME1A9FmPzrY8wzzLDg,3270
+arkindex_base_worker-0.4.0b1.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+arkindex_base_worker-0.4.0b1.dist-info/top_level.txt,sha256=58NuslgxQC2vT4DiqZEgO4JqJRrYa2yeNI9QvkbfGQU,40
+arkindex_base_worker-0.4.0b1.dist-info/RECORD,,

{arkindex_base_worker-0.4.0a1.dist-info → arkindex_base_worker-0.4.0b1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

arkindex_worker/worker/__init__.py

@@ -83,7 +83,20 @@ class ElementsWorker(
         """
         super().__init__(description, support_cache)
 
-        # Add mandatory argument to process elements
+        self.classes = {}
+
+        self.entity_types = {}
+        """Known and available entity types in processed corpus
+        """
+
+        self.corpus_types = {}
+        """Known and available element types in processed corpus
+        """
+
+        self._worker_version_cache = {}
+
+    def add_arguments(self):
+        """Define specific ``argparse`` arguments for this worker"""
         self.parser.add_argument(
             "--elements-list",
             help="JSON elements list to use",
@@ -97,14 +110,6 @@ class ElementsWorker(
             help="One or more Arkindex element ID",
         )
 
-        self.classes = {}
-
-        self.entity_types = {}
-        """Known and available entity types in processed corpus
-        """
-
-        self._worker_version_cache = {}
-
     def list_elements(self) -> Iterable[CachedElement] | list[str]:
         """
         List the elements to be processed, either from the CLI arguments or
@@ -222,7 +227,9 @@ class ElementsWorker(
                     element = item
                 else:
                     # Load element using the Arkindex API
-                    element = Element(**self.request("RetrieveElement", id=item))
+                    element = Element(
+                        **self.api_client.request("RetrieveElement", id=item)
+                    )
 
                 logger.info(f"Processing {element} ({i}/{count})")
 
@@ -301,7 +308,7 @@ class ElementsWorker(
         assert isinstance(state, ActivityState), "state should be an ActivityState"
 
         try:
-            self.request(
+            self.api_client.request(
                 "UpdateWorkerActivity",
                 id=self.worker_run_id,
                 body={
@@ -376,6 +383,8 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
         # Set as an instance variable as dataset workers might use it to easily extract its content
         self.downloaded_dataset_artifact: Path | None = None
 
+    def add_arguments(self):
+        """Define specific ``argparse`` arguments for this worker"""
         self.parser.add_argument(
             "--set",
             type=check_dataset_set,
@@ -472,7 +481,7 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
             # Retrieving dataset information is not already cached
             if dataset_id not in datasets:
                 datasets[dataset_id] = Dataset(
-                    **self.request("RetrieveDataset", id=dataset_id)
+                    **self.api_client.request("RetrieveDataset", id=dataset_id)
                 )
 
             yield Set(name=set_name, dataset=datasets[dataset_id])

arkindex_worker/worker/base.py

@@ -231,7 +231,7 @@ class BaseWorker:
             logger.debug("Debug output enabled")
 
         # Load worker run information
-        worker_run = self.request("RetrieveWorkerRun", id=self.worker_run_id)
+        worker_run = self.api_client.request("RetrieveWorkerRun", id=self.worker_run_id)
 
         # Load process information
         self.process_information = worker_run["process"]
@@ -290,7 +290,7 @@ class BaseWorker:
         if self.support_cache and self.args.database is not None:
             self.use_cache = True
         elif self.support_cache and self.task_id:
-            task = self.request("RetrieveTaskFromAgent", id=self.task_id)
+            task = self.api_client.request("RetrieveTask", id=self.task_id)
             self.task_parents = task["parents"]
             paths = self.find_parents_file_paths(Path("db.sqlite"))
             self.use_cache = len(paths) > 0
@@ -331,7 +331,7 @@ class BaseWorker:
 
         # Load from the backend
         try:
-            resp = self.request("RetrieveSecret", name=str(name))
+            resp = self.api_client.request("RetrieveSecret", name=str(name))
             secret = resp["content"]
             logging.info(f"Loaded API secret {name}")
         except ErrorResponse as e:
@@ -471,12 +471,6 @@ class BaseWorker:
             # Clean up
             shutil.rmtree(base_extracted_path)
 
-    def request(self, *args, **kwargs):
-        """
-        Wrapper around the ``ArkindexClient.request`` method.
-        """
-        return self.api_client.request(*args, **kwargs)
-
     def add_arguments(self):
         """Override this method to add ``argparse`` arguments to this worker"""
 

arkindex_worker/worker/classification.py

@@ -39,7 +39,7 @@ class ClassificationMixin:
         if ml_class_id is None:
             logger.info(f"Creating ML class {ml_class} on corpus {self.corpus_id}")
             try:
-                response = self.request(
+                response = self.api_client.request(
                     "CreateMLClass", id=self.corpus_id, body={"name": ml_class}
                 )
                 ml_class_id = self.classes[ml_class] = response["id"]
@@ -119,7 +119,7 @@ class ClassificationMixin:
             )
             return
         try:
-            created = self.request(
+            created = self.api_client.request(
                 "CreateClassification",
                 body={
                     "element": str(element.id),
@@ -220,7 +220,7 @@ class ClassificationMixin:
             )
             return
 
-        created_cls = self.request(
+        created_cls = self.api_client.request(
             "CreateClassifications",
             body={
                 "parent": str(element.id),

arkindex_worker/worker/corpus.py

@@ -63,7 +63,9 @@ class CorpusMixin:
         # Download latest export
         export_id: str = exports[0]["id"]
         logger.info(f"Downloading export ({export_id})...")
-        export: _TemporaryFileWrapper = self.request("DownloadExport", id=export_id)
+        export: _TemporaryFileWrapper = self.api_client.request(
+            "DownloadExport", id=export_id
+        )
         logger.info(f"Downloaded export ({export_id}) @ `{export.name}`")
 
         return export

arkindex_worker/worker/dataset.py

@@ -93,7 +93,7 @@ class DatasetMixin:
             logger.warning("Cannot update dataset as this worker is in read-only mode")
             return
 
-        updated_dataset = self.request(
+        updated_dataset = self.api_client.request(
             "PartialUpdateDataset",
             id=dataset.id,
             body={"state": state.value},

arkindex_worker/worker/element.py

@@ -31,6 +31,21 @@ class MissingTypeError(Exception):
 
 
 class ElementMixin:
+    def list_corpus_types(self):
+        """
+        Loads available element types in corpus.
+        """
+        self.corpus_types = {
+            element_type["slug"]: element_type
+            for element_type in self.api_client.request(
+                "RetrieveCorpus", id=self.corpus_id
+            )["types"]
+        }
+        count = len(self.corpus_types)
+        logger.info(
+            f'Loaded {count} element type{"s"[:count>1]} in corpus ({self.corpus_id}).'
+        )
+
     @unsupported_cache
     def create_required_types(self, element_types: list[ElementType]):
         """Creates given element types in the corpus.
@@ -38,7 +53,7 @@ class ElementMixin:
         :param element_types: The missing element types to create.
         """
         for element_type in element_types:
-            self.request(
+            self.api_client.request(
                 "CreateElementType",
                 body={
                     "slug": element_type.slug,
@@ -66,10 +81,10 @@ class ElementMixin:
             isinstance(slug, str) for slug in type_slugs
         ), "Element type slugs must be strings."
 
-        corpus = self.request("RetrieveCorpus", id=self.corpus_id)
-        available_slugs = {element_type["slug"] for element_type in corpus["types"]}
-        missing_slugs = set(type_slugs) - available_slugs
+        if not self.corpus_types:
+            self.list_corpus_types()
 
+        missing_slugs = set(type_slugs) - set(self.corpus_types)
         if missing_slugs:
             if create_missing:
                 self.create_required_types(
@@ -79,7 +94,7 @@ class ElementMixin:
                 )
             else:
                 raise MissingTypeError(
-                    f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
+                    f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in corpus ({self.corpus_id}).'
                 )
 
         return True
@@ -145,7 +160,7 @@ class ElementMixin:
             logger.warning("Cannot create element as this worker is in read-only mode")
             return
 
-        sub_element = self.request(
+        sub_element = self.api_client.request(
             "CreateElement",
             body={
                 "type": type,
@@ -243,7 +258,7 @@ class ElementMixin:
             logger.warning("Cannot create elements as this worker is in read-only mode")
             return
 
-        created_ids = self.request(
+        created_ids = self.api_client.request(
             "CreateElements",
             id=parent.id,
             body={
@@ -311,7 +326,7 @@ class ElementMixin:
             logger.warning("Cannot link elements as this worker is in read-only mode")
             return
 
-        return self.request(
+        return self.api_client.request(
             "CreateElementParent",
             parent=parent.id,
             child=child.id,
@@ -383,7 +398,7 @@ class ElementMixin:
             logger.warning("Cannot update element as this worker is in read-only mode")
             return
 
-        updated_element = self.request(
+        updated_element = self.api_client.request(
             "PartialUpdateElement",
             id=element.id,
             body=kwargs,

arkindex_worker/worker/entity.py

@@ -48,6 +48,7 @@ class EntityMixin:
         if not self.entity_types:
             # Load entity_types of corpus
             self.list_corpus_entity_types()
+
         for entity_type in entity_types:
             # Do nothing if type already exists
             if entity_type in self.entity_types:
@@ -60,7 +61,7 @@ class EntityMixin:
                 )
 
             # Create type if non-existent
-            self.entity_types[entity_type] = self.request(
+            self.entity_types[entity_type] = self.api_client.request(
                 "CreateEntityType",
                 body={
                     "name": entity_type,
@@ -106,7 +107,7 @@ class EntityMixin:
         entity_type_id = self.entity_types.get(type)
         assert entity_type_id, f"Entity type `{type}` not found in the corpus."
 
-        entity = self.request(
+        entity = self.api_client.request(
             "CreateEntity",
             body={
                 "name": name,
@@ -188,7 +189,7 @@ class EntityMixin:
         if confidence is not None:
             body["confidence"] = confidence
 
-        transcription_ent = self.request(
+        transcription_ent = self.api_client.request(
             "CreateTranscriptionEntity",
             id=transcription.id,
             body=body,
@@ -289,7 +290,7 @@ class EntityMixin:
             )
             return
 
-        created_ids = self.request(
+        created_ids = self.api_client.request(
             "CreateTranscriptionEntities",
             id=transcription.id,
             body={
@@ -385,9 +386,7 @@ class EntityMixin:
             f'Loaded {count} entit{"ies" if count > 1 else "y"} in corpus ({self.corpus_id})'
         )
 
-    def list_corpus_entity_types(
-        self,
-    ):
+    def list_corpus_entity_types(self):
         """
         Loads available entity types in corpus.
         """

arkindex_worker/worker/metadata.py

@@ -93,7 +93,7 @@ class MetaDataMixin:
             logger.warning("Cannot create metadata as this worker is in read-only mode")
             return
 
-        metadata = self.request(
+        metadata = self.api_client.request(
             "CreateMetaData",
             id=element.id,
             body={
@@ -168,7 +168,7 @@ class MetaDataMixin:
             logger.warning("Cannot create metadata as this worker is in read-only mode")
             return
 
-        created_metadata_list = self.request(
+        created_metadata_list = self.api_client.request(
             "CreateMetaDataBulk",
             id=element.id,
             body={

arkindex_worker/worker/task.py

@@ -22,7 +22,7 @@ class TaskMixin:
             task_id, uuid.UUID
         ), "task_id shouldn't be null and should be an UUID"
 
-        results = self.request("ListArtifacts", id=task_id)
+        results = self.api_client.request("ListArtifacts", id=task_id)
 
         return map(Artifact, results)
 
@@ -43,4 +43,6 @@ class TaskMixin:
             artifact, Artifact
         ), "artifact shouldn't be null and should be an Artifact"
 
-        return self.request("DownloadArtifact", id=task_id, path=artifact.path)
+        return self.api_client.request(
+            "DownloadArtifact", id=task_id, path=artifact.path
+        )

arkindex_worker/worker/training.py

@@ -185,7 +185,7 @@ class TrainingMixin:
         assert not self.model_version, "A model version has already been created."
 
         configuration = configuration or {}
-        self.model_version = self.request(
+        self.model_version = self.api_client.request(
             "CreateModelVersion",
             id=model_id,
             body=build_clean_payload(
@@ -217,7 +217,7 @@ class TrainingMixin:
         :param parent: ID of the parent model version
         """
         assert self.model_version, "No model version has been created yet."
-        self.model_version = self.request(
+        self.model_version = self.api_client.request(
             "UpdateModelVersion",
             id=self.model_version["id"],
             body=build_clean_payload(
@@ -273,41 +273,44 @@ class TrainingMixin:
         """
         assert self.model_version, "You must create the model version and upload its archive before validating it."
         try:
-            self.model_version = self.request(
-                "ValidateModelVersion",
+            self.model_version = self.api_client.request(
+                "PartialUpdateModelVersion",
                 id=self.model_version["id"],
                 body={
+                    "state": "available",
                     "size": size,
                     "hash": hash,
                     "archive_hash": archive_hash,
                 },
             )
         except ErrorResponse as e:
-            # Temporary fix while waiting for `ValidateModelVersion` refactoring as it can
-            # return errors even when the model version is properly validated
-            if e.status_code in [403, 500]:
-                logger.warning(
-                    f'An error occurred while validating model version {self.model_version["id"]}, please check its status.'
-                )
-                return
-
-            if e.status_code != 409:
+            model_version = e.content
+            if not model_version or "id" not in model_version:
                 raise e
 
             logger.warning(
                 f"An available model version exists with hash {hash}, using it instead of the pending version."
             )
             pending_version_id = self.model_version["id"]
-            self.model_version = getattr(e, "content", None)
-            assert self.model_version is not None, "An unexpected error occurred."
-
             logger.warning("Removing the pending model version.")
             try:
-                self.request("DestroyModelVersion", id=pending_version_id)
+                self.api_client.request("DestroyModelVersion", id=pending_version_id)
             except ErrorResponse as e:
                 msg = getattr(e, "content", str(e))
                 logger.error(
                     f"An error occurred removing the pending version {pending_version_id}: {msg}."
                 )
 
+            logger.info("Retrieving the existing model version.")
+            existing_version_id = model_version["id"].pop()
+            try:
+                self.model_version = self.api_client.request(
+                    "RetrieveModelVersion", id=existing_version_id
+                )
+            except ErrorResponse as e:
+                logger.error(
+                    f"An error occurred retrieving the existing version {existing_version_id}: {e.status_code} - {e.content}."
+                )
+                raise
+
         logger.info(f"Model version {self.model_version['id']} is now available.")

arkindex_worker/worker/transcription.py

@@ -77,7 +77,7 @@ class TranscriptionMixin:
             )
             return
 
-        created = self.request(
+        created = self.api_client.request(
             "CreateTranscription",
             id=element.id,
             body={
@@ -171,7 +171,7 @@ class TranscriptionMixin:
             )
             return
 
-        created_trs = self.request(
+        created_trs = self.api_client.request(
             "CreateTranscriptions",
             body={
                 "worker_run_id": self.worker_run_id,
@@ -291,7 +291,7 @@ class TranscriptionMixin:
             )
             return
 
-        annotations = self.request(
+        annotations = self.api_client.request(
             "CreateElementTranscriptions",
             id=element.id,
             body={

arkindex_worker/worker/version.py

@@ -34,7 +34,9 @@ class WorkerVersionMixin:
         if worker_version_id in self._worker_version_cache:
             return self._worker_version_cache[worker_version_id]
 
-        worker_version = self.request("RetrieveWorkerVersion", id=worker_version_id)
+        worker_version = self.api_client.request(
+            "RetrieveWorkerVersion", id=worker_version_id
+        )
         self._worker_version_cache[worker_version_id] = worker_version
 
         return worker_version

tests/test_base_worker.py

@@ -658,7 +658,7 @@ def test_find_extras_directory_not_found(monkeypatch, extras_path, exists, error
 def test_find_parents_file_paths(responses, mock_base_worker_with_cache, tmp_path):
     responses.add(
         responses.GET,
-        "http://testserver/api/v1/task/my_task/from-agent/",
+        "http://testserver/api/v1/task/my_task/",
         status=200,
         json={"parents": ["first", "second", "third"]},
     )

tests/test_elements_worker/test_elements.py

@@ -22,6 +22,24 @@ from tests import CORPUS_ID
 from . import BASE_API_CALLS
 
 
+def test_list_corpus_types(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/",
+        json={
+            "id": CORPUS_ID,
+            "types": [{"slug": "folder"}, {"slug": "page"}],
+        },
+    )
+
+    mock_elements_worker.list_corpus_types()
+
+    assert mock_elements_worker.corpus_types == {
+        "folder": {"slug": "folder"},
+        "page": {"slug": "page"},
+    }
+
+
 def test_check_required_types_argument_types(mock_elements_worker):
     with pytest.raises(
         AssertionError, match="At least one element type slug is required."
@@ -32,17 +50,11 @@ def test_check_required_types_argument_types(mock_elements_worker):
         mock_elements_worker.check_required_types("lol", 42)
 
 
-def test_check_required_types(responses, mock_elements_worker):
-    responses.add(
-        responses.GET,
-        f"http://testserver/api/v1/corpus/{CORPUS_ID}/",
-        json={
-            "id": CORPUS_ID,
-            "name": "Some Corpus",
-            "types": [{"slug": "folder"}, {"slug": "page"}],
-        },
-    )
-    mock_elements_worker.setup_api_client()
+def test_check_required_types(mock_elements_worker):
+    mock_elements_worker.corpus_types = {
+        "folder": {"slug": "folder"},
+        "page": {"slug": "page"},
+    }
 
     assert mock_elements_worker.check_required_types("page")
     assert mock_elements_worker.check_required_types("page", "folder")
@@ -50,22 +62,18 @@ def test_check_required_types(responses, mock_elements_worker):
     with pytest.raises(
         MissingTypeError,
         match=re.escape(
-            "Element type(s) act, text_line were not found in the Some Corpus corpus (11111111-1111-1111-1111-111111111111)."
+            "Element type(s) act, text_line were not found in corpus (11111111-1111-1111-1111-111111111111)."
         ),
     ):
         assert mock_elements_worker.check_required_types("page", "text_line", "act")
 
 
 def test_create_missing_types(responses, mock_elements_worker):
-    responses.add(
-        responses.GET,
-        f"http://testserver/api/v1/corpus/{CORPUS_ID}/",
-        json={
-            "id": CORPUS_ID,
-            "name": "Some Corpus",
-            "types": [{"slug": "folder"}, {"slug": "page"}],
-        },
-    )
+    mock_elements_worker.corpus_types = {
+        "folder": {"slug": "folder"},
+        "page": {"slug": "page"},
+    }
+
     responses.add(
         responses.POST,
         "http://testserver/api/v1/elements/type/",
@@ -94,7 +102,6 @@ def test_create_missing_types(responses, mock_elements_worker):
             )
         ],
     )
-    mock_elements_worker.setup_api_client()
 
     assert mock_elements_worker.check_required_types(
         "page", "text_line", "act", create_missing=True

tests/test_elements_worker/test_training.py

@@ -179,44 +179,12 @@ def test_validate_model_version_not_created(mock_training_worker):
         mock_training_worker.validate_model_version(hash="a", size=1, archive_hash="b")
 
 
-@pytest.mark.parametrize("status_code", [403, 500])
-def test_validate_model_version_catch_errors(
-    mocker, mock_training_worker, caplog, status_code
-):
-    mocker.patch(
-        "arkindex.client.ArkindexClient.request.retry.retry", return_value=False
-    )
-
-    mock_training_worker.model_version = {"id": "model_version_id"}
-    args = {
-        "hash": "hash",
-        "archive_hash": "archive_hash",
-        "size": 30,
-    }
-    mock_training_worker.api_client.add_error_response(
-        "ValidateModelVersion",
-        id="model_version_id",
-        status_code=status_code,
-        body=args,
-    )
-
-    mock_training_worker.validate_model_version(**args)
-    assert mock_training_worker.model_version == {"id": "model_version_id"}
-    assert [
-        (level, message)
-        for module, level, message in caplog.record_tuples
-        if module == "arkindex_worker"
-    ] == [
-        (
-            logging.WARNING,
-            "An error occurred while validating model version model_version_id, please check its status.",
-        ),
-    ]
-
-
 @pytest.mark.parametrize("deletion_failed", [True, False])
 def test_validate_model_version_hash_conflict(
-    mock_training_worker, default_model_version, caplog, deletion_failed
+    mock_training_worker,
+    default_model_version,
+    caplog,
+    deletion_failed,
 ):
     mock_training_worker.model_version = {"id": "another_id"}
     args = {
@@ -225,11 +193,11 @@ def test_validate_model_version_hash_conflict(
         "size": 30,
     }
     mock_training_worker.api_client.add_error_response(
-        "ValidateModelVersion",
+        "PartialUpdateModelVersion",
         id="another_id",
         status_code=409,
-        body=args,
-        content=default_model_version,
+        body={"state": "available", **args},
+        content={"id": ["model_version_id"]},
     )
     if deletion_failed:
         mock_training_worker.api_client.add_error_response(
@@ -244,6 +212,11 @@ def test_validate_model_version_hash_conflict(
             id="another_id",
             response="No content",
         )
+    mock_training_worker.api_client.add_response(
+        "RetrieveModelVersion",
+        id="model_version_id",
+        response=default_model_version,
+    )
 
     mock_training_worker.validate_model_version(**args)
     assert mock_training_worker.model_version == default_model_version
@@ -266,6 +239,7 @@ def test_validate_model_version_hash_conflict(
         ),
         (logging.WARNING, "Removing the pending model version."),
         *error_msg,
+        (logging.INFO, "Retrieving the existing model version."),
         (logging.INFO, "Model version model_version_id is now available."),
     ]
 
@@ -278,9 +252,9 @@ def test_validate_model_version(mock_training_worker, default_model_version, cap
         "size": 30,
     }
     mock_training_worker.api_client.add_response(
-        "ValidateModelVersion",
+        "PartialUpdateModelVersion",
         id="model_version_id",
-        body=args,
+        body={"state": "available", **args},
         response=default_model_version,
     )
 

tests/test_merge.py

@@ -161,7 +161,7 @@ def test_merge_from_worker(
     """
     responses.add(
         responses.GET,
-        "http://testserver/api/v1/task/my_task/from-agent/",
+        "http://testserver/api/v1/task/my_task/",
         status=200,
         json={"parents": ["first", "second"]},
     )