Support for microsoft/documentdb

.vscode/settings.json

@@ -48,5 +48,10 @@
     },
     "custom": true, // Enable the `custom` dictionary
     "internal-terms": true // Disable the `internal-terms` dictionary
-  }
+  },
+  "python.testing.pytestArgs": [
+    "tests"
+  ],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true
 }

graphrag/cache/factory.py

@@ -10,6 +10,7 @@
 from graphrag.config.enums import CacheType
 from graphrag.storage.blob_pipeline_storage import create_blob_storage
 from graphrag.storage.cosmosdb_pipeline_storage import create_cosmosdb_storage
+from graphrag.storage.documentdb_pipeline_storage import create_documentdb_storage
 from graphrag.storage.file_pipeline_storage import FilePipelineStorage
 
 if TYPE_CHECKING:
@@ -56,6 +57,8 @@ def create_cache(
                 return JsonPipelineCache(create_blob_storage(**kwargs))
             case CacheType.cosmosdb:
                 return JsonPipelineCache(create_cosmosdb_storage(**kwargs))
+            case CacheType.documentdb:
+                return JsonPipelineCache(create_documentdb_storage(**kwargs))
             case _:
                 if cache_type in cls.cache_types:
                     return cls.cache_types[cache_type](**kwargs)

graphrag/config/enums.py

@@ -21,6 +21,8 @@ class CacheType(str, Enum):
     """The blob cache configuration type."""
     cosmosdb = "cosmosdb"
     """The cosmosdb cache configuration type"""
+    documentdb = "documentdb"
+    """The documentdb cache configuration type"""
 
     def __repr__(self):
         """Get a string representation."""
@@ -64,6 +66,8 @@ class OutputType(str, Enum):
     """The blob output type."""
     cosmosdb = "cosmosdb"
     """The cosmosdb output type"""
+    documentdb = "documentdb"
+    """The documentdb output type"""
 
     def __repr__(self):
         """Get a string representation."""

graphrag/config/init_content.py

@@ -80,15 +80,15 @@
 ## connection_string and container_name must be provided
 
 cache:
-  type: {defs.CACHE_TYPE.value} # [file, blob, cosmosdb]
+  type: {defs.CACHE_TYPE.value} # [file, blob, cosmosdb, documentdb]
   base_dir: "{defs.CACHE_BASE_DIR}"
 
 reporting:
-  type: {defs.REPORTING_TYPE.value} # [file, blob, cosmosdb]
+  type: {defs.REPORTING_TYPE.value} # [file, blob, cosmosdb, documentdb]
   base_dir: "{defs.REPORTING_BASE_DIR}"
 
 output:
-  type: {defs.OUTPUT_TYPE.value} # [file, blob, cosmosdb]
+  type: {defs.OUTPUT_TYPE.value} # [file, blob, cosmosdb, documentdb]
   base_dir: "{defs.OUTPUT_BASE_DIR}"
 
 ### Workflow settings ###

graphrag/config/models/cache_config.py

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

graphrag/config/models/output_config.py

@@ -31,3 +31,6 @@ class OutputConfig(BaseModel):
     cosmosdb_account_url: str | None = Field(
         description="The cosmosdb account url to use.", default=None
     )
+    documentdb_account_url: str | None = Field(
+        description="The documentdb account url to use.", default=None
+    )

graphrag/config/models/vector_store_config.py

@@ -50,11 +50,17 @@ def _validate_url(self) -> None:
         ):
             msg = "vector_store.url is required when vector_store.type == cosmos_db. Please rerun `graphrag init` and select the correct vector store type."
             raise ValueError(msg)
+
+        if self.type == VectorStoreType.DocumentDB and (
+            self.url is None or self.url.strip() == ""
+        ):
+            msg = "vector_store.url is required when vector_store.type == document_db. Please rerun `graphrag init` and select the correct vector store type."
+            raise ValueError(msg)
 
         if self.type == VectorStoreType.LanceDB and (
             self.url is not None and self.url.strip() != ""
         ):
-            msg = "vector_store.url is only used when vector_store.type == azure_ai_search or vector_store.type == cosmos_db. Please rerun `graphrag init` and select the correct vector store type."
+            msg = "vector_store.url is only used when vector_store.type == azure_ai_search or vector_store.type == cosmos_db or vector_store.type == document_db. Please rerun `graphrag init` and select the correct vector store type."
             raise ValueError(msg)
 
     api_key: str | None = Field(
@@ -73,7 +79,7 @@ def _validate_url(self) -> None:
     )
 
     database_name: str | None = Field(
-        description="The database name to use when type == cosmos_db.", default=None
+        description="The database name to use when type == cosmos_db or document_db.", default=None
     )
 
     overwrite: bool = Field(