feature(asr-worker): index transcriptions

datashare-python/datashare_python/objects.py

@@ -179,6 +179,11 @@ class Translation(BaseModel):  # No camelcase here we don't know why
     content: Annotated[str, BeforeValidator(_from_sentences)]
 
 
+Routing = str
+DocID = str
+DocRoute = tuple[DocID, Routing]
+
+
 class Document(DatashareModel):
     id: str
     language: DatashareLanguage
@@ -254,6 +259,11 @@ def to_filesystem(self) -> FilesystemDocument:
             resource_name=resource_name,
         )
 
+    def to_route(self) -> DocRoute:
+        if self.root_document is not None:
+            return self.root_document, self.id
+        return self.id, self.id
+
 
 def _is_absolute_path(v: bytes | BytesIO | Path) -> Any:
     if isinstance(v, Path) and not v.is_absolute():

datashare-python/datashare_python/utils.py

@@ -3,6 +3,7 @@
 import fcntl
 import inspect
 import json
+import logging
 import os
 import shutil
 import sys
@@ -45,6 +46,8 @@
 from .objects import BaseModel, DocArtifact, DocumentLocation, FilesystemDocument
 from .types_ import RawAsyncProgressHandler
 
+logger = logging.getLogger(__name__)
+
 DependencyLabel = str | None
 DependencySetup = Callable[..., None]
 DependencyAsyncSetup = Callable[..., Coroutine[None, None, None]]
@@ -397,7 +400,7 @@ def write_artifact(
             )
         manifest[artifact.type] = manifest_entry
         manifest_path.write_text(json.dumps(manifest))
-        return artifact_path.relative_to(artif_dir)
+        return artifact_path.relative_to(root)
 
 
 @contextlib.contextmanager
@@ -556,6 +559,16 @@ def symlink_embedded_document_to_workdir(
             raise ValueError(f"unsupported location {doc.location}")
 
 
+def artifact_path(
+    doc_id: str,
+    artifact_type: type[DocArtifact],
+    *,
+    project: str,
+    root: Path,
+) -> Path:
+    return root / artifacts_dir(doc_id, project=project) / artifact_type.filename
+
+
 def read_jsonl(path: Path) -> Iterable[dict]:
     with path.open() as f:
         for line in f:
@@ -564,6 +577,34 @@ def read_jsonl(path: Path) -> Iterable[dict]:
                 yield json.loads(line)
 
 
+async def publish_and_consume(
+    publisher: asyncio.Task,
+    publisher_completion_callback: Callable[[], None],
+    *,
+    consumer: asyncio.Task,
+) -> tuple[Any, Any]:
+    # Publish and consume concurrently
+    logger.debug("starting publish and subscribe")
+    done, pending = await asyncio.wait(
+        [publisher, consumer], return_when=asyncio.FIRST_COMPLETED
+    )
+    for d in done:
+        # Stop everything case of exception
+        exc = d.exception()
+        if exc:
+            for p in pending:
+                p.cancel()
+            raise exc
+    # Wait for publish to be done and push the poison pill to stop consuming
+    p_res = await publisher
+    publisher_completion_callback()
+    logger.debug("done publishing, waiting for consumer to complete...")
+    # Wait for consumption to be done
+    c_res = await consumer
+    logger.debug("done consuming !")
+    return p_res, c_res
+
+
 class _PydanticPayloadConverter(CompositePayloadConverter):
     def __init__(self) -> None:
         json_payload_converter = PydanticJSONPlainPayloadConverter(

datashare-python/tests/test_utils.py

@@ -164,9 +164,13 @@ def test_write_artifact(tmp_path: Path) -> None:
         manifest_entry=manifest_entry,
     )
     # When
-    write_artifact(root_dir, artifact)
+    artifact_path = write_artifact(root_dir, artifact)
     # Then
-    artifact_dir = root_dir / TEST_PROJECT / "do" / "c_" / "doc_id"
+    expected_artifact_cache = (
+        Path(TEST_PROJECT) / "do" / "c_" / "doc_id" / "mocked-structure"
+    )
+    assert artifact_path == expected_artifact_cache
+    artifact_dir = (root_dir / artifact_path).parent
     assert artifact_dir.exists()
     assert artifact_dir.is_dir()
     manifest_path = artifact_dir / "manifest.json"
@@ -247,9 +251,13 @@ def test_write_artifact_with_existing_legacy_metadata(tmp_path: Path) -> None:
     meta_path = artifact_dir / "metadata.json"
     meta_path.write_text(json.dumps(existing_metadata))
     # When
-    write_artifact(root_dir, artifact)
+    artifact_path = write_artifact(root_dir, artifact)
     # Then
-    artifact_dir = root_dir / TEST_PROJECT / "do" / "c_" / "doc_id"
+    expected_artifact_path = (
+        Path(TEST_PROJECT) / "do" / "c_" / "doc_id" / "mocked-structure"
+    )
+    assert artifact_path == expected_artifact_path
+    artifact_dir = (root_dir / artifact_path).parent
     assert artifact_dir.exists()
     assert artifact_dir.is_dir()
     meta_path = artifact_dir / "metadata.json"