54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
class GenerationCache:
|
|
def __init__(self, cache_dir: Path) -> None:
|
|
self.cache_dir = cache_dir
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
self.index_path = cache_dir / "cache_index.json"
|
|
self._index: dict[str, str] = self._load_index()
|
|
|
|
def _load_index(self) -> dict[str, str]:
|
|
if self.index_path.exists():
|
|
with open(self.index_path) as f:
|
|
return json.load(f)
|
|
return {}
|
|
|
|
def _save_index(self) -> None:
|
|
with open(self.index_path, "w") as f:
|
|
json.dump(self._index, f, indent=2)
|
|
|
|
def _hash_key(self, text: str, voice_ref: str | None, instruct: str | None) -> str:
|
|
key_data = f"{text}|{voice_ref}|{instruct}"
|
|
return hashlib.sha256(key_data.encode()).hexdigest()[:16]
|
|
|
|
def get(self, text: str, voice_ref: str | None, instruct: str | None) -> Path | None:
|
|
key = self._hash_key(text, voice_ref, instruct)
|
|
if key in self._index:
|
|
cached_path = Path(self._index[key])
|
|
if cached_path.exists():
|
|
return cached_path
|
|
return None
|
|
|
|
def put(self, text: str, voice_ref: str | None, instruct: str | None, audio_path: Path) -> None:
|
|
key = self._hash_key(text, voice_ref, instruct)
|
|
self._index[key] = str(audio_path)
|
|
self._save_index()
|
|
|
|
def clear(self) -> None:
|
|
self._index = {}
|
|
self._save_index()
|
|
|
|
def stats(self) -> dict[str, int]:
|
|
return {
|
|
"cached_entries": len(self._index),
|
|
"cache_dir": str(self.cache_dir),
|
|
}
|