enpi_api.l2.client.api.enrichment_api

View Source

  1import os
  2import zipfile
  3from pathlib import Path
  4from uuid import UUID
  5
  6import pandas as pd
  7from loguru import logger
  8
  9from enpi_api.l1 import openapi_client
 10from enpi_api.l2.client.api.file_api import FileApi
 11from enpi_api.l2.events.workflow_execution_task_waitable import WorkflowExecutionTaskWaitable
 12from enpi_api.l2.types.api_error import ApiErrorContext
 13from enpi_api.l2.types.cluster import ClusterRunId
 14from enpi_api.l2.types.enrichment import (
 15    EnrichmentExportMode,
 16    EnrichmentRun,
 17    EnrichmentRunId,
 18    EnrichmentTemplate,
 19    EnrichmentTemplateId,
 20    EnrichmentTemplateOperation,
 21    EnrichmentWorkInput,
 22    SimplifiedEnrichmentTemplate,
 23    transform_operation,
 24    transform_operation_input,
 25)
 26from enpi_api.l2.types.execution import Execution
 27from enpi_api.l2.types.log import LogLevel
 28from enpi_api.l2.types.tag import TagId
 29from enpi_api.l2.types.task import TaskState
 30from enpi_api.l2.types.workflow import WorkflowExecutionId, WorkflowExecutionTaskId, WorkflowTaskTemplateName
 31from enpi_api.l2.util.file import unique_temp_dir
 32
 33
 34class EnrichmentApi:
 35    _inner_api_client: openapi_client.ApiClient
 36    _log_level: LogLevel
 37
 38    def __init__(self, inner_api_client: openapi_client.ApiClient, log_level: LogLevel):
 39        """@private"""
 40        self._inner_api_client = inner_api_client
 41        self._log_level = log_level
 42
 43    def get_runs(
 44        self,
 45    ) -> list[EnrichmentRun]:
 46        """Get all successful Enrichment Runs.
 47
 48        Returns:
 49            list[enpi_api.l2.types.enrichment.EnrichmentRun]: List of Enrichment Runs.
 50
 51        Raises:
 52            enpi_api.l2.types.api_error.ApiError: If API request fails.
 53
 54        Example:
 55            ```python
 56            with EnpiApiClient() as enpi_client:
 57                enrichment_runs = enpi_client.enrichment_api.get_runs()
 58            ```
 59        """
 60        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
 61
 62        with ApiErrorContext():
 63            data = enrichment_api_instance.get_enrichment_runs()
 64
 65        return [EnrichmentRun.from_raw(cr) for cr in data.runs]
 66
 67    def get_run(self, enrichment_run_id: EnrichmentRunId) -> EnrichmentRun:
 68        """Get a single Enrichment Run by its ID.
 69
 70        Args:
 71            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): ID of the Enrichment run to get.
 72
 73        Returns:
 74            enpi_api.l2.types.enrichment.EnrichmentRun: A successful Enrichment Run.
 75
 76        Raises:
 77            enpi_api.l2.types.api_error.ApiError: If API request fails.
 78
 79        Example:
 80            ```python
 81            with EnpiApiClient() as enpi_client:
 82                enrichment_run = enpi_client.enrichment_api.get_run(EnrichmentRunId(123))
 83            ```
 84        """
 85        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
 86
 87        with ApiErrorContext():
 88            data = enrichment_api_instance.get_enrichment_run(int(enrichment_run_id))
 89
 90        return EnrichmentRun.from_raw(data.run)
 91
 92    def get_run_by_task_id(self, task_id: WorkflowExecutionTaskId) -> EnrichmentRun:
 93        """Get a single Enrichment Run by its task ID.
 94
 95        Args:
 96            task_id (enpi_api.l2.types.task.TaskId): ID of a task linked to a successful Enrichment Run.
 97
 98        Returns:
 99            enpi_api.l2.types.enrichment.EnrichmentRun: Successful Enrichment Run linked to the provided task ID.
100
101        Raises:
102            enpi_api.l2.types.api_error.ApiError: If API request fails.
103
104        Example:
105            Fetch a Enrichment Run
106
107            ```python
108            with EnpiApiClient() as enpi_client:
109                enrichment_run = enpi_client.enrichment_api.get_run_by_task_id(TaskId(1234))
110            ```
111        """
112        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
113
114        with ApiErrorContext():
115            data = enrichment_api_instance.get_enrichment_run_by_task_id(task_id)
116
117        return EnrichmentRun.from_raw(data.run)
118
119    def get_templates(self) -> list[SimplifiedEnrichmentTemplate]:
120        """Get all available Enrichment templates.
121
122        Returns:
123            list[enpi_api.l2.types.enrichment.SimplifiedEnrichmentTemplate]: Available Enrichment templates.
124
125        Raises:
126            enpi_api.l2.types.api_error.ApiError: If API request fails.
127
128        Example:
129            ```python
130            with EnpiApiClient() as enpi_client:
131                templates = enpi_client.enrichment_api.get_templates()
132            ```
133        """
134        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
135
136        with ApiErrorContext():
137            data = enrichment_api_instance.get_enrichment_templates()
138
139        return [
140            SimplifiedEnrichmentTemplate(
141                id=EnrichmentTemplateId(str(d.id)),
142                name=str(d.name),
143                created_at=d.created_at,
144            )
145            for d in data.templates
146        ]
147
148    def get_template(self, enrichment_template_id: EnrichmentTemplateId) -> EnrichmentTemplate:
149        """Get a Enrichment template by its ID.
150
151        Args:
152            enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of a Enrichment template to get.
153
154        Returns:
155            enpi_api.l2.types.enrichment.EnrichmentTemplate: A Enrichment template matching the provided ID.
156
157        Raises:
158            enpi_api.l2.types.api_error.ApiError: If API request fails.
159
160        Example:
161            ```python
162            with EnpiApiClient() as enpi_client:
163                template = enpi_client.enrichment_api.get_template(EnrichmentTemplateId(24))
164            ```
165        """
166        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
167
168        with ApiErrorContext():
169            data = enrichment_api_instance.get_enrichment_template(UUID(enrichment_template_id))
170
171        return EnrichmentTemplate.from_raw(data.template)
172
173    def delete_template(self, enrichment_template_id: EnrichmentTemplateId) -> None:
174        """Delete a Enrichment Template by its ID.
175
176        Args:
177            enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of the deleted Enrichment template.
178
179        Raises:
180            enpi_api.l2.types.api_error.ApiError: If API request fails.
181
182        Example:
183            ```python
184            with EnpiApiClient() as enpi_client:
185                template = enpi_client.enrichment_api.delete_template(EnrichmentTemplateId(24))
186            ```
187        """
188        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
189
190        with ApiErrorContext():
191            enrichment_api_instance.delete_enrichment_template(UUID(enrichment_template_id), body={})
192
193    def create_template(
194        self,
195        name: str,
196        operations: list[EnrichmentTemplateOperation],
197    ) -> EnrichmentTemplate:
198        """Create a new Enrichment Template.
199
200        Args:
201            name (str): Enrichment template name.
202            operations (list[enpi_api.l2.types.enrichment.EnrichmentTemplateOperation]): Configs for the template's operations.
203
204        Returns:
205            enpi_api.l2.types.enrichment.EnrichmentTemplate: A newly created Enrichment template.
206
207        Raises:
208            enpi_api.l2.types.api_error.ApiError: If API request fails.
209
210        Example:
211            ```python
212            with EnpiApiClient() as enpi_client:
213                operations=[
214                    EnrichmentTemplateUnionOperation(
215                        name="Counter",
216                    ),
217                    EnrichmentTemplateIntersectionOperation(
218                        name="Target",
219                    ),
220                    EnrichmentTemplateDifferenceOperation(
221                        name="Result",
222                        input_operations=EnrichmentTemplateDifferenceInputs(
223                            remove_operation="Counter",
224                            from_operation="Target",
225                        ),
226                        annotations=[
227                            EnrichmentTemplateFoldChangeAnnotation(name="Target FC"),
228                        ],
229                    ),
230                ],
231                template = enpi_client.enrichment_api.create_template(name="my new template", operations=operations)
232            ```
233        """
234        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
235
236        payload = openapi_client.NewEnrichmentTemplate(name=name, saved=True, operations=[transform_operation(op) for op in operations])
237
238        with ApiErrorContext():
239            data = enrichment_api_instance.create_enrichment_template(payload)
240
241        return EnrichmentTemplate.from_raw(data.template)
242
243    def start(
244        self,
245        name: str,
246        enrichment_template: EnrichmentTemplate,
247        cluster_run_id: ClusterRunId,
248        inputs: list[EnrichmentWorkInput],
249    ) -> Execution[EnrichmentRun]:
250        """Start a new Enrichment run.
251
252        Args:
253            name (str):
254                Enrichment run name.
255            enrichment_template (enpi_api.l2.types.enrichment.EnrichmentTemplate):
256                Enrichment template ID.
257            cluster_run_id (enpi_api.l2.types.cluster.ClusterRunId):
258                Cluster run ID.
259            inputs (list[enpi_api.l2.types.enrichment.EnrichmentWorkInput]):
260                Configs for the template's operations.
261
262        Returns:
263            enpi_api.l2.types.execution.Execution[enpi_api.l2.types.enrichment.EnrichmentRun]: An awaitable that returns the new Enrichment Run.
264
265        Raises:
266            enpi_api.l2.types.api_error.ApiError: If API request fails.
267
268        Example:
269            Assuming that those are the collections meant for Enrichment run:
270                - Counter collection IDs: 1, 2
271                - Target collection IDs: 3, 4
272
273            ```python
274            with EnpiApiClient() as enpi_client:
275                enrichment_run = client.enrichment_api.start(
276                    name="Test Run",
277                    enrichment_template=template,
278                    cluster_run_id=cluster_run.id,
279                    inputs=[
280                        UnionOperationInput(
281                            name="Counter",
282                            input_collections=[CollectionSelector(value=CollectionId(1)),
283                                CollectionSelector(value=CollectionId(2))],
284                        ),
285                        IntersectionOperationInput(
286                            name="Target",
287                            input_collections=[CollectionSelector(value=CollectionId(3)),
288                                CollectionSelector(value=CollectionId(4))
289                            ],
290                        ),
291                        FoldChangeInput(
292                            name="Target FC",
293                            operation_name="Result"
294                            input_collections=FoldChangeInputCollections(
295                                from_collection=CollectionSelector(value=CollectionId(3)),
296                                to_collection=CollectionSelector(value=CollectionId(4)),
297                            ),
298                        ),
299                    ],
300                ).wait()
301            ```
302        """
303        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
304
305        payload = openapi_client.EnrichmentWork(
306            name=name,
307            template=openapi_client.EnrichmentTemplateIdVersion(id=str(enrichment_template.id), version=int(enrichment_template.version)),
308            cluster_run_id=cluster_run_id,
309            inputs=[transform_operation_input(x) for x in inputs],
310        )
311
312        with ApiErrorContext():
313            data = enrichment_api_instance.start_enrichment_run(payload)
314            assert data.workflow_execution_id is not None
315
316            workflow_execution_id = WorkflowExecutionId(int(data.workflow_execution_id))
317
318            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> EnrichmentRun:
319                enrichment_run = self.get_run_by_task_id(task_id)
320
321                logger.success(f"Enrichment run with task ID: {data.workflow_execution_id} has successfully finished.")
322
323                return enrichment_run
324
325            waitable = WorkflowExecutionTaskWaitable[EnrichmentRun](
326                workflow_execution_id=workflow_execution_id, task_template_name=WorkflowTaskTemplateName.ENPI_APP_ENRICHMENT, on_complete=on_complete
327            )
328            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
329
330    def export_as_tsv(
331        self,
332        enrichment_run_id: EnrichmentRunId,
333        operation: str,
334        mode: EnrichmentExportMode,
335        tag_ids: list[TagId],
336        limit: int | None = None,
337        output_directory: str | Path | None = None,
338    ) -> Execution[Path]:
339        """Run an export of a Enrichment operation result and download the TSV file.
340
341        Args:
342            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
343            operation (str): Name of the operation to export.
344            mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
345            tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
346            limit (int | None):
347                If specified, the export will contain only the first N clusters, N being the value passed as this param.
348            output_directory (str | Path | None): The directory where to download the TSV file. Defaults to a
349              unique temporary directory.
350
351        Returns:
352            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.
353
354        Raises:
355            enpi_api.l2.types.api_error.ApiError: If API request fails.
356
357        Example:
358            ```python
359            with EnpiApiClient() as enpi_client:
360                tsv_path = enpi_client.enrichment_api.export_as_tsv(
361                    enrichment_run_id=EnrichmentRunId(42),
362                    operation="Result",
363                    mode=EnrichmentExportMode.REPRESENTATIVES,
364                    tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
365                    limit=50,
366                    output_directory="/results",
367                ).wait()
368            ```
369        """
370
371        # Ensure that the directory exists
372        if output_directory is None:
373            output_directory = unique_temp_dir()
374
375        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
376
377        payload = openapi_client.EnrichmentExportPayload(
378            operation_name=operation,
379            enrichment_run_id=enrichment_run_id,
380            mode=openapi_client.ExportMode(mode.value),
381            tag_ids=[int(id) for id in tag_ids],
382            limit=limit,
383        )
384
385        with ApiErrorContext():
386            response = enrichment_api_instance.export_enrichment_results(payload)
387            assert response.workflow_execution_id is not None
388
389            workflow_execution_id = WorkflowExecutionId(response.workflow_execution_id)
390
391            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> Path:
392                file_api = FileApi(self._inner_api_client, self._log_level)
393                zip_path = file_api.download_export_by_workflow_execution_task_id(task_id=task_id, output_directory=output_directory)
394
395                with zipfile.ZipFile(zip_path, "r") as archive:
396                    names = archive.namelist()
397                    archive.extractall(output_directory)
398                paths = [Path(os.path.join(output_directory, name)) for name in names]
399
400                if len(paths) > 1:
401                    logger.warning(f"More than 1 file encountered in the zipped export: {','.join([str(path) for path in paths])}, only returning first file")
402                logger.success(f"Enrichment results export with task ID: {task_id} has successfully finished.")
403
404                return paths[0]
405
406            waitable = WorkflowExecutionTaskWaitable[Path](
407                workflow_execution_id=workflow_execution_id, on_complete=on_complete, task_template_name=WorkflowTaskTemplateName.ENPI_APP_ENRICHMENT_EXPORT
408            )
409            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
410
411    def export_as_df(
412        self,
413        enrichment_run_id: EnrichmentRunId,
414        operation: str,
415        mode: EnrichmentExportMode,
416        tag_ids: list[TagId],
417        limit: int | None = None,
418    ) -> Execution[pd.DataFrame]:
419        """Runs an Export of a Enrichment Operation results and loads it as a pandas DataFrame.
420
421        Args:
422            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
423            operation (str): Name of the operation to export.
424            mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
425            tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
426            limit (int | None):
427                If specified, the export will contain only the first N clusters, N being the value passed as this param.
428
429        Returns:
430            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.
431
432        Raises:
433            enpi_api.l2.types.api_error.ApiError: If API request fails.
434
435        Example:
436            ```python
437            with EnpiApiClient() as enpi_client:
438                tsv_path = enpi_client.enrichment_api.export_as_df(
439                    enrichment_run_id=EnrichmentRunId(42),
440                    operation="Result",
441                    mode=EnrichmentExportMode.REPRESENTATIVES,
442                    tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
443                    limit=50
444                ).wait()
445            ```
446        """
447
448        export_tsv = self.export_as_tsv(enrichment_run_id, operation, mode, tag_ids, limit)
449
450        def wait() -> pd.DataFrame:
451            file_path = export_tsv.wait()
452            return pd.read_csv(file_path, delimiter="\t")
453
454        return Execution(wait=wait, check_execution_state=export_tsv.check_execution_state)

