From 9b487b0c5d16b6eaa34d33f36611369822df6996 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Fri, 4 Sep 2026 12:19:48 +0200 Subject: [PATCH] chore(asr-worker): add both user and language --- datashare-python/datashare_python/objects.py | 10 ++++++++++ datashare-python/tests/test_utils.py | 6 +++--- workers/asr-worker/asr_worker/objects.py | 2 ++ workers/asr-worker/tests/conftest.py | 4 +++- workers/asr-worker/tests/test_activities.py | 6 ++++-- workers/asr-worker/tests/test_workflows.py | 3 +++ 6 files changed, 25 insertions(+), 6 deletions(-) diff --git a/datashare-python/datashare_python/objects.py b/datashare-python/datashare_python/objects.py index 62436dc9..8e74c3ec 100644 --- a/datashare-python/datashare_python/objects.py +++ b/datashare-python/datashare_python/objects.py @@ -366,7 +366,17 @@ class ManifestEntryStatus(StrEnum): PARTIAL = "partial" +class User(DatashareModel): + id: str + name: str | None = None + email: str | None = None + provider: str | None = None + details: dict[str, Any] = Field(default_factory=dict) + + class TaskArgs(DatashareModel, ABC): + user: User | None = None + def as_manifest_task_input(self) -> dict[str, Any]: # This is a base implementation, if the input is too large to be dumped, # override this and pop large keys diff --git a/datashare-python/tests/test_utils.py b/datashare-python/tests/test_utils.py index 6853d79f..2e52b52c 100644 --- a/datashare-python/tests/test_utils.py +++ b/datashare-python/tests/test_utils.py @@ -186,7 +186,7 @@ def test_write_artifact(tmp_path: Path) -> None: "structure": { "status": "complete", "taskId": None, - "taskInput": {"someValue": "value"}, + "taskInput": {"someValue": "value", "user": None}, "label": None, } } @@ -227,7 +227,7 @@ def test_write_artifact_with_existing_metadata(tmp_path: Path) -> None: "structure": { "status": "complete", "taskId": None, - "taskInput": {"someValue": "value"}, + "taskInput": {"someValue": "value", "user": None}, "label": None, }, "some": "value", @@ -312,7 +312,7 @@ def test_overwrite_artifact(tmp_path: Path) -> None: "structure": { "status": "complete", "taskId": None, - "taskInput": {"someValue": "value"}, + "taskInput": {"someValue": "value", "user": None}, "label": None, }, } diff --git a/workers/asr-worker/asr_worker/objects.py b/workers/asr-worker/asr_worker/objects.py index a0cee5a7..e8c149d0 100644 --- a/workers/asr-worker/asr_worker/objects.py +++ b/workers/asr-worker/asr_worker/objects.py @@ -7,6 +7,7 @@ from caul_core.objects import ASRLanguage, ASRModel from datashare_python.objects import ( ArtifactType, + DatashareLanguage, DatashareModel, DocArtifact, ManifestEntry, @@ -38,6 +39,7 @@ class ASRIndexingConfig(DatashareModel): class ASRArgs(TaskArgs): project: str docs: list[DocId] | DocumentSearchQuery + language: DatashareLanguage config: ASRPipelineConfig = Field(default_factory=ASRPipelineConfig.parakeet) batch_size: int indexing: ASRIndexingConfig = Field(default_factory=ASRIndexingConfig) diff --git a/workers/asr-worker/tests/conftest.py b/workers/asr-worker/tests/conftest.py index df92b56e..84fd7dab 100644 --- a/workers/asr-worker/tests/conftest.py +++ b/workers/asr-worker/tests/conftest.py @@ -27,12 +27,14 @@ test_worker_config, typer_asyncio_patch, ) -from datashare_python.objects import Document, WorkerPaths +from datashare_python.objects import DatashareLanguage, Document, WorkerPaths from datashare_python.utils import artifacts_dir from icij_common.es import ESClient from tests import AUDIOS_PATH +DS_ENGLISH = DatashareLanguage("ENGLISH") + @pytest.fixture(scope="session") def test_worker_config(tmp_path_factory: TempPathFactory) -> ASRWorkerConfig: # noqa: ANN001, ARG001, F811 diff --git a/workers/asr-worker/tests/test_activities.py b/workers/asr-worker/tests/test_activities.py index c71e8c5a..833076fc 100644 --- a/workers/asr-worker/tests/test_activities.py +++ b/workers/asr-worker/tests/test_activities.py @@ -43,6 +43,8 @@ from icij_common.iter_utils import batches from icij_common.registrable import RegistrableConfig +from .conftest import DS_ENGLISH + PREPROCESSED_INPUT_0 = PreprocessedInput( metadata=InputMetadata( input_ordering=0, @@ -357,7 +359,7 @@ async def test_infer_act(tmpdir: Path) -> None: def test_postprocess_act(tmpdir: Path) -> None: # Given - args = ASRArgs(project=TEST_PROJECT, docs=[], batch_size=2) + args = ASRArgs(project=TEST_PROJECT, docs=[], batch_size=2, language=DS_ENGLISH) postprocessor = MockPostprocessor() project = TEST_PROJECT artifacts_root = Path(tmpdir) @@ -460,7 +462,7 @@ def with_transcribed_docs( ) -> list[tuple[Document, str]]: artifacts_root = test_worker_config.paths.artifacts transcriptions = [] - args = ASRArgs(project=TEST_PROJECT, docs=[], batch_size=2) + args = ASRArgs(project=TEST_PROJECT, docs=[], batch_size=2, language=DS_ENGLISH) manifest_entry = TranscriptionManifestEntry.complete(args, confidence=1.0) for doc_i, doc in enumerate(with_audio_docs): artifact_factory = partial( diff --git a/workers/asr-worker/tests/test_workflows.py b/workers/asr-worker/tests/test_workflows.py index 8d7d5d8d..fac7e66a 100644 --- a/workers/asr-worker/tests/test_workflows.py +++ b/workers/asr-worker/tests/test_workflows.py @@ -29,6 +29,8 @@ from pydantic import TypeAdapter from temporalio.worker import Worker +from tests.conftest import DS_ENGLISH + _LIST_OF_PATH_ADAPTER = TypeAdapter(list[Path]) _MODEL_RESULT_0 = ASRResult( @@ -171,6 +173,7 @@ async def test_asr_workflow_e2e( # noqa: PLR0917 docs=doc_ids, config=ASRPipelineConfig.parakeet(), batch_size=batch_size, + language=DS_ENGLISH, ) workflow_id = f"asr-{uuid.uuid4().hex}"