updates

dworthen · dworthen · commit 3a472331ca75 · 2025-12-03T10:16:00.000-08:00
diff --git a/docs/config/yaml.md b/docs/config/yaml.md
@@ -81,7 +81,7 @@ Our pipeline can ingest .csv, .txt, or .json data from an input location. See th
 #### Fields
 
 - `storage` **StorageConfig**
-  - `type` **File|AzureBlob|AzureCosmos** - The storage type to use. Default=`file`
+  - `type` **file|memory|blob|cosmosdb** - The storage type to use. Default=`file`
   - `base_dir` **str** - The base directory to write output artifacts to, relative to the root.
   - `connection_string` **str** - (blob/cosmosdb only) The Azure Storage connection string.
   - `container_name` **str** - (blob/cosmosdb only) The Azure Storage container name.
@@ -115,7 +115,7 @@ This section controls the storage mechanism used by the pipeline used for export
 
 #### Fields
 
-- `type` **File|AzureBlob|AzureCosmos** - The storage type to use. Default=`file`
+- `type` **file|memory|blob|cosmosdb** - The storage type to use. Default=`file`
 - `base_dir` **str** - The base directory to write output artifacts to, relative to the root.
 - `connection_string` **str** - (blob/cosmosdb only) The Azure Storage connection string.
 - `container_name` **str** - (blob/cosmosdb only) The Azure Storage container name.
@@ -128,7 +128,7 @@ The section defines a secondary storage location for running incremental indexin
 
 #### Fields
 
-- `type` **File|AzureBlob|AzureCosmos** - The storage type to use. Default=`file`
+- `type` **file|memory|blob|cosmosdb** - The storage type to use. Default=`file`
 - `base_dir` **str** - The base directory to write output artifacts to, relative to the root.
 - `connection_string` **str** - (blob/cosmosdb only) The Azure Storage connection string.
 - `container_name` **str** - (blob/cosmosdb only) The Azure Storage container name.
diff --git a/packages/graphrag-storage/README.md b/packages/graphrag-storage/README.md
@@ -9,7 +9,7 @@ from graphrag_storage import StorageConfig, create_storage, StorageType
 async def run():
     storage = create_storage(
         StorageConfig(
-            type=StorageType.FILE
+            type=StorageType.File
             base_dir="output"
         )
     )
@@ -29,7 +29,7 @@ from typing import Any
 from graphrag_storage import Storage, StorageConfig, create_storage, register_storage
 
 class MyStorage(Storage):
-    def __init__(self, some_setting: str, optional_setting: str = "default setting"):
+    def __init__(self, some_setting: str, optional_setting: str = "default setting", **kwargs: Any):
         # Validate settings and initialize
         ...
 
@@ -55,3 +55,28 @@ async def run():
 
 if __name__ == "__main__":
     asyncio.run(run())
