microsoft--graphrag/packages/graphrag-cache

GraphRAG Cache

This package contains a collection of utilities to handle GraphRAG caching implementation.

Basic

This example shows how to create a JSON cache with file storage using the GraphRAG cache package's configuration system.

import asyncio
from graphrag_storage import StorageConfig, create_storage, StorageType
from graphrag_cache import CacheConfig, create_cache, CacheType, create_cache_key

async def run():
    cache = create_cache()

    # The above is equivalent to the following:
    cache = create_cache(
        CacheConfig(
            type=CacheType.Json,
            storage=StorageConfig(
                type=StorageType.File,
                base_dir="cache"
            )
        ),
    )

    await cache.set("my_key", {"some": "object to cache"})
    print(await cache.get("my_key"))

    # create cache key from data dict.
    cache_key = create_cache_key({
        "some_arg": "some_value",
        "something_else": 5
    })
    await cache.set(cache_key, {"some": "object to cache"})
    print(await cache.get(cache_key))

if __name__ == "__main__":
    asyncio.run(run())

Custom Cache

This 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.

import asyncio
from typing import Any
from graphrag_storage import Storage
from graphrag_cache import Cache, CacheConfig, create_cache, register_cache

class MyCache(Cache):
    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.
        ...

    #Implement rest of interface
    ...

register_cache("MyCache", MyCache)

async def run():
    cache = create_cache(
        CacheConfig(
            type="MyCache",
            some_setting="My Setting"
        )
    )

    # Or use the factory directly to instantiate with a dict instead of using
    # CacheConfig + create_factory
    # from graphrag_cache.cache_factory import cache_factory
    # cache = cache_factory.create(strategy="MyCache", init_args={"some_setting": "My Setting"})

    await cache.set("my_key", {"some": "object to cache"})
    print(await cache.get("my_key"))

if __name__ == "__main__":
    asyncio.run(run())

Details

By default, the create_cache comes with the following cache providers registered that correspond to the entries in the CacheType enum.

• JsonCache
• MemoryCache
• NoopCache

The preregistration happens dynamically, e.g., JsonCache is only imported and registered if you request a JsonCache with create_cache(CacheType.Json, ...). There is no need to manually import and register builtin cache providers when using create_cache.

If you want a clean factory with no preregistered cache providers then directly import cache_factory and bypass using create_cache. The downside is that cache_factory.create uses a dict for init args instead of the strongly typed CacheConfig used with create_cache.

from graphrag_cache.cache_factory import cache_factory
from graphrag_cache.json_cache import JsonCache

# cache_factory has no preregistered providers so you must register any
# providers you plan on using.
# May also register a custom implementation, see above for example.
cache_factory.register("my_cache_impl", JsonCache)

cache = cache_factory.create(strategy="my_cache_impl", init_args={"some_setting": "..."})

...