class EnrichmentApi: View Source

 35class EnrichmentApi:
 36    _inner_api_client: openapi_client.ApiClient
 37    _log_level: LogLevel
 38
 39    def __init__(self, inner_api_client: openapi_client.ApiClient, log_level: LogLevel):
 40        """@private"""
 41        self._inner_api_client = inner_api_client
 42        self._log_level = log_level
 43
 44    def get_runs(
 45        self,
 46    ) -> list[EnrichmentRun]:
 47        """Get all successful Enrichment Runs.
 48
 49        Returns:
 50            list[enpi_api.l2.types.enrichment.EnrichmentRun]: List of Enrichment Runs.
 51
 52        Raises:
 53            enpi_api.l2.types.api_error.ApiError: If API request fails.
 54
 55        Example:
 56            ```python
 57            with EnpiApiClient() as enpi_client:
 58                enrichment_runs = enpi_client.enrichment_api.get_runs()
 59            ```
 60        """
 61        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
 62
 63        with ApiErrorContext():
 64            data = enrichment_api_instance.get_enrichment_runs()
 65
 66        return [EnrichmentRun.from_raw(cr) for cr in data.runs]
 67
 68    def get_run(self, enrichment_run_id: EnrichmentRunId) -> EnrichmentRun:
 69        """Get a single Enrichment Run by its ID.
 70
 71        Args:
 72            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): ID of the Enrichment run to get.
 73
 74        Returns:
 75            enpi_api.l2.types.enrichment.EnrichmentRun: A successful Enrichment Run.
 76
 77        Raises:
 78            enpi_api.l2.types.api_error.ApiError: If API request fails.
 79
 80        Example:
 81            ```python
 82            with EnpiApiClient() as enpi_client:
 83                enrichment_run = enpi_client.enrichment_api.get_run(EnrichmentRunId(123))
 84            ```
 85        """
 86        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
 87
 88        with ApiErrorContext():
 89            data = enrichment_api_instance.get_enrichment_run(int(enrichment_run_id))
 90
 91        return EnrichmentRun.from_raw(data.run)
 92
 93    def get_run_by_task_id(self, task_id: WorkflowExecutionTaskId) -> EnrichmentRun:
 94        """Get a single Enrichment Run by its task ID.
 95
 96        Args:
 97            task_id (enpi_api.l2.types.task.TaskId): ID of a task linked to a successful Enrichment Run.
 98
 99        Returns:
100            enpi_api.l2.types.enrichment.EnrichmentRun: Successful Enrichment Run linked to the provided task ID.
101
102        Raises:
103            enpi_api.l2.types.api_error.ApiError: If API request fails.
104
105        Example:
106            Fetch a Enrichment Run
107
108            ```python
109            with EnpiApiClient() as enpi_client:
110                enrichment_run = enpi_client.enrichment_api.get_run_by_task_id(TaskId(1234))
111            ```
112        """
113        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
114
115        with ApiErrorContext():
116            data = enrichment_api_instance.get_enrichment_run_by_task_id(task_id)
117
118        return EnrichmentRun.from_raw(data.run)
119
120    def get_templates(self) -> list[SimplifiedEnrichmentTemplate]:
121        """Get all available Enrichment templates.
122
123        Returns:
124            list[enpi_api.l2.types.enrichment.SimplifiedEnrichmentTemplate]: Available Enrichment templates.
125
126        Raises:
127            enpi_api.l2.types.api_error.ApiError: If API request fails.
128
129        Example:
130            ```python
131            with EnpiApiClient() as enpi_client:
132                templates = enpi_client.enrichment_api.get_templates()
133            ```
134        """
135        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
136
137        with ApiErrorContext():
138            data = enrichment_api_instance.get_enrichment_templates()
139
140        return [
141            SimplifiedEnrichmentTemplate(
142                id=EnrichmentTemplateId(str(d.id)),
143                name=str(d.name),
144                created_at=d.created_at,
145            )
146            for d in data.templates
147        ]
148
149    def get_template(self, enrichment_template_id: EnrichmentTemplateId) -> EnrichmentTemplate:
150        """Get a Enrichment template by its ID.
151
152        Args:
153            enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of a Enrichment template to get.
154
155        Returns:
156            enpi_api.l2.types.enrichment.EnrichmentTemplate: A Enrichment template matching the provided ID.
157
158        Raises:
159            enpi_api.l2.types.api_error.ApiError: If API request fails.
160
161        Example:
162            ```python
163            with EnpiApiClient() as enpi_client:
164                template = enpi_client.enrichment_api.get_template(EnrichmentTemplateId(24))
165            ```
166        """
167        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
168
169        with ApiErrorContext():
170            data = enrichment_api_instance.get_enrichment_template(UUID(enrichment_template_id))
171
172        return EnrichmentTemplate.from_raw(data.template)
173
174    def delete_template(self, enrichment_template_id: EnrichmentTemplateId) -> None:
175        """Delete a Enrichment Template by its ID.
176
177        Args:
178            enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of the deleted Enrichment template.
179
180        Raises:
181            enpi_api.l2.types.api_error.ApiError: If API request fails.
182
183        Example:
184            ```python
185            with EnpiApiClient() as enpi_client:
186                template = enpi_client.enrichment_api.delete_template(EnrichmentTemplateId(24))
187            ```
188        """
189        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
190
191        with ApiErrorContext():
192            enrichment_api_instance.delete_enrichment_template(UUID(enrichment_template_id), body={})
193
194    def create_template(
195        self,
196        name: str,
197        operations: list[EnrichmentTemplateOperation],
198    ) -> EnrichmentTemplate:
199        """Create a new Enrichment Template.
200
201        Args:
202            name (str): Enrichment template name.
203            operations (list[enpi_api.l2.types.enrichment.EnrichmentTemplateOperation]): Configs for the template's operations.
204
205        Returns:
206            enpi_api.l2.types.enrichment.EnrichmentTemplate: A newly created Enrichment template.
207
208        Raises:
209            enpi_api.l2.types.api_error.ApiError: If API request fails.
210
211        Example:
212            ```python
213            with EnpiApiClient() as enpi_client:
214                operations=[
215                    EnrichmentTemplateUnionOperation(
216                        name="Counter",
217                    ),
218                    EnrichmentTemplateIntersectionOperation(
219                        name="Target",
220                    ),
221                    EnrichmentTemplateDifferenceOperation(
222                        name="Result",
223                        input_operations=EnrichmentTemplateDifferenceInputs(
224                            remove_operation="Counter",
225                            from_operation="Target",
226                        ),
227                        annotations=[
228                            EnrichmentTemplateFoldChangeAnnotation(name="Target FC"),
229                        ],
230                    ),
231                ],
232                template = enpi_client.enrichment_api.create_template(name="my new template", operations=operations)
233            ```
234        """
235        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
236
237        payload = openapi_client.NewEnrichmentTemplate(name=name, saved=True, operations=[transform_operation(op) for op in operations])
238
239        with ApiErrorContext():
240            data = enrichment_api_instance.create_enrichment_template(payload)
241
242        return EnrichmentTemplate.from_raw(data.template)
243
244    def start(
245        self,
246        name: str,
247        enrichment_template: EnrichmentTemplate,
248        cluster_run_id: ClusterRunId,
249        inputs: list[EnrichmentWorkInput],
250    ) -> Execution[EnrichmentRun]:
251        """Start a new Enrichment run.
252
253        Args:
254            name (str):
255                Enrichment run name.
256            enrichment_template (enpi_api.l2.types.enrichment.EnrichmentTemplate):
257                Enrichment template ID.
258            cluster_run_id (enpi_api.l2.types.cluster.ClusterRunId):
259                Cluster run ID.
260            inputs (list[enpi_api.l2.types.enrichment.EnrichmentWorkInput]):
261                Configs for the template's operations.
262
263        Returns:
264            enpi_api.l2.types.execution.Execution[enpi_api.l2.types.enrichment.EnrichmentRun]: An awaitable that returns the new Enrichment Run.
265
266        Raises:
267            enpi_api.l2.types.api_error.ApiError: If API request fails.
268
269        Example:
270            Assuming that those are the collections meant for Enrichment run:
271                - Counter collection IDs: 1, 2
272                - Target collection IDs: 3, 4
273
274            ```python
275            with EnpiApiClient() as enpi_client:
276                enrichment_run = client.enrichment_api.start(
277                    name="Test Run",
278                    enrichment_template=template,
279                    cluster_run_id=cluster_run.id,
280                    inputs=[
281                        UnionOperationInput(
282                            name="Counter",
283                            input_collections=[CollectionSelector(value=CollectionId(1)),
284                                CollectionSelector(value=CollectionId(2))],
285                        ),
286                        IntersectionOperationInput(
287                            name="Target",
288                            input_collections=[CollectionSelector(value=CollectionId(3)),
289                                CollectionSelector(value=CollectionId(4))
290                            ],
291                        ),
292                        FoldChangeInput(
293                            name="Target FC",
294                            operation_name="Result"
295                            input_collections=FoldChangeInputCollections(
296                                from_collection=CollectionSelector(value=CollectionId(3)),
297                                to_collection=CollectionSelector(value=CollectionId(4)),
298                            ),
299                        ),
300                    ],
301                ).wait()
302            ```
303        """
304        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
305
306        payload = openapi_client.EnrichmentWork(
307            name=name,
308            template=openapi_client.EnrichmentTemplateIdVersion(id=str(enrichment_template.id), version=int(enrichment_template.version)),
309            cluster_run_id=cluster_run_id,
310            inputs=[transform_operation_input(x) for x in inputs],
311        )
312
313        with ApiErrorContext():
314            data = enrichment_api_instance.start_enrichment_run(payload)
315            assert data.workflow_execution_id is not None
316
317            workflow_execution_id = WorkflowExecutionId(int(data.workflow_execution_id))
318
319            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> EnrichmentRun:
320                enrichment_run = self.get_run_by_task_id(task_id)
321
322                logger.success(f"Enrichment run with task ID: {data.workflow_execution_id} has successfully finished.")
323
324                return enrichment_run
325
326            waitable = WorkflowExecutionTaskWaitable[EnrichmentRun](
327                workflow_execution_id=workflow_execution_id, task_template_name=WorkflowTaskTemplateName.ENPI_APP_ENRICHMENT, on_complete=on_complete
328            )
329            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
330
331    def export_as_tsv(
332        self,
333        enrichment_run_id: EnrichmentRunId,
334        operation: str,
335        mode: EnrichmentExportMode,
336        tag_ids: list[TagId],
337        limit: int | None = None,
338        output_directory: str | Path | None = None,
339    ) -> Execution[Path]:
340        """Run an export of a Enrichment operation result and download the TSV file.
341
342        Args:
343            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
344            operation (str): Name of the operation to export.
345            mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
346            tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
347            limit (int | None):
348                If specified, the export will contain only the first N clusters, N being the value passed as this param.
349            output_directory (str | Path | None): The directory where to download the TSV file. Defaults to a
350              unique temporary directory.
351
352        Returns:
353            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.
354
355        Raises:
356            enpi_api.l2.types.api_error.ApiError: If API request fails.
357
358        Example:
359            ```python
360            with EnpiApiClient() as enpi_client:
361                tsv_path = enpi_client.enrichment_api.export_as_tsv(
362                    enrichment_run_id=EnrichmentRunId(42),
363                    operation="Result",
364                    mode=EnrichmentExportMode.REPRESENTATIVES,
365                    tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
366                    limit=50,
367                    output_directory="/results",
368                ).wait()
369            ```
370        """
371
372        # Ensure that the directory exists
373        if output_directory is None:
374            output_directory = unique_temp_dir()
375
376        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
377
378        payload = openapi_client.EnrichmentExportPayload(
379            operation_name=operation,
380            enrichment_run_id=enrichment_run_id,
381            mode=openapi_client.ExportMode(mode.value),
382            tag_ids=[int(id) for id in tag_ids],
383            limit=limit,
384        )
385
386        with ApiErrorContext():
387            response = enrichment_api_instance.export_enrichment_results(payload)
388            assert response.workflow_execution_id is not None
389
390            workflow_execution_id = WorkflowExecutionId(response.workflow_execution_id)
391
392            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> Path:
393                file_api = FileApi(self._inner_api_client, self._log_level)
394                zip_path = file_api.download_export_by_workflow_execution_task_id(task_id=task_id, output_directory=output_directory)
395
396                with zipfile.ZipFile(zip_path, "r") as archive:
397                    names = archive.namelist()
398                    archive.extractall(output_directory)
399                paths = [Path(os.path.join(output_directory, name)) for name in names]
400
401                if len(paths) > 1:
402                    logger.warning(f"More than 1 file encountered in the zipped export: {','.join([str(path) for path in paths])}, only returning first file")
403                logger.success(f"Enrichment results export with task ID: {task_id} has successfully finished.")
404
405                return paths[0]
406
407            waitable = WorkflowExecutionTaskWaitable[Path](
408                workflow_execution_id=workflow_execution_id, on_complete=on_complete, task_template_name=WorkflowTaskTemplateName.ENPI_APP_ENRICHMENT_EXPORT
409            )
410            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
411
412    def export_as_df(
413        self,
414        enrichment_run_id: EnrichmentRunId,
415        operation: str,
416        mode: EnrichmentExportMode,
417        tag_ids: list[TagId],
418        limit: int | None = None,
419    ) -> Execution[pd.DataFrame]:
420        """Runs an Export of a Enrichment Operation results and loads it as a pandas DataFrame.
421
422        Args:
423            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
424            operation (str): Name of the operation to export.
425            mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
426            tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
427            limit (int | None):
428                If specified, the export will contain only the first N clusters, N being the value passed as this param.
429
430        Returns:
431            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.
432
433        Raises:
434            enpi_api.l2.types.api_error.ApiError: If API request fails.
435
436        Example:
437            ```python
438            with EnpiApiClient() as enpi_client:
439                tsv_path = enpi_client.enrichment_api.export_as_df(
440                    enrichment_run_id=EnrichmentRunId(42),
441                    operation="Result",
442                    mode=EnrichmentExportMode.REPRESENTATIVES,
443                    tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
444                    limit=50
445                ).wait()
446            ```
447        """
448
449        export_tsv = self.export_as_tsv(enrichment_run_id, operation, mode, tag_ids, limit)
450
451        def wait() -> pd.DataFrame:
452            file_path = export_tsv.wait()
453            return pd.read_csv(file_path, delimiter="\t")
454
455        return Execution(wait=wait, check_execution_state=export_tsv.check_execution_state)

