2 changed files with 312 additions and 0 deletions
@ -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] |
||||
@ -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 |
||||
Loading…
Reference in new issue