azure-documentdb-haystack 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,77 @@
1
+ Metadata-Version: 2.5
2
+ Name: azure-documentdb-haystack
3
+ Version: 1.0.0
4
+ Summary: An integration of Azure DocumentDB with Haystack
5
+ Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations
6
+ Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/azure_documentdb/README.md
7
+ Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
8
+ Author-email: deepset GmbH <info@deepset.ai>
9
+ License-Expression: Apache-2.0
10
+ License-File: LICENSE.txt
11
+ Keywords: azure,documentdb,haystack,rag,vector-search
12
+ Classifier: Development Status :: 4 - Beta
13
+ Classifier: License :: OSI Approved :: Apache Software License
14
+ Classifier: Programming Language :: Python
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Programming Language :: Python :: 3.14
20
+ Classifier: Programming Language :: Python :: Implementation :: CPython
21
+ Classifier: Programming Language :: Python :: Implementation :: PyPy
22
+ Requires-Python: >=3.10
23
+ Requires-Dist: azure-identity>=1.15.0
24
+ Requires-Dist: haystack-ai>=2.28.0
25
+ Requires-Dist: pymongo[srv]>=4.13.0
26
+ Description-Content-Type: text/markdown
27
+
28
+ # azure-documentdb-haystack
29
+
30
+ [![PyPI - Version](https://img.shields.io/pypi/v/azure-documentdb-haystack.svg)](https://pypi.org/project/azure-documentdb-haystack)
31
+ [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/azure-documentdb-haystack.svg)](https://pypi.org/project/azure-documentdb-haystack)
32
+
33
+ Haystack document store and retrievers for [Azure DocumentDB](https://learn.microsoft.com/azure/documentdb/overview).
34
+
35
+ ## Installation
36
+
37
+ ```bash
38
+ pip install azure-documentdb-haystack
39
+ ```
40
+
41
+ ## Authentication
42
+
43
+ Microsoft Entra ID is the default and recommended authentication method. `DefaultAzureCredential` supports local Azure
44
+ CLI credentials, workload identity, and managed identity without changing application code. Assign the identity an
45
+ appropriate Azure DocumentDB data-plane role, then set the cluster name:
46
+
47
+ ```bash
48
+ AZURE_DOCUMENTDB_CLUSTER_NAME=my-cluster
49
+ ```
50
+
51
+ For local development and integration testing only, a connection string can be supplied through
52
+ `AZURE_DOCUMENTDB_CONNECTION_STRING`. The integration emits a warning whenever this fallback is used.
53
+
54
+ ## Usage
55
+
56
+ ```python
57
+ from haystack_integrations.document_stores.azure_documentdb import AzureDocumentDBDocumentStore
58
+
59
+ store = AzureDocumentDBDocumentStore(
60
+ database_name="haystack",
61
+ collection_name="documents",
62
+ )
63
+
64
+ # Create once. The collection must exist before initializing the store.
65
+ store.create_vector_index(dimensions=1536, kind="vector-hnsw", similarity="COS", m=16, efConstruction=64)
66
+ ```
67
+
68
+ Use `AzureDocumentDBEmbeddingRetriever` for `cosmosSearch` vector queries. `AzureDocumentDBFullTextRetriever` supports
69
+ Azure DocumentDB BM25 full-text search, which is currently a gated preview and must be enabled on the cluster.
70
+
71
+ - [Azure DocumentDB vector search](https://learn.microsoft.com/azure/documentdb/vector-search)
72
+ - [Azure DocumentDB full-text search](https://learn.microsoft.com/azure/documentdb/full-text-search-overview)
73
+ - [Changelog](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/azure_documentdb/CHANGELOG.md)
74
+
75
+ ## Contributing
76
+
77
+ Refer to the general [Contribution Guidelines](https://github.com/deepset-ai/haystack-core-integrations/blob/main/CONTRIBUTING.md).
@@ -0,0 +1,12 @@
1
+ haystack_integrations/components/retrievers/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
2
+ haystack_integrations/components/retrievers/azure_documentdb/__init__.py,sha256=B_XC2m3RSmH7v026Ue2YtjouucCjnxvXzoRbDRFm-7w,467
3
+ haystack_integrations/components/retrievers/azure_documentdb/embedding_retriever.py,sha256=QlCOMAt1leEnSivdBy7TiZYIpmz_lt9P15UHC3h-OEU,4915
4
+ haystack_integrations/components/retrievers/azure_documentdb/full_text_retriever.py,sha256=GmSZQ1rtQyUng4bL-AqneTueYq_T1P7h6OK3rirxRag,5178
5
+ haystack_integrations/document_stores/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
6
+ haystack_integrations/document_stores/azure_documentdb/__init__.py,sha256=Vqe31JYuSn1RPf58VVB2PcGmRVDt6b5Xj9IaYJjMkus,266
7
+ haystack_integrations/document_stores/azure_documentdb/document_store.py,sha256=UfOnxc6i8LvPN-soN6R1F5HhuNMqUlzmQlw82EKgLJ4,35492
8
+ haystack_integrations/document_stores/azure_documentdb/filters.py,sha256=gzP49vIVhxZc72BMJR3q5JuvWzPv0JB1IhT_PYrkIAc,3713
9
+ azure_documentdb_haystack-1.0.0.dist-info/METADATA,sha256=8XbKo8DPUJsUIZRaCpapFkdxh53awFjBpBNH21nKxgA,3494
10
+ azure_documentdb_haystack-1.0.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
11
+ azure_documentdb_haystack-1.0.0.dist-info/licenses/LICENSE.txt,sha256=ntqrug2pqJYB34vlUCCLoaSaMp20iNtHW3kVivo_jlI,10249
12
+ azure_documentdb_haystack-1.0.0.dist-info/RECORD,,
@@ -0,0 +1,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: hatchling 1.32.0
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
@@ -0,0 +1,73 @@
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
10
+
11
+ "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
12
+
13
+ "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
14
+
15
+ "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
16
+
17
+ "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
18
+
19
+ "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
20
+
21
+ "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
22
+
23
+ "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this definition, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of the Work and Derivative Works thereof.
24
+
25
+ "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
26
+
27
+ "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
28
+
29
+ 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
30
+
31
+ 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
32
+
33
+ 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
34
+
35
+ (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and
36
+
37
+ (b) You must cause any modified files to carry prominent notices stating that You changed the files; and
38
+
39
+ (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
40
+
41
+ (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
42
+
43
+ You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
44
+
45
+ 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
46
+
47
+ 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
48
+
49
+ 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
50
+
51
+ 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
52
+
53
+ 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
54
+
55
+ END OF TERMS AND CONDITIONS
56
+
57
+ APPENDIX: How to apply the Apache License to your work.
58
+
59
+ To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives.
60
+
61
+ Copyright 2026-present deepset GmbH
62
+
63
+ Licensed under the Apache License, Version 2.0 (the "License");
64
+ you may not use this file except in compliance with the License.
65
+ You may obtain a copy of the License at
66
+
67
+ http://www.apache.org/licenses/LICENSE-2.0
68
+
69
+ Unless required by applicable law or agreed to in writing, software
70
+ distributed under the License is distributed on an "AS IS" BASIS,
71
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
72
+ See the License for the specific language governing permissions and
73
+ limitations under the License.
@@ -0,0 +1,12 @@
1
+ # SPDX-FileCopyrightText: 2026-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from haystack_integrations.components.retrievers.azure_documentdb.embedding_retriever import (
6
+ AzureDocumentDBEmbeddingRetriever,
7
+ )
8
+ from haystack_integrations.components.retrievers.azure_documentdb.full_text_retriever import (
9
+ AzureDocumentDBFullTextRetriever,
10
+ )
11
+
12
+ __all__ = ["AzureDocumentDBEmbeddingRetriever", "AzureDocumentDBFullTextRetriever"]
@@ -0,0 +1,129 @@
1
+ # SPDX-FileCopyrightText: 2026-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from typing import Any
6
+
7
+ from haystack import component, default_from_dict, default_to_dict
8
+ from haystack.dataclasses import Document
9
+ from haystack.document_stores.types import FilterPolicy
10
+ from haystack.document_stores.types.filter_policy import apply_filter_policy
11
+
12
+ from haystack_integrations.document_stores.azure_documentdb import AzureDocumentDBDocumentStore
13
+
14
+
15
+ @component
16
+ class AzureDocumentDBEmbeddingRetriever:
17
+ """Retrieve documents from Azure DocumentDB using `cosmosSearch` vector similarity."""
18
+
19
+ def __init__(
20
+ self,
21
+ *,
22
+ document_store: AzureDocumentDBDocumentStore,
23
+ filters: dict[str, Any] | None = None,
24
+ top_k: int = 10,
25
+ filter_policy: str | FilterPolicy = FilterPolicy.REPLACE,
26
+ ) -> None:
27
+ """
28
+ Create the embedding retriever.
29
+
30
+ :param document_store: Azure DocumentDB document store to query.
31
+ :param filters: Default Haystack metadata filters.
32
+ :param top_k: Maximum number of documents to return.
33
+ :param filter_policy: Policy for combining initialization and runtime filters.
34
+ """
35
+ if not isinstance(document_store, AzureDocumentDBDocumentStore):
36
+ msg = "document_store must be an instance of AzureDocumentDBDocumentStore"
37
+ raise ValueError(msg)
38
+ if top_k <= 0:
39
+ msg = "top_k must be greater than zero"
40
+ raise ValueError(msg)
41
+ self.document_store = document_store
42
+ self.filters = filters or {}
43
+ self.top_k = top_k
44
+ self.filter_policy = (
45
+ filter_policy if isinstance(filter_policy, FilterPolicy) else FilterPolicy.from_str(filter_policy)
46
+ )
47
+
48
+ def to_dict(self) -> dict[str, Any]:
49
+ """
50
+ Serialize this component to a dictionary.
51
+
52
+ :returns: Serialized retriever configuration.
53
+ """
54
+ return default_to_dict(
55
+ self,
56
+ document_store=self.document_store.to_dict(),
57
+ filters=self.filters,
58
+ top_k=self.top_k,
59
+ filter_policy=self.filter_policy.value,
60
+ )
61
+
62
+ @classmethod
63
+ def from_dict(cls, data: dict[str, Any]) -> "AzureDocumentDBEmbeddingRetriever":
64
+ """
65
+ Deserialize this component from a dictionary.
66
+
67
+ :param data: Serialized retriever configuration.
68
+ :returns: The deserialized retriever.
69
+ """
70
+ data["init_parameters"]["document_store"] = AzureDocumentDBDocumentStore.from_dict(
71
+ data["init_parameters"]["document_store"]
72
+ )
73
+ if policy := data["init_parameters"].get("filter_policy"):
74
+ data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(policy)
75
+ return default_from_dict(cls, data)
76
+
77
+ def close(self) -> None:
78
+ """Release synchronous document-store resources."""
79
+ self.document_store.close()
80
+
81
+ async def close_async(self) -> None:
82
+ """Release asynchronous document-store resources."""
83
+ await self.document_store.close_async()
84
+
85
+ @component.output_types(documents=list[Document])
86
+ def run(
87
+ self,
88
+ query_embedding: list[float],
89
+ filters: dict[str, Any] | None = None,
90
+ top_k: int | None = None,
91
+ ) -> dict[str, list[Document]]:
92
+ """
93
+ Retrieve documents by vector similarity.
94
+
95
+ :param query_embedding: Query vector.
96
+ :param filters: Runtime Haystack metadata filters.
97
+ :param top_k: Runtime maximum number of documents.
98
+ :returns: A dictionary containing the retrieved `documents`.
99
+ """
100
+ effective_filters = apply_filter_policy(self.filter_policy, self.filters, filters)
101
+ documents = self.document_store._embedding_retrieval(
102
+ query_embedding=query_embedding,
103
+ filters=effective_filters,
104
+ top_k=self.top_k if top_k is None else top_k,
105
+ )
106
+ return {"documents": documents}
107
+
108
+ @component.output_types(documents=list[Document])
109
+ async def run_async(
110
+ self,
111
+ query_embedding: list[float],
112
+ filters: dict[str, Any] | None = None,
113
+ top_k: int | None = None,
114
+ ) -> dict[str, list[Document]]:
115
+ """
116
+ Asynchronously retrieve documents by vector similarity.
117
+
118
+ :param query_embedding: Query vector.
119
+ :param filters: Runtime Haystack metadata filters.
120
+ :param top_k: Runtime maximum number of documents.
121
+ :returns: A dictionary containing the retrieved `documents`.
122
+ """
123
+ effective_filters = apply_filter_policy(self.filter_policy, self.filters, filters)
124
+ documents = await self.document_store._embedding_retrieval_async(
125
+ query_embedding=query_embedding,
126
+ filters=effective_filters,
127
+ top_k=self.top_k if top_k is None else top_k,
128
+ )
129
+ return {"documents": documents}
@@ -0,0 +1,135 @@
1
+ # SPDX-FileCopyrightText: 2026-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from typing import Any
6
+
7
+ from haystack import component, default_from_dict, default_to_dict
8
+ from haystack.dataclasses import Document
9
+ from haystack.document_stores.types import FilterPolicy
10
+ from haystack.document_stores.types.filter_policy import apply_filter_policy
11
+
12
+ from haystack_integrations.document_stores.azure_documentdb import AzureDocumentDBDocumentStore
13
+
14
+
15
+ @component
16
+ class AzureDocumentDBFullTextRetriever:
17
+ """Retrieve documents using Azure DocumentDB BM25 full-text search, currently a gated preview."""
18
+
19
+ def __init__(
20
+ self,
21
+ *,
22
+ document_store: AzureDocumentDBDocumentStore,
23
+ filters: dict[str, Any] | None = None,
24
+ top_k: int = 10,
25
+ filter_policy: str | FilterPolicy = FilterPolicy.REPLACE,
26
+ ) -> None:
27
+ """
28
+ Create the full-text retriever.
29
+
30
+ :param document_store: Azure DocumentDB document store to query.
31
+ :param filters: Default Haystack metadata filters.
32
+ :param top_k: Maximum number of documents to return.
33
+ :param filter_policy: Policy for combining initialization and runtime filters.
34
+ """
35
+ if not isinstance(document_store, AzureDocumentDBDocumentStore):
36
+ msg = "document_store must be an instance of AzureDocumentDBDocumentStore"
37
+ raise ValueError(msg)
38
+ if top_k <= 0:
39
+ msg = "top_k must be greater than zero"
40
+ raise ValueError(msg)
41
+ self.document_store = document_store
42
+ self.filters = filters or {}
43
+ self.top_k = top_k
44
+ self.filter_policy = (
45
+ filter_policy if isinstance(filter_policy, FilterPolicy) else FilterPolicy.from_str(filter_policy)
46
+ )
47
+
48
+ def to_dict(self) -> dict[str, Any]:
49
+ """
50
+ Serialize this component to a dictionary.
51
+
52
+ :returns: Serialized retriever configuration.
53
+ """
54
+ return default_to_dict(
55
+ self,
56
+ document_store=self.document_store.to_dict(),
57
+ filters=self.filters,
58
+ top_k=self.top_k,
59
+ filter_policy=self.filter_policy.value,
60
+ )
61
+
62
+ @classmethod
63
+ def from_dict(cls, data: dict[str, Any]) -> "AzureDocumentDBFullTextRetriever":
64
+ """
65
+ Deserialize this component from a dictionary.
66
+
67
+ :param data: Serialized retriever configuration.
68
+ :returns: The deserialized retriever.
69
+ """
70
+ data["init_parameters"]["document_store"] = AzureDocumentDBDocumentStore.from_dict(
71
+ data["init_parameters"]["document_store"]
72
+ )
73
+ if policy := data["init_parameters"].get("filter_policy"):
74
+ data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(policy)
75
+ return default_from_dict(cls, data)
76
+
77
+ def close(self) -> None:
78
+ """Release synchronous document-store resources."""
79
+ self.document_store.close()
80
+
81
+ async def close_async(self) -> None:
82
+ """Release asynchronous document-store resources."""
83
+ await self.document_store.close_async()
84
+
85
+ @component.output_types(documents=list[Document])
86
+ def run(
87
+ self,
88
+ query: str | list[str],
89
+ fuzzy: dict[str, int] | None = None,
90
+ filters: dict[str, Any] | None = None,
91
+ top_k: int | None = None,
92
+ ) -> dict[str, list[Document]]:
93
+ """
94
+ Retrieve documents by BM25 keyword search.
95
+
96
+ :param query: Query string or strings.
97
+ :param fuzzy: Azure DocumentDB fuzzy-search options such as `maxEdits`.
98
+ :param filters: Runtime Haystack metadata filters.
99
+ :param top_k: Runtime maximum number of documents.
100
+ :returns: A dictionary containing the retrieved `documents`.
101
+ """
102
+ effective_filters = apply_filter_policy(self.filter_policy, self.filters, filters)
103
+ documents = self.document_store._full_text_retrieval(
104
+ query=query,
105
+ fuzzy=fuzzy,
106
+ filters=effective_filters,
107
+ top_k=self.top_k if top_k is None else top_k,
108
+ )
109
+ return {"documents": documents}
110
+
111
+ @component.output_types(documents=list[Document])
112
+ async def run_async(
113
+ self,
114
+ query: str | list[str],
115
+ fuzzy: dict[str, int] | None = None,
116
+ filters: dict[str, Any] | None = None,
117
+ top_k: int | None = None,
118
+ ) -> dict[str, list[Document]]:
119
+ """
120
+ Asynchronously retrieve documents by BM25 keyword search.
121
+
122
+ :param query: Query string or strings.
123
+ :param fuzzy: Azure DocumentDB fuzzy-search options such as `maxEdits`.
124
+ :param filters: Runtime Haystack metadata filters.
125
+ :param top_k: Runtime maximum number of documents.
126
+ :returns: A dictionary containing the retrieved `documents`.
127
+ """
128
+ effective_filters = apply_filter_policy(self.filter_policy, self.filters, filters)
129
+ documents = await self.document_store._full_text_retrieval_async(
130
+ query=query,
131
+ fuzzy=fuzzy,
132
+ filters=effective_filters,
133
+ top_k=self.top_k if top_k is None else top_k,
134
+ )
135
+ return {"documents": documents}
File without changes
@@ -0,0 +1,7 @@
1
+ # SPDX-FileCopyrightText: 2026-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from haystack_integrations.document_stores.azure_documentdb.document_store import AzureDocumentDBDocumentStore
6
+
7
+ __all__ = ["AzureDocumentDBDocumentStore"]
@@ -0,0 +1,765 @@
1
+ # SPDX-FileCopyrightText: 2026-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ import os
6
+ import re
7
+ from contextlib import suppress
8
+ from typing import Any, Literal
9
+
10
+ from azure.core.credentials import TokenCredential
11
+ from azure.identity import DefaultAzureCredential
12
+ from haystack import default_from_dict, default_to_dict, logging
13
+ from haystack.dataclasses import Document
14
+ from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
15
+ from haystack.document_stores.types import DuplicatePolicy
16
+ from haystack.utils import Secret, deserialize_secrets_inplace
17
+ from pymongo import AsyncMongoClient, InsertOne, MongoClient, ReplaceOne, UpdateOne
18
+ from pymongo.asynchronous.collection import AsyncCollection
19
+ from pymongo.auth_oidc import OIDCCallback, OIDCCallbackContext, OIDCCallbackResult
20
+ from pymongo.collection import Collection
21
+ from pymongo.driver_info import DriverInfo
22
+ from pymongo.errors import BulkWriteError
23
+
24
+ from haystack_integrations.document_stores.azure_documentdb.filters import _normalize_filters
25
+
26
+ logger = logging.getLogger(__name__)
27
+
28
+ _DOCUMENTDB_TOKEN_SCOPE = "https://ossrdbms-aad.database.windows.net/.default"
29
+ _DOCUMENTDB_HOST_SUFFIX = "global.mongocluster.cosmos.azure.com"
30
+ _DRIVER_INFO = DriverInfo(name="AzureDocumentDBHaystackIntegration")
31
+
32
+
33
+ class AzureIdentityTokenCallback(OIDCCallback):
34
+ """Fetch Microsoft Entra access tokens for PyMongo's OIDC authentication."""
35
+
36
+ def __init__(self, credential: TokenCredential) -> None:
37
+ self._credential = credential
38
+
39
+ def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult: # noqa: ARG002
40
+ """
41
+ Fetch an access token for Azure DocumentDB.
42
+
43
+ :param context: PyMongo OIDC callback context.
44
+ :returns: The OIDC callback result containing a Microsoft Entra access token.
45
+ """
46
+ token = self._credential.get_token(_DOCUMENTDB_TOKEN_SCOPE)
47
+ return OIDCCallbackResult(access_token=token.token)
48
+
49
+
50
+ class AzureDocumentDBDocumentStore:
51
+ """
52
+ A Haystack document store backed by Azure DocumentDB.
53
+
54
+ The default authentication mode uses Microsoft Entra ID through `DefaultAzureCredential`. Supply the Azure
55
+ DocumentDB cluster name with `cluster_name` or the `AZURE_DOCUMENTDB_CLUSTER_NAME` environment variable.
56
+
57
+ A connection string can be supplied through `mongo_connection_string` or
58
+ `AZURE_DOCUMENTDB_CONNECTION_STRING` for local development and integration tests. Connection strings can contain
59
+ credentials and aren't recommended for production workloads.
60
+
61
+ The collection must already exist. For embedding retrieval, create a `cosmosSearch` vector index by calling
62
+ `create_vector_index` or provisioning it separately. Filtered vector search also requires a regular index for
63
+ every filtered metadata field, such as `meta.category`. Values used with `>`, `>=`, `<`, or `<=` must be numbers
64
+ or ISO-formatted date strings.
65
+
66
+ Usage:
67
+
68
+ ```python
69
+ from haystack_integrations.document_stores.azure_documentdb import AzureDocumentDBDocumentStore
70
+
71
+ document_store = AzureDocumentDBDocumentStore(database_name="haystack", collection_name="documents")
72
+ document_store.create_vector_index(dimensions=1536, similarity="COS")
73
+ ```
74
+ """
75
+
76
+ def __init__(
77
+ self,
78
+ *,
79
+ database_name: str,
80
+ collection_name: str,
81
+ vector_search_index: str = "haystack_vector_index",
82
+ full_text_search_index: str | None = None,
83
+ cluster_name: str | None = None,
84
+ mongo_connection_string: Secret | None = Secret.from_env_var( # noqa: B008
85
+ "AZURE_DOCUMENTDB_CONNECTION_STRING", strict=False
86
+ ),
87
+ azure_token_credential: TokenCredential | None = None,
88
+ embedding_field: str = "embedding",
89
+ content_field: str = "content",
90
+ ) -> None:
91
+ """
92
+ Create an Azure DocumentDB document store.
93
+
94
+ :param database_name: Name of the existing database.
95
+ :param collection_name: Name of the existing collection.
96
+ :param vector_search_index: Name used when creating the vector index. Azure DocumentDB selects vector indexes
97
+ by path at query time, so this name is not included in vector search queries.
98
+ :param full_text_search_index: Name of an Azure DocumentDB full-text search index. Full-text search is currently
99
+ a gated preview and must be enabled on the cluster before using the full-text retriever.
100
+ :param cluster_name: Azure DocumentDB cluster name. If omitted, `AZURE_DOCUMENTDB_CLUSTER_NAME` is used.
101
+ :param mongo_connection_string: Optional MongoDB connection string intended only for local development and
102
+ integration tests. Microsoft Entra authentication is used when this value is absent.
103
+ :param azure_token_credential: Azure credential used for Microsoft Entra authentication. If omitted,
104
+ `DefaultAzureCredential` is used.
105
+ :param embedding_field: Field containing document embeddings.
106
+ :param content_field: Field containing document content.
107
+ :raises ValueError: If database, collection, or field names are invalid.
108
+ """
109
+ for name, value in (("database_name", database_name), ("collection_name", collection_name)):
110
+ if not value or not re.fullmatch(r"[A-Za-z0-9_-]+", value):
111
+ msg = f"Invalid {name}: {value!r}. It can only contain letters, numbers, hyphens, or underscores."
112
+ raise ValueError(msg)
113
+ if not embedding_field or embedding_field.startswith("$"):
114
+ msg = "embedding_field must be a non-empty MongoDB field path and cannot start with '$'."
115
+ raise ValueError(msg)
116
+ if not content_field or content_field.startswith("$"):
117
+ msg = "content_field must be a non-empty MongoDB field path and cannot start with '$'."
118
+ raise ValueError(msg)
119
+
120
+ self.database_name = database_name
121
+ self.collection_name = collection_name
122
+ self.vector_search_index = vector_search_index
123
+ self.full_text_search_index = full_text_search_index
124
+ self.cluster_name = cluster_name
125
+ self.mongo_connection_string = mongo_connection_string
126
+ self.azure_token_credential = azure_token_credential
127
+ self.embedding_field = embedding_field
128
+ self.content_field = content_field
129
+
130
+ self._connection: MongoClient | None = None
131
+ self._connection_async: AsyncMongoClient | None = None
132
+ self._collection: Collection | None = None
133
+ self._collection_async: AsyncCollection | None = None
134
+ self._credential: TokenCredential | None = None
135
+
136
+ def _client_kwargs(self) -> tuple[str, dict[str, Any]]:
137
+ connection_string = self.mongo_connection_string.resolve_value() if self.mongo_connection_string else None
138
+ if connection_string:
139
+ logger.warning(
140
+ "Azure DocumentDB is using connection-string authentication. This fallback is intended only for "
141
+ "local development and integration tests. Use Microsoft Entra managed identity in production."
142
+ )
143
+ return connection_string, {"retryWrites": False, "driver": _DRIVER_INFO}
144
+
145
+ cluster_name = self.cluster_name or os.getenv("AZURE_DOCUMENTDB_CLUSTER_NAME")
146
+ if not cluster_name:
147
+ msg = (
148
+ "Azure DocumentDB cluster name is required. Set `cluster_name` or the "
149
+ "AZURE_DOCUMENTDB_CLUSTER_NAME environment variable."
150
+ )
151
+ raise DocumentStoreError(msg)
152
+
153
+ if self._credential is None:
154
+ self._credential = self.azure_token_credential or DefaultAzureCredential()
155
+ callback = AzureIdentityTokenCallback(self._credential)
156
+ uri = f"mongodb+srv://{cluster_name}.{_DOCUMENTDB_HOST_SUFFIX}/"
157
+ return uri, {
158
+ "authMechanism": "MONGODB-OIDC",
159
+ "authMechanismProperties": {"OIDC_CALLBACK": callback},
160
+ "retryWrites": False,
161
+ "tls": True,
162
+ "driver": _DRIVER_INFO,
163
+ }
164
+
165
+ def _connection_is_valid(self, connection: MongoClient) -> bool:
166
+ try:
167
+ connection.admin.command("ping")
168
+ return True
169
+ except Exception as error:
170
+ logger.error("Connection to Azure DocumentDB failed: {error}", error=error)
171
+ return False
172
+
173
+ async def _connection_is_valid_async(self, connection: AsyncMongoClient) -> bool:
174
+ try:
175
+ await connection.admin.command("ping")
176
+ return True
177
+ except Exception as error:
178
+ logger.error("Connection to Azure DocumentDB failed: {error}", error=error)
179
+ return False
180
+
181
+ def _ensure_connection_setup(self) -> None:
182
+ if self._connection is None:
183
+ uri, kwargs = self._client_kwargs()
184
+ self._connection = MongoClient(uri, **kwargs)
185
+ if not self._connection_is_valid(self._connection):
186
+ msg = "Connection to Azure DocumentDB failed."
187
+ raise DocumentStoreError(msg)
188
+ database = self._connection[self.database_name]
189
+ if self.collection_name not in database.list_collection_names():
190
+ msg = f"Collection '{self.collection_name}' does not exist in database '{self.database_name}'."
191
+ raise DocumentStoreError(msg)
192
+ self._collection = database[self.collection_name]
193
+ self._collection.create_index("id", unique=True)
194
+
195
+ async def _ensure_connection_setup_async(self) -> None:
196
+ if self._connection_async is None:
197
+ uri, kwargs = self._client_kwargs()
198
+ self._connection_async = AsyncMongoClient(uri, **kwargs)
199
+ if not await self._connection_is_valid_async(self._connection_async):
200
+ msg = "Connection to Azure DocumentDB failed."
201
+ raise DocumentStoreError(msg)
202
+ database = self._connection_async[self.database_name]
203
+ if self.collection_name not in await database.list_collection_names():
204
+ msg = f"Collection '{self.collection_name}' does not exist in database '{self.database_name}'."
205
+ raise DocumentStoreError(msg)
206
+ self._collection_async = database[self.collection_name]
207
+ await self._collection_async.create_index("id", unique=True)
208
+
209
+ @property
210
+ def connection(self) -> MongoClient | AsyncMongoClient:
211
+ """
212
+ Return the active Azure DocumentDB client.
213
+
214
+ :returns: The synchronous or asynchronous PyMongo client.
215
+ :raises DocumentStoreError: If no connection has been established.
216
+ """
217
+ if self._connection is not None:
218
+ return self._connection
219
+ if self._connection_async is not None:
220
+ return self._connection_async
221
+ msg = "The connection is not established yet."
222
+ raise DocumentStoreError(msg)
223
+
224
+ @property
225
+ def collection(self) -> Collection | AsyncCollection:
226
+ """
227
+ Return the active Azure DocumentDB collection.
228
+
229
+ :returns: The synchronous or asynchronous PyMongo collection.
230
+ :raises DocumentStoreError: If no collection has been initialized.
231
+ """
232
+ if self._collection is not None:
233
+ return self._collection
234
+ if self._collection_async is not None:
235
+ return self._collection_async
236
+ msg = "The collection is not established yet."
237
+ raise DocumentStoreError(msg)
238
+
239
+ def close(self) -> None:
240
+ """Release synchronous client resources."""
241
+ if self._connection is not None:
242
+ with suppress(Exception):
243
+ self._connection.close()
244
+ self._connection = None
245
+ self._collection = None
246
+
247
+ async def close_async(self) -> None:
248
+ """Release asynchronous client resources."""
249
+ if self._connection_async is not None:
250
+ with suppress(Exception):
251
+ await self._connection_async.close()
252
+ self._connection_async = None
253
+ self._collection_async = None
254
+
255
+ def to_dict(self) -> dict[str, Any]:
256
+ """
257
+ Serialize this document store to a dictionary.
258
+
259
+ :returns: Serialized document-store configuration.
260
+ """
261
+ if self.azure_token_credential is not None:
262
+ logger.warning(
263
+ "AzureDocumentDBDocumentStore was initialized with `azure_token_credential`, which cannot be "
264
+ "serialized and must be provided again after deserialization."
265
+ )
266
+ return default_to_dict(
267
+ self,
268
+ database_name=self.database_name,
269
+ collection_name=self.collection_name,
270
+ vector_search_index=self.vector_search_index,
271
+ full_text_search_index=self.full_text_search_index,
272
+ cluster_name=self.cluster_name,
273
+ mongo_connection_string=self.mongo_connection_string.to_dict() if self.mongo_connection_string else None,
274
+ embedding_field=self.embedding_field,
275
+ content_field=self.content_field,
276
+ )
277
+
278
+ @classmethod
279
+ def from_dict(cls, data: dict[str, Any]) -> "AzureDocumentDBDocumentStore":
280
+ """
281
+ Deserialize this document store from a dictionary.
282
+
283
+ :param data: Serialized document-store configuration.
284
+ :returns: The deserialized document store.
285
+ """
286
+ deserialize_secrets_inplace(data["init_parameters"], keys=["mongo_connection_string"])
287
+ return default_from_dict(cls, data)
288
+
289
+ def count_documents(self) -> int:
290
+ """
291
+ Return the number of documents in the store.
292
+
293
+ :returns: The number of documents.
294
+ """
295
+ self._ensure_connection_setup()
296
+ assert self._collection is not None
297
+ return self._collection.count_documents({})
298
+
299
+ async def count_documents_async(self) -> int:
300
+ """
301
+ Asynchronously return the number of documents in the store.
302
+
303
+ :returns: The number of documents.
304
+ """
305
+ await self._ensure_connection_setup_async()
306
+ assert self._collection_async is not None
307
+ return await self._collection_async.count_documents({})
308
+
309
+ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]:
310
+ """
311
+ Return documents matching Haystack metadata filters.
312
+
313
+ :param filters: Haystack metadata filters. Strings in ordered comparisons must be ISO-formatted dates.
314
+ :returns: Documents matching the filters.
315
+ """
316
+ self._ensure_connection_setup()
317
+ assert self._collection is not None
318
+ query = _normalize_filters(filters) if filters else {}
319
+ return [self._mongo_doc_to_haystack_doc(doc) for doc in self._collection.find(query)]
320
+
321
+ async def filter_documents_async(self, filters: dict[str, Any] | None = None) -> list[Document]:
322
+ """
323
+ Asynchronously return documents matching Haystack metadata filters.
324
+
325
+ :param filters: Haystack metadata filters. Strings in ordered comparisons must be ISO-formatted dates.
326
+ :returns: Documents matching the filters.
327
+ """
328
+ await self._ensure_connection_setup_async()
329
+ assert self._collection_async is not None
330
+ query = _normalize_filters(filters) if filters else {}
331
+ documents = await self._collection_async.find(query).to_list(length=None)
332
+ return [self._mongo_doc_to_haystack_doc(doc) for doc in documents]
333
+
334
+ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE) -> int:
335
+ """
336
+ Write documents to Azure DocumentDB using the requested duplicate policy.
337
+
338
+ :param documents: Documents to write.
339
+ :param policy: How to handle documents whose IDs already exist.
340
+ :returns: The number of documents written.
341
+ :raises ValueError: If `documents` contains an object that is not a `Document`.
342
+ :raises DuplicateDocumentError: If a duplicate ID is written with `DuplicatePolicy.FAIL`.
343
+ """
344
+ self._ensure_connection_setup()
345
+ assert self._collection is not None
346
+ return self._write_documents(self._collection, documents, policy)
347
+
348
+ def _write_documents(self, collection: Collection, documents: list[Document], policy: DuplicatePolicy) -> int:
349
+ existing = 0
350
+ if policy == DuplicatePolicy.SKIP:
351
+ existing = collection.count_documents({"id": {"$in": [document.id for document in documents]}})
352
+ operations, written = self._prepare_write_operations(documents, policy, existing)
353
+ if not operations:
354
+ return 0
355
+ try:
356
+ collection.bulk_write(operations)
357
+ except BulkWriteError as error:
358
+ details = error.details.get("writeErrors", []) if error.details else []
359
+ msg = f"Duplicate documents found: {details}"
360
+ raise DuplicateDocumentError(msg) from error
361
+ return written
362
+
363
+ def _prepare_write_operations(
364
+ self, documents: list[Document], policy: DuplicatePolicy, existing: int = 0
365
+ ) -> tuple[list[InsertOne[dict[str, Any]] | ReplaceOne[dict[str, Any]] | UpdateOne], int]:
366
+ if any(not isinstance(document, Document) for document in documents):
367
+ msg = "param 'documents' must contain a list of objects of type Document"
368
+ raise ValueError(msg)
369
+ if not documents:
370
+ return [], 0
371
+ if policy == DuplicatePolicy.NONE:
372
+ policy = DuplicatePolicy.FAIL
373
+ mongo_documents = [self._haystack_doc_to_mongo_doc(document) for document in documents]
374
+ operations: list[InsertOne[dict[str, Any]] | ReplaceOne[dict[str, Any]] | UpdateOne]
375
+ if policy == DuplicatePolicy.SKIP:
376
+ operations = [UpdateOne({"id": doc["id"]}, {"$setOnInsert": doc}, upsert=True) for doc in mongo_documents]
377
+ written = len(documents) - existing
378
+ elif policy == DuplicatePolicy.FAIL:
379
+ operations = [InsertOne(doc) for doc in mongo_documents]
380
+ written = len(documents)
381
+ else:
382
+ operations = [ReplaceOne({"id": doc["id"]}, doc, upsert=True) for doc in mongo_documents]
383
+ written = len(documents)
384
+ return operations, written
385
+
386
+ async def write_documents_async(
387
+ self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE
388
+ ) -> int:
389
+ """
390
+ Asynchronously write documents using the requested duplicate policy.
391
+
392
+ :param documents: Documents to write.
393
+ :param policy: How to handle documents whose IDs already exist.
394
+ :returns: The number of documents written.
395
+ :raises ValueError: If `documents` contains an object that is not a `Document`.
396
+ :raises DuplicateDocumentError: If a duplicate ID is written with `DuplicatePolicy.FAIL`.
397
+ """
398
+ await self._ensure_connection_setup_async()
399
+ assert self._collection_async is not None
400
+ existing = 0
401
+ if policy == DuplicatePolicy.SKIP:
402
+ existing = await self._collection_async.count_documents(
403
+ {"id": {"$in": [document.id for document in documents]}}
404
+ )
405
+ operations, written = self._prepare_write_operations(documents, policy, existing)
406
+ if not operations:
407
+ return 0
408
+ try:
409
+ await self._collection_async.bulk_write(operations)
410
+ except BulkWriteError as error:
411
+ details = error.details.get("writeErrors", []) if error.details else []
412
+ msg = f"Duplicate documents found: {details}"
413
+ raise DuplicateDocumentError(msg) from error
414
+ return written
415
+
416
+ def delete_documents(self, document_ids: list[str]) -> None:
417
+ """
418
+ Delete documents with matching Haystack IDs.
419
+
420
+ :param document_ids: IDs of documents to delete.
421
+ """
422
+ self._ensure_connection_setup()
423
+ assert self._collection is not None
424
+ if document_ids:
425
+ self._collection.delete_many({"id": {"$in": document_ids}})
426
+
427
+ async def delete_documents_async(self, document_ids: list[str]) -> None:
428
+ """
429
+ Asynchronously delete documents with matching Haystack IDs.
430
+
431
+ :param document_ids: IDs of documents to delete.
432
+ """
433
+ await self._ensure_connection_setup_async()
434
+ assert self._collection_async is not None
435
+ if document_ids:
436
+ await self._collection_async.delete_many({"id": {"$in": document_ids}})
437
+
438
+ def delete_by_filter(self, filters: dict[str, Any]) -> int:
439
+ """
440
+ Delete documents matching filters.
441
+
442
+ :param filters: Haystack metadata filters selecting documents to delete.
443
+ :returns: The number of documents deleted.
444
+ """
445
+ self._ensure_connection_setup()
446
+ assert self._collection is not None
447
+ result = self._collection.delete_many(_normalize_filters(filters))
448
+ return result.deleted_count
449
+
450
+ async def delete_by_filter_async(self, filters: dict[str, Any]) -> int:
451
+ """
452
+ Asynchronously delete documents matching filters.
453
+
454
+ :param filters: Haystack metadata filters selecting documents to delete.
455
+ :returns: The number of documents deleted.
456
+ """
457
+ await self._ensure_connection_setup_async()
458
+ assert self._collection_async is not None
459
+ result = await self._collection_async.delete_many(_normalize_filters(filters))
460
+ return result.deleted_count
461
+
462
+ def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int:
463
+ """
464
+ Update metadata on documents matching filters.
465
+
466
+ :param filters: Haystack metadata filters selecting documents to update.
467
+ :param meta: Metadata fields and values to set.
468
+ :returns: The number of documents updated.
469
+ """
470
+ self._ensure_connection_setup()
471
+ assert self._collection is not None
472
+ update = {"$set": {f"meta.{key}": value for key, value in meta.items()}}
473
+ result = self._collection.update_many(_normalize_filters(filters), update)
474
+ return result.modified_count
475
+
476
+ async def update_by_filter_async(self, filters: dict[str, Any], meta: dict[str, Any]) -> int:
477
+ """
478
+ Asynchronously update metadata on documents matching filters.
479
+
480
+ :param filters: Haystack metadata filters selecting documents to update.
481
+ :param meta: Metadata fields and values to set.
482
+ :returns: The number of documents updated.
483
+ """
484
+ await self._ensure_connection_setup_async()
485
+ assert self._collection_async is not None
486
+ update = {"$set": {f"meta.{key}": value for key, value in meta.items()}}
487
+ result = await self._collection_async.update_many(_normalize_filters(filters), update)
488
+ return result.modified_count
489
+
490
+ def delete_all_documents(self, *, recreate_collection: bool = False) -> None:
491
+ """
492
+ Delete all documents, optionally recreating the collection.
493
+
494
+ :param recreate_collection: Drop and recreate the collection instead of deleting documents individually.
495
+ """
496
+ self._ensure_connection_setup()
497
+ assert self._collection is not None
498
+ assert self._connection is not None
499
+ if recreate_collection:
500
+ database = self._connection[self.database_name]
501
+ collection_info = database.list_collections(filter={"name": self.collection_name})
502
+ options = next(collection_info, {}).get("options", {})
503
+ indexes = [index for index in self._collection.list_indexes() if index["name"] != "_id_"]
504
+ self._collection.drop()
505
+ database.create_collection(self.collection_name, **options)
506
+ self._collection = database[self.collection_name]
507
+ for index in indexes:
508
+ keys = list(index["key"].items())
509
+ index_options = {key: value for key, value in index.items() if key not in {"key", "v", "ns"}}
510
+ self._collection.create_index(keys, **index_options)
511
+ else:
512
+ self._collection.delete_many({})
513
+
514
+ async def delete_all_documents_async(self, *, recreate_collection: bool = False) -> None:
515
+ """
516
+ Asynchronously delete all documents, optionally recreating the collection.
517
+
518
+ :param recreate_collection: Drop and recreate the collection instead of deleting documents individually.
519
+ """
520
+ await self._ensure_connection_setup_async()
521
+ assert self._collection_async is not None
522
+ assert self._connection_async is not None
523
+ if recreate_collection:
524
+ database = self._connection_async[self.database_name]
525
+ collection_info = await database.list_collections(filter={"name": self.collection_name})
526
+ collection_info_list = await collection_info.to_list(length=1)
527
+ options = collection_info_list[0].get("options", {}) if collection_info_list else {}
528
+ indexes_cursor = await self._collection_async.list_indexes()
529
+ indexes = [index for index in await indexes_cursor.to_list(length=None) if index["name"] != "_id_"]
530
+ await self._collection_async.drop()
531
+ await database.create_collection(self.collection_name, **options)
532
+ self._collection_async = database[self.collection_name]
533
+ for index in indexes:
534
+ keys = list(index["key"].items())
535
+ index_options = {key: value for key, value in index.items() if key not in {"key", "v", "ns"}}
536
+ await self._collection_async.create_index(keys, **index_options)
537
+ else:
538
+ await self._collection_async.delete_many({})
539
+
540
+ def create_vector_index(
541
+ self,
542
+ *,
543
+ dimensions: int,
544
+ similarity: Literal["COS", "L2", "IP"] = "COS",
545
+ kind: Literal["vector-ivf", "vector-hnsw", "vector-diskann"] = "vector-hnsw",
546
+ **index_options: Any,
547
+ ) -> None:
548
+ """
549
+ Create the configured Azure DocumentDB `cosmosSearch` vector index.
550
+
551
+ :param dimensions: Number of dimensions in each embedding.
552
+ :param similarity: Similarity metric: cosine (`COS`), Euclidean (`L2`), or inner product (`IP`).
553
+ :param kind: Vector index algorithm.
554
+ :param index_options: Algorithm-specific Azure DocumentDB index options.
555
+ :raises ValueError: If `dimensions` is not positive.
556
+ :raises DocumentStoreError: If index creation fails.
557
+ """
558
+ if dimensions <= 0:
559
+ msg = "dimensions must be greater than zero"
560
+ raise ValueError(msg)
561
+ self._ensure_connection_setup()
562
+ assert self._connection is not None
563
+ options = {"kind": kind, "dimensions": dimensions, "similarity": similarity, **index_options}
564
+ command = {
565
+ "createIndexes": self.collection_name,
566
+ "indexes": [
567
+ {
568
+ "name": self.vector_search_index,
569
+ "key": {self.embedding_field: "cosmosSearch"},
570
+ "cosmosSearchOptions": options,
571
+ }
572
+ ],
573
+ }
574
+ try:
575
+ self._connection[self.database_name].command(command)
576
+ except Exception as error:
577
+ msg = f"Failed to create Azure DocumentDB vector index: {error}"
578
+ raise DocumentStoreError(msg) from error
579
+
580
+ async def create_vector_index_async(
581
+ self,
582
+ *,
583
+ dimensions: int,
584
+ similarity: Literal["COS", "L2", "IP"] = "COS",
585
+ kind: Literal["vector-ivf", "vector-hnsw", "vector-diskann"] = "vector-hnsw",
586
+ **index_options: Any,
587
+ ) -> None:
588
+ """
589
+ Asynchronously create the configured `cosmosSearch` vector index.
590
+
591
+ :param dimensions: Number of dimensions in each embedding.
592
+ :param similarity: Similarity metric: cosine (`COS`), Euclidean (`L2`), or inner product (`IP`).
593
+ :param kind: Vector index algorithm.
594
+ :param index_options: Algorithm-specific Azure DocumentDB index options.
595
+ :raises ValueError: If `dimensions` is not positive.
596
+ :raises DocumentStoreError: If index creation fails.
597
+ """
598
+ if dimensions <= 0:
599
+ msg = "dimensions must be greater than zero"
600
+ raise ValueError(msg)
601
+ await self._ensure_connection_setup_async()
602
+ assert self._connection_async is not None
603
+ options = {"kind": kind, "dimensions": dimensions, "similarity": similarity, **index_options}
604
+ command = {
605
+ "createIndexes": self.collection_name,
606
+ "indexes": [
607
+ {
608
+ "name": self.vector_search_index,
609
+ "key": {self.embedding_field: "cosmosSearch"},
610
+ "cosmosSearchOptions": options,
611
+ }
612
+ ],
613
+ }
614
+ try:
615
+ await self._connection_async[self.database_name].command(command)
616
+ except Exception as error:
617
+ msg = f"Failed to create Azure DocumentDB vector index: {error}"
618
+ raise DocumentStoreError(msg) from error
619
+
620
+ def _embedding_pipeline(
621
+ self, query_embedding: list[float], filters: dict[str, Any] | None, top_k: int
622
+ ) -> list[dict[str, Any]]:
623
+ if not query_embedding:
624
+ msg = "Query embedding must not be empty"
625
+ raise ValueError(msg)
626
+ if top_k <= 0:
627
+ msg = "top_k must be greater than zero"
628
+ raise ValueError(msg)
629
+ cosmos_search: dict[str, Any] = {
630
+ "vector": query_embedding,
631
+ "path": self.embedding_field,
632
+ "k": top_k,
633
+ }
634
+ if filters:
635
+ cosmos_search["filter"] = _normalize_filters(filters)
636
+ return [
637
+ {"$search": {"cosmosSearch": cosmos_search, "returnStoredSource": True}},
638
+ {"$project": {"document": "$$ROOT", "score": {"$meta": "searchScore"}}},
639
+ ]
640
+
641
+ def _embedding_retrieval(
642
+ self, query_embedding: list[float], filters: dict[str, Any] | None = None, top_k: int = 10
643
+ ) -> list[Document]:
644
+ """Retrieve documents by Azure DocumentDB vector similarity."""
645
+ pipeline = self._embedding_pipeline(query_embedding, filters, top_k)
646
+ self._ensure_connection_setup()
647
+ assert self._collection is not None
648
+ try:
649
+ results = list(self._collection.aggregate(pipeline))
650
+ except Exception as error:
651
+ msg = f"Vector retrieval from Azure DocumentDB failed: {error}"
652
+ raise DocumentStoreError(msg) from error
653
+ return [self._search_result_to_haystack_doc(result) for result in results]
654
+
655
+ async def _embedding_retrieval_async(
656
+ self, query_embedding: list[float], filters: dict[str, Any] | None = None, top_k: int = 10
657
+ ) -> list[Document]:
658
+ """Asynchronously retrieve documents by Azure DocumentDB vector similarity."""
659
+ pipeline = self._embedding_pipeline(query_embedding, filters, top_k)
660
+ await self._ensure_connection_setup_async()
661
+ assert self._collection_async is not None
662
+ try:
663
+ cursor = await self._collection_async.aggregate(pipeline)
664
+ results = await cursor.to_list(length=None)
665
+ except Exception as error:
666
+ msg = f"Vector retrieval from Azure DocumentDB failed: {error}"
667
+ raise DocumentStoreError(msg) from error
668
+ return [self._search_result_to_haystack_doc(result) for result in results]
669
+
670
+ def _full_text_pipeline(
671
+ self, query: str | list[str], fuzzy: dict[str, int] | None, filters: dict[str, Any] | None, top_k: int
672
+ ) -> list[dict[str, Any]]:
673
+ if not query:
674
+ msg = "Argument query must not be empty."
675
+ raise ValueError(msg)
676
+ if not self.full_text_search_index:
677
+ msg = "full_text_search_index must be configured to use full-text retrieval."
678
+ raise ValueError(msg)
679
+ if top_k <= 0:
680
+ msg = "top_k must be greater than zero"
681
+ raise ValueError(msg)
682
+ text_search: dict[str, Any] = {"query": query, "path": self.content_field}
683
+ if fuzzy:
684
+ text_search["fuzzy"] = fuzzy
685
+ pipeline: list[dict[str, Any]] = [
686
+ {"$search": {"index": self.full_text_search_index, "text": text_search}},
687
+ ]
688
+ if filters:
689
+ pipeline.append({"$match": _normalize_filters(filters)})
690
+ pipeline.extend(
691
+ [
692
+ {"$limit": top_k},
693
+ {"$addFields": {"score": {"$meta": "searchScore"}}},
694
+ {"$project": {"_id": 0}},
695
+ ]
696
+ )
697
+ return pipeline
698
+
699
+ def _full_text_retrieval(
700
+ self,
701
+ query: str | list[str],
702
+ fuzzy: dict[str, int] | None = None,
703
+ filters: dict[str, Any] | None = None,
704
+ top_k: int = 10,
705
+ ) -> list[Document]:
706
+ """Retrieve documents with Azure DocumentDB BM25 full-text search (gated preview)."""
707
+ pipeline = self._full_text_pipeline(query, fuzzy, filters, top_k)
708
+ self._ensure_connection_setup()
709
+ assert self._collection is not None
710
+ try:
711
+ results = list(self._collection.aggregate(pipeline))
712
+ except Exception as error:
713
+ msg = f"Full-text retrieval from Azure DocumentDB failed: {error}"
714
+ raise DocumentStoreError(msg) from error
715
+ return [self._mongo_doc_to_haystack_doc(result) for result in results]
716
+
717
+ async def _full_text_retrieval_async(
718
+ self,
719
+ query: str | list[str],
720
+ fuzzy: dict[str, int] | None = None,
721
+ filters: dict[str, Any] | None = None,
722
+ top_k: int = 10,
723
+ ) -> list[Document]:
724
+ """Asynchronously retrieve documents with Azure DocumentDB BM25 search (gated preview)."""
725
+ pipeline = self._full_text_pipeline(query, fuzzy, filters, top_k)
726
+ await self._ensure_connection_setup_async()
727
+ assert self._collection_async is not None
728
+ try:
729
+ cursor = await self._collection_async.aggregate(pipeline)
730
+ results = await cursor.to_list(length=None)
731
+ except Exception as error:
732
+ msg = f"Full-text retrieval from Azure DocumentDB failed: {error}"
733
+ raise DocumentStoreError(msg) from error
734
+ return [self._mongo_doc_to_haystack_doc(result) for result in results]
735
+
736
+ def _search_result_to_haystack_doc(self, result: dict[str, Any]) -> Document:
737
+ document = dict(result.get("document", result))
738
+ if "score" in result:
739
+ document["score"] = result["score"]
740
+ return self._mongo_doc_to_haystack_doc(document)
741
+
742
+ def _mongo_doc_to_haystack_doc(self, mongo_doc: dict[str, Any]) -> Document:
743
+ document = dict(mongo_doc)
744
+ document.pop("_id", None)
745
+ if self.content_field != "content":
746
+ document["content"] = document.pop(self.content_field, None)
747
+ if self.embedding_field != "embedding":
748
+ document["embedding"] = document.pop(self.embedding_field, None)
749
+ return Document.from_dict(document)
750
+
751
+ def _haystack_doc_to_mongo_doc(self, haystack_doc: Document) -> dict[str, Any]:
752
+ document = haystack_doc.to_dict(flatten=False)
753
+ if self.content_field != "content":
754
+ document[self.content_field] = document.pop("content", None)
755
+ if self.embedding_field != "embedding":
756
+ document[self.embedding_field] = document.pop("embedding", None)
757
+ sparse_embedding = document.pop("sparse_embedding", None)
758
+ if sparse_embedding:
759
+ logger.warning(
760
+ "Document {id} has a sparse embedding, but Azure DocumentDB integration does not support sparse "
761
+ "embeddings. The field will be ignored.",
762
+ id=haystack_doc.id,
763
+ )
764
+ document.pop("_id", None)
765
+ return document
@@ -0,0 +1,100 @@
1
+ # SPDX-FileCopyrightText: 2026-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from collections.abc import Callable
6
+ from datetime import datetime
7
+ from typing import Any
8
+
9
+ from haystack.errors import FilterError
10
+
11
+
12
+ def _normalize_filters(filters: dict[str, Any]) -> dict[str, Any]:
13
+ """Convert Haystack filters to MongoDB-compatible Azure DocumentDB filters."""
14
+ if not isinstance(filters, dict):
15
+ msg = "Filters must be a dictionary"
16
+ raise FilterError(msg)
17
+ if "field" in filters:
18
+ return _parse_comparison_condition(filters)
19
+ return _parse_logical_condition(filters)
20
+
21
+
22
+ def _parse_logical_condition(condition: dict[str, Any]) -> dict[str, Any]:
23
+ if "operator" not in condition:
24
+ msg = f"'operator' key missing in {condition}"
25
+ raise FilterError(msg)
26
+ if "conditions" not in condition:
27
+ msg = f"'conditions' key missing in {condition}"
28
+ raise FilterError(msg)
29
+ conditions = [
30
+ _parse_comparison_condition(item) if "field" in item else _parse_logical_condition(item)
31
+ for item in condition["conditions"]
32
+ ]
33
+ operator = condition["operator"]
34
+ if operator == "AND":
35
+ return {"$and": conditions}
36
+ if operator == "OR":
37
+ return {"$or": conditions}
38
+ if operator == "NOT":
39
+ return {"$nor": [{"$and": conditions}]}
40
+ msg = f"Unknown logical operator '{operator}'. Valid operators are: 'AND', 'OR', 'NOT'"
41
+ raise FilterError(msg)
42
+
43
+
44
+ def _parse_comparison_condition(condition: dict[str, Any]) -> dict[str, Any]:
45
+ for key in ("field", "operator", "value"):
46
+ if key not in condition:
47
+ msg = f"'{key}' key missing in {condition}"
48
+ raise FilterError(msg)
49
+ operator = condition["operator"]
50
+ if operator not in COMPARISON_OPERATORS:
51
+ msg = f"Unknown comparison operator '{operator}'"
52
+ raise FilterError(msg)
53
+ return COMPARISON_OPERATORS[operator](condition["field"], condition["value"])
54
+
55
+
56
+ def _comparison(mongo_operator: str) -> Callable[[str, Any], dict[str, Any]]:
57
+ def convert(field: str, value: Any) -> dict[str, Any]:
58
+ return {field: {mongo_operator: value}}
59
+
60
+ return convert
61
+
62
+
63
+ def _ordered_comparison(mongo_operator: str) -> Callable[[str, Any], dict[str, Any]]:
64
+ def convert(field: str, value: Any) -> dict[str, Any]:
65
+ if isinstance(value, list):
66
+ msg = f"Can't compare {type(value)} using ordered comparison operators."
67
+ raise FilterError(msg)
68
+ if isinstance(value, str):
69
+ try:
70
+ datetime.fromisoformat(value)
71
+ except (TypeError, ValueError) as error:
72
+ msg = "Strings are only comparable if they are ISO formatted dates."
73
+ raise FilterError(msg) from error
74
+ if value is None and mongo_operator in {"$gte", "$lte"}:
75
+ return {field: {"$gt" if mongo_operator == "$gte" else "$lt": None}}
76
+ return {field: {mongo_operator: value}}
77
+
78
+ return convert
79
+
80
+
81
+ def _membership(mongo_operator: str) -> Callable[[str, Any], dict[str, Any]]:
82
+ def convert(field: str, value: Any) -> dict[str, Any]:
83
+ if not isinstance(value, list):
84
+ msg = f"{field}'s value must be a list when using a membership comparator"
85
+ raise FilterError(msg)
86
+ return {field: {mongo_operator: value}}
87
+
88
+ return convert
89
+
90
+
91
+ COMPARISON_OPERATORS: dict[str, Callable[[str, Any], dict[str, Any]]] = {
92
+ "==": _comparison("$eq"),
93
+ "!=": _comparison("$ne"),
94
+ ">": _ordered_comparison("$gt"),
95
+ ">=": _ordered_comparison("$gte"),
96
+ "<": _ordered_comparison("$lt"),
97
+ "<=": _ordered_comparison("$lte"),
98
+ "in": _membership("$in"),
99
+ "not in": _membership("$nin"),
100
+ }
File without changes