yujunjun/graphrag
1
0
Fork 0
mirror of https://github.com/microsoft/graphrag.git synced 2025-12-27 15:10:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

collapsed cosmosdb schema to use minimal containers and databases

Browse Source
This commit is contained in:
Kenny Zhang 2024-11-19 15:30:59 -05:00
parent c5281bb79a
commit 76511d0180
6 changed files with 29 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
graphrag/index/cache/load_cache.py vendored
View File

@ -54,7 +54,7 @@ def load_cache(config: PipelineCacheConfig | None, root_dir: str | None):
storage = create_cosmosdb_storage(
cosmosdb_account_url=config.cosmosdb_account_url,
connection_string=config.connection_string,
container_name=None,
container_name=config.container_name,
base_dir=config.base_dir,
)
return JsonPipelineCache(storage)

9
graphrag/index/config/cache.py
View File

@ -80,9 +80,14 @@ class PipelineCosmosDBCacheConfig(PipelineCacheConfig[Literal[CacheType.cosmosdb
"""The type of cache."""
base_dir: str = pydantic_Field(
description="The base directory for the cache.", default=None
description="The cosmosdb database name for the cache.", default=None
)
"""The base directory for the cache."""
"""The cosmosdb database name for the cache."""
container_name: str = pydantic_Field(
description="The container name for cache.", default=None
)
"""The container name for cache."""
connection_string: str | None = pydantic_Field(
description="The cosmosdb primary key for the cache.", default=None

5
graphrag/index/config/storage.py
View File

@ -77,6 +77,11 @@ class PipelineCosmosDBStorageConfig(PipelineStorageConfig[Literal[StorageType.co
)
"""The cosmosdb storage primary key for the storage."""
container_name: str = pydantic_Field(
description="The container name for storage", default=None
)
"""The container name for storage."""
base_dir: str = pydantic_Field(
description="The base directory for the storage.", default=None
)

12
graphrag/index/create_pipeline_config.py
View File

@ -421,16 +421,21 @@ def _get_storage_config(
cosmosdb_account_url = storage_settings.cosmosdb_account_url
connection_string = storage_settings.connection_string
base_dir = storage_settings.base_dir
container_name = storage_settings.container_name
if cosmosdb_account_url is None:
msg = "CosmosDB account url must be provided for cosmosdb storage."
raise ValueError(msg)
if base_dir is None:
msg = "Base directory must be provided for cosmosdb storage."
raise ValueError(msg)
if container_name is None:
msg = "Container name must be provided for cosmosdb storage."
raise ValueError(msg)
return PipelineCosmosDBStorageConfig(
cosmosdb_account_url=cosmosdb_account_url,
connection_string=connection_string,
base_dir=storage_settings.base_dir,
container_name=container_name,
)
case _:
# relative to the root_dir
@ -473,16 +478,21 @@ def _get_cache_config(
cosmosdb_account_url = settings.cache.cosmosdb_account_url
connection_string = settings.cache.connection_string
base_dir = settings.cache.base_dir
container_name = settings.cache.container_name
if base_dir is None:
msg = "Base directory must be provided for cosmosdb cache."
raise ValueError(msg)
if container_name is None:
msg = "Container name must be provided for cosmosdb cache."
raise ValueError(msg)
if connection_string is None and cosmosdb_account_url is None:
msg = "Connection string or cosmosDB account url must be provided for cosmosdb cache."
raise ValueError(msg)
return PipelineCosmosDBCacheConfig(
cosmosdb_account_url=cosmosdb_account_url,
connection_string=connection_string,
base_dir=settings.cache.base_dir,
base_dir=base_dir,
container_name=container_name,
)
case _:
# relative to root dir

10
graphrag/index/storage/cosmosdb_pipeline_storage.py
View File

@ -254,15 +254,7 @@ class CosmosDBPipelineStorage(PipelineStorage):
def child(self, name: str | None) -> "PipelineStorage":
"""Create a child storage instance."""
if name is None:
return self
return CosmosDBPipelineStorage(
cosmosdb_account_url=self._cosmosdb_account_url,
connection_string=self._connection_string,
database_name=self._database_name,
encoding=self._encoding,
current_container=name,
)
return self
def create_container(self) -> None:
"""Create a container for the current container name if it doesn't exist."""

8
graphrag/index/storage/load_storage.py
View File

@ -36,10 +36,10 @@ def load_storage(config: PipelineStorageConfig):
case StorageType.cosmosdb:
config = cast(PipelineCosmosDBStorageConfig, config)
return create_cosmosdb_storage(
config.cosmosdb_account_url,
config.connection_string,
config.base_dir,
config.base_dir,
connection_string=config.connection_string,
cosmosdb_account_url=config.cosmosdb_account_url,
container_name=config.container_name,
base_dir=config.base_dir,
)
case StorageType.file:
config = cast(PipelineFileStorageConfig, config)