chore: import upstream snapshot with attribution

packages/graphrag-cache/graphrag_cache/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""The GraphRAG Cache package."""
+
+from graphrag_cache.cache import Cache
+from graphrag_cache.cache_config import CacheConfig
+from graphrag_cache.cache_factory import create_cache, register_cache
+from graphrag_cache.cache_key import CacheKeyCreator, create_cache_key
+from graphrag_cache.cache_type import CacheType
+
+__all__ = [
+    "Cache",
+    "CacheConfig",
+    "CacheKeyCreator",
+    "CacheType",
+    "create_cache",
+    "create_cache_key",
+    "register_cache",
+]

packages/graphrag-cache/graphrag_cache/cache.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Abstract base class for cache."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from graphrag_storage import Storage
+
+
+class Cache(ABC):
+    """Provide a cache interface for the pipeline."""
+
+    @abstractmethod
+    def __init__(self, *, storage: Storage | None, **kwargs: Any) -> None:
+        """Create a cache instance."""
+
+    @abstractmethod
+    async def get(self, key: str) -> Any:
+        """Get the value for the given key.
+
+        Args:
+            - key - The key to get the value for.
+            - as_bytes - Whether or not to return the value as bytes.
+
+        Returns
+        -------
+            - output - The value for the given key.
+        """
+
+    @abstractmethod
+    async def set(self, key: str, value: Any, debug_data: dict | None = None) -> None:
+        """Set the value for the given key.
+
+        Args:
+            - key - The key to set the value for.
+            - value - The value to set.
+        """
+
+    @abstractmethod
+    async def has(self, key: str) -> bool:
+        """Return True if the given key exists in the cache.
+
+        Args:
+            - key - The key to check for.
+
+        Returns
+        -------
+            - output - True if the key exists in the cache, False otherwise.
+        """
+
+    @abstractmethod
+    async def delete(self, key: str) -> None:
+        """Delete the given key from the cache.
+
+        Args:
+            - key - The key to delete.
+        """
+
+    @abstractmethod
+    async def clear(self) -> None:
+        """Clear the cache."""
+
+    @abstractmethod
+    def child(self, name: str) -> Cache:
+        """Create a child cache with the given name.
+
+        Args:
+            - name - The name to create the sub cache with.
+        """

packages/graphrag-cache/graphrag_cache/cache_config.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Cache configuration model."""
+
+from graphrag_storage import StorageConfig, StorageType
+from pydantic import BaseModel, ConfigDict, Field
+
+from graphrag_cache.cache_type import CacheType
+
+
+class CacheConfig(BaseModel):
+    """The configuration section for cache."""
+
+    model_config = ConfigDict(extra="allow")
+    """Allow extra fields to support custom cache implementations."""
+
+    type: str = Field(
+        description="The cache type to use. Builtin types include 'Json', 'Memory', and 'Noop'.",
+        default=CacheType.Json,
+    )
+
+    storage: StorageConfig | None = Field(
+        description="The storage configuration to use for file-based caches such as 'Json'.",
+        default_factory=lambda: StorageConfig(type=StorageType.File, base_dir="cache"),
+    )

