Add Cosmos DB storage/cache option

.github/workflows/python-integration-tests.yml

@@ -23,7 +23,7 @@ permissions:
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
-  # Only run the for the latest commit
+  # only run the for the latest commit
   cancel-in-progress: true
 
 env:
@@ -37,7 +37,7 @@ jobs:
       matrix:
         python-version: ["3.10"]
         os: [ubuntu-latest, windows-latest]
-      fail-fast: false # Continue running all jobs even if one fails
+      fail-fast: false # continue running all jobs even if one fails
     env:
       DEBUG: 1
 
@@ -84,6 +84,17 @@ jobs:
         id: azuright
         uses: potatoqualitee/[email protected]
 
+      # For more information on installation/setup of Azure Cosmos DB Emulator
+      # https://learn.microsoft.com/en-us/azure/cosmos-db/how-to-develop-emulator?tabs=docker-linux%2Cpython&pivots=api-nosql
+      # Note: the emulator is only available on Windows runners. It can take longer than the default to initially startup so we increase the default timeout.
+      # If a job fails due to timeout, restarting the cicd job usually resolves the problem.
+      - name: Install Azure Cosmos DB Emulator
+        if: runner.os == 'Windows'
+        run: |
+          Write-Host "Launching Cosmos DB Emulator"
+          Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator"
+          Start-CosmosDbEmulator -Timeout 500
+
       - name: Integration Test
         run: |
           poetry run poe test_integration

.semversioner/next-release/patch-20241121202210026640.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Implement cosmosdb storage option for cache and output"
+}

DEVELOPING.md

@@ -45,7 +45,6 @@ graphrag
 ├── config          # configuration management
 ├── index           # indexing engine
 |    └─ run/run.py  #  main entrypoint to build an index
-├── llm             # generic llm interfaces
 ├── logger          # logger module supporting several options
 │    └─ factory.py  #  └─ main entrypoint to create a logger
 ├── model           # data model definitions associated with the knowledge graph

dictionary.txt

@@ -28,6 +28,7 @@ ints
 
 # Azure
 abfs
+cosmosdb
 Hnsw
 odata
 

graphrag/cache/factory.py

@@ -9,6 +9,7 @@
 
 from graphrag.config.enums import CacheType
 from graphrag.storage.blob_pipeline_storage import BlobPipelineStorage
+from graphrag.storage.cosmosdb_pipeline_storage import create_cosmosdb_storage
 from graphrag.storage.file_pipeline_storage import FilePipelineStorage
 
 if TYPE_CHECKING:
@@ -50,6 +51,8 @@ def create_cache(
                 )
             case CacheType.blob:
                 return JsonPipelineCache(BlobPipelineStorage(**kwargs))
+            case CacheType.cosmosdb:
+                return JsonPipelineCache(create_cosmosdb_storage(**kwargs))
             case _:
                 if cache_type in cls.cache_types:
                     return cls.cache_types[cache_type](**kwargs)

graphrag/config/create_graphrag_config.py

