Update cache config to support storage.

Author: dworthen

packages/graphrag-cache/README.md

@@ -8,19 +8,14 @@ from graphrag_storage import StorageConfig, create_storage, StorageType
 from graphrag_cache import CacheConfig, create_cache, CacheType
 
 async def run():
-    # Json cache requires a storage implementation.
-    storage = create_storage(
-        StorageConfig(
-            type=StorageType.File
-            base_dir="output"
-        )
-    )
-
     cache = create_cache(
         CacheConfig(
             type=CacheType.Json
+            storage=StorageConfig(
+                type=StorageType.File
+                base_dir="cache"
+            )
         ),
-        storage=storage
     )
 
     await cache.set("my_key", {"some": "object to cache"})
@@ -39,8 +34,9 @@ from graphrag_storage import Storage
 from graphrag_cache import Cache, CacheConfig, create_cache, register_cache
 
 class MyCache(Cache):
-    def __init__(self, storage: Storage, some_setting: str, optional_setting: str = "default setting", **kwargs: Any):
+    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
@@ -54,13 +50,12 @@ async def run():
             type="MyCache"
             some_setting="My Setting"
         )
-        # if your cache relies on a storage implementation you can pass that here
-        # storage=some_storage
     )
+
     # 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={"storage": storage_implementation, "some_setting": "My Setting"})
+    # 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"))

packages/graphrag-cache/graphrag_cache/cache_config.py

@@ -3,6 +3,7 @@
 
 """Cache configuration model."""
 
+from graphrag_storage import StorageConfig
 from pydantic import BaseModel, ConfigDict, Field
 
 from graphrag_cache.cache_type import CacheType
@@ -19,12 +20,7 @@ class CacheConfig(BaseModel):
         default=CacheType.Json,
     )
 