packages/graphrag-cache/graphrag_cache/cache_factory.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+
+"""Cache factory implementation."""
+
+from collections.abc import Callable
+
+from graphrag_common.factory import Factory, ServiceScope
+from graphrag_storage import Storage, create_storage
+
+from graphrag_cache.cache import Cache
+from graphrag_cache.cache_config import CacheConfig
+from graphrag_cache.cache_type import CacheType
+
+
+class CacheFactory(Factory["Cache"]):
+    """A factory class for cache implementations."""
+
+
+cache_factory = CacheFactory()
+
+
+def register_cache(
+    cache_type: str,
+    cache_initializer: Callable[..., Cache],
+    scope: ServiceScope = "transient",
+) -> None:
+    """Register a custom cache implementation.
+
+    Args
+    ----
+        - cache_type: str
+            The cache id to register.
+        - cache_initializer: Callable[..., Cache]
+            The cache initializer to register.
+    """
+    cache_factory.register(cache_type, cache_initializer, scope)
+
+
+def create_cache(
+    config: CacheConfig | None = None, storage: Storage | None = None
+) -> "Cache":
+    """Create a cache implementation based on the given configuration.
+
+    Args
+    ----
+        - config: CacheConfig
+            The cache configuration to use.
+        - storage: Storage | None
+            The storage implementation to use for file-based caches such as 'Json'.
+
+    Returns
+    -------
+        Cache
+            The created cache implementation.
+    """
+    config = config or CacheConfig()
+    config_model = config.model_dump()
+    cache_strategy = config.type
+
+    if not storage and config.storage:
+        storage = create_storage(config.storage)
+
+    if cache_strategy not in cache_factory:
+        match cache_strategy:
+            case CacheType.Json:
+                from graphrag_cache.json_cache import JsonCache
+
+                register_cache(CacheType.Json, JsonCache)
+
+            case CacheType.Memory:
+                from graphrag_cache.memory_cache import MemoryCache
+
+                register_cache(CacheType.Memory, MemoryCache)
+
+            case CacheType.Noop:
+                from graphrag_cache.noop_cache import NoopCache
+
+                register_cache(CacheType.Noop, NoopCache)
+
+            case _:
+                msg = f"CacheConfig.type '{cache_strategy}' is not registered in the CacheFactory. Registered types: {', '.join(cache_factory.keys())}."
+                raise ValueError(msg)
+
+    if storage:
+        config_model["storage"] = storage
+
+    return cache_factory.create(strategy=cache_strategy, init_args=config_model)

