diff --git a/.semversioner/next-release/patch-20250306224613336336.json b/.semversioner/next-release/patch-20250306224613336336.json new file mode 100644 index 0000000000..087d1d4aa3 --- /dev/null +++ b/.semversioner/next-release/patch-20250306224613336336.json @@ -0,0 +1,4 @@ +{ + "type": "patch", + "description": "add support for cosmosdb to use url+AccountKey based authentication" +} diff --git a/graphrag/vector_stores/cosmosdb.py b/graphrag/vector_stores/cosmosdb.py index 9c736076bf..8c5ede12b7 100644 --- a/graphrag/vector_stores/cosmosdb.py +++ b/graphrag/vector_stores/cosmosdb.py @@ -40,9 +40,15 @@ def connect(self, **kwargs: Any) -> Any: if not url: msg = "Either connection_string or url must be provided." raise ValueError(msg) - self._cosmos_client = CosmosClient( - url=url, credential=DefaultAzureCredential() - ) + access_key = kwargs.get("api_key") + if access_key: + self._cosmos_client = CosmosClient( + url=url, credential=access_key + ) + else: + self._cosmos_client = CosmosClient( + url=url, credential=DefaultAzureCredential() + ) database_name = kwargs.get("database_name") if database_name is None: