software/src/evaluation/storage.py

1"""Storage utilities for translation evaluations."""23from __future__ import annotations45import json6from datetime import datetime7from pathlib import Path8from typing import Any910from src.models.evaluation import (11    BookEvaluationSummary,12    ChapterEvaluation,13    EvaluationScore,14    VerseEvaluation,15)16from src.models.verse import VerseReference171819class EvaluationStorage:20    """Manages persistence of translation evaluations."""2122    def __init__(self, data_dir: Path | str) -> None:23        """Initialize the evaluation storage.2425        Args:26            data_dir: Path to the data directory.27        """28        self.data_dir = Path(data_dir)29        self.output_dir = self.data_dir / "output"3031    def _get_eval_path(self, book: str, chapter: int) -> Path:32        """Get path to evaluation file for a chapter."""33        return self.output_dir / book / f"chapter_{chapter:02d}_eval.json"3435    def _serialize_datetime(self, obj: Any) -> Any:36        """Serialize datetime objects for JSON."""37        if isinstance(obj, datetime):38            return obj.isoformat()39        raise TypeError(f"Object of type {type(obj)} is not JSON serializable")4041    def save_chapter_evaluation(self, evaluation: ChapterEvaluation) -> Path:42        """Save a chapter evaluation to JSON.4344        Args:45            evaluation: ChapterEvaluation to save.4647        Returns:48            Path to the saved file.49        """50        path = self._get_eval_path(evaluation.book, evaluation.chapter)51        path.parent.mkdir(parents=True, exist_ok=True)5253        # Convert to dict for JSON serialization54        data = {55            "book": evaluation.book,56            "chapter": evaluation.chapter,57            "evaluator_model": evaluation.evaluator_model,58            "evaluated_at": evaluation.evaluated_at.isoformat(),59            "verses": {},60            "aggregate": {61                "scores": evaluation.aggregate_scores,62                "overall": evaluation.overall_score,63                "flagged_count": evaluation.flagged_count,64                "flagged_verses": evaluation.flagged_verses,65            },66        }6768        for verse_num, verse_eval in evaluation.verses.items():69            data["verses"][str(verse_num)] = {70                "scores": {71                    dim: {72                        "score": score.score,73                        "confidence": score.confidence,74                        "reasoning": score.reasoning,75                    }76                    for dim, score in verse_eval.scores.items()77                },78                "overall_score": verse_eval.overall_score,79                "flags": verse_eval.flags,80                "notes": verse_eval.notes,81                "timestamp": verse_eval.timestamp.isoformat(),82            }8384        with path.open("w", encoding="utf-8") as f:85            json.dump(data, f, ensure_ascii=False, indent=2)8687        return path8889    def load_chapter_evaluation(90        self, book: str, chapter: int91    ) -> ChapterEvaluation | None:92        """Load a chapter evaluation from JSON.9394        Args:95            book: Book name.96            chapter: Chapter number.9798        Returns:99            ChapterEvaluation if found, None otherwise.100        """101        path = self._get_eval_path(book, chapter)102        if not path.exists():103            return None104105        with path.open("r", encoding="utf-8") as f:106            data = json.load(f)107108        verses = {}109        for verse_num_str, verse_data in data.get("verses", {}).items():110            verse_num = int(verse_num_str)111            scores = {}112            for dim, score_data in verse_data.get("scores", {}).items():113                scores[dim] = EvaluationScore(114                    dimension=dim,115                    score=score_data["score"],116                    confidence=score_data["confidence"],117                    reasoning=score_data["reasoning"],118                )119120            verses[verse_num] = VerseEvaluation(121                verse_number=verse_num,122                evaluator_model=data.get("evaluator_model", "unknown"),123                timestamp=datetime.fromisoformat(124                    verse_data.get("timestamp", datetime.utcnow().isoformat())125                ),126                scores=scores,127                flags=verse_data.get("flags", []),128                notes=verse_data.get("notes"),129            )130131        return ChapterEvaluation(132            book=data["book"],133            chapter=data["chapter"],134            evaluator_model=data.get("evaluator_model", "unknown"),135            evaluated_at=datetime.fromisoformat(136                data.get("evaluated_at", datetime.utcnow().isoformat())137            ),138            verses=verses,139        )140141    def get_flagged_verses(self, book: str) -> list[VerseReference]:142        """Get all flagged verses for a book.143144        Args:145            book: Book name.146147        Returns:148            List of VerseReference objects for flagged verses.149        """150        flagged = []151        book_dir = self.output_dir / book152153        if not book_dir.exists():154            return flagged155156        for eval_file in sorted(book_dir.glob("chapter_*_eval.json")):157            chapter_eval = self.load_chapter_evaluation(158                book, int(eval_file.stem.split("_")[1])159            )160            if chapter_eval:161                for verse_num in chapter_eval.flagged_verses:162                    flagged.append(163                        VerseReference(164                            book=book, chapter=chapter_eval.chapter, verse=verse_num165                        )166                    )167168        return flagged169170    def get_book_summary(self, book: str) -> BookEvaluationSummary | None:171        """Generate a summary of evaluations for an entire book.172173        Args:174            book: Book name.175176        Returns:177            BookEvaluationSummary if evaluations exist, None otherwise.178        """179        book_dir = self.output_dir / book180        if not book_dir.exists():181            return None182183        eval_files = list(book_dir.glob("chapter_*_eval.json"))184        if not eval_files:185            return None186187        all_scores: dict[str, list[float]] = {}188        total_verses = 0189        flagged_refs: list[str] = []190191        for eval_file in sorted(eval_files):192            chapter_num = int(eval_file.stem.split("_")[1])193            chapter_eval = self.load_chapter_evaluation(book, chapter_num)194195            if chapter_eval:196                total_verses += len(chapter_eval.verses)197198                # Aggregate dimension scores199                for dim, score in chapter_eval.aggregate_scores.items():200                    if dim not in all_scores:201                        all_scores[dim] = []202                    all_scores[dim].append(score)203204                # Collect flagged verses205                for verse_num in chapter_eval.flagged_verses:206                    flagged_refs.append(f"{book}:{chapter_num}:{verse_num}")207208        if not all_scores:209            return None210211        aggregate = {212            dim: round(sum(scores) / len(scores), 1)213            for dim, scores in all_scores.items()214        }215        overall = round(sum(aggregate.values()) / len(aggregate), 1)216217        return BookEvaluationSummary(218            book=book,219            total_chapters=len(eval_files),220            total_verses=total_verses,221            aggregate_scores=aggregate,222            overall_score=overall,223            flagged_verses=flagged_refs,224            evaluated_at=datetime.utcnow(),225        )226227    def is_chapter_evaluated(self, book: str, chapter: int) -> bool:228        """Check if a chapter has been evaluated.229230        Args:231            book: Book name.232            chapter: Chapter number.233234        Returns:235            True if evaluation exists.236        """237        return self._get_eval_path(book, chapter).exists()238239    def save_verse_evaluation(240        self, book: str, chapter: int, evaluation: VerseEvaluation241    ) -> Path:242        """Save a single verse evaluation.243244        Loads existing chapter evaluation if it exists, updates the verse,245        and saves back.246247        Args:248            book: Book name.249            chapter: Chapter number.250            evaluation: The verse evaluation to save.251252        Returns:253            Path to the saved file.254        """255        # Load existing or create new chapter evaluation256        chapter_eval = self.load_chapter_evaluation(book, chapter)257        if chapter_eval is None:258            chapter_eval = ChapterEvaluation(259                book=book,260                chapter=chapter,261                evaluator_model=evaluation.evaluator_model,262                verses={},263            )264265        # Update the verse266        chapter_eval.verses[evaluation.verse_number] = evaluation267268        # Save back269        return self.save_chapter_evaluation(chapter_eval)270