@@ -362,6 +362,7 @@ def hydrate_parallelization_params(
                 storage_account_blob_url=reader.str(Fragment.storage_account_blob_url),
                 container_name=reader.str(Fragment.container_name),
                 base_dir=reader.str(Fragment.base_dir) or defs.CACHE_BASE_DIR,
+                cosmosdb_account_url=reader.str(Fragment.cosmosdb_account_url),
             )
         with (
             reader.envvar_prefix(Section.reporting),
@@ -383,6 +384,7 @@ def hydrate_parallelization_params(
                 storage_account_blob_url=reader.str(Fragment.storage_account_blob_url),
                 container_name=reader.str(Fragment.container_name),
                 base_dir=reader.str(Fragment.base_dir) or defs.STORAGE_BASE_DIR,
+                cosmosdb_account_url=reader.str(Fragment.cosmosdb_account_url),
             )
 
         with (
@@ -667,6 +669,7 @@ class Fragment(str, Enum):
     concurrent_requests = "CONCURRENT_REQUESTS"
     conn_string = "CONNECTION_STRING"
     container_name = "CONTAINER_NAME"
+    cosmosdb_account_url = "COSMOSDB_ACCOUNT_URL"
     deployment_name = "DEPLOYMENT_NAME"
     description = "DESCRIPTION"
     enabled = "ENABLED"

graphrag/config/enums.py

@@ -19,6 +19,8 @@ class CacheType(str, Enum):
     """The none cache configuration type."""
     blob = "blob"
     """The blob cache configuration type."""
+    cosmosdb = "cosmosdb"
+    """The cosmosdb cache configuration type"""
 
     def __repr__(self):
         """Get a string representation."""
@@ -60,6 +62,8 @@ class StorageType(str, Enum):
     """The memory storage type."""
     blob = "blob"
     """The blob storage type."""
+    cosmosdb = "cosmosdb"
+    """The cosmosdb storage type"""
 
     def __repr__(self):
         """Get a string representation."""

graphrag/config/init_content.py

@@ -63,15 +63,15 @@
 ## connection_string and container_name must be provided
 
 cache:
-  type: {defs.CACHE_TYPE.value} # or blob
+  type: {defs.CACHE_TYPE.value} # one of [blob, cosmosdb, file]
   base_dir: "{defs.CACHE_BASE_DIR}"
 
 reporting:
   type: {defs.REPORTING_TYPE.value} # or console, blob
   base_dir: "{defs.REPORTING_BASE_DIR}"
 
 storage:
-  type: {defs.STORAGE_TYPE.value} # or blob
+  type: {defs.STORAGE_TYPE.value} # one of [blob, cosmosdb, file]
   base_dir: "{defs.STORAGE_BASE_DIR}"
 
 ## only turn this on if running `graphrag index` with custom settings

graphrag/config/input_models/cache_config_input.py

@@ -16,3 +16,4 @@ class CacheConfigInput(TypedDict):
     connection_string: NotRequired[str | None]
     container_name: NotRequired[str | None]
     storage_account_blob_url: NotRequired[str | None]
+    cosmosdb_account_url: NotRequired[str | None]

graphrag/config/input_models/storage_config_input.py

@@ -16,3 +16,4 @@ class StorageConfigInput(TypedDict):
     connection_string: NotRequired[str | None]
     container_name: NotRequired[str | None]
     storage_account_blob_url: NotRequired[str | None]
+    cosmosdb_account_url: NotRequired[str | None]

graphrag/config/models/cache_config.py

@@ -27,3 +27,6 @@ class CacheConfig(BaseModel):
     storage_account_blob_url: str | None = Field(
         description="The storage account blob url to use.", default=None
     )
+    cosmosdb_account_url: str | None = Field(
+        description="The cosmosdb account url to use.", default=None
+    )

graphrag/config/models/storage_config.py

@@ -28,3 +28,6 @@ class StorageConfig(BaseModel):
     storage_account_blob_url: str | None = Field(
         description="The storage account blob url to use.", default=None
     )
+    cosmosdb_account_url: str | None = Field(
+        description="The cosmosdb account url to use.", default=None
+    )

graphrag/index/config/cache.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2024 Microsoft Corporation.
 # Licensed under the MIT License
 
-"""A module containing 'PipelineCacheConfig', 'PipelineFileCacheConfig', 'PipelineMemoryCacheConfig', 'PipelineBlobCacheConfig' models."""
+"""A module containing 'PipelineCacheConfig', 'PipelineFileCacheConfig', 'PipelineMemoryCacheConfig', 'PipelineBlobCacheConfig', 'PipelineCosmosDBCacheConfig' models."""
 
 from __future__ import annotations
 
@@ -71,9 +71,35 @@ class PipelineBlobCacheConfig(PipelineCacheConfig[Literal[CacheType.blob]]):
     """The storage account blob url for cache"""
 
 
+class PipelineCosmosDBCacheConfig(PipelineCacheConfig[Literal[CacheType.cosmosdb]]):
+    """Represents the cosmosdb cache configuration for the pipeline."""
+
+    type: Literal[CacheType.cosmosdb] = CacheType.cosmosdb
+    """The type of cache."""
+
+    base_dir: str | None = Field(
+        description="The cosmosdb database name for the cache.", default=None
+    )
+    """The cosmosdb database name for the cache."""
+
+    container_name: str = Field(description="The container name for cache.", default="")
+    """The container name for cache."""
+
+    connection_string: str | None = Field(
+        description="The cosmosdb primary key for the cache.", default=None
+    )
+    """The cosmosdb primary key for the cache."""
+
+    cosmosdb_account_url: str | None = Field(
+        description="The cosmosdb account url for cache", default=None
+    )
+    """The cosmosdb account url for cache"""
+
+
 PipelineCacheConfigTypes = (
     PipelineFileCacheConfig
     | PipelineMemoryCacheConfig
     | PipelineBlobCacheConfig
     | PipelineNoneCacheConfig
+    | PipelineCosmosDBCacheConfig
 )

graphrag/index/config/storage.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2024 Microsoft Corporation.
 # Licensed under the MIT License
 
-"""A module containing 'PipelineStorageConfig', 'PipelineFileStorageConfig' and 'PipelineMemoryStorageConfig' models."""
+"""A module containing 'PipelineStorageConfig', 'PipelineFileStorageConfig','PipelineMemoryStorageConfig', 'PipelineBlobStorageConfig', and 'PipelineCosmosDBStorageConfig' models."""
 
 from __future__ import annotations
 
@@ -66,6 +66,38 @@ class PipelineBlobStorageConfig(PipelineStorageConfig[Literal[StorageType.blob]]
     """The storage account blob url."""
 
 
+class PipelineCosmosDBStorageConfig(
+    PipelineStorageConfig[Literal[StorageType.cosmosdb]]
+):
+    """Represents the cosmosdb storage configuration for the pipeline."""
+
+    type: Literal[StorageType.cosmosdb] = StorageType.cosmosdb
+    """The type of storage."""
+
+    connection_string: str | None = Field(
+        description="The cosmosdb storage primary key for the storage.", default=None
+    )
+    """The cosmosdb storage primary key for the storage."""
+
+    container_name: str = Field(
+        description="The container name for storage", default=""
+    )
+    """The container name for storage."""
+
+    base_dir: str | None = Field(
+        description="The base directory for the storage.", default=None
+    )
+    """The base directory for the storage."""
+
+    cosmosdb_account_url: str | None = Field(
+        description="The cosmosdb account url.", default=None
+    )
+    """The cosmosdb account url."""
+
+
 PipelineStorageConfigTypes = (
-    PipelineFileStorageConfig | PipelineMemoryStorageConfig | PipelineBlobStorageConfig
+    PipelineFileStorageConfig
+    | PipelineMemoryStorageConfig
+    | PipelineBlobStorageConfig
+    | PipelineCosmosDBStorageConfig
 )

graphrag/index/create_pipeline_config.py

@@ -20,6 +20,7 @@
 from graphrag.index.config.cache import (
     PipelineBlobCacheConfig,
     PipelineCacheConfigTypes,
+    PipelineCosmosDBCacheConfig,
     PipelineFileCacheConfig,
     PipelineMemoryCacheConfig,
     PipelineNoneCacheConfig,
@@ -44,6 +45,7 @@
 )
 from graphrag.index.config.storage import (
     PipelineBlobStorageConfig,
+    PipelineCosmosDBStorageConfig,
     PipelineFileStorageConfig,
     PipelineMemoryStorageConfig,
     PipelineStorageConfigTypes,
@@ -420,6 +422,26 @@ def _get_storage_config(
                 base_dir=storage_settings.base_dir,
                 storage_account_blob_url=storage_account_blob_url,
             )
+        case StorageType.cosmosdb:
+            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
             base_dir = storage_settings.base_dir
@@ -457,6 +479,26 @@ def _get_cache_config(
                 base_dir=settings.cache.base_dir,
                 storage_account_blob_url=storage_account_blob_url,
             )
+        case CacheType.cosmosdb:
+            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=base_dir,
+                container_name=container_name,
+            )
         case _:
             # relative to root dir
             return PipelineFileCacheConfig(base_dir="./cache")

graphrag/index/run/workflow.py

@@ -95,8 +95,12 @@ async def _process_workflow(
         return None
 
     context.stats.workflows[workflow_name] = {"overall": 0.0}
+
     await _inject_workflow_data_dependencies(
-        workflow, workflow_dependencies, dataset, context.storage
+        workflow,
+        workflow_dependencies,
+        dataset,
+        context.storage,
     )
 
     workflow_start_time = time.time()

graphrag/query/input/loaders/utils.py

@@ -165,8 +165,7 @@ def to_optional_float(data: pd.Series, column_name: str | None) -> float | None:
         if value is None:
             return None
         if not isinstance(value, float):
-            msg = f"value is not a float: {value} ({type(value)})"
-            raise ValueError(msg)
+            return float(value)
     else:
         msg = f"Column {column_name} not found in data"
         raise ValueError(msg)