engines/cache.py

from __future__ import annotations

import hashlib
import json
import re
from pathlib import Path
from typing import Any, Mapping


def _slugify(value: str) -> str:
    slug = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")
    return slug or "value"


def build_cache_key(request: Mapping[str, Any]) -> str:
    custom_prompt = str(request.get("custom_prompt") or "").strip()
    prompt_marker = request.get("preset")
    if custom_prompt:
        prompt_hash = hashlib.sha1(custom_prompt.encode("utf-8")).hexdigest()[:10]
        prompt_marker = f"custom-{prompt_hash}"
    elif not prompt_marker:
        prompt = str(request.get("custom_prompt") or "")
        prompt_hash = hashlib.sha1(prompt.encode("utf-8")).hexdigest()[:10]
        prompt_marker = f"custom-{prompt_hash}"
    parts = (
        str(request.get("model") or ""),
        str(prompt_marker),
        str(request.get("context_length") or ""),
        str(request.get("mode") or ""),
        f"p{request.get('page_size')}",
        f"bk{request.get('bits_k')}",
        f"bv{request.get('bits_v')}",
        f"rw{request.get('recent_window')}",
        f"sw{request.get('sink_window')}",
        str(request.get("shortlist_policy") or ""),
    )
    return "__".join(_slugify(part) for part in parts) + ".json"


class ResultCache:
    def __init__(self, root: Path) -> None:
        self.root = root
        self.precomputed_dir = root / "data" / "preset_runs"
        self.runtime_cache_dir = root / "data" / "runtime_cache"
        self.runtime_cache_dir.mkdir(parents=True, exist_ok=True)

    def load_precomputed(self, request: Mapping[str, Any]) -> dict[str, Any] | None:
        path = self.precomputed_dir / build_cache_key(request)
        if not path.exists():
            return None
        return json.loads(path.read_text(encoding="utf-8"))

    def load_runtime_cache(self, request: Mapping[str, Any]) -> dict[str, Any] | None:
        path = self.runtime_cache_dir / build_cache_key(request)
        if not path.exists():
            return None
        return json.loads(path.read_text(encoding="utf-8"))

    def write_runtime_cache(self, request: Mapping[str, Any], payload: Mapping[str, Any]) -> Path:
        path = self.runtime_cache_dir / build_cache_key(request)
        path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
        return path