software/src/data/storage.py

1"""Storage utilities for translations and progress tracking."""23from __future__ import annotations45import json6from datetime import datetime7from pathlib import Path8from typing import TYPE_CHECKING, Any910from src.models.verse import TranslationOutput, VerseReference1112if TYPE_CHECKING:13    from src.orchestration.provenance import TranslationProvenance14    from src.protos.converters import (15        TranslationPromptInputProto,16        TranslationResultProto,17    )181920class Storage:21    """Manages storage of translations and progress tracking."""2223    def __init__(24        self,25        data_dir: Path | str,26        enable_archive: bool = False,27    ) -> None:28        """Initialize storage.2930        Args:31            data_dir: Path to the data directory.32            enable_archive: Whether to enable proto archive storage.33        """34        self.data_dir = Path(data_dir)35        self.output_dir = self.data_dir / "output"36        self.progress_dir = self.data_dir / "progress"37        self.archive_dir = self.data_dir / "archive"38        self._enable_archive = enable_archive39        self._ensure_dirs()4041    @property42    def enable_archive(self) -> bool:43        """Get archive enabled status."""44        return self._enable_archive4546    @enable_archive.setter47    def enable_archive(self, value: bool) -> None:48        """Set archive enabled status."""49        self._enable_archive = value50        if value:51            self.archive_dir.mkdir(parents=True, exist_ok=True)5253    def _ensure_dirs(self) -> None:54        """Ensure required directories exist."""55        self.output_dir.mkdir(parents=True, exist_ok=True)56        self.progress_dir.mkdir(parents=True, exist_ok=True)57        if self._enable_archive:58            self.archive_dir.mkdir(parents=True, exist_ok=True)5960    def _load_json(self, path: Path) -> dict[str, Any]:61        """Load a JSON file."""62        if not path.exists():63            return {}64        with path.open("r", encoding="utf-8") as f:65            return json.load(f)6667    def _save_json(self, path: Path, data: dict[str, Any]) -> None:68        """Save data to a JSON file."""69        path.parent.mkdir(parents=True, exist_ok=True)70        with path.open("w", encoding="utf-8") as f:71            json.dump(data, f, ensure_ascii=False, indent=2, default=str)7273    def save_translation(self, output: TranslationOutput) -> None:74        """Save a translation output.7576        Args:77            output: The translation output to save.78        """79        book = output.reference.book80        chapter = output.reference.chapter81        verse = output.reference.verse8283        # Load existing chapter file or create new84        chapter_file = self.output_dir / book / f"chapter_{chapter:02d}.json"85        chapter_data = self._load_json(chapter_file)8687        if "verses" not in chapter_data:88            chapter_data["verses"] = {}89            chapter_data["book"] = book90            chapter_data["chapter"] = chapter9192        verse_key = str(verse)93        chapter_data["verses"][verse_key] = {94            "text": output.text,95            "text_with_furigana": output.text_with_furigana,96            "furigana_annotations": output.furigana_annotations,97            "model": output.model,98            "timestamp": output.timestamp.isoformat(),99            "prompt_version": output.prompt_version,100            "sources_used": output.sources_used,101            "glossary_terms_applied": output.glossary_terms_applied,102        }103104        self._save_json(chapter_file, chapter_data)105        self._update_progress(book, chapter, verse)106107    def load_translation(self, reference: VerseReference) -> TranslationOutput | None:108        """Load a saved translation.109110        Args:111            reference: The verse reference to load.112113        Returns:114            TranslationOutput if found, None otherwise.115        """116        chapter_file = (117            self.output_dir / reference.book / f"chapter_{reference.chapter:02d}.json"118        )119        chapter_data = self._load_json(chapter_file)120121        verses = chapter_data.get("verses", {})122        verse_data = verses.get(str(reference.verse))123        if not verse_data:124            return None125126        return TranslationOutput(127            reference=reference,128            text=verse_data["text"],129            text_with_furigana=verse_data.get("text_with_furigana", ""),130            furigana_annotations=verse_data.get("furigana_annotations", []),131            model=verse_data["model"],132            timestamp=datetime.fromisoformat(verse_data["timestamp"]),133            prompt_version=verse_data.get("prompt_version", "v1.0"),134            sources_used=verse_data.get("sources_used", []),135            glossary_terms_applied=verse_data.get("glossary_terms_applied", []),136        )137138    def load_chapter_translations(139        self, book: str, chapter: int140    ) -> dict[int, TranslationOutput]:141        """Load all translations for a chapter.142143        Returns:144            Dictionary mapping verse numbers to TranslationOutput.145        """146        chapter_file = self.output_dir / book / f"chapter_{chapter:02d}.json"147        chapter_data = self._load_json(chapter_file)148149        translations = {}150        verses = chapter_data.get("verses", {})151        for verse_num_str, verse_data in verses.items():152            verse_num = int(verse_num_str)153            ref = VerseReference(book=book, chapter=chapter, verse=verse_num)154            translations[verse_num] = TranslationOutput(155                reference=ref,156                text=verse_data["text"],157                text_with_furigana=verse_data.get("text_with_furigana", ""),158                furigana_annotations=verse_data.get("furigana_annotations", []),159                model=verse_data["model"],160                timestamp=datetime.fromisoformat(verse_data["timestamp"]),161                prompt_version=verse_data.get("prompt_version", "v1.0"),162                sources_used=verse_data.get("sources_used", []),163                glossary_terms_applied=verse_data.get("glossary_terms_applied", []),164            )165166        return translations167168    def _update_progress(self, book: str, chapter: int, verse: int) -> None:169        """Update progress tracking for a book."""170        progress_file = self.progress_dir / f"{book}_progress.json"171        progress = self._load_json(progress_file)172173        if "completed" not in progress:174            progress["completed"] = {}175            progress["book"] = book176            progress["last_updated"] = None177178        chapter_key = str(chapter)179        if chapter_key not in progress["completed"]:180            progress["completed"][chapter_key] = []181182        if verse not in progress["completed"][chapter_key]:183            progress["completed"][chapter_key].append(verse)184            progress["completed"][chapter_key].sort()185186        progress["last_updated"] = datetime.utcnow().isoformat()187        self._save_json(progress_file, progress)188189    def get_progress(self, book: str) -> dict[str, Any]:190        """Get progress for a book.191192        Returns:193            Progress dictionary with completed verses per chapter.194        """195        progress_file = self.progress_dir / f"{book}_progress.json"196        return self._load_json(progress_file)197198    def is_verse_translated(self, reference: VerseReference) -> bool:199        """Check if a verse has already been translated."""200        progress = self.get_progress(reference.book)201        completed = progress.get("completed", {})202        chapter_verses = completed.get(str(reference.chapter), [])203        return reference.verse in chapter_verses204205    def get_completion_stats(self, book: str, total_verses: int) -> dict[str, Any]:206        """Get completion statistics for a book.207208        Args:209            book: Book name.210            total_verses: Total verses expected in the book.211212        Returns:213            Dictionary with completion statistics.214        """215        progress = self.get_progress(book)216        completed = progress.get("completed", {})217218        completed_count = sum(len(verses) for verses in completed.values())219        percentage = (completed_count / total_verses * 100) if total_verses > 0 else 0220221        return {222            "book": book,223            "completed": completed_count,224            "total": total_verses,225            "percentage": round(percentage, 2),226            "chapters_started": len(completed),227        }228229    def save_translation_with_provenance(  # pragma: no cover230        self,231        output: TranslationOutput,232        provenance: TranslationProvenance,233    ) -> None:234        """Save a translation output with full provenance tracking.235236        Args:237            output: The translation output to save.238            provenance: Full provenance record for this translation.239        """240        book = output.reference.book241        chapter = output.reference.chapter242        verse = output.reference.verse243244        # Load existing chapter file or create new245        chapter_file = self.output_dir / book / f"chapter_{chapter:02d}.json"246        chapter_data = self._load_json(chapter_file)247248        if "verses" not in chapter_data:249            chapter_data["verses"] = {}250            chapter_data["book"] = book251            chapter_data["chapter"] = chapter252253        verse_key = str(verse)254        chapter_data["verses"][verse_key] = {255            "text": output.text,256            "text_with_furigana": output.text_with_furigana,257            "furigana_annotations": output.furigana_annotations,258            "model": output.model,259            "timestamp": output.timestamp.isoformat(),260            "prompt_version": output.prompt_version,261            "sources_used": output.sources_used,262            "glossary_terms_applied": output.glossary_terms_applied,263            "provenance": provenance.to_dict(),264        }265266        self._save_json(chapter_file, chapter_data)267        self._update_progress(book, chapter, verse)268269    def load_provenance(  # pragma: no cover270        self, reference: VerseReference271    ) -> TranslationProvenance | None:272        """Load provenance for a saved translation.273274        Args:275            reference: The verse reference to load provenance for.276277        Returns:278            TranslationProvenance if found, None otherwise.279        """280        from src.orchestration.provenance import TranslationProvenance281282        chapter_file = (283            self.output_dir / reference.book / f"chapter_{reference.chapter:02d}.json"284        )285        chapter_data = self._load_json(chapter_file)286287        verses = chapter_data.get("verses", {})288        verse_data = verses.get(str(reference.verse))289        if not verse_data:290            return None291292        provenance_data = verse_data.get("provenance")293        if not provenance_data:294            return None295296        return TranslationProvenance.from_dict(provenance_data)297298    # =========================================================================299    # Proto Archive Methods300    # =========================================================================301302    def _get_archive_path(self, reference: VerseReference) -> Path:303        """Get the archive directory path for a verse.304305        Args:306            reference: The verse reference.307308        Returns:309            Path to the archive directory for this verse.310        """311        return (312            self.archive_dir313            / reference.book314            / f"chapter_{reference.chapter:02d}"315        )316317    def save_translation_archive(318        self,319        reference: VerseReference,320        input_proto: TranslationPromptInputProto,321        output_proto: TranslationResultProto,322    ) -> Path:323        """Save translation input/output protos to archive.324325        Args:326            reference: The verse reference.327            input_proto: The input proto with all translation inputs.328            output_proto: The output proto with translation results.329330        Returns:331            Path to the archive directory where files were saved.332        """333        from src.protos.serialization import serialize_proto334335        archive_path = self._get_archive_path(reference)336        archive_path.mkdir(parents=True, exist_ok=True)337338        verse_num = reference.verse339        input_file = archive_path / f"verse_{verse_num:02d}_input.msgpack"340        output_file = archive_path / f"verse_{verse_num:02d}_output.msgpack"341342        input_bytes = serialize_proto(input_proto)343        output_bytes = serialize_proto(output_proto)344345        with input_file.open("wb") as f:346            f.write(input_bytes)347348        with output_file.open("wb") as f:349            f.write(output_bytes)350351        return archive_path352353    def load_translation_archive(354        self,355        reference: VerseReference,356    ) -> tuple[TranslationPromptInputProto, TranslationResultProto] | None:357        """Load translation input/output protos from archive.358359        Args:360            reference: The verse reference.361362        Returns:363            Tuple of (input_proto, output_proto) if found, None otherwise.364        """365        from src.protos.converters import (366            TranslationPromptInputProto,367            TranslationResultProto,368        )369        from src.protos.serialization import deserialize_proto370371        archive_path = self._get_archive_path(reference)372        verse_num = reference.verse373        input_file = archive_path / f"verse_{verse_num:02d}_input.msgpack"374        output_file = archive_path / f"verse_{verse_num:02d}_output.msgpack"375376        if not input_file.exists() or not output_file.exists():377            return None378379        with input_file.open("rb") as f:380            input_bytes = f.read()381382        with output_file.open("rb") as f:383            output_bytes = f.read()384385        input_proto = deserialize_proto(input_bytes, TranslationPromptInputProto)386        output_proto = deserialize_proto(output_bytes, TranslationResultProto)387388        return input_proto, output_proto389