-    encoding: str | None = Field(
-        description="The encoding to use for file-based caching.",
-        default=None,
-    )
-
-    name: str | None = Field(
-        description="The name to use for the cache instance.",
+    storage: StorageConfig | None = Field(
+        description="The storage configuration to use for file-based caches such as 'Json'.",
         default=None,
     )

packages/graphrag-cache/graphrag_cache/cache_factory.py

@@ -77,6 +77,7 @@ def create_cache(config: CacheConfig, storage: Storage | None = None) -> Cache:
                 msg = f"CacheConfig.type '{cache_strategy}' is not registered in the CacheFactory. Registered types: {', '.join(cache_factory.keys())}."
                 raise ValueError(msg)
 
-    return cache_factory.create(
-        strategy=cache_strategy, init_args={"storage": storage, **config_model}
-    )
+    if storage:
+        config_model["storage"] = storage
+
+    return cache_factory.create(strategy=cache_strategy, init_args=config_model)

packages/graphrag-cache/graphrag_cache/cache_type.py

@@ -12,4 +12,4 @@ class CacheType(StrEnum):
 
     Json = "json"
     Memory = "memory"
-    Noop = "noop"
+    Noop = "none"

packages/graphrag-cache/graphrag_cache/json_cache.py

@@ -6,7 +6,7 @@
 import json
 from typing import Any
 
-from graphrag_storage import Storage
+from graphrag_storage import Storage, StorageConfig, create_storage
 
 from graphrag_cache.cache import Cache
 
@@ -15,18 +15,26 @@ class JsonCache(Cache):
     """File pipeline cache class definition."""
 
     _storage: Storage
-    _encoding: str
 
-    def __init__(self, storage: Storage, encoding="utf-8", **kwargs: Any) -> None:
+    def __init__(
+        self,
+        storage: Storage | dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
         """Init method definition."""
-        self._storage = storage
-        self._encoding = encoding
+        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, encoding=self._encoding)
+                data = await self._storage.get(key)
                 data = json.loads(data)
             except UnicodeDecodeError:
                 await self._storage.delete(key)
@@ -44,9 +52,7 @@ async def set(self, key: str, value: Any, debug_data: dict | None = None) -> Non
         if value is None:
             return
         data = {"result": value, **(debug_data or {})}
-        await self._storage.set(
-            key, json.dumps(data, ensure_ascii=False), encoding=self._encoding
-        )
+        await self._storage.set(key, json.dumps(data, ensure_ascii=False))
 
     async def has(self, key: str) -> bool:
         """Has method definition."""
@@ -63,4 +69,4 @@ async def clear(self) -> None:
 
     def child(self, name: str) -> "Cache":
         """Child method definition."""
-        return JsonCache(self._storage.child(name), encoding=self._encoding)
+        return JsonCache(storage=self._storage.child(name))

packages/graphrag-cache/graphrag_cache/memory_cache.py

@@ -14,10 +14,9 @@ class MemoryCache(Cache):
     _cache: dict[str, Any]
     _name: str
 
-    def __init__(self, name: str | None = None, **kwargs: Any) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         """Init method definition."""
         self._cache = {}
-        self._name = name or ""
 
     async def get(self, key: str) -> Any:
         """Get the value for the given key.
@@ -30,7 +29,6 @@ async def get(self, key: str) -> Any:
         -------
             - output - The value for the given key.
         """
-        key = self._create_cache_key(key)
         return self._cache.get(key)
 
     async def set(self, key: str, value: Any, debug_data: dict | None = None) -> None:
@@ -40,7 +38,6 @@ async def set(self, key: str, value: Any, debug_data: dict | None = None) -> Non
             - key - The key to set the value for.
             - value - The value to set.
         """
-        key = self._create_cache_key(key)
         self._cache[key] = value
 
     async def has(self, key: str) -> bool:
@@ -53,7 +50,6 @@ async def has(self, key: str) -> bool:
         -------
             - output - True if the key exists in the storage, False otherwise.
         """
-        key = self._create_cache_key(key)
         return key in self._cache
 
     async def delete(self, key: str) -> None:
@@ -62,7 +58,6 @@ async def delete(self, key: str) -> None:
         Args:
             - key - The key to delete.
         """
-        key = self._create_cache_key(key)
         del self._cache[key]
 
     async def clear(self) -> None:
@@ -71,8 +66,4 @@ async def clear(self) -> None:
 
     def child(self, name: str) -> "Cache":
         """Create a sub cache with the given name."""
-        return MemoryCache(name)
-
-    def _create_cache_key(self, key: str) -> str:
-        """Create a cache key for the given key."""
-        return f"{self._name}{key}"
+        return MemoryCache()

packages/graphrag/graphrag/config/defaults.py

@@ -56,15 +56,6 @@ class BasicSearchDefaults:
     embedding_model_id: str = DEFAULT_EMBEDDING_MODEL_ID
 
 
-@dataclass
-class CacheDefaults:
-    """Default values for cache."""
-
-    type: CacheType = CacheType.Json
-    encoding: str | None = None
-    name: str | None = None
-
-
 @dataclass
 class ChunksDefaults:
     """Default values for chunks."""
@@ -238,13 +229,6 @@ class StorageDefaults:
     azure_cosmosdb_account_url: None = None
 
 
-@dataclass
-class CacheStorageDefaults(StorageDefaults):
-    """Default values for cache storage."""
-
-    base_dir: str | None = DEFAULT_CACHE_BASE_DIR
-
-
 @dataclass
 class InputStorageDefaults(StorageDefaults):
     """Default values for input storage."""
@@ -265,6 +249,21 @@ class InputDefaults:
     metadata: None = None
 
 
+@dataclass
+class CacheStorageDefaults(StorageDefaults):
+    """Default values for cache storage."""
+
+    base_dir: str | None = DEFAULT_CACHE_BASE_DIR
+
+
+@dataclass
+class CacheDefaults:
+    """Default values for cache."""
+
+    type: CacheType = CacheType.Json
+    storage: CacheStorageDefaults = field(default_factory=CacheStorageDefaults)
+
+
 @dataclass
 class LanguageModelDefaults:
     """Default values for language model."""
@@ -401,7 +400,6 @@ class GraphRagConfigDefaults:
         default_factory=UpdateIndexOutputDefaults
     )
     cache: CacheDefaults = field(default_factory=CacheDefaults)
-    cache_storage: CacheStorageDefaults = field(default_factory=CacheStorageDefaults)
     input: InputDefaults = field(default_factory=InputDefaults)
     embed_text: EmbedTextDefaults = field(default_factory=EmbedTextDefaults)
     chunks: ChunksDefaults = field(default_factory=ChunksDefaults)

packages/graphrag/graphrag/config/init_content.py

@@ -50,7 +50,7 @@
 
 input:
   storage:
-    type: {graphrag_config_defaults.input.storage.type} # or blob, cosmosdb
+    type: {graphrag_config_defaults.input.storage.type} # [file, blob, cosmosdb]
     base_dir: "{graphrag_config_defaults.input.storage.base_dir}"
   file_type: {graphrag_config_defaults.input.file_type.value} # [csv, text, json]
 
@@ -63,15 +63,14 @@
 ## connection_string and container_name must be provided
 
 output:
-  type: {graphrag_config_defaults.output.type} # or blob, cosmosdb
+  type: {graphrag_config_defaults.output.type} # [file, blob, cosmosdb]
   base_dir: "{graphrag_config_defaults.output.base_dir}"
-    
-cache_storage:
-  type: {graphrag_config_defaults.cache_storage.type} # [file, blob, cosmosdb]
-  base_dir: "{graphrag_config_defaults.cache_storage.base_dir}"
 
 cache:
-  type: {graphrag_config_defaults.cache.type} # [json, memory, noop]
+  type: {graphrag_config_defaults.cache.type} # [json, memory, none]
+  storage:
+    type: {graphrag_config_defaults.cache.storage.type} # [file, blob, cosmosdb]
+    base_dir: "{graphrag_config_defaults.cache.storage.base_dir}"
 
 reporting:
   type: {graphrag_config_defaults.reporting.type.value} # [file, blob]

packages/graphrag/graphrag/config/models/graph_rag_config.py

@@ -165,14 +165,6 @@ def _validate_update_index_output_base_dir(self) -> None:
                 Path(self.update_index_output.base_dir).resolve()
             )
 
-    cache_storage: StorageConfig | None = Field(
-        description="The cache storage configuration.",
-        default=StorageConfig(
-            **asdict(graphrag_config_defaults.cache_storage),
-        ),
-    )
-    """The cache storage configuration."""
-
     cache: CacheConfig = Field(
         description="The cache configuration.",
         default=CacheConfig(**asdict(graphrag_config_defaults.cache)),

packages/graphrag/graphrag/index/run/run_pipeline.py

@@ -37,10 +37,7 @@ async def run_pipeline(
     """Run all workflows using a simplified pipeline."""
     input_storage = create_storage(config.input.storage)
     output_storage = create_storage(config.output)
-    cache_storage: Storage | None = None
-    if config.cache_storage:
-        cache_storage = create_storage(config.cache_storage)
-    cache = create_cache(config.cache, storage=cache_storage)
+    cache = create_cache(config.cache)
 
     # load existing state in case any workflows are stateful
     state_json = await output_storage.get("context.json")