update storage factory design

Author: jgbradley1

graphrag/storage/blob_pipeline_storage.py

@@ -29,15 +29,20 @@ class BlobPipelineStorage(PipelineStorage):
     _encoding: str
     _storage_account_blob_url: str | None
 
-    def __init__(
-        self,
-        connection_string: str | None,
-        container_name: str,
-        encoding: str = "utf-8",
-        path_prefix: str | None = None,
-        storage_account_blob_url: str | None = None,
-    ):
+    def __init__(self, **kwargs: Any) -> None:
         """Create a new BlobStorage instance."""
+        connection_string = kwargs.get("connection_string")
+        storage_account_blob_url = kwargs.get("storage_account_blob_url")
+        path_prefix = kwargs.get("base_dir")
+        container_name = kwargs["container_name"]
+        if container_name is None:
+            msg = "No container name provided for blob storage."
+            raise ValueError(msg)
+        if connection_string is None and storage_account_blob_url is None:
+            msg = "No storage account blob url provided for blob storage."
+            raise ValueError(msg)
+
+        logger.info("Creating blob storage at %s", container_name)
         if connection_string:
             self._blob_service_client = BlobServiceClient.from_connection_string(
                 connection_string
@@ -51,7 +56,7 @@ def __init__(
                 account_url=storage_account_blob_url,
                 credential=DefaultAzureCredential(),
             )
-        self._encoding = encoding
+        self._encoding = kwargs.get("encoding", "utf-8")
         self._container_name = container_name
         self._connection_string = connection_string
         self._path_prefix = path_prefix or ""
@@ -308,27 +313,6 @@ async def get_creation_date(self, key: str) -> str:
             return ""
 
 
-def create_blob_storage(**kwargs: Any) -> PipelineStorage:
-    """Create a blob based storage."""
-    connection_string = kwargs.get("connection_string")
-    storage_account_blob_url = kwargs.get("storage_account_blob_url")
-    base_dir = kwargs.get("base_dir")
-    container_name = kwargs["container_name"]
-    logger.info("Creating blob storage at %s", container_name)
-    if container_name is None:
-        msg = "No container name provided for blob storage."
-        raise ValueError(msg)
-    if connection_string is None and storage_account_blob_url is None:
-        msg = "No storage account blob url provided for blob storage."
-        raise ValueError(msg)
-    return BlobPipelineStorage(
-        connection_string=connection_string,
-        container_name=container_name,
-        path_prefix=base_dir,
-        storage_account_blob_url=storage_account_blob_url,
-    )
-
-
 def validate_blob_container_name(container_name: str):
     """
     Check if the provided blob container name is valid based on Azure rules.

graphrag/storage/cosmosdb_pipeline_storage.py

@@ -39,15 +39,20 @@ class CosmosDBPipelineStorage(PipelineStorage):
     _encoding: str
     _no_id_prefixes: list[str]
 
-    def __init__(
-        self,
-        database_name: str,
-        container_name: str,
-        cosmosdb_account_url: str | None = None,
-        connection_string: str | None = None,
-        encoding: str = "utf-8",
-    ):
-        """Initialize the CosmosDB Storage."""
+    def __init__(self, **kwargs: Any) -> None:
+        """Create a CosmosDB storage instance."""
+        logger.info("Creating cosmosdb storage")
+        cosmosdb_account_url = kwargs.get("cosmosdb_account_url")
+        connection_string = kwargs.get("connection_string")
+        database_name = kwargs["base_dir"]
+        container_name = kwargs["container_name"]
+        if not database_name:
+            msg = "No base_dir provided for database name"
+            raise ValueError(msg)
+        if connection_string is None and cosmosdb_account_url is None:
+            msg = "connection_string or cosmosdb_account_url is required."
+            raise ValueError(msg)
+
         if connection_string:
             self._cosmos_client = CosmosClient.from_connection_string(connection_string)
         else:
@@ -60,7 +65,7 @@ def __init__(
                 url=cosmosdb_account_url,
                 credential=DefaultAzureCredential(),
             )
-        self._encoding = encoding
+        self._encoding = kwargs.get("encoding", "utf-8")
         self._database_name = database_name
         self._connection_string = connection_string
         self._cosmosdb_account_url = cosmosdb_account_url
@@ -348,29 +353,6 @@ async def get_creation_date(self, key: str) -> str:
             return ""
 
 
-# TODO remove this helper function and have the factory instantiate the class directly
-# once the new config system is in place and will enforce the correct types/existence of certain fields
-def create_cosmosdb_storage(**kwargs: Any) -> PipelineStorage:
-    """Create a CosmosDB storage instance."""
-    logger.info("Creating cosmosdb storage")
-    cosmosdb_account_url = kwargs.get("cosmosdb_account_url")
-    connection_string = kwargs.get("connection_string")
-    base_dir = kwargs["base_dir"]
-    container_name = kwargs["container_name"]
-    if not base_dir:
-        msg = "No base_dir provided for database name"
-        raise ValueError(msg)
-    if connection_string is None and cosmosdb_account_url is None:
-        msg = "connection_string or cosmosdb_account_url is required."
-        raise ValueError(msg)
-    return CosmosDBPipelineStorage(
-        cosmosdb_account_url=cosmosdb_account_url,
-        connection_string=connection_string,
-        database_name=base_dir,
-        container_name=container_name,
-    )
-
-
 def _create_progress_status(
     num_loaded: int, num_filtered: int, num_total: int
 ) -> Progress:

graphrag/storage/factory.py

@@ -8,9 +8,9 @@
 from typing import TYPE_CHECKING, ClassVar
 
 from graphrag.config.enums import StorageType
-from graphrag.storage.blob_pipeline_storage import create_blob_storage
-from graphrag.storage.cosmosdb_pipeline_storage import create_cosmosdb_storage
-from graphrag.storage.file_pipeline_storage import create_file_storage
+from graphrag.storage.blob_pipeline_storage import BlobPipelineStorage
+from graphrag.storage.cosmosdb_pipeline_storage import CosmosDBPipelineStorage
+from graphrag.storage.file_pipeline_storage import FilePipelineStorage
 from graphrag.storage.memory_pipeline_storage import MemoryPipelineStorage
 
 if TYPE_CHECKING:
@@ -28,7 +28,7 @@ class StorageFactory:
     for individual enforcement of required/optional arguments.
     """
 
-    _storage_registry: ClassVar[dict[str, Callable[..., PipelineStorage]]] = {}
+    _registry: ClassVar[dict[str, Callable[..., PipelineStorage]]] = {}
 
     @classmethod
     def register(
@@ -44,10 +44,7 @@ def register(
         ------
             TypeError: If creator is a class type instead of a factory function.
         """
-        if isinstance(creator, type):
-            msg = "Registering classes directly is no longer supported. Please provide a factory function instead."
-            raise TypeError(msg)
-        cls._storage_registry[storage_type] = creator
+        cls._registry[storage_type] = creator
 
     @classmethod
     def create_storage(
@@ -67,31 +64,31 @@ def create_storage(
         ------
             ValueError: If the storage type is not registered.
         """
-        storage_type_str = (
+        type_str = (
             storage_type.value
             if isinstance(storage_type, StorageType)
             else storage_type
         )
 
-        if storage_type_str not in cls._storage_registry:
+        if type_str not in cls._storage_registry:
             msg = f"Unknown storage type: {storage_type}"
             raise ValueError(msg)
 
-        return cls._storage_registry[storage_type_str](**kwargs)
+        return cls._storage_registry[type_str](**kwargs)
 
     @classmethod
     def get_storage_types(cls) -> list[str]:
         """Get the registered storage implementations."""
-        return list(cls._storage_registry.keys())
+        return list(cls._registry.keys())
 
     @classmethod
-    def is_supported_storage_type(cls, storage_type: str) -> bool:
+    def is_supported_type(cls, storage_type: str) -> bool:
         """Check if the given storage type is supported."""
-        return storage_type in cls._storage_registry
+        return storage_type in cls._registry
 
 
-# --- Register default implementations ---
-StorageFactory.register(StorageType.blob.value, create_blob_storage)
-StorageFactory.register(StorageType.cosmosdb.value, create_cosmosdb_storage)
-StorageFactory.register(StorageType.file.value, create_file_storage)
-StorageFactory.register(StorageType.memory.value, lambda **_: MemoryPipelineStorage())
+# --- register built-in storage implementations ---
+StorageFactory.register(StorageType.blob.value, BlobPipelineStorage)
+StorageFactory.register(StorageType.cosmosdb.value, CosmosDBPipelineStorage)
+StorageFactory.register(StorageType.file.value, FilePipelineStorage)
+StorageFactory.register(StorageType.memory.value, MemoryPipelineStorage)

graphrag/storage/file_pipeline_storage.py

@@ -30,10 +30,11 @@ class FilePipelineStorage(PipelineStorage):
     _root_dir: str
     _encoding: str
 
-    def __init__(self, root_dir: str = "", encoding: str = "utf-8"):
-        """Init method definition."""
-        self._root_dir = root_dir
-        self._encoding = encoding
+    def __init__(self, **kwargs: Any) -> None:
+        """Create a file based storage."""
+        self._root_dir = kwargs["base_dir"]
+        self._encoding = kwargs.get("encoding", "utf-8")
+        logger.info("Creating file storage at %s", self._root_dir)
         Path(self._root_dir).mkdir(parents=True, exist_ok=True)
 
     def find(
@@ -167,10 +168,3 @@ async def get_creation_date(self, key: str) -> str:
 def join_path(file_path: str, file_name: str) -> Path:
     """Join a path and a file. Independent of the OS."""
     return Path(file_path) / Path(file_name).parent / Path(file_name).name
-
-
-def create_file_storage(**kwargs: Any) -> PipelineStorage:
-    """Create a file based storage."""
-    base_dir = kwargs["base_dir"]
-    logger.info("Creating file storage at %s", base_dir)
-    return FilePipelineStorage(root_dir=base_dir)

graphrag/vector_stores/factory.py

@@ -8,6 +8,10 @@
 from enum import Enum
 from typing import TYPE_CHECKING, ClassVar
 
+from graphrag.vector_stores.azure_ai_search import AzureAISearchVectorStore
+from graphrag.vector_stores.cosmosdb import CosmosDBVectorStore
+from graphrag.vector_stores.lancedb import LanceDBVectorStore
+
 if TYPE_CHECKING:
     from collections.abc import Callable
 
@@ -31,7 +35,7 @@ class VectorStoreFactory:
     for individual enforcement of required/optional arguments.
     """
 
-    _vector_store_registry: ClassVar[dict[str, Callable[..., BaseVectorStore]]] = {}
+    _registry: ClassVar[dict[str, Callable[..., BaseVectorStore]]] = {}
 
     @classmethod
     def register(
@@ -47,10 +51,7 @@ def register(
         ------
             TypeError: If creator is a class type instead of a factory function.
         """
-        if isinstance(creator, type):
-            msg = "Registering classes directly is no longer supported. Please provide a factory function instead."
-            raise TypeError(msg)
-        cls._vector_store_registry[vector_store_type] = creator
+        cls._registry[vector_store_type] = creator
 
     @classmethod
     def create_vector_store(
@@ -70,56 +71,32 @@ def create_vector_store(
         ------
             ValueError: If the vector store type is not registered.
         """
-        vector_store_type_str = (
+        type_str = (
             vector_store_type.value
             if isinstance(vector_store_type, VectorStoreType)
             else vector_store_type
         )
 
-        if vector_store_type_str not in cls._vector_store_registry:
+        if type_str not in cls._registry:
             msg = f"Unknown vector store type: {vector_store_type}"
             raise ValueError(msg)
 
-        return cls._vector_store_registry[vector_store_type_str](**kwargs)
+        return cls._registry[type_str](**kwargs)
 
     @classmethod
     def get_vector_store_types(cls) -> list[str]:
         """Get the registered vector store implementations."""
-        return list(cls._vector_store_registry.keys())
+        return list(cls._registry.keys())
 
     @classmethod
-    def is_supported_vector_store_type(cls, vector_store_type: str) -> bool:
+    def is_supported_type(cls, vector_store_type: str) -> bool:
         """Check if the given vector store type is supported."""
-        return vector_store_type in cls._vector_store_registry
-
-
-# --- Factory functions for built-in vector stores ---
-def create_lancedb_vector_store(**kwargs) -> BaseVectorStore:
-    """Create a LanceDB vector store."""
-    from graphrag.vector_stores.lancedb import LanceDBVectorStore
-
-    return LanceDBVectorStore(**kwargs)
-
-
-def create_azure_ai_search_vector_store(**kwargs) -> BaseVectorStore:
-    """Create an Azure AI Search vector store."""
-    from graphrag.vector_stores.azure_ai_search import AzureAISearchVectorStore
+        return vector_store_type in cls._registry
 
-    return AzureAISearchVectorStore(**kwargs)
 
-
-def create_cosmosdb_vector_store(**kwargs) -> BaseVectorStore:
-    """Create a CosmosDB vector store."""
-    from graphrag.vector_stores.cosmosdb import CosmosDBVectorStore
-
-    return CosmosDBVectorStore(**kwargs)
-
-
-# --- register default implementations ---
-VectorStoreFactory.register(VectorStoreType.LanceDB.value, create_lancedb_vector_store)
-VectorStoreFactory.register(
-    VectorStoreType.AzureAISearch.value, create_azure_ai_search_vector_store
-)
+# --- register built-in vector store implementations ---
+VectorStoreFactory.register(VectorStoreType.LanceDB.value, LanceDBVectorStore)
 VectorStoreFactory.register(
-    VectorStoreType.CosmosDB.value, create_cosmosdb_vector_store
+    VectorStoreType.AzureAISearch.value, AzureAISearchVectorStore
 )
+VectorStoreFactory.register(VectorStoreType.CosmosDB.value, CosmosDBVectorStore)

tests/integration/storage/test_factory.py

@@ -84,7 +84,7 @@ def test_register_and_create_custom_storage():
 
     # Check if it's in the list of registered storage types
     assert "custom" in StorageFactory.get_storage_types()
-    assert StorageFactory.is_supported_storage_type("custom")
+    assert StorageFactory.is_supported_type("custom")
 
 
 def test_get_storage_types():

tests/integration/vector_stores/test_factory.py

@@ -75,7 +75,7 @@ def test_register_and_create_custom_vector_store():
 
     # Check if it's in the list of registered vector store types
     assert "custom" in VectorStoreFactory.get_vector_store_types()
-    assert VectorStoreFactory.is_supported_vector_store_type("custom")
+    assert VectorStoreFactory.is_supported_type("custom")
 
 
 def test_get_vector_store_types():
@@ -91,20 +91,14 @@ def test_create_unknown_vector_store():
         VectorStoreFactory.create_vector_store("unknown", {})
 
 
-def test_is_supported_vector_store_type():
+def test_is_supported_type():
     # Test built-in types
-    assert VectorStoreFactory.is_supported_vector_store_type(
-        VectorStoreType.LanceDB.value
-    )
-    assert VectorStoreFactory.is_supported_vector_store_type(
-        VectorStoreType.AzureAISearch.value
-    )
-    assert VectorStoreFactory.is_supported_vector_store_type(
-        VectorStoreType.CosmosDB.value
-    )
+    assert VectorStoreFactory.is_supported_type(VectorStoreType.LanceDB.value)
+    assert VectorStoreFactory.is_supported_type(VectorStoreType.AzureAISearch.value)
+    assert VectorStoreFactory.is_supported_type(VectorStoreType.CosmosDB.value)
 
     # Test unknown type
-    assert not VectorStoreFactory.is_supported_vector_store_type("unknown")
+    assert not VectorStoreFactory.is_supported_type("unknown")
 
 
 def test_enum_and_string_compatibility():