def get_runs(self) -> list[enpi_api.l2.types.enrichment.EnrichmentRun]: View Source

44    def get_runs(
45        self,
46    ) -> list[EnrichmentRun]:
47        """Get all successful Enrichment Runs.
48
49        Returns:
50            list[enpi_api.l2.types.enrichment.EnrichmentRun]: List of Enrichment Runs.
51
52        Raises:
53            enpi_api.l2.types.api_error.ApiError: If API request fails.
54
55        Example:
56            ```python
57            with EnpiApiClient() as enpi_client:
58                enrichment_runs = enpi_client.enrichment_api.get_runs()
59            ```
60        """
61        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
62
63        with ApiErrorContext():
64            data = enrichment_api_instance.get_enrichment_runs()
65
66        return [EnrichmentRun.from_raw(cr) for cr in data.runs]

Get all successful Enrichment Runs.

Returns:

list[enpi_api.l2.types.enrichment.EnrichmentRun]: List of Enrichment Runs.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    enrichment_runs = enpi_client.enrichment_api.get_runs()

def get_run( self, enrichment_run_id: enpi_api.l2.types.enrichment.EnrichmentRunId) -> enpi_api.l2.types.enrichment.EnrichmentRun: View Source

68    def get_run(self, enrichment_run_id: EnrichmentRunId) -> EnrichmentRun:
69        """Get a single Enrichment Run by its ID.
70
71        Args:
72            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): ID of the Enrichment run to get.
73
74        Returns:
75            enpi_api.l2.types.enrichment.EnrichmentRun: A successful Enrichment Run.
76
77        Raises:
78            enpi_api.l2.types.api_error.ApiError: If API request fails.
79
80        Example:
81            ```python
82            with EnpiApiClient() as enpi_client:
83                enrichment_run = enpi_client.enrichment_api.get_run(EnrichmentRunId(123))
84            ```
85        """
86        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
87
88        with ApiErrorContext():
89            data = enrichment_api_instance.get_enrichment_run(int(enrichment_run_id))
90
91        return EnrichmentRun.from_raw(data.run)

