tts-voice-creator-clone-and-design-2/core/vision_cache.py

"""Record/replay cache for vision-model (OCR) responses, so OCR/layout A/B runs are comparable.

A live vision model is not repeatable (model swaps, server-side sampling, load-dependent output), so two
imports of one PDF differ even when only the post-processing changed. This module freezes the raw
completion text per request:

* ``live``   (default) -- no cache; behaviour is exactly as before.
* ``record`` -- call the model, store the response, return it.
* ``replay`` -- never call the model; a request with no stored response raises ``VisionCacheMiss``
                (fail loudly rather than silently mixing frozen and fresh output).

Select with ``MAGAZINE_VISION_CACHE=live|record|replay``; the directory is
``MAGAZINE_VISION_CACHE_DIR`` (default ``<CONFIG_DIR or ./config>/vision_cache``; ``config/`` is git-ignored).

The key is sha256 over (model, prompt text, sha256 of the image bytes, all other request parameters such as
temperature, max_tokens, seed and vendor fields). Endpoint URL and timeouts are not part of it (same model
behind another address is the same model). Each entry stores the key inputs, the response and a timestamp.
Every non-live run keeps ``last_run.json`` (mode, models, hits/recorded/misses) in the cache directory for
comparison reports.
"""
from __future__ import annotations

import hashlib
import json
import os
import threading
import time
from pathlib import Path
from typing import Any, Callable

MODES = ("live", "record", "replay")
ENV_MODE = "MAGAZINE_VISION_CACHE"
ENV_DIR = "MAGAZINE_VISION_CACHE_DIR"


class VisionCacheMiss(RuntimeError):
    """Replay mode asked for a response that was never recorded."""


def get_mode() -> str:
    mode = (os.environ.get(ENV_MODE) or "live").strip().lower()
    if mode not in MODES:
        raise ValueError(f"{ENV_MODE}={mode!r} is invalid; use one of {', '.join(MODES)}")
    return mode


def cache_dir() -> Path:
    env = os.environ.get(ENV_DIR)
    if env:
        return Path(env)
    base = os.environ.get("CONFIG_DIR") or str(Path(__file__).resolve().parents[1] / "config")
    return Path(base) / "vision_cache"


def make_key(model: str, prompt: str, image_bytes: bytes, params: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    inputs = {
        "model": model,
        "prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(),
        "image_sha256": hashlib.sha256(image_bytes).hexdigest(),
        "params": params,
    }
    blob = json.dumps(inputs, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
    return hashlib.sha256(blob.encode("utf-8")).hexdigest(), inputs


class _Stats:
    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.models: set[str] = set()
        self.hits = self.recorded = self.misses = 0
        self.started = time.strftime("%Y-%m-%dT%H:%M:%S%z")

    def snapshot(self, mode: str) -> dict[str, Any]:
        return {"mode": mode, "models": sorted(self.models), "hits": self.hits, "recorded": self.recorded,
                "misses": self.misses, "started": self.started}


_stats = _Stats()


def run_summary() -> dict[str, Any]:
    """What this process did with the cache so far (for reports)."""
    with _stats.lock:
        return _stats.snapshot(get_mode())


def _write_manifest(directory: Path, mode: str) -> None:
    try:
        directory.mkdir(parents=True, exist_ok=True)
        tmp = directory / "last_run.json.tmp"
        tmp.write_text(json.dumps(_stats.snapshot(mode), indent=1), "utf-8")
        tmp.replace(directory / "last_run.json")
    except OSError:
        pass  # the manifest is a convenience; never fail an OCR call over it


def cached_completion(model: str, prompt: str, image_bytes: bytes, params: dict[str, Any],
                      fetch: Callable[[], str | None]) -> str | None:
    """Return the completion text for this request according to the cache mode.

    ``fetch`` performs the real call (live/record only) and returns the text, or None on failure;
    failures are never stored.
    """
    mode = get_mode()
    if mode == "live":
        return fetch()
    directory = cache_dir()
    key, inputs = make_key(model, prompt, image_bytes, params)
    path = directory / f"{key}.json"
    with _stats.lock:
        _stats.models.add(model)
    if mode == "replay":
        try:
            entry = json.loads(path.read_text("utf-8"))
            content = entry["response"]
        except (OSError, ValueError, KeyError):
            with _stats.lock:
                _stats.misses += 1
                _write_manifest(directory, mode)
            raise VisionCacheMiss(
                f"no recorded vision response for model {model!r} (key {key[:12]}, image "
                f"{inputs['image_sha256'][:12]}) in {directory}; record it first with {ENV_MODE}=record")
        with _stats.lock:
            _stats.hits += 1
            _write_manifest(directory, mode)
        return content
    content = fetch()
    if content is not None:
        directory.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps({"key_inputs": inputs, "response": content,
                                   "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z")},
                                  ensure_ascii=False, indent=1), "utf-8")
        tmp.replace(path)
        with _stats.lock:
            _stats.recorded += 1
    with _stats.lock:
        _write_manifest(directory, mode)
    return content