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.
- azure_documentdb_haystack-1.0.0.dist-info/METADATA +77 -0
- azure_documentdb_haystack-1.0.0.dist-info/RECORD +12 -0
- azure_documentdb_haystack-1.0.0.dist-info/WHEEL +4 -0
- azure_documentdb_haystack-1.0.0.dist-info/licenses/LICENSE.txt +73 -0
- haystack_integrations/components/retrievers/azure_documentdb/__init__.py +12 -0
- haystack_integrations/components/retrievers/azure_documentdb/embedding_retriever.py +129 -0
- haystack_integrations/components/retrievers/azure_documentdb/full_text_retriever.py +135 -0
- haystack_integrations/components/retrievers/py.typed +0 -0
- haystack_integrations/document_stores/azure_documentdb/__init__.py +7 -0
- haystack_integrations/document_stores/azure_documentdb/document_store.py +765 -0
- haystack_integrations/document_stores/azure_documentdb/filters.py +100 -0
- haystack_integrations/document_stores/py.typed +0 -0
|
@@ -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
|
+
[](https://pypi.org/project/azure-documentdb-haystack)
|
|
31
|
+
[](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,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
|