Get a single Enrichment Run by its ID.

Arguments:

enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): ID of the Enrichment run to get.

Returns:

enpi_api.l2.types.enrichment.EnrichmentRun: A successful Enrichment Run.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    enrichment_run = enpi_client.enrichment_api.get_run(EnrichmentRunId(123))

def get_run_by_task_id( self, task_id: enpi_api.l2.types.workflow.WorkflowExecutionTaskId) -> enpi_api.l2.types.enrichment.EnrichmentRun: View Source

 93    def get_run_by_task_id(self, task_id: WorkflowExecutionTaskId) -> EnrichmentRun:
 94        """Get a single Enrichment Run by its task ID.
 95
 96        Args:
 97            task_id (enpi_api.l2.types.task.TaskId): ID of a task linked to a successful Enrichment Run.
 98
 99        Returns:
100            enpi_api.l2.types.enrichment.EnrichmentRun: Successful Enrichment Run linked to the provided task ID.
101
102        Raises:
103            enpi_api.l2.types.api_error.ApiError: If API request fails.
104
105        Example:
106            Fetch a Enrichment Run
107
108            ```python
109            with EnpiApiClient() as enpi_client:
110                enrichment_run = enpi_client.enrichment_api.get_run_by_task_id(TaskId(1234))
111            ```
112        """
113        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
114
115        with ApiErrorContext():
116            data = enrichment_api_instance.get_enrichment_run_by_task_id(task_id)
117
118        return EnrichmentRun.from_raw(data.run)