+```
+
+### Information
+
+By default, the `create_storage` comes with the following storage providers registered that correspond to the entries in the `StorageType` enum.
+
+- `FileStorage`
+- `AzureBlobStorage`
+- `AzureCosmosStorage`
+- `MemoryStorage`
+
+You can directly import `storage_factory` if you want a clean factory with no preregistered storage providers.
+
+```python
+from graphrag_storage.storage_factory import storage_factory
+from graphrag_storage.file_storage import FileStorage
+
+# Or register a custom implementation, see above for example.
+storage_factory.register("my_storage_key", FileStorage)
+
+storage = storage_factory.create(strategy="my_storage_key", init_args={"base_dir": "...", "other_settings": "...",})
+
+...
+
+```
diff --git a/packages/graphrag-storage/graphrag_storage/azure_blob_storage.py b/packages/graphrag-storage/graphrag_storage/azure_blob_storage.py
@@ -38,6 +38,7 @@ def __init__(
         azure_connection_string: str | None = None,
         base_dir: str | None = None,
         encoding: str = "utf-8",
+        **kwargs: Any,
     ) -> None:
         """Create a new BlobStorage instance."""
         if azure_connection_string is not None and azure_account_url is not None:
diff --git a/packages/graphrag-storage/graphrag_storage/azure_cosmos_storage.py b/packages/graphrag-storage/graphrag_storage/azure_cosmos_storage.py
@@ -46,6 +46,7 @@ def __init__(
         azure_connection_string: str | None = None,
         azure_account_url: str | None = None,
         encoding: str = "utf-8",
+        **kwargs: Any,
     ) -> None:
         """Create a CosmosDB storage instance."""
         logger.info("Creating cosmosdb storage")
diff --git a/packages/graphrag-storage/graphrag_storage/file_storage.py b/packages/graphrag-storage/graphrag_storage/file_storage.py
@@ -30,7 +30,7 @@ class FileStorage(Storage):
     _base_dir: Path
     _encoding: str
 
-    def __init__(self, base_dir: str, encoding: str = "utf-8") -> None:
+    def __init__(self, base_dir: str, encoding: str = "utf-8", **kwargs: Any) -> None:
         """Create a file based storage."""
         self._base_dir = Path(base_dir).resolve()
         self._encoding = encoding
diff --git a/packages/graphrag-storage/graphrag_storage/storage.py b/packages/graphrag-storage/graphrag_storage/storage.py
@@ -13,6 +13,10 @@
 class Storage(ABC):
     """Provide a storage interface."""
 
+    @abstractmethod
+    def __init__(self, **kwargs: Any) -> None:
+        """Create a storage instance."""
+
     @abstractmethod
     def find(
         self,
diff --git a/packages/graphrag-storage/graphrag_storage/storage_config.py b/packages/graphrag-storage/graphrag_storage/storage_config.py
@@ -16,7 +16,7 @@ class StorageConfig(BaseModel):
 
     type: str = Field(
         description="The storage type to use. Builtin types include 'File', 'AzureBlob', and 'AzureCosmos'.",
-        default=StorageType.FILE,
+        default=StorageType.File,
     )
 
     encoding: str | None = Field(
diff --git a/packages/graphrag-storage/graphrag_storage/storage_factory.py b/packages/graphrag-storage/graphrag_storage/storage_factory.py
@@ -13,11 +13,11 @@
 from graphrag_storage.storage_type import StorageType
 
 
-class _StorageFactory(Factory[Storage]):
+class StorageFactory(Factory[Storage]):
     """A factory class for storage implementations."""
 
 
-storage_factory = _StorageFactory()
+storage_factory = StorageFactory()
 
 
 def register_storage(
@@ -49,31 +49,26 @@ def create_storage(config: StorageConfig) -> Storage:
             The created storage implementation.
     """
     config_model = config.model_dump()
-    storage_strategy = config_model.pop("type")
-
-    # Check storage_strategy is a string
-    if not isinstance(storage_strategy, str):
-        msg = f"StorageConfig.type must be a string, got {type(storage_strategy)}"
-        raise TypeError(msg)
+    storage_strategy = config.type
 
     if storage_strategy not in storage_factory:
         match storage_strategy:
-            case StorageType.FILE:
+            case StorageType.File:
                 from graphrag_storage.file_storage import FileStorage
 
-                register_storage(StorageType.FILE, FileStorage)
-            case StorageType.MEMORY:
+                register_storage(StorageType.File, FileStorage)
+            case StorageType.Memory:
                 from graphrag_storage.memory_storage import MemoryStorage
 
-                register_storage(StorageType.MEMORY, MemoryStorage)
-            case StorageType.AZURE_BLOB:
+                register_storage(StorageType.Memory, MemoryStorage)
+            case StorageType.AzureBlob:
                 from graphrag_storage.azure_blob_storage import AzureBlobStorage
 
-                register_storage(StorageType.AZURE_BLOB, AzureBlobStorage)
-            case StorageType.AZURE_COSMOS:
+                register_storage(StorageType.AzureBlob, AzureBlobStorage)
+            case StorageType.AzureCosmos:
                 from graphrag_storage.azure_cosmos_storage import AzureCosmosStorage
 
-                register_storage(StorageType.AZURE_COSMOS, AzureCosmosStorage)
+                register_storage(StorageType.AzureCosmos, AzureCosmosStorage)
             case _:
                 msg = f"StorageConfig.type '{storage_strategy}' is not registered in the StorageFactory. Registered types: {', '.join(storage_factory.keys())}."
                 raise ValueError(msg)
diff --git a/packages/graphrag-storage/graphrag_storage/storage_type.py b/packages/graphrag-storage/graphrag_storage/storage_type.py
@@ -10,7 +10,7 @@
 class StorageType(StrEnum):
     """Enum for storage types."""
 
-    FILE = "File"
-    MEMORY = "Memory"
-    AZURE_BLOB = "AzureBlob"
-    AZURE_COSMOS = "AzureCosmos"
+    File = "file"
+    Memory = "memory"
+    AzureBlob = "blob"
+    AzureCosmos = "cosmosdb"
diff --git a/packages/graphrag/graphrag/cache/factory.py b/packages/graphrag/graphrag/cache/factory.py
@@ -29,7 +29,6 @@ def create_file_cache(**kwargs) -> PipelineCache:
     """Create a file-based cache implementation."""
     from graphrag_storage.file_storage import FileStorage
 
-    kwargs.pop("type", None)
     storage = FileStorage(**kwargs)
     return JsonPipelineCache(storage)
 
@@ -38,7 +37,6 @@ def create_blob_cache(**kwargs) -> PipelineCache:
     """Create a blob storage-based cache implementation."""
     from graphrag_storage.azure_blob_storage import AzureBlobStorage
 
-    kwargs.pop("type", None)
     storage = AzureBlobStorage(**kwargs)
     return JsonPipelineCache(storage)
 
@@ -47,7 +45,6 @@ def create_cosmosdb_cache(**kwargs) -> PipelineCache:
     """Create a CosmosDB-based cache implementation."""
     from graphrag_storage.azure_cosmos_storage import AzureCosmosStorage
 
-    kwargs.pop("type", None)
     storage = AzureCosmosStorage(**kwargs)
     return JsonPipelineCache(storage)
 
diff --git a/packages/graphrag/graphrag/config/defaults.py b/packages/graphrag/graphrag/config/defaults.py
@@ -231,7 +231,7 @@ class GlobalSearchDefaults:
 class StorageDefaults:
     """Default values for storage."""
 
-    type: str = StorageType.FILE
+    type: str = StorageType.File
     encoding: str | None = None
     base_dir: str | None = None
     azure_connection_string: None = None
diff --git a/packages/graphrag/graphrag/config/init_content.py b/packages/graphrag/graphrag/config/init_content.py
@@ -50,7 +50,7 @@
 
 input:
   storage:
-    type: {graphrag_config_defaults.input.storage.type} # or AzureBlob, AzureCosmos
+    type: {graphrag_config_defaults.input.storage.type} # or blob, cosmosdb
     base_dir: "{graphrag_config_defaults.input.storage.base_dir}"
   file_type: {graphrag_config_defaults.input.file_type.value} # [csv, text, json]
 
@@ -63,7 +63,7 @@
 ## connection_string and container_name must be provided
 
 output:
-  type: {graphrag_config_defaults.output.type} # or AzureBlob, AzureCosmos
+  type: {graphrag_config_defaults.output.type} # or blob, cosmosdb
   base_dir: "{graphrag_config_defaults.output.base_dir}"
     
 cache:
diff --git a/packages/graphrag/graphrag/config/models/graph_rag_config.py b/packages/graphrag/graphrag/config/models/graph_rag_config.py
@@ -116,7 +116,7 @@ def _validate_input_pattern(self) -> None:
 
     def _validate_input_base_dir(self) -> None:
         """Validate the input base directory."""
-        if self.input.storage.type == StorageType.FILE:
+        if self.input.storage.type == StorageType.File:
             if not self.input.storage.base_dir:
                 msg = "input storage base directory is required for file input storage. Please rerun `graphrag init` and set the input storage configuration."
                 raise ValueError(msg)
@@ -140,7 +140,7 @@ def _validate_input_base_dir(self) -> None:
 
     def _validate_output_base_dir(self) -> None:
         """Validate the output base directory."""
-        if self.output.type == StorageType.FILE:
+        if self.output.type == StorageType.File:
             if not self.output.base_dir:
                 msg = "output base directory is required for file output. Please rerun `graphrag init` and set the output configuration."
                 raise ValueError(msg)
@@ -156,7 +156,7 @@ def _validate_output_base_dir(self) -> None:
 
     def _validate_update_index_output_base_dir(self) -> None:
         """Validate the update index output base directory."""
-        if self.update_index_output.type == StorageType.FILE:
+        if self.update_index_output.type == StorageType.File:
             if not self.update_index_output.base_dir:
                 msg = "update_index_output base directory is required for file output. Please rerun `graphrag init` and set the update_index_output configuration."
                 raise ValueError(msg)
diff --git a/tests/integration/storage/test_factory.py b/tests/integration/storage/test_factory.py
@@ -29,7 +29,7 @@
 @pytest.mark.skip(reason="Blob storage emulator is not available in this environment")
 def test_create_blob_storage():
     config = StorageConfig(
-        type=StorageType.AZURE_BLOB,
+        type=StorageType.AzureBlob,
         azure_connection_string=WELL_KNOWN_BLOB_STORAGE_KEY,
         base_dir="testbasedir",
         azure_container_name="testcontainer",
@@ -44,7 +44,7 @@ def test_create_blob_storage():
 )
 def test_create_cosmosdb_storage():
     config = StorageConfig(
-        type=StorageType.AZURE_COSMOS,
+        type=StorageType.AzureCosmos,
         azure_connection_string=WELL_KNOWN_COSMOS_CONNECTION_STRING,
         azure_cosmosdb_database_name="testdatabase",
         azure_container_name="testcontainer",
@@ -55,7 +55,7 @@ def test_create_cosmosdb_storage():
 
 def test_create_file():
     config = StorageConfig(
-        type=StorageType.FILE,
+        type=StorageType.File,
         base_dir="/tmp/teststorage",
     )
     storage = create_storage(config)
@@ -65,7 +65,7 @@ def test_create_file():
 def test_create_memory_storage():
     config = StorageConfig(
         base_dir="",
-        type=StorageType.MEMORY,
+        type=StorageType.Memory,
     )
     storage = create_storage(config)
     assert isinstance(storage, MemoryStorage)

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ class StorageConfig(BaseModel):`
`16`	`16`
`17`	`17`	`type: str = Field(`
`18`	`18`	`description="The storage type to use. Builtin types include 'File', 'AzureBlob', and 'AzureCosmos'.",`
`19`		`- default=StorageType.FILE,`
	`19`	`+ default=StorageType.File,`
`20`	`20`	`)`
`21`	`21`
`22`	`22`	`encoding: str \| None = Field(`