fix: missing pipeline arg

ICIJ · Nov 28, 2024 · cbac9db · cbac9db
1 parent 582cd76
commit cbac9db
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/datashare_spacy_worker/app.py b/datashare_spacy_worker/app.py
@@ -1,3 +1,5 @@
+from typing import Literal
+
 from icij_worker import AsyncApp
 from icij_worker.app import TaskGroup
 from icij_worker.typing_ import RateProgress
@@ -7,7 +9,7 @@
 
 app = AsyncApp("spacy", dependencies=APP_LIFESPAN_DEPS)
 PYTHON_TASK_GROUP = "PYTHON"
-
+_SPACY_PIPELINE = "SPACY"
 
 @app.task(name="BatchNlp", group=TaskGroup(name=PYTHON_TASK_GROUP))
 async def spacy_ner(
@@ -17,7 +19,10 @@ async def spacy_ner(
     model_size: str | None = None,
     max_length: int,
     progress: RateProgress | None = None,
+    pipeline: Literal["SPACY"],
 ) -> int:
+    if pipeline != _SPACY_PIPELINE:
+        raise ValueError(f"invalid pipeline: {pipeline} expected {_SPACY_PIPELINE}")
     return await spacy_ner_(
         docs,
         categories=categories,