Get a single Enrichment Run by its task ID.

Arguments:

task_id (enpi_api.l2.types.task.TaskId): ID of a task linked to a successful Enrichment Run.

Returns:

enpi_api.l2.types.enrichment.EnrichmentRun: Successful Enrichment Run linked to the provided task ID.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

Fetch a Enrichment Run

with EnpiApiClient() as enpi_client:
    enrichment_run = enpi_client.enrichment_api.get_run_by_task_id(TaskId(1234))

def get_templates(self) -> list[enpi_api.l2.types.enrichment.SimplifiedEnrichmentTemplate]: View Source

120    def get_templates(self) -> list[SimplifiedEnrichmentTemplate]:
121        """Get all available Enrichment templates.
122
123        Returns:
124            list[enpi_api.l2.types.enrichment.SimplifiedEnrichmentTemplate]: Available Enrichment templates.
125
126        Raises:
127            enpi_api.l2.types.api_error.ApiError: If API request fails.
128
129        Example:
130            ```python
131            with EnpiApiClient() as enpi_client:
132                templates = enpi_client.enrichment_api.get_templates()
133            ```
134        """
135        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
136
137        with ApiErrorContext():
138            data = enrichment_api_instance.get_enrichment_templates()
139
140        return [
141            SimplifiedEnrichmentTemplate(
142                id=EnrichmentTemplateId(str(d.id)),
143                name=str(d.name),
144                created_at=d.created_at,
145            )
146            for d in data.templates
147        ]

