feature(asr-worker): add user and language to asr

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

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,
         },
     }

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)

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

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(

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}"