software/src/context/translation_memory.py

1"""Translation Memory management for segment storage and retrieval."""23from __future__ import annotations45import fcntl6import json7import logging8from contextlib import contextmanager9from datetime import datetime, timezone10from pathlib import Path11from typing import Iterator1213from src.models.translation_memory import (14    MatchType,15    SegmentLevel,16    TMMatch,17    TMSegment,18    TMStatistics,19    calculate_similarity,20    normalize_text,21)22from src.models.verse import VerseReference2324logger = logging.getLogger(__name__)2526#: Default store location. Kept as a module constant so the test suite can27#: redirect it — a test that constructs a manager with no path would otherwise28#: read and write the real, irreplaceable production memory.29DEFAULT_STORAGE_PATH = Path("data/memory/segments.json")303132class TranslationMemoryManager:33    """Manages translation memory segments for reuse.3435    Provides storage, retrieval, and fuzzy matching of36    previously translated segments.37    """3839    def __init__(self, storage_path: Path | str | None = None) -> None:40        """Initialize the translation memory manager.4142        Args:43            storage_path: Path to the storage file. If None, uses default.44        """45        if storage_path is None:46            self._storage_path = DEFAULT_STORAGE_PATH47        else:48            self._storage_path = Path(storage_path)4950        self._segments: dict[str, TMSegment] = {}51        self._by_reference: dict[str, list[str]] = {}  # ref_str -> [segment_ids]52        self._loaded = False5354    @property55    def storage_path(self) -> Path:56        """Get the storage file path."""57        return self._storage_path5859    def _ensure_loaded(self) -> None:60        """Ensure segments are loaded from storage."""61        if not self._loaded:62            self.load()6364    def add_segment(self, segment: TMSegment) -> None:65        """Add a segment to the translation memory.6667        Args:68            segment: The segment to add.69        """70        self._ensure_loaded()7172        self._segments[segment.id] = segment7374        # Index by reference75        ref_str = str(segment.reference)76        if ref_str not in self._by_reference:77            self._by_reference[ref_str] = []78        if segment.id not in self._by_reference[ref_str]:79            self._by_reference[ref_str].append(segment.id)8081    def add_from_translation(82        self,83        source_text: str,84        target_text: str,85        reference: VerseReference,86        level: SegmentLevel = SegmentLevel.VERSE,87        metadata: dict | None = None,88    ) -> TMSegment:89        """Create and add a segment from a translation result.9091        Args:92            source_text: Original source text.93            target_text: Translated text.94            reference: Verse reference.95            level: Segment level.96            metadata: Additional metadata.9798        Returns:99            The created TMSegment.100        """101        segment = TMSegment.create(102            source_text=source_text,103            target_text=target_text,104            reference=reference,105            level=level,106            metadata=metadata,107        )108        self.add_segment(segment)109        return segment110111    def get_segment(self, segment_id: str) -> TMSegment | None:112        """Get a segment by ID.113114        Args:115            segment_id: The segment ID.116117        Returns:118            The segment if found, None otherwise.119        """120        self._ensure_loaded()121        return self._segments.get(segment_id)122123    def get_by_reference(self, reference: VerseReference) -> list[TMSegment]:124        """Get all segments for a verse reference.125126        Args:127            reference: The verse reference.128129        Returns:130            List of segments for that reference.131        """132        self._ensure_loaded()133134        ref_str = str(reference)135        segment_ids = self._by_reference.get(ref_str, [])136        return [self._segments[sid] for sid in segment_ids if sid in self._segments]137138    def find_matches(139        self,140        source_text: str,141        threshold: float = 0.7,142        max_results: int = 5,143    ) -> list[TMMatch]:144        """Find matching segments for a source text.145146        Args:147            source_text: The source text to match.148            threshold: Minimum similarity threshold (0.0-1.0).149            max_results: Maximum number of results to return.150151        Returns:152            List of TMMatch objects sorted by similarity (highest first).153        """154        self._ensure_loaded()155156        if not self._segments:157            return []158159        normalized = normalize_text(source_text)160        matches = []161162        for segment in self._segments.values():163            similarity = calculate_similarity(normalized, segment.source_normalized)164165            if similarity >= threshold:166                if similarity == 1.0:167                    match_type = MatchType.EXACT168                elif segment.level == SegmentLevel.PHRASE:169                    match_type = MatchType.PHRASE170                else:171                    match_type = MatchType.FUZZY172173                matches.append(174                    TMMatch(175                        segment=segment,176                        similarity=similarity,177                        match_type=match_type,178                    )179                )180181        # Sort by similarity (highest first)182        matches.sort(key=lambda m: m.similarity, reverse=True)183184        return matches[:max_results]185186    def find_exact_match(self, source_text: str) -> TMSegment | None:187        """Find an exact match for a source text.188189        Args:190            source_text: The source text to match.191192        Returns:193            The matching segment if found, None otherwise.194        """195        matches = self.find_matches(source_text, threshold=1.0, max_results=1)196        if matches and matches[0].is_exact:197            return matches[0].segment198        return None199200    def increment_reuse(self, segment_id: str) -> bool:201        """Increment the reuse count for a segment.202203        Args:204            segment_id: The segment ID.205206        Returns:207            True if successful, False if segment not found.208        """209        self._ensure_loaded()210211        segment = self._segments.get(segment_id)212        if segment:213            segment.increment_reuse()214            return True215        return False216217    def get_statistics(self) -> TMStatistics:218        """Get statistics about the translation memory.219220        Returns:221            TMStatistics object.222        """223        self._ensure_loaded()224        return TMStatistics.from_segments(list(self._segments.values()))225226    def get_all_segments(self) -> list[TMSegment]:227        """Get all segments in the translation memory.228229        Returns:230            List of all TMSegment objects.231        """232        self._ensure_loaded()233        return list(self._segments.values())234235    def remove_segment(self, segment_id: str) -> bool:236        """Remove a segment from the translation memory.237238        Args:239            segment_id: The segment ID to remove.240241        Returns:242            True if removed, False if not found.243        """244        self._ensure_loaded()245246        segment = self._segments.pop(segment_id, None)247        if segment:248            ref_str = str(segment.reference)249            if ref_str in self._by_reference:250                self._by_reference[ref_str] = [251                    sid for sid in self._by_reference[ref_str] if sid != segment_id252                ]253            return True254        return False255256    def clear(self) -> None:257        """Clear all segments from memory (not from storage)."""258        self._segments.clear()259        self._by_reference.clear()260261    def save(self) -> None:262        """Merge in-memory segments into the stored file, atomically.263264        Parallel book workers each hold their own manager over the same file,265        each loaded when its book started and each saving hours later. A plain266        overwrite therefore discards every other book's segments — the last267        writer wins and days of paid translation vanish silently.268269        So the write is both:270271        - **atomic** (temp file then rename), so a crash cannot leave a torn272          file that poisons every later run; and273        - **merging** (re-read under an exclusive lock, combine, then write),274          so a slow worker cannot roll back a fast one's work.275276        On conflict the in-memory segment wins, being the more recent277        translation of that identical source text.278        """279        self._ensure_loaded()280        self._storage_path.parent.mkdir(parents=True, exist_ok=True)281282        with self._exclusive_lock():283            merged = self._read_disk_segments()284            merged.update(self._segments)285286            data = {287                "version": "1.0",288                "updated_at": datetime.now(timezone.utc).isoformat(),289                "segments": [290                    {291                        "id": s.id,292                        "source_text": s.source_text,293                        "source_normalized": s.source_normalized,294                        "target_text": s.target_text,295                        "reference": str(s.reference),296                        "level": s.level.value,297                        "created_at": s.created_at.isoformat(),298                        "reuse_count": s.reuse_count,299                        "metadata": s.metadata,300                    }301                    for s in merged.values()302                ],303            }304305            import tempfile306307            tmp_fd, tmp_path = tempfile.mkstemp(308                dir=self._storage_path.parent,309                suffix=".tmp",310            )311            try:312                with open(tmp_fd, "w", encoding="utf-8") as f:313                    json.dump(data, f, ensure_ascii=False, indent=2)314                Path(tmp_path).replace(self._storage_path)315            except BaseException:316                Path(tmp_path).unlink(missing_ok=True)317                raise318319            self._segments = merged320            self._rebuild_reference_index()321322    @contextmanager323    def _exclusive_lock(self) -> Iterator[None]:324        """Hold an exclusive lock over the store for the duration of the block.325326        Uses a sidecar lock file rather than the store itself, so the atomic327        replace cannot swap the locked file out from under a held lock.328329        Yields:330            None, once the lock is held.331        """332        lock_path = self._storage_path.with_suffix(333            self._storage_path.suffix + ".lock"334        )335        lock_path.parent.mkdir(parents=True, exist_ok=True)336        handle = lock_path.open("w")337        try:338            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)339            yield340        finally:341            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)342            handle.close()343344    def _rebuild_reference_index(self) -> None:345        """Rebuild the reference index from the current segments."""346        self._by_reference.clear()347        for segment in self._segments.values():348            ref_str = str(segment.reference)349            self._by_reference.setdefault(ref_str, []).append(segment.id)350351    def _quarantine_corrupt_file(self) -> Path:352        """Move an unreadable store aside so a later save cannot overwrite it.353354        Returns:355            The path the corrupt file was moved to.356        """357        stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")358        target = self._storage_path.with_name(359            f"{self._storage_path.stem}.corrupt.{stamp}.json"360        )361        self._storage_path.replace(target)362        return target363364    def _read_disk_segments(self) -> dict[str, TMSegment]:365        """Read stored segments without disturbing in-memory state.366367        A corrupt file is quarantined rather than ignored. Ignoring it meant368        the next save silently replaced days of accumulated translation with369        whatever the current process happened to hold.370371        Returns:372            The segments currently on disk, keyed by id.373        """374        if not self._storage_path.exists():375            return {}376377        try:378            with self._storage_path.open("r", encoding="utf-8") as f:379                data = json.load(f)380            if not isinstance(data, dict):381                # Valid JSON of the wrong shape is just as unreadable; without382                # this it crashed load/save with an AttributeError instead.383                raise json.JSONDecodeError("not a JSON object", "", 0)384        except json.JSONDecodeError:385            quarantined = self._quarantine_corrupt_file()386            logger.error(387                "Translation memory file was unreadable; preserved at %s. "388                "Starting from empty rather than overwriting it.",389                quarantined,390            )391            return {}392393        segments: dict[str, TMSegment] = {}394        for seg_data in data.get("segments", []):395            reference = VerseReference.from_string(seg_data["reference"])396            segment = TMSegment(397                id=seg_data["id"],398                source_text=seg_data["source_text"],399                source_normalized=seg_data["source_normalized"],400                target_text=seg_data["target_text"],401                reference=reference,402                level=SegmentLevel(seg_data["level"]),403                created_at=datetime.fromisoformat(seg_data["created_at"]),404                reuse_count=seg_data.get("reuse_count", 0),405                metadata=seg_data.get("metadata", {}),406            )407            segments[segment.id] = segment408        return segments409410    def load(self) -> None:411        """Load segments from the storage file."""412        self._loaded = True413        self._segments = self._read_disk_segments()414        self._rebuild_reference_index()415416    def reload(self) -> None:417        """Force reload from storage, discarding in-memory changes."""418        self._loaded = False419        self.load()420421    def __len__(self) -> int:422        """Get the number of segments in memory."""423        self._ensure_loaded()424        return len(self._segments)425426    def format_matches_for_prompt(427        self, matches: list[TMMatch], max_display: int = 3428    ) -> str:429        """Format TM matches for inclusion in a translation prompt.430431        Args:432            matches: List of TM matches.433            max_display: Maximum number of matches to display.434435        Returns:436            Formatted string for prompt inclusion.437        """438        if not matches:439            return "No similar translations found in memory."440441        lines = ["Similar translations from memory:"]442        for i, match in enumerate(matches[:max_display]):443            similarity_pct = int(match.similarity * 100)444            seg = match.segment445            lines.append(446                f"  [{similarity_pct}% match] {seg.reference}: {seg.target_text}"447            )448449        if len(matches) > max_display:450            lines.append(f"  ... and {len(matches) - max_display} more")451452        return "\n".join(lines)453