Get all available Enrichment templates.

Returns:

list[enpi_api.l2.types.enrichment.SimplifiedEnrichmentTemplate]: Available Enrichment templates.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    templates = enpi_client.enrichment_api.get_templates()

def get_template( self, enrichment_template_id: enpi_api.l2.types.enrichment.EnrichmentTemplateId) -> enpi_api.l2.types.enrichment.EnrichmentTemplate: View Source

149    def get_template(self, enrichment_template_id: EnrichmentTemplateId) -> EnrichmentTemplate:
150        """Get a Enrichment template by its ID.
151
152        Args:
153            enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of a Enrichment template to get.
154
155        Returns:
156            enpi_api.l2.types.enrichment.EnrichmentTemplate: A Enrichment template matching the provided ID.
157
158        Raises:
159            enpi_api.l2.types.api_error.ApiError: If API request fails.
160
161        Example:
162            ```python
163            with EnpiApiClient() as enpi_client:
164                template = enpi_client.enrichment_api.get_template(EnrichmentTemplateId(24))
165            ```
166        """
167        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
168
169        with ApiErrorContext():
170            data = enrichment_api_instance.get_enrichment_template(UUID(enrichment_template_id))
171
172        return EnrichmentTemplate.from_raw(data.template)

Get a Enrichment template by its ID.

Arguments:

enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of a Enrichment template to get.

Returns:

enpi_api.l2.types.enrichment.EnrichmentTemplate: A Enrichment template matching the provided ID.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    template = enpi_client.enrichment_api.get_template(EnrichmentTemplateId(24))

def delete_template( self, enrichment_template_id: enpi_api.l2.types.enrichment.EnrichmentTemplateId) -> None: View Source

174    def delete_template(self, enrichment_template_id: EnrichmentTemplateId) -> None:
175        """Delete a Enrichment Template by its ID.
176
177        Args:
178            enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of the deleted Enrichment template.
179
180        Raises:
181            enpi_api.l2.types.api_error.ApiError: If API request fails.
182
183        Example:
184            ```python
185            with EnpiApiClient() as enpi_client:
186                template = enpi_client.enrichment_api.delete_template(EnrichmentTemplateId(24))
187            ```
188        """
189        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
190
191        with ApiErrorContext():
192            enrichment_api_instance.delete_enrichment_template(UUID(enrichment_template_id), body={})

Delete a Enrichment Template by its ID.

Arguments:

enrichment_template_id (enpi_api.l2.types.enrichment.EnrichmentTemplateId): ID of the deleted Enrichment template.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    template = enpi_client.enrichment_api.delete_template(EnrichmentTemplateId(24))

def create_template( self, name: str, operations: list[enpi_api.l2.types.enrichment.EnrichmentTemplateUnionOperation | enpi_api.l2.types.enrichment.EnrichmentTemplateIntersectionOperation | enpi_api.l2.types.enrichment.EnrichmentTemplateDifferenceOperation]) -> enpi_api.l2.types.enrichment.EnrichmentTemplate: View Source

194    def create_template(
195        self,
196        name: str,
197        operations: list[EnrichmentTemplateOperation],
198    ) -> EnrichmentTemplate:
199        """Create a new Enrichment Template.
200
201        Args:
202            name (str): Enrichment template name.
203            operations (list[enpi_api.l2.types.enrichment.EnrichmentTemplateOperation]): Configs for the template's operations.
204
205        Returns:
206            enpi_api.l2.types.enrichment.EnrichmentTemplate: A newly created Enrichment template.
207
208        Raises:
209            enpi_api.l2.types.api_error.ApiError: If API request fails.
210
211        Example:
212            ```python
213            with EnpiApiClient() as enpi_client:
214                operations=[
215                    EnrichmentTemplateUnionOperation(
216                        name="Counter",
217                    ),
218                    EnrichmentTemplateIntersectionOperation(
219                        name="Target",
220                    ),
221                    EnrichmentTemplateDifferenceOperation(
222                        name="Result",
223                        input_operations=EnrichmentTemplateDifferenceInputs(
224                            remove_operation="Counter",
225                            from_operation="Target",
226                        ),
227                        annotations=[
228                            EnrichmentTemplateFoldChangeAnnotation(name="Target FC"),
229                        ],
230                    ),
231                ],
232                template = enpi_client.enrichment_api.create_template(name="my new template", operations=operations)
233            ```
234        """
235        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
236
237        payload = openapi_client.NewEnrichmentTemplate(name=name, saved=True, operations=[transform_operation(op) for op in operations])
238
239        with ApiErrorContext():
240            data = enrichment_api_instance.create_enrichment_template(payload)
241
242        return EnrichmentTemplate.from_raw(data.template)