packages/graphrag-cache/graphrag_cache/cache_key.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Create cache key."""
+
+from typing import Any, Protocol, runtime_checkable
+
+from graphrag_common.hasher import hash_data
+
+
+@runtime_checkable
+class CacheKeyCreator(Protocol):
+    """Create cache key function protocol.
+
+    Args
+    ----
+        input_args: dict[str, Any]
+            The input arguments for creating the cache key.
+
+    Returns
+    -------
+        str
+            The generated cache key.
+    """
+
+    def __call__(
+        self,
+        input_args: dict[str, Any],
+    ) -> str:
+        """Create cache key."""
+        ...
+
+
+def create_cache_key(input_args: dict[str, Any]) -> str:
+    """Create a cache key based on the input arguments."""
+    return hash_data(input_args)

packages/graphrag-cache/graphrag_cache/cache_type.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+
+"""Builtin cache implementation types."""
+
+from enum import StrEnum
+
+
+class CacheType(StrEnum):
+    """Enum for cache types."""
+
+    Json = "json"
+    Memory = "memory"
+    Noop = "none"

packages/graphrag-cache/graphrag_cache/json_cache.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'JsonCache' model."""
+
+import json
+from typing import Any
+
+from graphrag_storage import Storage, StorageConfig, create_storage
+
+from graphrag_cache.cache import Cache
+
+
+class JsonCache(Cache):
+    """File pipeline cache class definition."""
+
+    _storage: Storage
+
+    def __init__(
+        self,
+        storage: Storage | dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Init method definition."""
+        if storage is None:
+            msg = "JsonCache requires either a Storage instance to be provided or a StorageConfig to create one."
+            raise ValueError(msg)
+        if isinstance(storage, Storage):
+            self._storage = storage
+        else:
+            self._storage = create_storage(StorageConfig(**storage))
+
+    async def get(self, key: str) -> Any | None:
+        """Get method definition."""
+        if await self.has(key):
+            try:
+                data = await self._storage.get(key)
+                data = json.loads(data)
+            except UnicodeDecodeError:
+                await self._storage.delete(key)
+                return None
+            except json.decoder.JSONDecodeError:
+                await self._storage.delete(key)
+                return None
+            else:
+                return data.get("result")
+
+        return None
+
+    async def set(self, key: str, value: Any, debug_data: dict | None = None) -> None:
+        """Set method definition."""
+        if value is None:
+            return
+        data = {"result": value, **(debug_data or {})}
+        await self._storage.set(key, json.dumps(data, ensure_ascii=False))
+
+    async def has(self, key: str) -> bool:
+        """Has method definition."""
+        return await self._storage.has(key)
+
+    async def delete(self, key: str) -> None:
+        """Delete method definition."""
+        if await self.has(key):
+            await self._storage.delete(key)
+
+    async def clear(self) -> None:
+        """Clear method definition."""
+        await self._storage.clear()
+
+    def child(self, name: str) -> "Cache":
+        """Child method definition."""
+        return JsonCache(storage=self._storage.child(name))

packages/graphrag-cache/graphrag_cache/memory_cache.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""MemoryCache implementation."""
+
+from typing import Any
+
+from graphrag_cache.cache import Cache
+
+
+class MemoryCache(Cache):
+    """In memory cache class definition."""
+
+    _cache: dict[str, Any]
+    _name: str
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Init method definition."""
+        self._cache = {}
+
+    async def get(self, key: str) -> Any:
+        """Get the value for the given key.
+
+        Args:
+            - key - The key to get the value for.
+            - as_bytes - Whether or not to return the value as bytes.
+
+        Returns
+        -------
+            - output - The value for the given key.
+        """
+        return self._cache.get(key)
+
+    async def set(self, key: str, value: Any, debug_data: dict | None = None) -> None:
+        """Set the value for the given key.
+
+        Args:
+            - key - The key to set the value for.
+            - value - The value to set.
+        """
+        self._cache[key] = value
+
+    async def has(self, key: str) -> bool:
+        """Return True if the given key exists in the storage.
+
+        Args:
+            - key - The key to check for.
+
+        Returns
+        -------
+            - output - True if the key exists in the storage, False otherwise.
+        """
+        return key in self._cache
+
+    async def delete(self, key: str) -> None:
+        """Delete the given key from the storage.
+
+        Args:
+            - key - The key to delete.
+        """
+        del self._cache[key]
+
+    async def clear(self) -> None:
+        """Clear the storage."""
+        self._cache.clear()
+
+    def child(self, name: str) -> "Cache":
+        """Create a sub cache with the given name."""
+        return MemoryCache()

packages/graphrag-cache/graphrag_cache/noop_cache.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""NoopCache implementation."""
+
+from typing import Any
+
+from graphrag_cache.cache import Cache
+
+
+class NoopCache(Cache):
+    """A no-op implementation of Cache, usually useful for testing."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Init method definition."""
+
+    async def get(self, key: str) -> Any:
+        """Get the value for the given key.
+
+        Args:
+            - key - The key to get the value for.
+            - as_bytes - Whether or not to return the value as bytes.
+
+        Returns
+        -------
+            - output - The value for the given key.
+        """
+        return None
+
+    async def set(
+        self, key: str, value: str | bytes | None, debug_data: dict | None = None
+    ) -> None:
+        """Set the value for the given key.
+
+        Args:
+            - key - The key to set the value for.
+            - value - The value to set.
+        """
+
+    async def has(self, key: str) -> bool:
+        """Return True if the given key exists in the cache.
+
+        Args:
+            - key - The key to check for.
+
+        Returns
+        -------
+            - output - True if the key exists in the cache, False otherwise.
+        """
+        return False
+
+    async def delete(self, key: str) -> None:
+        """Delete the given key from the cache.
+
+        Args:
+            - key - The key to delete.
+        """
+
+    async def clear(self) -> None:
+        """Clear the cache."""
+
+    def child(self, name: str) -> "Cache":
+        """Create a child cache with the given name.
+
+        Args:
+            - name - The name to create the sub cache with.
+        """
+        return self