microsoft--graphrag/packages/graphrag-llm/graphrag_llm/cache/create_cache_key.py

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License

"""Create cache key."""

from typing import Any

from graphrag_cache import create_cache_key as default_create_cache_key

_CACHE_VERSION = 4
"""
If there's a breaking change in what we cache, we should increment this version number to invalidate existing caches.

fnllm was on cache version 2 and though we generate
similar cache keys, the objects stored in cache by fnllm and litellm are different.
Using litellm model providers will not be able to reuse caches generated by fnllm
thus we start with version 3 for litellm.

graphrag-llm package is now on version 4.
This is to account for changes to the ModelConfig that affect the cache key and
occurred when pulling this package out of graphrag.
graphrag-llm, now that is supports metrics, also caches metrics which were not cached before.
"""


def create_cache_key(
    input_args: dict[str, Any],
) -> str:
    """Generate a cache key based on the model configuration and input arguments.

    Args
    ____
        input_args: dict[str, Any]
            The input arguments for the model call.

    Returns
    -------
        str
            The generated cache key in the format
            `{prefix}_{data_hash}_v{version}` if prefix is provided.
    """
    cache_key_parameters = _get_parameters(
        input_args=input_args,
    )
    return default_create_cache_key(cache_key_parameters)


def _get_parameters(
    # model_config: "ModelConfig",
    input_args: dict[str, Any],
) -> dict[str, Any]:
    """Pluck out the parameters that define a cache key."""
    excluded_keys = [
        "metrics",
        "stream",
        "stream_options",
        "mock_response",
        "timeout",
        "base_url",
        "api_base",
        "api_version",
        "api_key",
        "azure_ad_token_provider",
        "drop_params",
    ]

    parameters: dict[str, Any] = {
        k: v for k, v in input_args.items() if k not in excluded_keys
    }

    return parameters