Create a new Enrichment Template.

Arguments:

name (str): Enrichment template name.
operations (list[enpi_api.l2.types.enrichment.EnrichmentTemplateOperation]): Configs for the template's operations.

Returns:

enpi_api.l2.types.enrichment.EnrichmentTemplate: A newly created Enrichment template.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    operations=[
        EnrichmentTemplateUnionOperation(
            name="Counter",
        ),
        EnrichmentTemplateIntersectionOperation(
            name="Target",
        ),
        EnrichmentTemplateDifferenceOperation(
            name="Result",
            input_operations=EnrichmentTemplateDifferenceInputs(
                remove_operation="Counter",
                from_operation="Target",
            ),
            annotations=[
                EnrichmentTemplateFoldChangeAnnotation(name="Target FC"),
            ],
        ),
    ],
    template = enpi_client.enrichment_api.create_template(name="my new template", operations=operations)

def start( self, name: str, enrichment_template: enpi_api.l2.types.enrichment.EnrichmentTemplate, cluster_run_id: enpi_api.l2.types.cluster.ClusterRunId, inputs: list[enpi_api.l2.types.enrichment.UnionOperationInput | enpi_api.l2.types.enrichment.IntersectionOperationInput | enpi_api.l2.types.enrichment.DifferenceOperationInput | enpi_api.l2.types.enrichment.FoldChangeInput]) -> enpi_api.l2.types.execution.Execution[EnrichmentRun]: View Source

244    def start(
245        self,
246        name: str,
247        enrichment_template: EnrichmentTemplate,
248        cluster_run_id: ClusterRunId,
249        inputs: list[EnrichmentWorkInput],
250    ) -> Execution[EnrichmentRun]:
251        """Start a new Enrichment run.
252
253        Args:
254            name (str):
255                Enrichment run name.
256            enrichment_template (enpi_api.l2.types.enrichment.EnrichmentTemplate):
257                Enrichment template ID.
258            cluster_run_id (enpi_api.l2.types.cluster.ClusterRunId):
259                Cluster run ID.
260            inputs (list[enpi_api.l2.types.enrichment.EnrichmentWorkInput]):
261                Configs for the template's operations.
262
263        Returns:
264            enpi_api.l2.types.execution.Execution[enpi_api.l2.types.enrichment.EnrichmentRun]: An awaitable that returns the new Enrichment Run.
265
266        Raises:
267            enpi_api.l2.types.api_error.ApiError: If API request fails.
268
269        Example:
270            Assuming that those are the collections meant for Enrichment run:
271                - Counter collection IDs: 1, 2
272                - Target collection IDs: 3, 4
273
274            ```python
275            with EnpiApiClient() as enpi_client:
276                enrichment_run = client.enrichment_api.start(
277                    name="Test Run",
278                    enrichment_template=template,
279                    cluster_run_id=cluster_run.id,
280                    inputs=[
281                        UnionOperationInput(
282                            name="Counter",
283                            input_collections=[CollectionSelector(value=CollectionId(1)),
284                                CollectionSelector(value=CollectionId(2))],
285                        ),
286                        IntersectionOperationInput(
287                            name="Target",
288                            input_collections=[CollectionSelector(value=CollectionId(3)),
289                                CollectionSelector(value=CollectionId(4))
290                            ],
291                        ),
292                        FoldChangeInput(
293                            name="Target FC",
294                            operation_name="Result"
295                            input_collections=FoldChangeInputCollections(
296                                from_collection=CollectionSelector(value=CollectionId(3)),
297                                to_collection=CollectionSelector(value=CollectionId(4)),
298                            ),
299                        ),
300                    ],
301                ).wait()
302            ```
303        """
304        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
305
306        payload = openapi_client.EnrichmentWork(
307            name=name,
308            template=openapi_client.EnrichmentTemplateIdVersion(id=str(enrichment_template.id), version=int(enrichment_template.version)),
309            cluster_run_id=cluster_run_id,
310            inputs=[transform_operation_input(x) for x in inputs],
311        )
312
313        with ApiErrorContext():
314            data = enrichment_api_instance.start_enrichment_run(payload)
315            assert data.workflow_execution_id is not None
316
317            workflow_execution_id = WorkflowExecutionId(int(data.workflow_execution_id))
318
319            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> EnrichmentRun:
320                enrichment_run = self.get_run_by_task_id(task_id)
321
322                logger.success(f"Enrichment run with task ID: {data.workflow_execution_id} has successfully finished.")
323
324                return enrichment_run
325
326            waitable = WorkflowExecutionTaskWaitable[EnrichmentRun](
327                workflow_execution_id=workflow_execution_id, task_template_name=WorkflowTaskTemplateName.ENPI_APP_ENRICHMENT, on_complete=on_complete
328            )
329            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)

Start a new Enrichment run.

Arguments:

name (str): Enrichment run name.
enrichment_template (enpi_api.l2.types.enrichment.EnrichmentTemplate): Enrichment template ID.
cluster_run_id (enpi_api.l2.types.cluster.ClusterRunId): Cluster run ID.
inputs (list[enpi_api.l2.types.enrichment.EnrichmentWorkInput]): Configs for the template's operations.

Returns:

enpi_api.l2.types.execution.Execution[enpi_api.l2.types.enrichment.EnrichmentRun]: An awaitable that returns the new Enrichment Run.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

Assuming that those are the collections meant for Enrichment run: - Counter collection IDs: 1, 2 - Target collection IDs: 3, 4

with EnpiApiClient() as enpi_client:
    enrichment_run = client.enrichment_api.start(
        name="Test Run",
        enrichment_template=template,
        cluster_run_id=cluster_run.id,
        inputs=[
            UnionOperationInput(
                name="Counter",
                input_collections=[CollectionSelector(value=CollectionId(1)),
                    CollectionSelector(value=CollectionId(2))],
            ),
            IntersectionOperationInput(
                name="Target",
                input_collections=[CollectionSelector(value=CollectionId(3)),
                    CollectionSelector(value=CollectionId(4))
                ],
            ),
            FoldChangeInput(
                name="Target FC",
                operation_name="Result"
                input_collections=FoldChangeInputCollections(
                    from_collection=CollectionSelector(value=CollectionId(3)),
                    to_collection=CollectionSelector(value=CollectionId(4)),
                ),
            ),
        ],
    ).wait()

