diff --git a/src/voice_transcriptor/services/job_manifest.py b/src/voice_transcriptor/services/job_manifest.py new file mode 100644 index 0000000..f8271b4 --- /dev/null +++ b/src/voice_transcriptor/services/job_manifest.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import json +from datetime import datetime, timezone +from enum import StrEnum +from pathlib import Path, PurePosixPath +from typing import Iterable, Sequence + + +class JobManifestError(RuntimeError): + pass + + +class ChunkStatus(StrEnum): + PENDING = "pending" + PROCESSING = "processing" + COMPLETED = "completed" + FAILED = "failed" + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +class JobManifestRepository: + def create( + self, + job_directory: Path, + source: dict, + settings: dict, + chunks: Sequence[dict], + ) -> dict: + job_directory = job_directory.resolve() + job_directory.mkdir(parents=True, exist_ok=True) + (job_directory / "chunks").mkdir(exist_ok=True) + timestamp = _now() + items = [self._new_chunk(item) for item in chunks] + manifest_path = job_directory / "manifest.json" + manifest = { + "schema_version": 2, + "job_id": job_directory.name, + "state": "preprocessing", + "source": dict(source), + "settings": dict(settings), + "created_at": timestamp, + "updated_at": timestamp, + "started_at": None, + "completed_at": None, + "chunks": items, + "completed_chunks": 0, + "total_chunks": len(items), + "error": None, + } + self.save(manifest_path, manifest) + return {**manifest, "manifest_path": str(manifest_path)} + + def load(self, path: Path) -> dict: + try: + manifest = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise JobManifestError("Could not load the job manifest.") from exc + self._validate(manifest) + return manifest + + def save(self, path: Path, manifest: dict) -> None: + self._validate(manifest) + manifest["updated_at"] = _now() + temporary = path.with_suffix(".tmp") + try: + temporary.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8") + temporary.replace(path) + except OSError as exc: + raise JobManifestError("Could not save the job manifest.") from exc + + def mark_processing(self, path: Path, index: int) -> dict: + manifest = self.load(path) + item = self._item(manifest, index) + if item["status"] == ChunkStatus.COMPLETED: + raise JobManifestError("A completed chunk cannot be processed again.") + item["status"] = ChunkStatus.PROCESSING + item["attempt_count"] += 1 + item["last_error"] = None + if manifest["started_at"] is None: + manifest["started_at"] = _now() + manifest["state"] = "transcribing" + self.save(path, manifest) + return manifest + + def mark_completed(self, path: Path, index: int, text: str) -> dict: + manifest = self.load(path) + item = self._item(manifest, index) + item["status"] = ChunkStatus.COMPLETED + item["transcript"] = text + item["last_error"] = None + item["completed_at"] = _now() + self._recount(manifest) + self.save(path, manifest) + return manifest + + def mark_failed(self, path: Path, index: int, message: str) -> dict: + manifest = self.load(path) + item = self._item(manifest, index) + if item["status"] == ChunkStatus.COMPLETED: + raise JobManifestError("A completed chunk cannot be marked failed.") + item["status"] = ChunkStatus.FAILED + item["last_error"] = self._sanitize(message) + manifest["state"] = "failed" + manifest["error"] = item["last_error"] + self.save(path, manifest) + return manifest + + def mark_job_state(self, path: Path, state: str, error: str | None = None) -> dict: + manifest = self.load(path) + manifest["state"] = state + manifest["error"] = self._sanitize(error) if error else None + if state == "completed": + manifest["completed_at"] = _now() + self.save(path, manifest) + return manifest + + def recover_for_resume(self, path: Path) -> dict: + manifest = self.load(path) + for item in manifest["chunks"]: + if item["status"] in (ChunkStatus.PROCESSING, ChunkStatus.FAILED): + item["status"] = ChunkStatus.PENDING + item["last_error"] = None + manifest["state"] = "transcribing" + manifest["error"] = None + self._recount(manifest) + self.save(path, manifest) + return manifest + + def reset_completed(self, path: Path, indexes: Iterable[int] | None = None) -> dict: + manifest = self.load(path) + selected = set(indexes) if indexes is not None else {item["index"] for item in manifest["chunks"]} + for item in manifest["chunks"]: + if item["index"] in selected and item["status"] == ChunkStatus.COMPLETED: + item.update(status=ChunkStatus.PENDING, transcript=None, completed_at=None, last_error=None) + self._recount(manifest) + self.save(path, manifest) + return manifest + + def assemble_transcript(self, manifest_path: Path) -> Path: + manifest = self.load(manifest_path) + texts = [ + item["transcript"] + for item in sorted(manifest["chunks"], key=lambda value: value["index"]) + if item["status"] == ChunkStatus.COMPLETED and item["transcript"] is not None + ] + target = manifest_path.parent / "transcript.txt" + temporary = target.with_suffix(".tmp") + temporary.write_text("".join(f"{text}\n" for text in texts), encoding="utf-8") + temporary.replace(target) + return target + + @staticmethod + def _new_chunk(source: dict) -> dict: + relative = PurePosixPath(str(source.get("path", ""))) + if relative.is_absolute() or ".." in relative.parts or not relative.parts or relative.parts[0] != "chunks": + raise JobManifestError("Invalid chunk path in job manifest.") + return { + **source, + "path": relative.as_posix(), + "status": ChunkStatus.PENDING, + "attempt_count": 0, + "transcript": None, + "last_error": None, + "completed_at": None, + } + + @staticmethod + def _validate(manifest: object) -> None: + if not isinstance(manifest, dict) or manifest.get("schema_version") != 2: + raise JobManifestError("Unsupported job manifest schema.") + chunks = manifest.get("chunks") + if not isinstance(chunks, list): + raise JobManifestError("Invalid chunks in job manifest.") + for item in chunks: + if not isinstance(item, dict): + raise JobManifestError("Invalid chunk in job manifest.") + relative = PurePosixPath(str(item.get("path", ""))) + if relative.is_absolute() or ".." in relative.parts or not relative.parts or relative.parts[0] != "chunks": + raise JobManifestError("Invalid chunk path in job manifest.") + + @staticmethod + def _item(manifest: dict, index: int) -> dict: + for item in manifest["chunks"]: + if item.get("index") == index: + return item + raise JobManifestError("Unknown chunk index.") + + @staticmethod + def _recount(manifest: dict) -> None: + manifest["completed_chunks"] = sum(item["status"] == ChunkStatus.COMPLETED for item in manifest["chunks"]) + + @staticmethod + def _sanitize(message: str) -> str: + value = str(message).replace("\r", " ").replace("\n", " ") + return value[:500] diff --git a/tests/test_job_manifest.py b/tests/test_job_manifest.py new file mode 100644 index 0000000..5931221 --- /dev/null +++ b/tests/test_job_manifest.py @@ -0,0 +1,113 @@ +import json +from pathlib import Path + +import pytest + +from voice_transcriptor.services.job_manifest import JobManifestError, JobManifestRepository + + +def chunks() -> list[dict]: + return [ + { + "index": 0, + "path": "chunks/chunk-00000.m4a", + "source_start_seconds": "0", + "source_end_seconds": "900.125", + "duration_seconds": "900.125", + }, + { + "index": 1, + "path": "chunks/chunk-00001.m4a", + "source_start_seconds": "885.125", + "source_end_seconds": "1000.5", + "duration_seconds": "115.375", + }, + ] + + +def create_manifest(tmp_path: Path) -> tuple[JobManifestRepository, Path]: + job = tmp_path / "job" + repository = JobManifestRepository() + manifest = repository.create( + job, + source={"path": str(tmp_path / "source.mp3"), "size_bytes": 10}, + settings={"model": "gpt-transcribe", "language": "pt-BR"}, + chunks=chunks(), + ) + return repository, Path(manifest["manifest_path"]) + + +def test_create_persists_schema_timing_and_pending_statuses(tmp_path: Path) -> None: + repository, path = create_manifest(tmp_path) + + manifest = repository.load(path) + + assert manifest["schema_version"] == 2 + assert manifest["state"] == "preprocessing" + assert manifest["completed_chunks"] == 0 + assert manifest["total_chunks"] == 2 + assert [item["status"] for item in manifest["chunks"]] == ["pending", "pending"] + assert manifest["chunks"][1]["source_start_seconds"] == "885.125" + assert not path.with_suffix(".tmp").exists() + + +def test_create_rejects_chunk_path_outside_job(tmp_path: Path) -> None: + invalid = chunks() + invalid[0]["path"] = "../secret.m4a" + + with pytest.raises(JobManifestError, match="chunk path"): + JobManifestRepository().create(tmp_path / "job", {}, {}, invalid) + + +def test_completed_chunk_is_written_immediately_and_cannot_be_reprocessed(tmp_path: Path) -> None: + repository, path = create_manifest(tmp_path) + + repository.mark_processing(path, 0) + repository.mark_completed(path, 0, "Primeiro trecho") + on_disk = json.loads(path.read_text(encoding="utf-8")) + + assert on_disk["chunks"][0]["status"] == "completed" + assert on_disk["chunks"][0]["transcript"] == "Primeiro trecho" + assert on_disk["completed_chunks"] == 1 + with pytest.raises(JobManifestError, match="completed"): + repository.mark_processing(path, 0) + + +def test_resume_recovers_processing_and_failed_but_preserves_completed(tmp_path: Path) -> None: + repository, path = create_manifest(tmp_path) + repository.mark_processing(path, 0) + repository.mark_completed(path, 0, "Feito") + repository.mark_processing(path, 1) + repository.mark_failed(path, 1, "temporary failure") + + manifest = repository.recover_for_resume(path) + + assert manifest["chunks"][0]["status"] == "completed" + assert manifest["chunks"][0]["transcript"] == "Feito" + assert manifest["chunks"][1]["status"] == "pending" + assert manifest["chunks"][1]["last_error"] is None + + +def test_assemble_transcript_uses_only_completed_chunks_in_index_order(tmp_path: Path) -> None: + repository, path = create_manifest(tmp_path) + repository.mark_processing(path, 1) + repository.mark_completed(path, 1, "Segundo") + repository.mark_processing(path, 0) + repository.mark_completed(path, 0, "Primeiro") + + transcript_path = repository.assemble_transcript(path) + + assert transcript_path.read_text(encoding="utf-8") == "Primeiro\nSegundo\n" + assert repository.load(path)["chunks"][1]["source_end_seconds"] == "1000.5" + + +def test_explicit_reset_is_required_to_clear_completed_work(tmp_path: Path) -> None: + repository, path = create_manifest(tmp_path) + repository.mark_processing(path, 0) + repository.mark_completed(path, 0, "Feito") + + repository.reset_completed(path, indexes=[0]) + + item = repository.load(path)["chunks"][0] + assert item["status"] == "pending" + assert item["transcript"] is None