utils/cache_utils.py

"""Utilities for caching async pipeline calls."""

from __future__ import annotations

import functools
import hashlib
import json
from collections.abc import Callable, Coroutine
from typing import Protocol, TypeVar, runtime_checkable

from pydantic import BaseModel, JsonValue

from parsing_pipeline.db.models import LLMCacheEntry
from utils.constants import FormatStrings


class SupportsStr(Protocol):
    """Protocol for values that provide a string representation."""

    def __str__(self) -> str:
        """Return a stable string representation."""
        ...


T = TypeVar("T")


@runtime_checkable
class LLMCacheManager(Protocol):
    """Protocol for cache-aware DatabaseManager APIs."""

    async def get_llm_cache_entry(self, cache_key_hash: str) -> LLMCacheEntry | None:
        """Return cached LLM entry for the given key, if any."""
        ...

    async def upsert_llm_cache_entry(self, entry: LLMCacheEntry) -> None:
        """Insert or update a cached LLM entry."""
        ...


def serialize_for_cache(obj: SupportsStr) -> JsonValue | SupportsStr:
    """Deterministically serialise an object for use in cache keys.

    If *obj* is a Pydantic **model class** (e.g. ChunkAnalysis), return
    "ClassName@<12-hex-digest>" where the digest is SHA-256 of its
    `model_json_schema()`. Any structural change to the model therefore
    invalidates previous cache entries.

    Otherwise return the object unchanged.
    """
    if isinstance(obj, type) and issubclass(obj, BaseModel):
        schema_json = json.dumps(obj.model_json_schema(), sort_keys=True)
        digest = hashlib.sha256(schema_json.encode()).hexdigest()[:12]
        return f"{obj.__qualname__}@{digest}"
    return obj


def make_cache_key(
    func: Callable[..., SupportsStr],
    args: tuple[SupportsStr, ...],
    kwargs: dict[str, SupportsStr],
) -> str:
    """Create a deterministic cache key hash from function name and arguments."""
    key_data = {
        "func": func.__qualname__,
        "args": [serialize_for_cache(a) for a in args],
        "kwargs": {k: serialize_for_cache(v) for k, v in kwargs.items()},
    }
    key_str = json.dumps(key_data, sort_keys=True, default=str)
    return hashlib.sha256(key_str.encode(FormatStrings.ENCODING_UTF8)).hexdigest()


async def cache_wrapper_impl(
    func: Callable[..., Coroutine[None, None, T]],
    db_manager_arg_index: int,
    serialize_cached: Callable[[T], str],
    deserialize_cached: Callable[[str], T],
    args: tuple[SupportsStr, ...],
    kwargs: dict[str, SupportsStr],
) -> T:
    """Implementation of cache wrapper logic - extracted to avoid nested function."""
    db_manager = args[db_manager_arg_index] if db_manager_arg_index < len(args) else None
    cache_key = make_cache_key(func, args, kwargs)

    if isinstance(db_manager, LLMCacheManager):
        cached_entry = await db_manager.get_llm_cache_entry(cache_key)
        if cached_entry is not None:
            return deserialize_cached(cached_entry.response_json)

    result = await func(*args, **kwargs)

    if isinstance(db_manager, LLMCacheManager):
        entry = LLMCacheEntry(
            cache_key_hash=cache_key,
            function_name=func.__qualname__,
            response_json=serialize_cached(result),
        )
        await db_manager.upsert_llm_cache_entry(entry)

    return result


def cache_decorator_with_db(  # nosemgrep: python-no-nested-functions -- Decorator pattern requires closure to wrap the decorated function
    db_manager_arg_index: int,
    serialize_cached: Callable[[T], str],
    deserialize_cached: Callable[[str], T],
) -> Callable[
    [Callable[..., Coroutine[None, None, T]]],
    Callable[..., Coroutine[None, None, T]],
]:
    """Decorator factory for async function caching with DatabaseManager.

    Creates a decorator that caches async function results using the DatabaseManager
    passed as a positional argument at the specified index.

    Args:
        db_manager_arg_index: Index of the DatabaseManager in the function's
            positional args.
        serialize_cached: Serialize a cached response to JSON text.
        deserialize_cached: Restore a cached response from JSON text.

    Example:
        def serialize_result(result: dict[str, JsonValue]) -> str:
            return json.dumps(result, default=str, ensure_ascii=False)

        def deserialize_result(payload: str) -> dict[str, JsonValue]:
            return json.loads(payload)

        @cache_decorator_with_db(2, serialize_result, deserialize_result)
        async def fetch_data(
            query: str, pref: str, db_manager: DatabaseManager
        ) -> dict:
            ...

        result = await fetch_data("query", "pref", db_manager)
    """

    # nosemgrep: python-no-nested-functions -- Decorator captures db_manager_arg_index.
    def decorator(
        func: Callable[..., Coroutine[None, None, T]],
    ) -> Callable[..., Coroutine[None, None, T]]:
        @functools.wraps(func)
        async def wrapper(
            *args: SupportsStr,
            **kwargs: SupportsStr,
        ) -> T:  # nosemgrep: python-no-nested-functions -- Decorator pattern requires closure to wrap the decorated function
            return await cache_wrapper_impl(
                func,
                db_manager_arg_index,
                serialize_cached,
                deserialize_cached,
                args,
                kwargs,
            )

        return wrapper

    return decorator