def export_as_tsv( self, enrichment_run_id: enpi_api.l2.types.enrichment.EnrichmentRunId, operation: str, mode: enpi_api.l2.types.enrichment.EnrichmentExportMode, tag_ids: list[enpi_api.l2.types.tag.TagId], limit: int | None = None, output_directory: str | pathlib.Path | None = None) -> enpi_api.l2.types.execution.Execution[Path]: View Source

331    def export_as_tsv(
332        self,
333        enrichment_run_id: EnrichmentRunId,
334        operation: str,
335        mode: EnrichmentExportMode,
336        tag_ids: list[TagId],
337        limit: int | None = None,
338        output_directory: str | Path | None = None,
339    ) -> Execution[Path]:
340        """Run an export of a Enrichment operation result and download the TSV file.
341
342        Args:
343            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
344            operation (str): Name of the operation to export.
345            mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
346            tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
347            limit (int | None):
348                If specified, the export will contain only the first N clusters, N being the value passed as this param.
349            output_directory (str | Path | None): The directory where to download the TSV file. Defaults to a
350              unique temporary directory.
351
352        Returns:
353            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.
354
355        Raises:
356            enpi_api.l2.types.api_error.ApiError: If API request fails.
357
358        Example:
359            ```python
360            with EnpiApiClient() as enpi_client:
361                tsv_path = enpi_client.enrichment_api.export_as_tsv(
362                    enrichment_run_id=EnrichmentRunId(42),
363                    operation="Result",
364                    mode=EnrichmentExportMode.REPRESENTATIVES,
365                    tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
366                    limit=50,
367                    output_directory="/results",
368                ).wait()
369            ```
370        """
371
372        # Ensure that the directory exists
373        if output_directory is None:
374            output_directory = unique_temp_dir()
375
376        enrichment_api_instance = openapi_client.EnrichmentApi(self._inner_api_client)
377
378        payload = openapi_client.EnrichmentExportPayload(
379            operation_name=operation,
380            enrichment_run_id=enrichment_run_id,
381            mode=openapi_client.ExportMode(mode.value),
382            tag_ids=[int(id) for id in tag_ids],
383            limit=limit,
384        )
385
386        with ApiErrorContext():
387            response = enrichment_api_instance.export_enrichment_results(payload)
388            assert response.workflow_execution_id is not None
389
390            workflow_execution_id = WorkflowExecutionId(response.workflow_execution_id)
391
392            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> Path:
393                file_api = FileApi(self._inner_api_client, self._log_level)
394                zip_path = file_api.download_export_by_workflow_execution_task_id(task_id=task_id, output_directory=output_directory)
395
396                with zipfile.ZipFile(zip_path, "r") as archive:
397                    names = archive.namelist()
398                    archive.extractall(output_directory)
399                paths = [Path(os.path.join(output_directory, name)) for name in names]
400
401                if len(paths) > 1:
402                    logger.warning(f"More than 1 file encountered in the zipped export: {','.join([str(path) for path in paths])}, only returning first file")
403                logger.success(f"Enrichment results export with task ID: {task_id} has successfully finished.")
404
405                return paths[0]
406
407            waitable = WorkflowExecutionTaskWaitable[Path](
408                workflow_execution_id=workflow_execution_id, on_complete=on_complete, task_template_name=WorkflowTaskTemplateName.ENPI_APP_ENRICHMENT_EXPORT
409            )
410            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)

Run an export of a Enrichment operation result and download the TSV file.

Arguments:

enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
operation (str): Name of the operation to export.
mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
limit (int | None): If specified, the export will contain only the first N clusters, N being the value passed as this param.
output_directory (str | Path | None): The directory where to download the TSV file. Defaults to a unique temporary directory.

Returns:

enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    tsv_path = enpi_client.enrichment_api.export_as_tsv(
        enrichment_run_id=EnrichmentRunId(42),
        operation="Result",
        mode=EnrichmentExportMode.REPRESENTATIVES,
        tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
        limit=50,
        output_directory="/results",
    ).wait()

def export_as_df( self, enrichment_run_id: enpi_api.l2.types.enrichment.EnrichmentRunId, operation: str, mode: enpi_api.l2.types.enrichment.EnrichmentExportMode, tag_ids: list[enpi_api.l2.types.tag.TagId], limit: int | None = None) -> enpi_api.l2.types.execution.Execution[DataFrame]: View Source

412    def export_as_df(
413        self,
414        enrichment_run_id: EnrichmentRunId,
415        operation: str,
416        mode: EnrichmentExportMode,
417        tag_ids: list[TagId],
418        limit: int | None = None,
419    ) -> Execution[pd.DataFrame]:
420        """Runs an Export of a Enrichment Operation results and loads it as a pandas DataFrame.
421
422        Args:
423            enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
424            operation (str): Name of the operation to export.
425            mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
426            tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
427            limit (int | None):
428                If specified, the export will contain only the first N clusters, N being the value passed as this param.
429
430        Returns:
431            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.
432
433        Raises:
434            enpi_api.l2.types.api_error.ApiError: If API request fails.
435
436        Example:
437            ```python
438            with EnpiApiClient() as enpi_client:
439                tsv_path = enpi_client.enrichment_api.export_as_df(
440                    enrichment_run_id=EnrichmentRunId(42),
441                    operation="Result",
442                    mode=EnrichmentExportMode.REPRESENTATIVES,
443                    tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
444                    limit=50
445                ).wait()
446            ```
447        """
448
449        export_tsv = self.export_as_tsv(enrichment_run_id, operation, mode, tag_ids, limit)
450
451        def wait() -> pd.DataFrame:
452            file_path = export_tsv.wait()
453            return pd.read_csv(file_path, delimiter="\t")
454
455        return Execution(wait=wait, check_execution_state=export_tsv.check_execution_state)

Runs an Export of a Enrichment Operation results and loads it as a pandas DataFrame.

Arguments:

enrichment_run_id (enpi_api.l2.types.enrichment.EnrichmentRunId): Enrichment run ID.
operation (str): Name of the operation to export.
mode (enpi_api.l2.types.enrichment.EnrichmentExportMode): Mode in which export will be run.
tag_ids (list[enpi_api.l2.types.tag.TagId]): List of tags to be included in the export.
limit (int | None): If specified, the export will contain only the first N clusters, N being the value passed as this param.

Returns:

enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns the path to the downloaded TSV file containing the sequences.

Raises:

enpi_api.l2.types.api_error.ApiError: If API request fails.

Example:

with EnpiApiClient() as enpi_client:
    tsv_path = enpi_client.enrichment_api.export_as_df(
        enrichment_run_id=EnrichmentRunId(42),
        operation="Result",
        mode=EnrichmentExportMode.REPRESENTATIVES,
        tag_ids=[SequenceTags.Cdr3AminoAcids, SequenceTags.Chain],
        limit=50
    ).wait()