Back to Graphrag

Copyright (c) 2026 Microsoft Corporation.

packages/graphrag-cache/example_notebooks/custom_cache_example.ipynb

3.0.92.4 KB
Original Source
python
# Copyright (c) 2026 Microsoft Corporation.
# Licensed under the MIT License.

Custom cache example

This example demonstrates how to create a custom cache implementation by extending the base Cache class and registering it with the GraphRAG cache system. Once registered, the custom cache can be instantiated through the factory pattern using either CacheConfig or directly via cache_factory, allowing for extensible caching solutions tailored to specific needs.

python
from typing import Any

from graphrag_cache import Cache, CacheConfig, create_cache, register_cache


class MyCache(Cache):
    """Custom cache implementation for storing and retrieving cached data."""

    def __init__(
        self,
        some_setting: str,
        optional_setting: str = "default setting",
        **kwargs: Any,
    ):
        # Validate settings and initialize
        # View the JsonCache implementation to see how to create a cache that relies on a Storage provider.
        self.some_setting = some_setting
        self.optional_setting = optional_setting
        self._cache: dict[str, Any] = {}
        self._child: Cache = self

    async def get(self, key: str) -> Any:
        """Retrieve a value from the cache by key."""
        return self._cache.get(key)

    async def set(self, key: str, value: Any, debug_data: Any = None) -> None:
        """Store a value in the cache with the specified key."""
        self._cache[key] = value

    async def has(self, key: str) -> bool:
        """Check if a key exists in the cache."""
        return key in self._cache

    async def delete(self, key: str) -> None:
        """Remove a key and its value from the cache."""
        self._cache.pop(key, None)

    async def clear(self) -> None:
        """Clear all items from the cache."""
        self._cache.clear()

    def child(self, name: str) -> Cache:
        """Create or access a child cache (not implemented in this example)."""
        return self._child


register_cache("MyCache", MyCache)


async def run():
    """Demonstrate usage of the custom cache implementation."""
    cache = create_cache(
        CacheConfig(
            type="MyCache",
            some_setting="important setting value",  # type: ignore
        )
    )

    await cache.set("my_key", {"k1": "object to cache"})
    print("Value stored in cache for 'my_key':")
    print(await cache.get("my_key"))


if __name__ == "__main__":
    await run()