From 2fe0a6ea7cbcfa747017cce23b89e5e04d9c0990 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Mon, 13 Apr 2026 14:16:10 +0200 Subject: [PATCH 01/17] feat: add Amazon S3 Vectors document store integration Implements issue #2110 - Amazon S3 Vectors document store integration with: - S3VectorsDocumentStore: full DocumentStore protocol (count, write, filter, delete) - S3VectorsEmbeddingRetriever: embedding-based retrieval with metadata filtering - Filter conversion from Haystack format to S3 Vectors filter syntax - Auto-creation of vector buckets and indexes - AWS credential support via Secret (or default credential chain) - 49 unit tests covering store, retriever, filters, and serialization - README with usage examples and known limitations --- .github/labeler.yml | 5 + .github/workflows/CI_coverage_comment.yml | 1 + .github/workflows/amazon_s3_vectors.yml | 139 ++++ README.md | 1 + integrations/amazon_s3_vectors/LICENSE.txt | 201 ++++++ integrations/amazon_s3_vectors/README.md | 207 ++++++ .../amazon_s3_vectors/examples/example.py | 108 ++++ .../pydoc/config_docusaurus.yml | 14 + integrations/amazon_s3_vectors/pyproject.toml | 171 +++++ .../haystack_integrations/components/py.typed | 0 .../retrievers/amazon_s3_vectors/__init__.py | 7 + .../amazon_s3_vectors/embedding_retriever.py | 128 ++++ .../amazon_s3_vectors/__init__.py | 7 + .../amazon_s3_vectors/document_store.py | 592 ++++++++++++++++++ .../amazon_s3_vectors/filters.py | 144 +++++ .../document_stores/py.typed | 0 .../amazon_s3_vectors/tests/__init__.py | 3 + .../tests/test_document_store.py | 296 +++++++++ .../tests/test_embedding_retriever.py | 112 ++++ .../amazon_s3_vectors/tests/test_filters.py | 113 ++++ .../tests/test_integration.py | 220 +++++++ 21 files changed, 2469 insertions(+) create mode 100644 .github/workflows/amazon_s3_vectors.yml create mode 100644 integrations/amazon_s3_vectors/LICENSE.txt create mode 100644 integrations/amazon_s3_vectors/README.md create mode 100644 integrations/amazon_s3_vectors/examples/example.py create mode 100644 integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml create mode 100644 integrations/amazon_s3_vectors/pyproject.toml create mode 100644 integrations/amazon_s3_vectors/src/haystack_integrations/components/py.typed create mode 100644 integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py create mode 100644 integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py create mode 100644 integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py create mode 100644 integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py create mode 100644 integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/filters.py create mode 100644 integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/py.typed create mode 100644 integrations/amazon_s3_vectors/tests/__init__.py create mode 100644 integrations/amazon_s3_vectors/tests/test_document_store.py create mode 100644 integrations/amazon_s3_vectors/tests/test_embedding_retriever.py create mode 100644 integrations/amazon_s3_vectors/tests/test_filters.py create mode 100644 integrations/amazon_s3_vectors/tests/test_integration.py diff --git a/.github/labeler.yml b/.github/labeler.yml index c4eb00b1ad..5ff72a907a 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -19,6 +19,11 @@ integration:amazon-bedrock: - any-glob-to-any-file: "integrations/amazon_bedrock/**/*" - any-glob-to-any-file: ".github/workflows/amazon_bedrock.yml" +integration:amazon-s3-vectors: + - changed-files: + - any-glob-to-any-file: "integrations/amazon_s3_vectors/**/*" + - any-glob-to-any-file: ".github/workflows/amazon_s3_vectors.yml" + integration:amazon-sagemaker: - changed-files: - any-glob-to-any-file: "integrations/amazon_sagemaker/**/*" diff --git a/.github/workflows/CI_coverage_comment.yml b/.github/workflows/CI_coverage_comment.yml index c1e5d36477..954a2cb756 100644 --- a/.github/workflows/CI_coverage_comment.yml +++ b/.github/workflows/CI_coverage_comment.yml @@ -7,6 +7,7 @@ on: - "Test / aimlapi" - "Test / alloydb" - "Test / amazon-bedrock" + - "Test / amazon-s3-vectors" - "Test / amazon-sagemaker" - "Test / amazon-textract" - "Test / anthropic" diff --git a/.github/workflows/amazon_s3_vectors.yml b/.github/workflows/amazon_s3_vectors.yml new file mode 100644 index 0000000000..703d462bd5 --- /dev/null +++ b/.github/workflows/amazon_s3_vectors.yml @@ -0,0 +1,139 @@ +# This workflow comes from https://github.com/ofek/hatch-mypyc +# https://github.com/ofek/hatch-mypyc/blob/5a198c0ba8660494d02716cfc9d79ce4adfb1442/.github/workflows/test.yml +name: Test / amazon-s3-vectors + +on: + schedule: + - cron: "0 0 * * *" + pull_request: + paths: + - "integrations/amazon_s3_vectors/**" + - "!integrations/amazon_s3_vectors/*.md" + - ".github/workflows/amazon_s3_vectors.yml" + push: + branches: + - main + paths: + - "integrations/amazon_s3_vectors/**" + - "!integrations/amazon_s3_vectors/*.md" + - ".github/workflows/amazon_s3_vectors.yml" + +defaults: + run: + working-directory: integrations/amazon_s3_vectors + +concurrency: + group: amazon_s3_vectors-${{ github.head_ref || github.sha }} + cancel-in-progress: true + +env: + PYTHONUNBUFFERED: "1" + FORCE_COLOR: "1" + TEST_MATRIX_OS: '["ubuntu-latest", "windows-latest", "macos-latest"]' + TEST_MATRIX_PYTHON: '["3.10", "3.14"]' + +jobs: + compute-test-matrix: + runs-on: ubuntu-slim + defaults: + run: + working-directory: . + outputs: + os: ${{ steps.set.outputs.os }} + python-version: ${{ steps.set.outputs.python-version }} + steps: + - id: set + run: | + echo 'os=${{ github.event_name == 'push' && '["ubuntu-latest"]' || env.TEST_MATRIX_OS }}' >> $GITHUB_OUTPUT + echo 'python-version=${{ github.event_name == 'push' && '["3.10"]' || env.TEST_MATRIX_PYTHON }}' >> $GITHUB_OUTPUT + + run: + name: Python ${{ matrix.python-version }} on ${{ startsWith(matrix.os, 'macos-') && 'macOS' || startsWith(matrix.os, 'windows-') && 'Windows' || 'Linux' }} + needs: compute-test-matrix + permissions: + contents: write + pull-requests: write + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: ${{ fromJSON(needs.compute-test-matrix.outputs.os) }} + python-version: ${{ fromJSON(needs.compute-test-matrix.outputs.python-version) }} + + steps: + - name: Support longpaths + if: matrix.os == 'windows-latest' + working-directory: . + run: git config --system core.longpaths true + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: ${{ matrix.python-version }} + + - name: Install Hatch + run: pip install --upgrade hatch + - name: Lint + if: matrix.python-version == '3.10' && runner.os == 'Linux' + run: hatch run fmt-check && hatch run test:types + + - name: Run unit tests + run: hatch run test:unit-cov-retry + + # On PR: posts coverage comment (directly on same-repo PRs; via artifact for fork PRs). On push to main: stores coverage baseline on data branch. + - name: Store unit tests coverage + id: coverage_comment + if: matrix.python-version == '3.10' && runner.os == 'Linux' && github.event_name != 'schedule' + uses: py-cov-action/python-coverage-comment-action@7188638f871f721a365d644f505d1ff3df20d683 # v3.40 + with: + GITHUB_TOKEN: ${{ github.token }} + COVERAGE_PATH: integrations/amazon_s3_vectors + SUBPROJECT_ID: amazon_s3_vectors + MINIMUM_GREEN: 90 + MINIMUM_ORANGE: 60 + + - name: Upload coverage comment to be posted + if: matrix.python-version == '3.10' && runner.os == 'Linux' && github.event_name == 'pull_request' && steps.coverage_comment.outputs.COMMENT_FILE_WRITTEN == 'true' + uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + with: + name: coverage-comment-amazon_s3_vectors + path: python-coverage-comment-action-amazon_s3_vectors.txt + + - name: Run integration tests + run: hatch run test:integration-cov-append-retry + + - name: Store combined coverage + if: github.event_name == 'push' + uses: py-cov-action/python-coverage-comment-action@7188638f871f721a365d644f505d1ff3df20d683 # v3.40 + with: + GITHUB_TOKEN: ${{ github.token }} + COVERAGE_PATH: integrations/amazon_s3_vectors + SUBPROJECT_ID: amazon_s3_vectors-combined + MINIMUM_GREEN: 90 + MINIMUM_ORANGE: 60 + + - name: Run unit tests with lowest direct dependencies + if: github.event_name != 'push' + run: | + hatch run uv pip compile pyproject.toml --resolution lowest-direct --output-file requirements_lowest_direct.txt + hatch -e test env run -- uv pip install -r requirements_lowest_direct.txt + hatch run test:unit + + - name: Nightly - run unit tests with Haystack main branch + if: github.event_name == 'schedule' + run: | + hatch env prune + hatch -e test env run -- uv pip install git+https://github.com/deepset-ai/haystack.git@main + hatch run test:unit + + + notify-slack-on-failure: + needs: run + if: failure() && github.event_name == 'schedule' + runs-on: ubuntu-slim + steps: + - uses: deepset-ai/notify-slack-action@3cda73b77a148f16f703274198e7771340cf862b # v1 + with: + slack-webhook-url: ${{ secrets.SLACK_WEBHOOK_URL_NOTIFICATIONS }} diff --git a/README.md b/README.md index 0a3c07f06d..20a5c0c374 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ Please check out our [Contribution Guidelines](CONTRIBUTING.md) for all the deta | [aimlapi-haystack](integrations/aimlapi/) | Generator | [![PyPI - Version](https://img.shields.io/pypi/v/aimlapi-haystack.svg)](https://pypi.org/project/aimlapi-haystack) | [![Test / aimlapi](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/aimlapi.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/aimlapi.yml) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-aimlapi/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-aimlapi/htmlcov/index.html) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-aimlapi-combined/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-aimlapi-combined/htmlcov/index.html) | | [alloydb-haystack](integrations/alloydb/) | Document Store | [![PyPI - Version](https://img.shields.io/pypi/v/alloydb-haystack.svg)](https://pypi.org/project/alloydb-haystack) | [![Test / alloydb](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/alloydb.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/alloydb.yml) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-alloydb/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-alloydb/htmlcov/index.html) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-alloydb-combined/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-alloydb-combined/htmlcov/index.html) | | [amazon-bedrock-haystack](integrations/amazon_bedrock/) | Embedder, Generator, Ranker, Downloader | [![PyPI - Version](https://img.shields.io/pypi/v/amazon-bedrock-haystack.svg)](https://pypi.org/project/amazon-bedrock-haystack) | [![Test / amazon_bedrock](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_bedrock.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_bedrock.yml) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_bedrock/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_bedrock/htmlcov/index.html) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_bedrock-combined/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_bedrock-combined/htmlcov/index.html) | +| [amazon-s3-vectors-haystack](integrations/amazon_s3_vectors/) | Document Store | [![PyPI - Version](https://img.shields.io/pypi/v/amazon-s3-vectors-haystack.svg)](https://pypi.org/project/amazon-s3-vectors-haystack) | [![Test / amazon_s3_vectors](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_s3_vectors.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_s3_vectors.yml) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_s3_vectors/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_s3_vectors/htmlcov/index.html) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_s3_vectors-combined/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_s3_vectors-combined/htmlcov/index.html) | | [amazon-sagemaker-haystack](integrations/amazon_sagemaker/) | Generator | [![PyPI - Version](https://img.shields.io/pypi/v/amazon-sagemaker-haystack.svg)](https://pypi.org/project/amazon-sagemaker-haystack) | [![Test / amazon_sagemaker](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_sagemaker.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_sagemaker.yml) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_sagemaker/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_sagemaker/htmlcov/index.html) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_sagemaker-combined/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_sagemaker-combined/htmlcov/index.html) | | [amazon-textract-haystack](integrations/amazon_textract/) | Converter | [![PyPI - Version](https://img.shields.io/pypi/v/amazon-textract-haystack.svg)](https://pypi.org/project/amazon-textract-haystack) | [![Test / amazon_textract](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_textract.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/amazon_textract.yml) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_textract/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_textract/htmlcov/index.html) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-amazon_textract-combined/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-amazon_textract-combined/htmlcov/index.html) | | [anthropic-haystack](integrations/anthropic/) | Generator | [![PyPI - Version](https://img.shields.io/pypi/v/anthropic-haystack.svg)](https://pypi.org/project/anthropic-haystack) | [![Test / anthropic](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/anthropic.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/anthropic.yml) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-anthropic/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-anthropic/htmlcov/index.html) | [![Coverage badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/deepset-ai/haystack-core-integrations/python-coverage-comment-action-data-anthropic-combined/endpoint.json&label=)](https://htmlpreview.github.io/?https://github.com/deepset-ai/haystack-core-integrations/blob/python-coverage-comment-action-data-anthropic-combined/htmlcov/index.html) | diff --git a/integrations/amazon_s3_vectors/LICENSE.txt b/integrations/amazon_s3_vectors/LICENSE.txt new file mode 100644 index 0000000000..6134ab324f --- /dev/null +++ b/integrations/amazon_s3_vectors/LICENSE.txt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "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. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "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. + + "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). + + "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 License, 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. + + "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." + + "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. + + 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. + + 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. + + 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: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (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 + + (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. + + 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. + + 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. + + 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. + + 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. + + 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. + + 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. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + 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. + + Copyright 2023-present deepset GmbH + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/integrations/amazon_s3_vectors/README.md b/integrations/amazon_s3_vectors/README.md new file mode 100644 index 0000000000..8fd50ff2d1 --- /dev/null +++ b/integrations/amazon_s3_vectors/README.md @@ -0,0 +1,207 @@ +# amazon-s3-vectors-haystack + +[![PyPI - Version](https://img.shields.io/pypi/v/amazon-s3-vectors-haystack.svg)](https://pypi.org/project/amazon-s3-vectors-haystack) +[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/amazon-s3-vectors-haystack.svg)](https://pypi.org/project/amazon-s3-vectors-haystack) + +--- + +A [Haystack](https://haystack.deepset.ai/) integration for [Amazon S3 Vectors](https://aws.amazon.com/s3/features/vectors/), providing a Document Store and Embedding Retriever backed by native vector storage in Amazon S3. + +## Installation + +```bash +pip install amazon-s3-vectors-haystack +``` + +## Usage + +### Document Store + +```python +from haystack.dataclasses import Document +from haystack.document_stores.types import DuplicatePolicy +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + +document_store = S3VectorsDocumentStore( + vector_bucket_name="my-vectors", + index_name="my-index", + dimension=768, + distance_metric="cosine", # or "euclidean" + region_name="us-east-1", +) + +# Write documents (embeddings are required) +docs = [ + Document(id="1", content="First document", embedding=[0.1] * 768, meta={"category": "news"}), + Document(id="2", content="Second document", embedding=[0.2] * 768, meta={"category": "sports"}), +] +document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) + +# Count documents +print(document_store.count_documents()) + +# Delete documents +document_store.delete_documents(["1", "2"]) +``` + +### Embedding Retriever in a Pipeline + +```python +from haystack import Pipeline +from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder +from haystack.document_stores.types import DuplicatePolicy +from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + +document_store = S3VectorsDocumentStore( + vector_bucket_name="my-vectors", + index_name="my-index", + dimension=768, +) + +# Index documents +doc_embedder = SentenceTransformersDocumentEmbedder() +doc_embedder.warm_up() +# ... embed and write documents ... + +# Query pipeline +pipeline = Pipeline() +pipeline.add_component("embedder", SentenceTransformersTextEmbedder()) +pipeline.add_component("retriever", S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5)) +pipeline.connect("embedder.embedding", "retriever.query_embedding") + +result = pipeline.run({"embedder": {"text": "What is the latest news?"}}) +print(result["retriever"]["documents"]) +``` + +### Filtering + +The retriever supports Haystack metadata filters, which are converted to S3 Vectors filter syntax: + +```python +# Filter during retrieval +result = pipeline.run({ + "embedder": {"text": "sports news"}, + "retriever": { + "filters": { + "operator": "AND", + "conditions": [ + {"field": "meta.category", "operator": "==", "value": "sports"}, + {"field": "meta.year", "operator": ">=", "value": 2024}, + ], + } + }, +}) +``` + +**Supported filter operators:** `==`, `!=`, `>`, `>=`, `<`, `<=`, `in`, `not in`, `AND`, `OR` + +### AWS Authentication + +The document store uses the standard boto3 credential chain by default. You can also pass credentials explicitly: + +```python +from haystack.utils.auth import Secret + +document_store = S3VectorsDocumentStore( + vector_bucket_name="my-vectors", + index_name="my-index", + dimension=768, + region_name="us-east-1", + aws_access_key_id=Secret.from_env_var("AWS_ACCESS_KEY_ID"), + aws_secret_access_key=Secret.from_env_var("AWS_SECRET_ACCESS_KEY"), +) +``` + +## Configuration Reference + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `vector_bucket_name` | `str` | *required* | Name of the S3 vector bucket | +| `index_name` | `str` | *required* | Name of the vector index within the bucket | +| `dimension` | `int` | *required* | Dimensionality of the embeddings (e.g. 768, 1536) | +| `distance_metric` | `str` | `"cosine"` | `"cosine"` or `"euclidean"` | +| `region_name` | `str \| None` | `None` | AWS region (uses default from env if not set) | +| `aws_access_key_id` | `Secret \| None` | `None` | AWS access key ID | +| `aws_secret_access_key` | `Secret \| None` | `None` | AWS secret access key | +| `aws_session_token` | `Secret \| None` | `None` | AWS session token (for temporary credentials) | +| `create_bucket_and_index` | `bool` | `True` | Auto-create bucket and index if they don't exist | +| `non_filterable_metadata_keys` | `list[str] \| None` | `None` | Additional metadata keys to mark as non-filterable | + +## Known Limitations & Considerations + +### No Native Document Count API +S3 Vectors does not provide a dedicated count endpoint. `count_documents()` paginates through all +vector keys, which can be slow for large indexes (millions of vectors). + +### `filter_documents()` Is Expensive +S3 Vectors only supports metadata filtering **during vector similarity queries** — there is no +standalone "list documents matching filter" API. As a result, `filter_documents()` must: +1. List all vectors with their data and metadata (paginated) +2. Apply filters client-side in memory + +**For filtered retrieval, always prefer `S3VectorsEmbeddingRetriever` with filters**, which uses +the native `query_vectors` API with server-side filtering. + +### Embedding Required +Every document written to the store **must** have an embedding. Documents without embeddings +will be rejected. This is a fundamental constraint of S3 Vectors as a pure vector store. + +### No Keyword / BM25 Retrieval +S3 Vectors only supports dense vector similarity search. There is no keyword or BM25 search +capability. If you need hybrid search, consider pairing this with Amazon OpenSearch. + +### Vector Data Type +Only `float32` vectors are supported. Higher-precision values are automatically downcast. + +### Metadata Size Limits +- **Total metadata per vector: 40 KB** (filterable + non-filterable combined) +- **Filterable metadata per vector: 2 KB** — user `meta` fields used in filters must fit in this budget +- **Non-filterable metadata keys per index: 10** — the integration reserves 4 internal keys + (`_content`, `_blob_data`, `_blob_meta`, `_blob_mime_type`), leaving 6 for user-defined keys +- Keys are set at index creation and **cannot be changed later** + +Large content (e.g. full document text) is stored as non-filterable metadata automatically. +If you store additional large metadata fields, declare them via `non_filterable_metadata_keys`. + +### Strict API Limits +- `put_vectors`: up to 500 vectors per call (handled automatically by the integration) +- `get_vectors`: up to 100 keys per call (handled automatically) +- `delete_vectors`: up to 500 keys per call (handled automatically) +- **`query_vectors`: maximum 100 results per query** — this is the hard cap on `top_k`. + If you need more than 100 results, you must implement pagination or use a different store. +- Combined PutVectors + DeleteVectors: up to 1,000 requests/second per index + +### Distance Metrics and Scoring +Only `cosine` and `euclidean` are supported. The metric is set at index creation time and cannot +be changed afterward. + +S3 Vectors returns raw **distances** (lower = more similar). The integration converts these to +Haystack-convention **scores** (higher = more similar): +- **Cosine:** `score = 1.0 - distance` (1.0 = identical, 0.0 = orthogonal) +- **Euclidean:** `score = -distance` (0.0 = identical, more negative = further) + +### No Embeddings Returned from Queries +The `query_vectors` API does not support returning the stored vector data alongside results. +Documents retrieved via `S3VectorsEmbeddingRetriever` will have `embedding=None`. If you need +the embedding vectors, use `filter_documents()` or fetch them separately via the boto3 client. + +### Eventual Consistency +Newly written vectors may not be immediately visible in query results. S3 Vectors provides +eventual consistency for write-then-read operations. + +## Running Tests + +```bash +cd integrations/amazon_s3_vectors + +# Unit tests (no AWS credentials required) +hatch run test:unit + +# Integration tests (requires AWS credentials and S3 Vectors access) +hatch run test:integration +``` + +## License + +`amazon-s3-vectors-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license. diff --git a/integrations/amazon_s3_vectors/examples/example.py b/integrations/amazon_s3_vectors/examples/example.py new file mode 100644 index 0000000000..67d3a67c20 --- /dev/null +++ b/integrations/amazon_s3_vectors/examples/example.py @@ -0,0 +1,108 @@ +# Amazon S3 Vectors — Haystack Integration Example +# +# A serverless vector store on AWS: no database to provision, no cluster to manage. +# Just a bucket name and AWS credentials. +# +# This example indexes documents, runs unfiltered and filtered queries, and cleans up. +# +# Prerequisites: +# pip install amazon-s3-vectors-haystack "sentence-transformers>=2.2.0" +# AWS credentials configured (env vars, ~/.aws/credentials, or IAM role) + +import time + +from haystack import Document, Pipeline +from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder +from haystack.components.writers import DocumentWriter +from haystack.document_stores.types import DuplicatePolicy + +from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + +MODEL = "sentence-transformers/all-MiniLM-L6-v2" + +# --- Create the document store (bucket + index are created automatically) --- +document_store = S3VectorsDocumentStore( + vector_bucket_name="haystack-example", + index_name="products", + dimension=384, # all-MiniLM-L6-v2 produces 384-dim embeddings + distance_metric="cosine", + region_name="us-east-1", +) + +# --- Index documents with metadata --- +products = [ + Document( + content="Lightweight running shoes with breathable mesh", + meta={"category": "shoes", "price": 89.99, "in_stock": True}, + ), + Document( + content="Waterproof hiking boots with ankle support", + meta={"category": "shoes", "price": 149.99, "in_stock": True}, + ), + Document( + content="Classic leather dress shoes", + meta={"category": "shoes", "price": 199.99, "in_stock": False}, + ), + Document( + content="Insulated winter jacket with down filling", + meta={"category": "jackets", "price": 249.99, "in_stock": True}, + ), + Document( + content="Lightweight windbreaker for running", + meta={"category": "jackets", "price": 79.99, "in_stock": True}, + ), +] + +indexing = Pipeline() +indexing.add_component("embedder", SentenceTransformersDocumentEmbedder(model=MODEL)) +indexing.add_component("writer", DocumentWriter(document_store=document_store, policy=DuplicatePolicy.OVERWRITE)) +indexing.connect("embedder", "writer") + +print("Indexing products...") +indexing.run({"embedder": {"documents": products}}) +print(f"Indexed {document_store.count_documents()} products.\n") + +# --- Build the query pipeline --- +querying = Pipeline() +querying.add_component("embedder", SentenceTransformersTextEmbedder(model=MODEL)) +querying.add_component("retriever", S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5)) +querying.connect("embedder", "retriever") + + +def search(query: str, filters: dict | None = None) -> None: + print(f"Query: '{query}'" + (f" Filter: {filters}" if filters else "")) + start = time.time() + result = querying.run({"embedder": {"text": query}, "retriever": {"filters": filters}}) + elapsed = (time.time() - start) * 1000 + for i, doc in enumerate(result["retriever"]["documents"], 1): + meta = doc.meta + print(f" {i}. [{doc.score:.3f}] {doc.content}") + print(f" category={meta['category']} price=${meta['price']} in_stock={meta['in_stock']}") + print(f" ({elapsed:.0f}ms)\n") + + +# Unfiltered: returns shoes AND jackets +search("something for running") + +# Filtered to shoes only — server-side, inside the vector search +search("something for running", filters={"field": "meta.category", "operator": "==", "value": "shoes"}) + +# Combined filter: in-stock AND under $100 +search( + "something for running", + filters={ + "operator": "AND", + "conditions": [ + {"field": "meta.in_stock", "operator": "==", "value": True}, + {"field": "meta.price", "operator": "<", "value": 100.0}, + ], + }, +) + +# --- Cleanup --- +document_store.delete_documents([doc.id for doc in products]) +client = document_store._get_client() +client.delete_index(vectorBucketName="haystack-example", indexName="products") +client.delete_vector_bucket(vectorBucketName="haystack-example") +print("Cleaned up all AWS resources.") diff --git a/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml b/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml new file mode 100644 index 0000000000..97cb55f004 --- /dev/null +++ b/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml @@ -0,0 +1,14 @@ +loaders: + - modules: + - haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever + - haystack_integrations.document_stores.amazon_s3_vectors.document_store + search_path: [../src] +processors: + - type: filter + documented_only: true + skip_empty_modules: true +renderer: + description: Amazon S3 Vectors integration for Haystack + id: integrations-amazon_s3_vectors + filename: amazon-s3-vectors.md + title: Amazon S3 Vectors diff --git a/integrations/amazon_s3_vectors/pyproject.toml b/integrations/amazon_s3_vectors/pyproject.toml new file mode 100644 index 0000000000..3a8fa55a65 --- /dev/null +++ b/integrations/amazon_s3_vectors/pyproject.toml @@ -0,0 +1,171 @@ +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" + +[project] +name = "amazon-s3-vectors-haystack" +dynamic = ["version"] +description = "Haystack integration for amazon_s3_vectors" +readme = "README.md" +requires-python = ">=3.10" +license = "Apache-2.0" +keywords = ["aws", "s3", "vectors", "document-store", "haystack", "vector-search"] +authors = [{ name = "deepset GmbH", email = "info@deepset.ai" }] +classifiers = [ + "License :: OSI Approved :: Apache Software License", + "Development Status :: 4 - Beta", + "Programming Language :: Python", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", +] +dependencies = ["haystack-ai", "boto3"] + +[project.urls] +Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_s3_vectors#readme" +Issues = "https://github.com/deepset-ai/haystack-core-integrations/issues" +Source = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_s3_vectors" + +[tool.hatch.build.targets.wheel] +packages = ["src/haystack_integrations"] + +[tool.hatch.version] +source = "vcs" +tag-pattern = 'integrations\/amazon_s3_vectors-v(?P.*)' + +[tool.hatch.version.raw-options] +root = "../.." +git_describe_command = 'git describe --tags --match="integrations/amazon_s3_vectors-v[0-9]*"' + +[tool.hatch.envs.default] +installer = "uv" +dependencies = ["haystack-pydoc-tools", "ruff"] + +[tool.hatch.envs.default.scripts] +docs = ["haystack-pydoc pydoc/config_docusaurus.yml"] +fmt = "ruff check --fix {args}; ruff format {args}" +fmt-check = "ruff check {args} && ruff format --check {args}" + +[tool.hatch.envs.test] +dependencies = [ + "pytest", + "pytest-asyncio", + "pytest-cov", + "pytest-rerunfailures", + "mypy", + "pip", +] + +[tool.hatch.envs.test.scripts] +unit = 'pytest -m "not integration" {args:tests}' +integration = 'pytest -m "integration" {args:tests}' +all = 'pytest {args:tests}' +unit-cov-retry = 'pytest --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x -m "not integration" {args:tests}' +integration-cov-append-retry = 'pytest --cov=haystack_integrations --cov-append --reruns 3 --reruns-delay 30 -x -m "integration" {args:tests}' +types = """mypy -p haystack_integrations.document_stores.amazon_s3_vectors \ +-p haystack_integrations.components.retrievers.amazon_s3_vectors {args}""" + +[tool.mypy] +install_types = true +non_interactive = true +check_untyped_defs = true +disallow_incomplete_defs = true + +[[tool.mypy.overrides]] +module = [ +"boto3.*", +"botocore.*", +] +ignore_missing_imports = true + +[tool.ruff] +line-length = 120 + +[tool.ruff.lint] +select = [ + "A", + "ANN", + "ARG", + "B", + "C", + "D102", # Missing docstring in public method + "D103", # Missing docstring in public function + "D205", # 1 blank line required between summary line and description + "D209", # Closing triple quotes go to new line + "D213", # summary lines must be positioned on the second physical line of the docstring + "D417", # Missing argument descriptions in the docstring + "D419", # Docstring is empty + "DTZ", + "E", + "EM", + "F", + "I", + "ICN", + "ISC", + "N", + "PLC", + "PLE", + "PLR", + "PLW", + "Q", + "RUF", + "S", + "T", + "TID", + "UP", + "W", + "YTT", +] +ignore = [ + # Allow non-abstract empty methods in abstract base classes + "B027", + # Allow function calls in argument defaults (common Haystack pattern for Secret.from_env_var) + "B008", + # Ignore checks for possible passwords + "S105", + "S106", + "S107", + # Ignore complexity + "C901", + "PLR0911", + "PLR0912", + "PLR0913", + "PLR0915", + # Allow `Any` type - used legitimately for dynamic types and SDK boundaries + "ANN401", +] + +[tool.ruff.lint.isort] +known-first-party = ["haystack_integrations"] + +[tool.ruff.lint.flake8-tidy-imports] +ban-relative-imports = "parents" + +[tool.ruff.lint.per-file-ignores] +# Tests can use magic values, assertions, relative imports, and don't need type annotations +"tests/**/*" = ["PLC0415", "PLR2004", "S101", "S110", "TID252", "D", "ANN"] +# Examples can contain print statements and long lines +"examples/**/*" = ["D", "T201", "E501"] + +[tool.coverage.run] +source = ["haystack_integrations"] +branch = true +parallel = false +relative_files = true + +[tool.coverage.report] +omit = ["*/tests/*", "*/__init__.py"] +show_missing = true +exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"] + +[tool.pytest.ini_options] +addopts = "--strict-markers" +markers = [ + "integration: integration tests", +] +log_cli = true +asyncio_default_fixture_loop_scope = "function" diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/py.typed b/integrations/amazon_s3_vectors/src/haystack_integrations/components/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py new file mode 100644 index 0000000000..39ebcd2f58 --- /dev/null +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from .embedding_retriever import S3VectorsEmbeddingRetriever + +__all__ = ["S3VectorsEmbeddingRetriever"] diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py new file mode 100644 index 0000000000..e04ad20990 --- /dev/null +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack import component, default_from_dict, default_to_dict +from haystack.dataclasses import Document +from haystack.document_stores.types import FilterPolicy +from haystack.document_stores.types.filter_policy import apply_filter_policy + +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + + +@component +class S3VectorsEmbeddingRetriever: + """ + Retrieve documents from an ``S3VectorsDocumentStore`` based on their dense embeddings. + + Usage example: + ```python + from haystack import Document, Pipeline + from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder + from haystack.document_stores.types import DuplicatePolicy + from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever + from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + + document_store = S3VectorsDocumentStore( + vector_bucket_name="my-vectors", + index_name="my-index", + dimension=768, + ) + + documents = [ + Document(content="There are over 7,000 languages spoken around the world today."), + Document(content="Elephants have been observed to behave in a way that indicates..."), + Document(content="In certain places, you can witness the phenomenon of bioluminescent waves."), + ] + + document_embedder = SentenceTransformersDocumentEmbedder() + document_embedder.warm_up() + documents_with_embeddings = document_embedder.run(documents) + + document_store.write_documents(documents_with_embeddings.get("documents"), policy=DuplicatePolicy.OVERWRITE) + + query_pipeline = Pipeline() + query_pipeline.add_component("text_embedder", SentenceTransformersTextEmbedder()) + query_pipeline.add_component("retriever", S3VectorsEmbeddingRetriever(document_store=document_store)) + query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding") + + query = "How many languages are there?" + res = query_pipeline.run({"text_embedder": {"text": query}}) + ``` + """ + + def __init__( + self, + *, + document_store: S3VectorsDocumentStore, + filters: dict[str, Any] | None = None, + top_k: int = 10, + filter_policy: str | FilterPolicy = FilterPolicy.REPLACE, + ) -> None: + """ + Initialize the S3VectorsEmbeddingRetriever. + + :param document_store: An instance of ``S3VectorsDocumentStore``. + :param filters: Filters applied to the retrieved Documents. + :param top_k: Maximum number of Documents to return. + :param filter_policy: Policy to determine how filters are applied. + :raises ValueError: If ``document_store`` is not an ``S3VectorsDocumentStore``. + """ + if not isinstance(document_store, S3VectorsDocumentStore): + msg = "document_store must be an instance of S3VectorsDocumentStore" + raise ValueError(msg) + + self.document_store = document_store + self.filters = filters or {} + self.top_k = top_k + self.filter_policy = ( + filter_policy if isinstance(filter_policy, FilterPolicy) else FilterPolicy.from_str(filter_policy) + ) + + def to_dict(self) -> dict[str, Any]: + """Serialize this component to a dictionary.""" + return default_to_dict( + self, + filters=self.filters, + top_k=self.top_k, + filter_policy=self.filter_policy.value, + document_store=self.document_store.to_dict(), + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "S3VectorsEmbeddingRetriever": + """Deserialize the component from a dictionary.""" + data["init_parameters"]["document_store"] = S3VectorsDocumentStore.from_dict( + data["init_parameters"]["document_store"] + ) + if filter_policy := data["init_parameters"].get("filter_policy"): + data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(filter_policy) + return default_from_dict(cls, data) + + @component.output_types(documents=list[Document]) + def run( + self, + query_embedding: list[float], + filters: dict[str, Any] | None = None, + top_k: int | None = None, + ) -> dict[str, list[Document]]: + """ + Retrieve documents from the S3VectorsDocumentStore based on dense embeddings. + + :param query_embedding: Embedding of the query. + :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on + the ``filter_policy`` chosen at retriever initialization. + :param top_k: Maximum number of Documents to return. + :returns: A dictionary with key ``"documents"`` containing the retrieved Documents. + """ + filters = apply_filter_policy(self.filter_policy, self.filters, filters) + top_k = top_k or self.top_k + + docs = self.document_store._embedding_retrieval( + query_embedding=query_embedding, + filters=filters, + top_k=top_k, + ) + return {"documents": docs} diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py new file mode 100644 index 0000000000..4e6beddd97 --- /dev/null +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from .document_store import S3VectorsDocumentStore + +__all__ = ["S3VectorsDocumentStore"] diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py new file mode 100644 index 0000000000..01c2c9ceb7 --- /dev/null +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -0,0 +1,592 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import base64 +import json +from typing import Any, Literal + +import boto3 +from botocore.exceptions import ClientError +from haystack import default_from_dict, default_to_dict, logging +from haystack.dataclasses import ByteStream, Document +from haystack.document_stores.errors import DocumentStoreError +from haystack.document_stores.types import DuplicatePolicy +from haystack.utils.auth import Secret, deserialize_secrets_inplace + +from .filters import _normalize_filters, _validate_filters + +logger = logging.getLogger(__name__) + +# S3 Vectors allows up to 500 vectors per put_vectors call +_WRITE_BATCH_SIZE = 500 + +# S3 Vectors allows up to 100 keys per get_vectors call, 500 per delete_vectors call +_GET_BATCH_SIZE = 100 +_DELETE_BATCH_SIZE = 500 + +# Maximum number of results from query_vectors +_MAX_TOP_K = 100 + +# S3 Vectors metadata limits +_MAX_TOTAL_METADATA_BYTES = 40 * 1024 # 40 KB total + +# Reserved metadata keys used to store Haystack Document fields +_CONTENT_KEY = "_content" +_BLOB_DATA_KEY = "_blob_data" +_BLOB_META_KEY = "_blob_meta" +_BLOB_MIME_TYPE_KEY = "_blob_mime_type" + +_RESERVED_META_KEYS = {_CONTENT_KEY, _BLOB_DATA_KEY, _BLOB_META_KEY, _BLOB_MIME_TYPE_KEY} + +# These keys are stored but should not be used in query filters. +# They are configured as nonFilterableMetadataKeys on the index. +_NON_FILTERABLE_KEYS = [_CONTENT_KEY, _BLOB_DATA_KEY, _BLOB_META_KEY, _BLOB_MIME_TYPE_KEY] + + +class S3VectorsDocumentStore: + """ + A Document Store using [Amazon S3 Vectors](https://aws.amazon.com/s3/features/vectors/). + + Amazon S3 Vectors provides native vector storage and similarity search within Amazon S3. + This document store stores Haystack `Document` objects as vectors with associated metadata + in an S3 vector bucket and index. + + Usage example: + ```python + from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + + document_store = S3VectorsDocumentStore( + vector_bucket_name="my-vectors", + index_name="my-index", + dimension=768, + ) + ``` + """ + + def __init__( + self, + *, + vector_bucket_name: str, + index_name: str, + dimension: int, + distance_metric: Literal["cosine", "euclidean"] = "cosine", + region_name: str | None = None, + aws_access_key_id: Secret | None = None, + aws_secret_access_key: Secret | None = None, + aws_session_token: Secret | None = None, + create_bucket_and_index: bool = True, + non_filterable_metadata_keys: list[str] | None = None, + ) -> None: + """ + Create an S3VectorsDocumentStore instance. + + :param vector_bucket_name: Name of the S3 vector bucket. + :param index_name: Name of the vector index within the bucket. + :param dimension: Dimensionality of the embeddings (e.g. 768, 1536). + :param distance_metric: Distance metric for similarity search: ``"cosine"`` or ``"euclidean"``. + :param region_name: AWS region. If not provided, uses the default from the environment/config. + :param aws_access_key_id: AWS access key ID. If not provided, uses the default credential chain. + :param aws_secret_access_key: AWS secret access key. + :param aws_session_token: AWS session token for temporary credentials. + :param create_bucket_and_index: Whether to automatically create the vector bucket and index + if they do not exist. Defaults to ``True``. + :param non_filterable_metadata_keys: Additional metadata keys to mark as non-filterable + on the index (beyond the internal keys used for Document content/blob storage). + """ + self.vector_bucket_name = vector_bucket_name + self.index_name = index_name + self.dimension = dimension + self.distance_metric = distance_metric + self.region_name = region_name + self.aws_access_key_id = aws_access_key_id + self.aws_secret_access_key = aws_secret_access_key + self.aws_session_token = aws_session_token + self.create_bucket_and_index = create_bucket_and_index + self.non_filterable_metadata_keys = non_filterable_metadata_keys or [] + + self._client: Any = None + + def _get_client(self) -> Any: + """Lazily create and return the boto3 s3vectors client.""" + if self._client is not None: + return self._client + + kwargs: dict[str, Any] = {"service_name": "s3vectors"} + if self.region_name: + kwargs["region_name"] = self.region_name + if self.aws_access_key_id: + kwargs["aws_access_key_id"] = self.aws_access_key_id.resolve_value() + if self.aws_secret_access_key: + kwargs["aws_secret_access_key"] = self.aws_secret_access_key.resolve_value() + if self.aws_session_token: + kwargs["aws_session_token"] = self.aws_session_token.resolve_value() + + self._client = boto3.client(**kwargs) + + if self.create_bucket_and_index: + self._ensure_bucket_and_index() + + return self._client + + def _ensure_bucket_and_index(self) -> None: + """Create the vector bucket and index if they don't already exist.""" + client = self._client + + # Ensure bucket exists + try: + client.get_vector_bucket(vectorBucketName=self.vector_bucket_name) + logger.info("Using existing vector bucket '{bucket}'.", bucket=self.vector_bucket_name) + except ClientError as e: + if e.response["Error"]["Code"] == "NotFoundException": + logger.info("Creating vector bucket '{bucket}'.", bucket=self.vector_bucket_name) + client.create_vector_bucket(vectorBucketName=self.vector_bucket_name) + else: + raise + + # Ensure index exists + all_non_filterable = list(set(_NON_FILTERABLE_KEYS + self.non_filterable_metadata_keys)) + try: + client.get_index(vectorBucketName=self.vector_bucket_name, indexName=self.index_name) + logger.info( + "Using existing index '{index}' in bucket '{bucket}'. " + "`dimension`, `distance_metric`, and `non_filterable_metadata_keys` will be ignored.", + index=self.index_name, + bucket=self.vector_bucket_name, + ) + except ClientError as e: + if e.response["Error"]["Code"] == "NotFoundException": + logger.info( + "Creating index '{index}' in bucket '{bucket}' (dimension={dim}, metric={metric}).", + index=self.index_name, + bucket=self.vector_bucket_name, + dim=self.dimension, + metric=self.distance_metric, + ) + client.create_index( + vectorBucketName=self.vector_bucket_name, + indexName=self.index_name, + dataType="float32", + dimension=self.dimension, + distanceMetric=self.distance_metric, + metadataConfiguration={"nonFilterableMetadataKeys": all_non_filterable}, + ) + else: + raise + + def to_dict(self) -> dict[str, Any]: + """Serialize this document store to a dictionary.""" + return default_to_dict( + self, + vector_bucket_name=self.vector_bucket_name, + index_name=self.index_name, + dimension=self.dimension, + distance_metric=self.distance_metric, + region_name=self.region_name, + aws_access_key_id=self.aws_access_key_id.to_dict() if self.aws_access_key_id else None, + aws_secret_access_key=self.aws_secret_access_key.to_dict() if self.aws_secret_access_key else None, + aws_session_token=self.aws_session_token.to_dict() if self.aws_session_token else None, + create_bucket_and_index=self.create_bucket_and_index, + non_filterable_metadata_keys=self.non_filterable_metadata_keys, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "S3VectorsDocumentStore": + """Deserialize a document store from a dictionary.""" + deserialize_secrets_inplace( + data["init_parameters"], + keys=["aws_access_key_id", "aws_secret_access_key", "aws_session_token"], + ) + return default_from_dict(cls, data) + + def count_documents(self) -> int: + """ + Return the number of documents in the document store. + + .. note:: + + S3 Vectors does not provide a dedicated count API. This method lists all vector keys + via pagination, which can be slow for large indexes. + """ + client = self._get_client() + count = 0 + next_token = None + while True: + kwargs: dict[str, Any] = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + } + if next_token: + kwargs["nextToken"] = next_token + response = client.list_vectors(**kwargs) + count += len(response.get("vectors", [])) + next_token = response.get("nextToken") + if not next_token: + break + return count + + def write_documents(self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.OVERWRITE) -> int: + """ + Write Documents to the S3 Vectors index. + + S3 Vectors ``put_vectors`` is an upsert operation by default, so ``DuplicatePolicy.OVERWRITE`` is the + natural behavior. ``DuplicatePolicy.SKIP`` will check for existing documents first (slower). + ``DuplicatePolicy.NONE`` will raise an error if a document already exists. + + :param documents: A list of Documents to write. + :param policy: The duplicate policy. Defaults to ``DuplicatePolicy.OVERWRITE``. + :returns: The number of documents written. + """ + if len(documents) == 0: + return 0 + + client = self._get_client() + written = 0 + + for i in range(0, len(documents), _WRITE_BATCH_SIZE): + batch = documents[i : i + _WRITE_BATCH_SIZE] + + # Validate embeddings upfront + for doc in batch: + if doc.embedding is None: + msg = f"Document '{doc.id}' has no embedding. S3VectorsDocumentStore requires embeddings." + raise DocumentStoreError(msg) + + # Batch-check for existing documents when needed + existing_ids: set[str] = set() + if policy in (DuplicatePolicy.SKIP, DuplicatePolicy.NONE): + batch_ids = [doc.id for doc in batch] + for j in range(0, len(batch_ids), _GET_BATCH_SIZE): + id_chunk = batch_ids[j : j + _GET_BATCH_SIZE] + response = client.get_vectors( + vectorBucketName=self.vector_bucket_name, + indexName=self.index_name, + keys=id_chunk, + ) + for v in response.get("vectors", []): + existing_ids.add(v["key"]) + + if policy == DuplicatePolicy.NONE and existing_ids: + msg = ( + f"Document(s) {sorted(existing_ids)} already exist in the document store. " + "Use DuplicatePolicy.OVERWRITE or DuplicatePolicy.SKIP." + ) + raise DocumentStoreError(msg) + + vectors_to_write = [] + for doc in batch: + if policy == DuplicatePolicy.SKIP and doc.id in existing_ids: + continue + vectors_to_write.append(self._document_to_s3_vector(doc)) + + if vectors_to_write: + client.put_vectors( + vectorBucketName=self.vector_bucket_name, + indexName=self.index_name, + vectors=vectors_to_write, + ) + written += len(vectors_to_write) + + return written + + def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]: + """ + Return documents matching the provided filters. + + .. warning:: + + S3 Vectors only supports metadata filtering during vector similarity queries, not as a + standalone operation. This method lists all vectors and applies filters client-side, + which can be very slow for large indexes. For filtered retrieval, prefer using + ``S3VectorsEmbeddingRetriever`` with filters instead. + + :param filters: Haystack-format filters to apply. + :returns: A list of matching Documents. + """ + logger.warning( + "S3 Vectors does not support standalone filtered listing. " + "filter_documents() will fetch ALL vectors and apply filters client-side, " + "which can be very slow for large indexes. " + "Prefer using S3VectorsEmbeddingRetriever with filters for efficient filtered retrieval." + ) + + client = self._get_client() + + # list_vectors supports returnData and returnMetadata directly, + # so we can read documents in a single paginated pass without + # a separate get_vectors round-trip. + documents: list[Document] = [] + next_token = None + while True: + kwargs: dict[str, Any] = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "returnData": True, + "returnMetadata": True, + } + if next_token: + kwargs["nextToken"] = next_token + response = client.list_vectors(**kwargs) + for v in response.get("vectors", []): + documents.append(self._s3_vector_to_document(v)) + next_token = response.get("nextToken") + if not next_token: + break + + if filters: + _validate_filters(filters) + documents = _apply_filters_in_memory(documents, filters) + + return documents + + def delete_documents(self, document_ids: list[str]) -> None: + """ + Delete documents by their IDs. + + :param document_ids: List of document IDs to delete. + """ + if not document_ids: + return + + client = self._get_client() + for i in range(0, len(document_ids), _DELETE_BATCH_SIZE): + batch = document_ids[i : i + _DELETE_BATCH_SIZE] + client.delete_vectors( + vectorBucketName=self.vector_bucket_name, + indexName=self.index_name, + keys=batch, + ) + + def _embedding_retrieval( + self, + query_embedding: list[float], + *, + filters: dict[str, Any] | None = None, + top_k: int = 10, + ) -> list[Document]: + """ + Retrieve documents most similar to the query embedding. + + This method is not part of the public interface. + Use ``S3VectorsEmbeddingRetriever`` instead. + + :param query_embedding: The query embedding vector. + :param filters: Optional Haystack-format metadata filters. + :param top_k: Maximum number of results to return. + :returns: List of Documents sorted by similarity. + """ + if not query_embedding: + msg = "query_embedding must be a non-empty list of floats" + raise ValueError(msg) + + _validate_filters(filters) + s3_filter = _normalize_filters(filters) if filters else None + + if top_k > _MAX_TOP_K: + logger.warning( + "Requested top_k={top_k} exceeds S3 Vectors maximum of {max_k}. Results will be capped.", + top_k=top_k, + max_k=_MAX_TOP_K, + ) + + client = self._get_client() + + query_kwargs: dict[str, Any] = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "topK": min(top_k, _MAX_TOP_K), # S3 Vectors caps at 100 + "queryVector": {"float32": query_embedding}, + "returnMetadata": True, + "returnDistance": True, + } + if s3_filter: + query_kwargs["filter"] = s3_filter + + result = client.query_vectors(**query_kwargs) + + # Convert distance to a score. + # For cosine: S3 Vectors returns cosine *distance* (0 = identical, 2 = opposite). + # Haystack convention is higher score = more similar, so we convert: score = 1 - distance. + # For euclidean: we negate the distance so that closer vectors score higher. + distance_metric = result.get("distanceMetric", self.distance_metric) + + documents = [] + for v in result.get("vectors", []): + metadata = v.get("metadata", {}) + content = metadata.pop(_CONTENT_KEY, None) + + # Remove other internal keys from user-facing metadata + blob_data = metadata.pop(_BLOB_DATA_KEY, None) + blob_meta = metadata.pop(_BLOB_META_KEY, None) + blob_mime_type = metadata.pop(_BLOB_MIME_TYPE_KEY, None) + + blob = None + if blob_data is not None: + blob = ByteStream( + data=base64.b64decode(blob_data) if isinstance(blob_data, str) else blob_data, + meta=blob_meta or {}, + mime_type=blob_mime_type, + ) + + raw_distance = v.get("distance") + score = None + if raw_distance is not None: + if distance_metric == "cosine": + score = 1.0 - raw_distance + else: + # euclidean: negate so higher = more similar + score = -raw_distance + + doc = Document( + id=v["key"], + content=content, + meta=metadata, + blob=blob, + score=score, + ) + documents.append(doc) + + return documents + + @staticmethod + def _document_to_s3_vector(doc: Document) -> dict[str, Any]: + """Convert a Haystack Document to an S3 Vectors vector entry.""" + metadata: dict[str, Any] = {} + + # Store content as non-filterable metadata + if doc.content is not None: + metadata[_CONTENT_KEY] = doc.content + + # Store blob fields + if doc.blob is not None: + metadata[_BLOB_DATA_KEY] = base64.b64encode(doc.blob.data).decode("ascii") + if doc.blob.meta: + metadata[_BLOB_META_KEY] = doc.blob.meta + if doc.blob.mime_type: + metadata[_BLOB_MIME_TYPE_KEY] = doc.blob.mime_type + + # Store user metadata + if doc.meta: + for key, value in doc.meta.items(): + if key in _RESERVED_META_KEYS: + logger.warning( + "Metadata key '{key}' is reserved; user value will be ignored.", + key=key, + ) + continue + metadata[key] = value + + # Warn if metadata is likely too large + try: + meta_size = len(json.dumps(metadata).encode("utf-8")) + if meta_size > _MAX_TOTAL_METADATA_BYTES: + logger.warning( + "Document '{doc_id}' has ~{size} bytes of metadata, exceeding the S3 Vectors " + "limit of {limit} bytes. The put_vectors call may fail.", + doc_id=doc.id, + size=meta_size, + limit=_MAX_TOTAL_METADATA_BYTES, + ) + except (TypeError, ValueError): + pass # Best-effort size check + + return { + "key": doc.id, + "data": {"float32": doc.embedding}, + "metadata": metadata, + } + + @staticmethod + def _s3_vector_to_document(vector: dict[str, Any]) -> Document: + """Convert an S3 Vectors vector response to a Haystack Document.""" + metadata = dict(vector.get("metadata", {})) + content = metadata.pop(_CONTENT_KEY, None) + blob_data = metadata.pop(_BLOB_DATA_KEY, None) + blob_meta = metadata.pop(_BLOB_META_KEY, None) + blob_mime_type = metadata.pop(_BLOB_MIME_TYPE_KEY, None) + + blob = None + if blob_data is not None: + blob = ByteStream( + data=base64.b64decode(blob_data) if isinstance(blob_data, str) else blob_data, + meta=blob_meta or {}, + mime_type=blob_mime_type, + ) + + embedding = None + data = vector.get("data", {}) + if "float32" in data: + embedding = data["float32"] + + return Document( + id=vector["key"], + content=content, + meta=metadata, + embedding=embedding, + blob=blob, + ) + + +def _apply_filters_in_memory(documents: list[Document], filters: dict[str, Any]) -> list[Document]: + """ + Apply Haystack filters to a list of Documents in memory. + + This is used by ``filter_documents`` since S3 Vectors doesn't support + standalone metadata queries without a vector search. + """ + return [doc for doc in documents if _document_matches(doc, filters)] + + +def _document_matches(doc: Document, filters: dict[str, Any]) -> bool: + """Check if a single Document matches the given Haystack filter.""" + if "operator" in filters and "conditions" in filters: + operator = filters["operator"] + conditions = filters["conditions"] + if operator == "AND": + return all(_document_matches(doc, c) for c in conditions) + if operator == "OR": + return any(_document_matches(doc, c) for c in conditions) + msg = f"Unknown logical operator '{operator}'" + raise ValueError(msg) + + if "field" in filters: + field = filters["field"] + operator = filters["operator"] + value = filters["value"] + + # Resolve the field value from the Document + if field.startswith("meta."): + field_name = field[5:] + doc_value = doc.meta.get(field_name) if doc.meta else None + elif field == "content": + doc_value = doc.content + elif field == "id": + doc_value = doc.id + else: + doc_value = doc.meta.get(field) if doc.meta else None + + return _compare(doc_value, operator, value) + + return True + + +def _compare(doc_value: Any, operator: str, filter_value: Any) -> bool: + """Perform a comparison between a document field value and a filter value.""" + if operator == "==": + return doc_value == filter_value + if operator == "!=": + return doc_value != filter_value + if operator == ">": + return doc_value is not None and doc_value > filter_value + if operator == ">=": + return doc_value is not None and doc_value >= filter_value + if operator == "<": + return doc_value is not None and doc_value < filter_value + if operator == "<=": + return doc_value is not None and doc_value <= filter_value + if operator == "in": + return doc_value is not None and doc_value in filter_value + if operator == "not in": + return doc_value is not None and doc_value not in filter_value + msg = f"Unknown comparison operator '{operator}'" + raise ValueError(msg) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/filters.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/filters.py new file mode 100644 index 0000000000..3141868f12 --- /dev/null +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/filters.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.errors import FilterError + + +def _normalize_filters(filters: dict[str, Any]) -> dict[str, Any]: + """ + Convert Haystack filters to Amazon S3 Vectors compatible filters. + + Reference: https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-vectors-metadata-filtering.html + """ + if not isinstance(filters, dict): + msg = "Filters must be a dictionary" + raise FilterError(msg) + + if "field" in filters: + return _parse_comparison_condition(filters) + return _parse_logical_condition(filters) + + +def _parse_logical_condition(condition: dict[str, Any]) -> dict[str, Any]: + if "operator" not in condition: + msg = f"'operator' key missing in {condition}" + raise FilterError(msg) + if "conditions" not in condition: + msg = f"'conditions' key missing in {condition}" + raise FilterError(msg) + + operator = condition["operator"] + conditions = [_normalize_filters(c) for c in condition["conditions"]] + + if operator in LOGICAL_OPERATORS: + return {LOGICAL_OPERATORS[operator]: conditions} + + msg = f"Unknown logical operator '{operator}'" + raise FilterError(msg) + + +def _parse_comparison_condition(condition: dict[str, Any]) -> dict[str, Any]: + if "field" not in condition: + return _parse_logical_condition(condition) + + field: str = condition["field"] + if "operator" not in condition: + msg = f"'operator' key missing in {condition}" + raise FilterError(msg) + if "value" not in condition: + msg = f"'value' key missing in {condition}" + raise FilterError(msg) + + operator: str = condition["operator"] + value: Any = condition["value"] + + # Strip the "meta." prefix — metadata is stored flat in S3 Vectors + if field.startswith("meta."): + field = field[5:] + + if operator not in COMPARISON_OPERATORS: + msg = f"Unknown comparison operator '{operator}'" + raise FilterError(msg) + + return COMPARISON_OPERATORS[operator](field, value) + + +def _equal(field: str, value: Any) -> dict[str, Any]: + _assert_supported_type(value, (str, int, float, bool), "equal") + return {field: {"$eq": value}} + + +def _not_equal(field: str, value: Any) -> dict[str, Any]: + _assert_supported_type(value, (str, int, float, bool), "not equal") + return {field: {"$ne": value}} + + +def _greater_than(field: str, value: Any) -> dict[str, Any]: + _assert_supported_type(value, (int, float), "greater than") + return {field: {"$gt": value}} + + +def _greater_than_equal(field: str, value: Any) -> dict[str, Any]: + _assert_supported_type(value, (int, float), "greater than equal") + return {field: {"$gte": value}} + + +def _less_than(field: str, value: Any) -> dict[str, Any]: + _assert_supported_type(value, (int, float), "less than") + return {field: {"$lt": value}} + + +def _less_than_equal(field: str, value: Any) -> dict[str, Any]: + _assert_supported_type(value, (int, float), "less than equal") + return {field: {"$lte": value}} + + +def _in(field: str, value: Any) -> dict[str, Any]: + if not isinstance(value, list): + msg = f"{field}'s value must be a list when using 'in' comparator" + raise FilterError(msg) + for v in value: + _assert_supported_type(v, (str, int, float), "in") + return {field: {"$in": value}} + + +def _not_in(field: str, value: Any) -> dict[str, Any]: + if not isinstance(value, list): + msg = f"{field}'s value must be a list when using 'not in' comparator" + raise FilterError(msg) + for v in value: + _assert_supported_type(v, (str, int, float), "not in") + return {field: {"$nin": value}} + + +def _assert_supported_type(value: Any, supported_types: tuple[type, ...], operator_name: str) -> None: + if not isinstance(value, supported_types): + msg = ( + f"Unsupported type for '{operator_name}' comparison: {type(value)}. " + f"Types supported by S3 Vectors are: {supported_types}" + ) + raise FilterError(msg) + + +COMPARISON_OPERATORS = { + "==": _equal, + "!=": _not_equal, + ">": _greater_than, + ">=": _greater_than_equal, + "<": _less_than, + "<=": _less_than_equal, + "in": _in, + "not in": _not_in, +} + +LOGICAL_OPERATORS = {"AND": "$and", "OR": "$or"} + + +def _validate_filters(filters: dict[str, Any] | None) -> None: + """Validate Haystack filter syntax.""" + if filters and "operator" not in filters and "conditions" not in filters: + msg = "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details." + raise ValueError(msg) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/py.typed b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/integrations/amazon_s3_vectors/tests/__init__.py b/integrations/amazon_s3_vectors/tests/__init__.py new file mode 100644 index 0000000000..c1764a6e03 --- /dev/null +++ b/integrations/amazon_s3_vectors/tests/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py new file mode 100644 index 0000000000..adf1133af4 --- /dev/null +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import MagicMock, patch + +import pytest +from haystack.dataclasses import Document +from haystack.document_stores.errors import DocumentStoreError +from haystack.document_stores.types import DuplicatePolicy + +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + + +@pytest.fixture +def mock_client(): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.list_vectors.return_value = {"vectors": []} + client.put_vectors.return_value = {} + client.delete_vectors.return_value = {} + client.get_vectors.return_value = {"vectors": []} + client.query_vectors.return_value = {"vectors": []} + return client + + +@pytest.fixture +def doc_store(mock_client): + with patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") as mock_boto3: + mock_boto3.client.return_value = mock_client + store = S3VectorsDocumentStore( + vector_bucket_name="test-bucket", + index_name="test-index", + dimension=768, + region_name="us-east-1", + ) + # Force initialization + store._get_client() + return store + + +class TestInit: + def test_default_params(self): + store = S3VectorsDocumentStore( + vector_bucket_name="my-bucket", + index_name="my-index", + dimension=768, + ) + assert store.vector_bucket_name == "my-bucket" + assert store.index_name == "my-index" + assert store.dimension == 768 + assert store.distance_metric == "cosine" + assert store.create_bucket_and_index is True + assert store.non_filterable_metadata_keys == [] + + +class TestSerialization: + def test_to_dict(self): + store = S3VectorsDocumentStore( + vector_bucket_name="my-bucket", + index_name="my-index", + dimension=768, + distance_metric="euclidean", + region_name="us-west-2", + create_bucket_and_index=False, + ) + d = store.to_dict() + assert d["init_parameters"]["vector_bucket_name"] == "my-bucket" + assert d["init_parameters"]["index_name"] == "my-index" + assert d["init_parameters"]["dimension"] == 768 + assert d["init_parameters"]["distance_metric"] == "euclidean" + assert d["init_parameters"]["region_name"] == "us-west-2" + assert d["init_parameters"]["create_bucket_and_index"] is False + + def test_from_dict(self): + store = S3VectorsDocumentStore( + vector_bucket_name="my-bucket", + index_name="my-index", + dimension=768, + ) + d = store.to_dict() + restored = S3VectorsDocumentStore.from_dict(d) + assert restored.vector_bucket_name == "my-bucket" + assert restored.index_name == "my-index" + assert restored.dimension == 768 + + +class TestCountDocuments: + def test_empty_index(self, doc_store, mock_client): + mock_client.list_vectors.return_value = {"vectors": []} + assert doc_store.count_documents() == 0 + + def test_with_documents(self, doc_store, mock_client): + mock_client.list_vectors.return_value = { + "vectors": [{"key": "1"}, {"key": "2"}, {"key": "3"}], + } + assert doc_store.count_documents() == 3 + + def test_pagination(self, doc_store, mock_client): + mock_client.list_vectors.side_effect = [ + {"vectors": [{"key": "1"}, {"key": "2"}], "nextToken": "token1"}, + {"vectors": [{"key": "3"}]}, + ] + assert doc_store.count_documents() == 3 + + +class TestWriteDocuments: + def test_write_documents(self, doc_store, mock_client): + docs = [ + Document(id="1", content="Hello", embedding=[0.1] * 768), + Document(id="2", content="World", embedding=[0.2] * 768), + ] + result = doc_store.write_documents(docs) + assert result == 2 + mock_client.put_vectors.assert_called_once() + + def test_write_empty_list(self, doc_store, mock_client): + result = doc_store.write_documents([]) + assert result == 0 + mock_client.put_vectors.assert_not_called() + + def test_write_without_embedding_raises(self, doc_store): + docs = [Document(id="1", content="Hello")] + with pytest.raises(DocumentStoreError, match="has no embedding"): + doc_store.write_documents(docs) + + def test_write_with_metadata(self, doc_store, mock_client): + docs = [ + Document( + id="1", + content="Hello", + embedding=[0.1] * 768, + meta={"category": "test", "year": 2024}, + ) + ] + doc_store.write_documents(docs) + call_args = mock_client.put_vectors.call_args + vectors = call_args[1]["vectors"] + assert len(vectors) == 1 + assert vectors[0]["key"] == "1" + assert vectors[0]["metadata"]["_content"] == "Hello" + assert vectors[0]["metadata"]["category"] == "test" + assert vectors[0]["metadata"]["year"] == 2024 + + def test_write_skip_existing(self, doc_store, mock_client): + mock_client.get_vectors.return_value = {"vectors": [{"key": "1"}]} + docs = [Document(id="1", content="Hello", embedding=[0.1] * 768)] + result = doc_store.write_documents(docs, policy=DuplicatePolicy.SKIP) + assert result == 0 + mock_client.put_vectors.assert_not_called() + + def test_write_none_policy_existing_raises(self, doc_store, mock_client): + mock_client.get_vectors.return_value = {"vectors": [{"key": "1"}]} + docs = [Document(id="1", content="Hello", embedding=[0.1] * 768)] + with pytest.raises(DocumentStoreError, match="already exist"): + doc_store.write_documents(docs, policy=DuplicatePolicy.NONE) + + +class TestDeleteDocuments: + def test_delete(self, doc_store, mock_client): + doc_store.delete_documents(["1", "2"]) + mock_client.delete_vectors.assert_called_once_with( + vectorBucketName="test-bucket", + indexName="test-index", + keys=["1", "2"], + ) + + def test_delete_empty_list(self, doc_store, mock_client): + doc_store.delete_documents([]) + mock_client.delete_vectors.assert_not_called() + + +class TestFilterDocuments: + def test_no_filters(self, doc_store, mock_client): + mock_client.list_vectors.return_value = { + "vectors": [ + { + "key": "1", + "data": {"float32": [0.1] * 768}, + "metadata": {"_content": "Hello", "category": "news"}, + } + ] + } + docs = doc_store.filter_documents() + assert len(docs) == 1 + assert docs[0].id == "1" + assert docs[0].content == "Hello" + assert docs[0].meta == {"category": "news"} + + def test_with_filters(self, doc_store, mock_client): + mock_client.list_vectors.return_value = { + "vectors": [ + { + "key": "1", + "data": {"float32": [0.1] * 768}, + "metadata": {"_content": "Hello", "category": "news"}, + }, + { + "key": "2", + "data": {"float32": [0.2] * 768}, + "metadata": {"_content": "World", "category": "sports"}, + }, + ] + } + filters = {"field": "meta.category", "operator": "==", "value": "news"} + docs = doc_store.filter_documents(filters=filters) + assert len(docs) == 1 + assert docs[0].id == "1" + + +class TestEmbeddingRetrieval: + def test_basic_retrieval(self, doc_store, mock_client): + mock_client.query_vectors.return_value = { + "vectors": [ + { + "key": "1", + "distance": 0.05, + "metadata": {"_content": "Hello", "category": "news"}, + } + ], + "distanceMetric": "cosine", + } + docs = doc_store._embedding_retrieval( + query_embedding=[0.1] * 768, + top_k=5, + ) + assert len(docs) == 1 + assert docs[0].id == "1" + assert docs[0].content == "Hello" + # cosine distance 0.05 → similarity score 0.95 + assert docs[0].score == pytest.approx(0.95) + assert docs[0].meta == {"category": "news"} + + def test_retrieval_with_filters(self, doc_store, mock_client): + mock_client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} + filters = { + "operator": "AND", + "conditions": [ + {"field": "meta.category", "operator": "==", "value": "news"}, + ], + } + doc_store._embedding_retrieval( + query_embedding=[0.1] * 768, + filters=filters, + top_k=5, + ) + call_args = mock_client.query_vectors.call_args[1] + assert "filter" in call_args + assert call_args["filter"] == {"$and": [{"category": {"$eq": "news"}}]} + + def test_empty_embedding_raises(self, doc_store): + with pytest.raises(ValueError, match="non-empty"): + doc_store._embedding_retrieval(query_embedding=[]) + + +class TestDocumentConversion: + def test_document_to_s3_vector(self): + doc = Document( + id="test-1", + content="Hello world", + embedding=[0.1, 0.2, 0.3], + meta={"category": "test", "year": 2024}, + ) + result = S3VectorsDocumentStore._document_to_s3_vector(doc) + assert result["key"] == "test-1" + assert result["data"] == {"float32": [0.1, 0.2, 0.3]} + assert result["metadata"]["_content"] == "Hello world" + assert result["metadata"]["category"] == "test" + assert result["metadata"]["year"] == 2024 + + def test_s3_vector_to_document(self): + vector = { + "key": "test-1", + "data": {"float32": [0.1, 0.2, 0.3]}, + "metadata": {"_content": "Hello world", "category": "test"}, + } + doc = S3VectorsDocumentStore._s3_vector_to_document(vector) + assert doc.id == "test-1" + assert doc.content == "Hello world" + assert doc.embedding == [0.1, 0.2, 0.3] + assert doc.meta == {"category": "test"} + + def test_roundtrip(self): + doc = Document( + id="test-1", + content="Hello world", + embedding=[0.1, 0.2, 0.3], + meta={"category": "test", "year": 2024}, + ) + vector = S3VectorsDocumentStore._document_to_s3_vector(doc) + restored = S3VectorsDocumentStore._s3_vector_to_document(vector) + assert restored.id == doc.id + assert restored.content == doc.content + assert restored.embedding == doc.embedding + assert restored.meta == doc.meta diff --git a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py new file mode 100644 index 0000000000..4bb86fdb1d --- /dev/null +++ b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import MagicMock, patch + +import pytest +from haystack.document_stores.types import FilterPolicy + +from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + + +@pytest.fixture +def mock_client(): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} + return client + + +@pytest.fixture +def doc_store(mock_client): + with patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") as mock_boto3: + mock_boto3.client.return_value = mock_client + store = S3VectorsDocumentStore( + vector_bucket_name="test-bucket", + index_name="test-index", + dimension=768, + region_name="us-east-1", + ) + store._get_client() + return store + + +class TestInit: + def test_valid_init(self, doc_store): + retriever = S3VectorsEmbeddingRetriever(document_store=doc_store) + assert retriever.document_store == doc_store + assert retriever.top_k == 10 + assert retriever.filters == {} + assert retriever.filter_policy == FilterPolicy.REPLACE + + def test_invalid_store_raises(self): + with pytest.raises(ValueError, match="must be an instance"): + S3VectorsEmbeddingRetriever(document_store="not a store") + + def test_custom_params(self, doc_store): + retriever = S3VectorsEmbeddingRetriever( + document_store=doc_store, + top_k=5, + filters={"operator": "AND", "conditions": []}, + filter_policy=FilterPolicy.MERGE, + ) + assert retriever.top_k == 5 + assert retriever.filter_policy == FilterPolicy.MERGE + + +class TestSerialization: + def test_to_dict(self, doc_store): + retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=5) + d = retriever.to_dict() + assert d["init_parameters"]["top_k"] == 5 + assert "document_store" in d["init_parameters"] + + def test_from_dict(self, doc_store): + retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=5) + d = retriever.to_dict() + with patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3"): + restored = S3VectorsEmbeddingRetriever.from_dict(d) + assert restored.top_k == 5 + + +class TestRun: + def test_run(self, doc_store, mock_client): + mock_client.query_vectors.return_value = { + "vectors": [ + { + "key": "1", + "distance": 0.05, + "metadata": {"_content": "Hello", "category": "news"}, + } + ], + "distanceMetric": "cosine", + } + retriever = S3VectorsEmbeddingRetriever(document_store=doc_store) + result = retriever.run(query_embedding=[0.1] * 768) + assert len(result["documents"]) == 1 + assert result["documents"][0].id == "1" + # cosine distance 0.05 -> similarity 0.95 + assert result["documents"][0].score == pytest.approx(0.95) + + def test_run_with_top_k_override(self, doc_store, mock_client): + mock_client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} + retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=10) + retriever.run(query_embedding=[0.1] * 768, top_k=3) + call_args = mock_client.query_vectors.call_args[1] + assert call_args["topK"] == 3 + + def test_run_with_filters(self, doc_store, mock_client): + mock_client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} + retriever = S3VectorsEmbeddingRetriever(document_store=doc_store) + filters = { + "operator": "AND", + "conditions": [ + {"field": "meta.category", "operator": "==", "value": "news"}, + ], + } + retriever.run(query_embedding=[0.1] * 768, filters=filters) + call_args = mock_client.query_vectors.call_args[1] + assert "filter" in call_args diff --git a/integrations/amazon_s3_vectors/tests/test_filters.py b/integrations/amazon_s3_vectors/tests/test_filters.py new file mode 100644 index 0000000000..88739f8837 --- /dev/null +++ b/integrations/amazon_s3_vectors/tests/test_filters.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from haystack.errors import FilterError + +from haystack_integrations.document_stores.amazon_s3_vectors.filters import _normalize_filters, _validate_filters + + +class TestNormalizeFilters: + def test_simple_equality(self): + filters = {"field": "meta.category", "operator": "==", "value": "news"} + result = _normalize_filters(filters) + assert result == {"category": {"$eq": "news"}} + + def test_not_equal(self): + filters = {"field": "meta.status", "operator": "!=", "value": "draft"} + result = _normalize_filters(filters) + assert result == {"status": {"$ne": "draft"}} + + def test_greater_than(self): + filters = {"field": "meta.year", "operator": ">", "value": 2020} + result = _normalize_filters(filters) + assert result == {"year": {"$gt": 2020}} + + def test_greater_than_equal(self): + filters = {"field": "meta.year", "operator": ">=", "value": 2020} + result = _normalize_filters(filters) + assert result == {"year": {"$gte": 2020}} + + def test_less_than(self): + filters = {"field": "meta.price", "operator": "<", "value": 50.0} + result = _normalize_filters(filters) + assert result == {"price": {"$lt": 50.0}} + + def test_less_than_equal(self): + filters = {"field": "meta.price", "operator": "<=", "value": 50.0} + result = _normalize_filters(filters) + assert result == {"price": {"$lte": 50.0}} + + def test_in_operator(self): + filters = {"field": "meta.genre", "operator": "in", "value": ["comedy", "drama"]} + result = _normalize_filters(filters) + assert result == {"genre": {"$in": ["comedy", "drama"]}} + + def test_not_in_operator(self): + filters = {"field": "meta.genre", "operator": "not in", "value": ["horror"]} + result = _normalize_filters(filters) + assert result == {"genre": {"$nin": ["horror"]}} + + def test_and_logical(self): + filters = { + "operator": "AND", + "conditions": [ + {"field": "meta.genre", "operator": "==", "value": "drama"}, + {"field": "meta.year", "operator": ">=", "value": 2020}, + ], + } + result = _normalize_filters(filters) + assert result == {"$and": [{"genre": {"$eq": "drama"}}, {"year": {"$gte": 2020}}]} + + def test_or_logical(self): + filters = { + "operator": "OR", + "conditions": [ + {"field": "meta.genre", "operator": "==", "value": "drama"}, + {"field": "meta.genre", "operator": "==", "value": "comedy"}, + ], + } + result = _normalize_filters(filters) + assert result == {"$or": [{"genre": {"$eq": "drama"}}, {"genre": {"$eq": "comedy"}}]} + + def test_field_without_meta_prefix(self): + filters = {"field": "category", "operator": "==", "value": "news"} + result = _normalize_filters(filters) + assert result == {"category": {"$eq": "news"}} + + def test_unsupported_type_raises(self): + filters = {"field": "meta.data", "operator": "==", "value": [1, 2, 3]} + with pytest.raises(FilterError): + _normalize_filters(filters) + + def test_in_with_non_list_raises(self): + filters = {"field": "meta.genre", "operator": "in", "value": "comedy"} + with pytest.raises(FilterError): + _normalize_filters(filters) + + def test_missing_operator_raises(self): + filters = {"field": "meta.genre", "value": "comedy"} + with pytest.raises(FilterError): + _normalize_filters(filters) + + def test_missing_value_raises(self): + filters = {"field": "meta.genre", "operator": "=="} + with pytest.raises(FilterError): + _normalize_filters(filters) + + def test_invalid_dict_raises(self): + with pytest.raises(FilterError): + _normalize_filters("not a dict") + + +class TestValidateFilters: + def test_valid_filters(self): + _validate_filters({"operator": "AND", "conditions": []}) + + def test_none_is_valid(self): + _validate_filters(None) + + def test_invalid_syntax_raises(self): + with pytest.raises(ValueError, match="Invalid filter syntax"): + _validate_filters({"field": "meta.x"}) diff --git a/integrations/amazon_s3_vectors/tests/test_integration.py b/integrations/amazon_s3_vectors/tests/test_integration.py new file mode 100644 index 0000000000..2ec71353f4 --- /dev/null +++ b/integrations/amazon_s3_vectors/tests/test_integration.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Integration tests for S3VectorsDocumentStore. + +Requires valid AWS credentials with s3vectors permissions. +Run with: hatch run test:integration +""" + +import os +import time +import uuid + +import pytest +from haystack.dataclasses import Document +from haystack.document_stores.errors import DocumentStoreError +from haystack.document_stores.types import DuplicatePolicy + +from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + + +def _aws_credentials_available() -> bool: + """Check if AWS credentials are available via any mechanism.""" + # Explicit env vars + if os.environ.get("AWS_ACCESS_KEY_ID") or os.environ.get("AWS_PROFILE") or os.environ.get("AWS_ROLE_ARN"): + return True + # Fallback: try to resolve credentials from the default chain + try: + import boto3 + + session = boto3.Session() + credentials = session.get_credentials() + return credentials is not None + except Exception: + return False + + +# Guard: skip all integration tests when AWS credentials are not available +pytestmark = pytest.mark.skipif( + not _aws_credentials_available(), + reason="AWS credentials not configured", +) + +# Use a small dimension to keep payloads light +DIMENSION = 4 +REGION = "us-east-1" + + +def _random_name(prefix: str = "haystack-test") -> str: + return f"{prefix}-{uuid.uuid4().hex[:8]}" + + +def _make_doc(doc_id: str, content: str, embedding: list[float] | None = None, meta: dict | None = None) -> Document: + return Document( + id=doc_id, + content=content, + embedding=embedding or [0.1] * DIMENSION, + meta=meta or {}, + ) + + +@pytest.fixture(scope="module") +def doc_store(): + """Create a real S3 Vectors document store with a unique bucket + index, tear down after.""" + bucket_name = _random_name("hs-integ") + index_name = _random_name("idx") + + store = S3VectorsDocumentStore( + vector_bucket_name=bucket_name, + index_name=index_name, + dimension=DIMENSION, + distance_metric="cosine", + region_name=REGION, + create_bucket_and_index=True, + non_filterable_metadata_keys=[], + ) + # Force initialization + store._get_client() + + yield store + + # Cleanup + client = store._client + try: + client.delete_index(vectorBucketName=bucket_name, indexName=index_name) + except Exception: + pass + try: + client.delete_vector_bucket(vectorBucketName=bucket_name) + except Exception: + pass + + +@pytest.mark.integration +class TestWriteAndCount: + def test_write_and_count(self, doc_store): + docs = [ + _make_doc("int-1", "First document"), + _make_doc("int-2", "Second document"), + _make_doc("int-3", "Third document"), + ] + written = doc_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) + assert written == 3 + + # S3 Vectors is eventually consistent — give it a moment + time.sleep(2) + + count = doc_store.count_documents() + assert count == 3 + + def test_overwrite(self, doc_store): + doc = _make_doc("int-1", "Updated first document", embedding=[0.9] * DIMENSION) + written = doc_store.write_documents([doc], policy=DuplicatePolicy.OVERWRITE) + assert written == 1 + + def test_skip_existing(self, doc_store): + doc = _make_doc("int-1", "Should be skipped") + written = doc_store.write_documents([doc], policy=DuplicatePolicy.SKIP) + assert written == 0 + + def test_none_policy_raises_on_existing(self, doc_store): + doc = _make_doc("int-1", "Should fail") + with pytest.raises(DocumentStoreError, match="already exist"): + doc_store.write_documents([doc], policy=DuplicatePolicy.NONE) + + def test_write_without_embedding_raises(self, doc_store): + doc = Document(id="no-emb", content="No embedding") + with pytest.raises(DocumentStoreError, match="has no embedding"): + doc_store.write_documents([doc]) + + +@pytest.mark.integration +class TestQuery: + def test_embedding_retrieval(self, doc_store): + # Query with an embedding close to [0.1, 0.1, 0.1, 0.1] + docs = doc_store._embedding_retrieval( + query_embedding=[0.1] * DIMENSION, + top_k=10, + ) + assert len(docs) > 0 + # All returned docs should have a score (cosine similarity) + for doc in docs: + assert doc.score is not None + assert doc.content is not None + + def test_embedding_retrieval_with_metadata_filter(self, doc_store): + # Write a doc with distinctive metadata + tagged_doc = _make_doc( + "int-tagged", + "Tagged document", + embedding=[0.5] * DIMENSION, + meta={"category": "special"}, + ) + doc_store.write_documents([tagged_doc], policy=DuplicatePolicy.OVERWRITE) + time.sleep(2) + + # Query with filter + filters = {"field": "meta.category", "operator": "==", "value": "special"} + docs = doc_store._embedding_retrieval( + query_embedding=[0.5] * DIMENSION, + filters=filters, + top_k=10, + ) + assert len(docs) >= 1 + assert all(d.meta.get("category") == "special" for d in docs) + + def test_retriever_component(self, doc_store): + retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=5) + result = retriever.run(query_embedding=[0.1] * DIMENSION) + assert "documents" in result + assert len(result["documents"]) > 0 + + +@pytest.mark.integration +class TestFilterDocuments: + def test_filter_documents_no_filter(self, doc_store): + docs = doc_store.filter_documents() + assert len(docs) > 0 + for doc in docs: + assert doc.id is not None + + def test_filter_documents_with_filter(self, doc_store): + filters = {"field": "meta.category", "operator": "==", "value": "special"} + docs = doc_store.filter_documents(filters=filters) + assert all(d.meta.get("category") == "special" for d in docs) + + +@pytest.mark.integration +class TestDelete: + def test_delete_documents(self, doc_store): + # Write a doc to delete + doc = _make_doc("int-to-delete", "Delete me") + doc_store.write_documents([doc], policy=DuplicatePolicy.OVERWRITE) + time.sleep(2) + + doc_store.delete_documents(["int-to-delete"]) + time.sleep(2) + + # Verify it's gone by trying to get it + client = doc_store._client + response = client.get_vectors( + vectorBucketName=doc_store.vector_bucket_name, + indexName=doc_store.index_name, + keys=["int-to-delete"], + ) + assert len(response.get("vectors", [])) == 0 + + +@pytest.mark.integration +class TestSerialization: + def test_roundtrip(self, doc_store): + d = doc_store.to_dict() + restored = S3VectorsDocumentStore.from_dict(d) + assert restored.vector_bucket_name == doc_store.vector_bucket_name + assert restored.index_name == doc_store.index_name + assert restored.dimension == doc_store.dimension + assert restored.distance_metric == doc_store.distance_metric From 3999b498aa4a90c67cbfe042f788dbccfd12fad8 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Mon, 13 Apr 2026 14:54:19 +0200 Subject: [PATCH 02/17] fix: pin boto3>=1.42.0, fix pydoc filename, quote GITHUB_OUTPUT in workflow - boto3 lower bound set to 1.42.0 (when s3vectors service was added) - pydoc filename changed to amazon_s3_vectors.md (underscores, matching folder name) - Quote $GITHUB_OUTPUT in workflow to fix shellcheck SC2086 --- .github/workflows/amazon_s3_vectors.yml | 4 ++-- integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml | 2 +- integrations/amazon_s3_vectors/pyproject.toml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/amazon_s3_vectors.yml b/.github/workflows/amazon_s3_vectors.yml index 703d462bd5..1fa282e190 100644 --- a/.github/workflows/amazon_s3_vectors.yml +++ b/.github/workflows/amazon_s3_vectors.yml @@ -44,8 +44,8 @@ jobs: steps: - id: set run: | - echo 'os=${{ github.event_name == 'push' && '["ubuntu-latest"]' || env.TEST_MATRIX_OS }}' >> $GITHUB_OUTPUT - echo 'python-version=${{ github.event_name == 'push' && '["3.10"]' || env.TEST_MATRIX_PYTHON }}' >> $GITHUB_OUTPUT + echo 'os=${{ github.event_name == 'push' && '["ubuntu-latest"]' || env.TEST_MATRIX_OS }}' >> "$GITHUB_OUTPUT" + echo 'python-version=${{ github.event_name == 'push' && '["3.10"]' || env.TEST_MATRIX_PYTHON }}' >> "$GITHUB_OUTPUT" run: name: Python ${{ matrix.python-version }} on ${{ startsWith(matrix.os, 'macos-') && 'macOS' || startsWith(matrix.os, 'windows-') && 'Windows' || 'Linux' }} diff --git a/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml b/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml index 97cb55f004..4320722107 100644 --- a/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml +++ b/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml @@ -10,5 +10,5 @@ processors: renderer: description: Amazon S3 Vectors integration for Haystack id: integrations-amazon_s3_vectors - filename: amazon-s3-vectors.md + filename: amazon_s3_vectors.md title: Amazon S3 Vectors diff --git a/integrations/amazon_s3_vectors/pyproject.toml b/integrations/amazon_s3_vectors/pyproject.toml index 3a8fa55a65..716ac88db8 100644 --- a/integrations/amazon_s3_vectors/pyproject.toml +++ b/integrations/amazon_s3_vectors/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["haystack-ai", "boto3"] +dependencies = ["haystack-ai", "boto3>=1.42.0"] [project.urls] Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_s3_vectors#readme" From aee963a8c54a8db7484a1a8e50e8e814f70c31dd Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Mon, 13 Apr 2026 14:56:02 +0200 Subject: [PATCH 03/17] fix: pin haystack-ai>=2.26.1 for FilterPolicy support --- integrations/amazon_s3_vectors/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integrations/amazon_s3_vectors/pyproject.toml b/integrations/amazon_s3_vectors/pyproject.toml index 716ac88db8..2aa84e5472 100644 --- a/integrations/amazon_s3_vectors/pyproject.toml +++ b/integrations/amazon_s3_vectors/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["haystack-ai", "boto3>=1.42.0"] +dependencies = ["haystack-ai>=2.26.1", "boto3>=1.42.0"] [project.urls] Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_s3_vectors#readme" From dcbf73dc839a3b06f904245cceeee3d716f38bcc Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Mon, 13 Apr 2026 15:03:07 +0200 Subject: [PATCH 04/17] refactor: align test style with repo conventions - Flatten test classes into standalone functions (matching pinecone/qdrant pattern) - Assert full serialized dict structure in to_dict/from_dict tests - Use Mock(spec=...) for retriever tests instead of MagicMock+patch - Verify _embedding_retrieval call args match exactly - Add test_from_dict_no_filter_policy (backward compat) - Add test_init_is_lazy --- integrations/amazon_s3_vectors/pyproject.toml | 2 +- .../tests/test_document_store.py | 570 +++++++++--------- .../tests/test_embedding_retriever.py | 249 +++++--- 3 files changed, 444 insertions(+), 377 deletions(-) diff --git a/integrations/amazon_s3_vectors/pyproject.toml b/integrations/amazon_s3_vectors/pyproject.toml index 2aa84e5472..c44b588f4b 100644 --- a/integrations/amazon_s3_vectors/pyproject.toml +++ b/integrations/amazon_s3_vectors/pyproject.toml @@ -147,7 +147,7 @@ ban-relative-imports = "parents" [tool.ruff.lint.per-file-ignores] # Tests can use magic values, assertions, relative imports, and don't need type annotations -"tests/**/*" = ["PLC0415", "PLR2004", "S101", "S110", "TID252", "D", "ANN"] +"tests/**/*" = ["E501", "PLC0415", "PLR2004", "S101", "S110", "TID252", "D", "ANN"] # Examples can contain print statements and long lines "examples/**/*" = ["D", "T201", "E501"] diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index adf1133af4..d9ac3d836c 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -12,285 +12,303 @@ from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore -@pytest.fixture -def mock_client(): +def test_init_is_lazy(): + store = S3VectorsDocumentStore( + vector_bucket_name="my-bucket", + index_name="my-index", + dimension=768, + ) + assert store._client is None + + +def test_init_default_params(): + store = S3VectorsDocumentStore( + vector_bucket_name="my-bucket", + index_name="my-index", + dimension=768, + ) + assert store.vector_bucket_name == "my-bucket" + assert store.index_name == "my-index" + assert store.dimension == 768 + assert store.distance_metric == "cosine" + assert store.create_bucket_and_index is True + assert store.non_filterable_metadata_keys == [] + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_to_dict(_mock_boto3): + store = S3VectorsDocumentStore( + vector_bucket_name="my-bucket", + index_name="my-index", + dimension=768, + distance_metric="euclidean", + region_name="us-west-2", + create_bucket_and_index=False, + ) + d = store.to_dict() + assert d == { + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "init_parameters": { + "vector_bucket_name": "my-bucket", + "index_name": "my-index", + "dimension": 768, + "distance_metric": "euclidean", + "region_name": "us-west-2", + "aws_access_key_id": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "create_bucket_and_index": False, + "non_filterable_metadata_keys": [], + }, + } + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_from_dict(_mock_boto3): + data = { + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "init_parameters": { + "vector_bucket_name": "my-bucket", + "index_name": "my-index", + "dimension": 768, + "distance_metric": "euclidean", + "region_name": "us-west-2", + "aws_access_key_id": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "create_bucket_and_index": False, + "non_filterable_metadata_keys": [], + }, + } + store = S3VectorsDocumentStore.from_dict(data) + assert store.vector_bucket_name == "my-bucket" + assert store.index_name == "my-index" + assert store.dimension == 768 + assert store.distance_metric == "euclidean" + assert store.region_name == "us-west-2" + assert store.create_bucket_and_index is False + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_count_documents_empty(mock_boto3): client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} client.list_vectors.return_value = {"vectors": []} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + assert store.count_documents() == 0 + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_count_documents_pagination(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.list_vectors.side_effect = [ + {"vectors": [{"key": "1"}, {"key": "2"}], "nextToken": "tok"}, + {"vectors": [{"key": "3"}]}, + ] + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + assert store.count_documents() == 3 + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.put_vectors.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + docs = [ + Document(id="1", content="Hello", embedding=[0.1] * 4), + Document(id="2", content="World", embedding=[0.2] * 4), + ] + assert store.write_documents(docs) == 2 + client.put_vectors.assert_called_once() + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents_empty(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + assert store.write_documents([]) == 0 + client.put_vectors.assert_not_called() + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents_no_embedding_raises(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + with pytest.raises(DocumentStoreError, match="has no embedding"): + store.write_documents([Document(id="1", content="Hello")]) + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents_skip_existing(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.get_vectors.return_value = {"vectors": [{"key": "1"}]} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + docs = [Document(id="1", content="Hello", embedding=[0.1] * 4)] + assert store.write_documents(docs, policy=DuplicatePolicy.SKIP) == 0 + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents_none_policy_raises(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.get_vectors.return_value = {"vectors": [{"key": "1"}]} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + docs = [Document(id="1", content="Hello", embedding=[0.1] * 4)] + with pytest.raises(DocumentStoreError, match="already exist"): + store.write_documents(docs, policy=DuplicatePolicy.NONE) + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents_metadata(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} client.put_vectors.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + docs = [Document(id="1", content="Hello", embedding=[0.1] * 4, meta={"category": "test", "year": 2024})] + store.write_documents(docs) + + vectors = client.put_vectors.call_args[1]["vectors"] + assert len(vectors) == 1 + assert vectors[0]["key"] == "1" + assert vectors[0]["metadata"]["_content"] == "Hello" + assert vectors[0]["metadata"]["category"] == "test" + assert vectors[0]["metadata"]["year"] == 2024 + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_delete_documents(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} client.delete_vectors.return_value = {} - client.get_vectors.return_value = {"vectors": []} - client.query_vectors.return_value = {"vectors": []} - return client - - -@pytest.fixture -def doc_store(mock_client): - with patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") as mock_boto3: - mock_boto3.client.return_value = mock_client - store = S3VectorsDocumentStore( - vector_bucket_name="test-bucket", - index_name="test-index", - dimension=768, - region_name="us-east-1", - ) - # Force initialization - store._get_client() - return store - - -class TestInit: - def test_default_params(self): - store = S3VectorsDocumentStore( - vector_bucket_name="my-bucket", - index_name="my-index", - dimension=768, - ) - assert store.vector_bucket_name == "my-bucket" - assert store.index_name == "my-index" - assert store.dimension == 768 - assert store.distance_metric == "cosine" - assert store.create_bucket_and_index is True - assert store.non_filterable_metadata_keys == [] - - -class TestSerialization: - def test_to_dict(self): - store = S3VectorsDocumentStore( - vector_bucket_name="my-bucket", - index_name="my-index", - dimension=768, - distance_metric="euclidean", - region_name="us-west-2", - create_bucket_and_index=False, - ) - d = store.to_dict() - assert d["init_parameters"]["vector_bucket_name"] == "my-bucket" - assert d["init_parameters"]["index_name"] == "my-index" - assert d["init_parameters"]["dimension"] == 768 - assert d["init_parameters"]["distance_metric"] == "euclidean" - assert d["init_parameters"]["region_name"] == "us-west-2" - assert d["init_parameters"]["create_bucket_and_index"] is False - - def test_from_dict(self): - store = S3VectorsDocumentStore( - vector_bucket_name="my-bucket", - index_name="my-index", - dimension=768, - ) - d = store.to_dict() - restored = S3VectorsDocumentStore.from_dict(d) - assert restored.vector_bucket_name == "my-bucket" - assert restored.index_name == "my-index" - assert restored.dimension == 768 - - -class TestCountDocuments: - def test_empty_index(self, doc_store, mock_client): - mock_client.list_vectors.return_value = {"vectors": []} - assert doc_store.count_documents() == 0 - - def test_with_documents(self, doc_store, mock_client): - mock_client.list_vectors.return_value = { - "vectors": [{"key": "1"}, {"key": "2"}, {"key": "3"}], - } - assert doc_store.count_documents() == 3 - - def test_pagination(self, doc_store, mock_client): - mock_client.list_vectors.side_effect = [ - {"vectors": [{"key": "1"}, {"key": "2"}], "nextToken": "token1"}, - {"vectors": [{"key": "3"}]}, - ] - assert doc_store.count_documents() == 3 - - -class TestWriteDocuments: - def test_write_documents(self, doc_store, mock_client): - docs = [ - Document(id="1", content="Hello", embedding=[0.1] * 768), - Document(id="2", content="World", embedding=[0.2] * 768), - ] - result = doc_store.write_documents(docs) - assert result == 2 - mock_client.put_vectors.assert_called_once() - - def test_write_empty_list(self, doc_store, mock_client): - result = doc_store.write_documents([]) - assert result == 0 - mock_client.put_vectors.assert_not_called() - - def test_write_without_embedding_raises(self, doc_store): - docs = [Document(id="1", content="Hello")] - with pytest.raises(DocumentStoreError, match="has no embedding"): - doc_store.write_documents(docs) - - def test_write_with_metadata(self, doc_store, mock_client): - docs = [ - Document( - id="1", - content="Hello", - embedding=[0.1] * 768, - meta={"category": "test", "year": 2024}, - ) - ] - doc_store.write_documents(docs) - call_args = mock_client.put_vectors.call_args - vectors = call_args[1]["vectors"] - assert len(vectors) == 1 - assert vectors[0]["key"] == "1" - assert vectors[0]["metadata"]["_content"] == "Hello" - assert vectors[0]["metadata"]["category"] == "test" - assert vectors[0]["metadata"]["year"] == 2024 - - def test_write_skip_existing(self, doc_store, mock_client): - mock_client.get_vectors.return_value = {"vectors": [{"key": "1"}]} - docs = [Document(id="1", content="Hello", embedding=[0.1] * 768)] - result = doc_store.write_documents(docs, policy=DuplicatePolicy.SKIP) - assert result == 0 - mock_client.put_vectors.assert_not_called() - - def test_write_none_policy_existing_raises(self, doc_store, mock_client): - mock_client.get_vectors.return_value = {"vectors": [{"key": "1"}]} - docs = [Document(id="1", content="Hello", embedding=[0.1] * 768)] - with pytest.raises(DocumentStoreError, match="already exist"): - doc_store.write_documents(docs, policy=DuplicatePolicy.NONE) - - -class TestDeleteDocuments: - def test_delete(self, doc_store, mock_client): - doc_store.delete_documents(["1", "2"]) - mock_client.delete_vectors.assert_called_once_with( - vectorBucketName="test-bucket", - indexName="test-index", - keys=["1", "2"], - ) - - def test_delete_empty_list(self, doc_store, mock_client): - doc_store.delete_documents([]) - mock_client.delete_vectors.assert_not_called() - - -class TestFilterDocuments: - def test_no_filters(self, doc_store, mock_client): - mock_client.list_vectors.return_value = { - "vectors": [ - { - "key": "1", - "data": {"float32": [0.1] * 768}, - "metadata": {"_content": "Hello", "category": "news"}, - } - ] - } - docs = doc_store.filter_documents() - assert len(docs) == 1 - assert docs[0].id == "1" - assert docs[0].content == "Hello" - assert docs[0].meta == {"category": "news"} - - def test_with_filters(self, doc_store, mock_client): - mock_client.list_vectors.return_value = { - "vectors": [ - { - "key": "1", - "data": {"float32": [0.1] * 768}, - "metadata": {"_content": "Hello", "category": "news"}, - }, - { - "key": "2", - "data": {"float32": [0.2] * 768}, - "metadata": {"_content": "World", "category": "sports"}, - }, - ] - } - filters = {"field": "meta.category", "operator": "==", "value": "news"} - docs = doc_store.filter_documents(filters=filters) - assert len(docs) == 1 - assert docs[0].id == "1" - - -class TestEmbeddingRetrieval: - def test_basic_retrieval(self, doc_store, mock_client): - mock_client.query_vectors.return_value = { - "vectors": [ - { - "key": "1", - "distance": 0.05, - "metadata": {"_content": "Hello", "category": "news"}, - } - ], - "distanceMetric": "cosine", - } - docs = doc_store._embedding_retrieval( - query_embedding=[0.1] * 768, - top_k=5, - ) - assert len(docs) == 1 - assert docs[0].id == "1" - assert docs[0].content == "Hello" - # cosine distance 0.05 → similarity score 0.95 - assert docs[0].score == pytest.approx(0.95) - assert docs[0].meta == {"category": "news"} - - def test_retrieval_with_filters(self, doc_store, mock_client): - mock_client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} - filters = { - "operator": "AND", - "conditions": [ - {"field": "meta.category", "operator": "==", "value": "news"}, - ], - } - doc_store._embedding_retrieval( - query_embedding=[0.1] * 768, - filters=filters, - top_k=5, - ) - call_args = mock_client.query_vectors.call_args[1] - assert "filter" in call_args - assert call_args["filter"] == {"$and": [{"category": {"$eq": "news"}}]} - - def test_empty_embedding_raises(self, doc_store): - with pytest.raises(ValueError, match="non-empty"): - doc_store._embedding_retrieval(query_embedding=[]) - - -class TestDocumentConversion: - def test_document_to_s3_vector(self): - doc = Document( - id="test-1", - content="Hello world", - embedding=[0.1, 0.2, 0.3], - meta={"category": "test", "year": 2024}, - ) - result = S3VectorsDocumentStore._document_to_s3_vector(doc) - assert result["key"] == "test-1" - assert result["data"] == {"float32": [0.1, 0.2, 0.3]} - assert result["metadata"]["_content"] == "Hello world" - assert result["metadata"]["category"] == "test" - assert result["metadata"]["year"] == 2024 - - def test_s3_vector_to_document(self): - vector = { - "key": "test-1", - "data": {"float32": [0.1, 0.2, 0.3]}, - "metadata": {"_content": "Hello world", "category": "test"}, - } - doc = S3VectorsDocumentStore._s3_vector_to_document(vector) - assert doc.id == "test-1" - assert doc.content == "Hello world" - assert doc.embedding == [0.1, 0.2, 0.3] - assert doc.meta == {"category": "test"} - - def test_roundtrip(self): - doc = Document( - id="test-1", - content="Hello world", - embedding=[0.1, 0.2, 0.3], - meta={"category": "test", "year": 2024}, - ) - vector = S3VectorsDocumentStore._document_to_s3_vector(doc) - restored = S3VectorsDocumentStore._s3_vector_to_document(vector) - assert restored.id == doc.id - assert restored.content == doc.content - assert restored.embedding == doc.embedding - assert restored.meta == doc.meta + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store.delete_documents(["1", "2"]) + client.delete_vectors.assert_called_once_with(vectorBucketName="b", indexName="i", keys=["1", "2"]) + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_delete_documents_empty(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store.delete_documents([]) + client.delete_vectors.assert_not_called() + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_embedding_retrieval(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.query_vectors.return_value = { + "vectors": [{"key": "1", "distance": 0.05, "metadata": {"_content": "Hello", "category": "news"}}], + "distanceMetric": "cosine", + } + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + docs = store._embedding_retrieval(query_embedding=[0.1] * 4, top_k=5) + assert len(docs) == 1 + assert docs[0].id == "1" + assert docs[0].content == "Hello" + assert docs[0].score == pytest.approx(0.95) # 1 - 0.05 + assert docs[0].meta == {"category": "news"} + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_embedding_retrieval_with_filters(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + filters = {"operator": "AND", "conditions": [{"field": "meta.category", "operator": "==", "value": "news"}]} + store._embedding_retrieval(query_embedding=[0.1] * 4, filters=filters, top_k=5) + + call_args = client.query_vectors.call_args[1] + assert call_args["filter"] == {"$and": [{"category": {"$eq": "news"}}]} + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_embedding_retrieval_empty_embedding_raises(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + with pytest.raises(ValueError, match="non-empty"): + store._embedding_retrieval(query_embedding=[]) + + +def test_document_to_s3_vector(): + doc = Document( + id="test-1", content="Hello world", embedding=[0.1, 0.2, 0.3], meta={"category": "test", "year": 2024} + ) + result = S3VectorsDocumentStore._document_to_s3_vector(doc) + assert result["key"] == "test-1" + assert result["data"] == {"float32": [0.1, 0.2, 0.3]} + assert result["metadata"]["_content"] == "Hello world" + assert result["metadata"]["category"] == "test" + assert result["metadata"]["year"] == 2024 + + +def test_s3_vector_to_document(): + vector = { + "key": "test-1", + "data": {"float32": [0.1, 0.2, 0.3]}, + "metadata": {"_content": "Hello world", "category": "test"}, + } + doc = S3VectorsDocumentStore._s3_vector_to_document(vector) + assert doc.id == "test-1" + assert doc.content == "Hello world" + assert doc.embedding == [0.1, 0.2, 0.3] + assert doc.meta == {"category": "test"} + + +def test_document_roundtrip(): + doc = Document( + id="test-1", content="Hello world", embedding=[0.1, 0.2, 0.3], meta={"category": "test", "year": 2024} + ) + vector = S3VectorsDocumentStore._document_to_s3_vector(doc) + restored = S3VectorsDocumentStore._s3_vector_to_document(vector) + assert restored.id == doc.id + assert restored.content == doc.content + assert restored.embedding == doc.embedding + assert restored.meta == doc.meta diff --git a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py index 4bb86fdb1d..a3566a5c1f 100644 --- a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py +++ b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py @@ -2,111 +2,160 @@ # # SPDX-License-Identifier: Apache-2.0 -from unittest.mock import MagicMock, patch +from unittest.mock import Mock, patch import pytest +from haystack.dataclasses import Document from haystack.document_stores.types import FilterPolicy from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore -@pytest.fixture -def mock_client(): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} - return client - - -@pytest.fixture -def doc_store(mock_client): - with patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") as mock_boto3: - mock_boto3.client.return_value = mock_client - store = S3VectorsDocumentStore( - vector_bucket_name="test-bucket", - index_name="test-index", - dimension=768, - region_name="us-east-1", - ) - store._get_client() - return store - - -class TestInit: - def test_valid_init(self, doc_store): - retriever = S3VectorsEmbeddingRetriever(document_store=doc_store) - assert retriever.document_store == doc_store - assert retriever.top_k == 10 - assert retriever.filters == {} - assert retriever.filter_policy == FilterPolicy.REPLACE - - def test_invalid_store_raises(self): - with pytest.raises(ValueError, match="must be an instance"): - S3VectorsEmbeddingRetriever(document_store="not a store") - - def test_custom_params(self, doc_store): - retriever = S3VectorsEmbeddingRetriever( - document_store=doc_store, - top_k=5, - filters={"operator": "AND", "conditions": []}, - filter_policy=FilterPolicy.MERGE, - ) - assert retriever.top_k == 5 - assert retriever.filter_policy == FilterPolicy.MERGE - - -class TestSerialization: - def test_to_dict(self, doc_store): - retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=5) - d = retriever.to_dict() - assert d["init_parameters"]["top_k"] == 5 - assert "document_store" in d["init_parameters"] - - def test_from_dict(self, doc_store): - retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=5) - d = retriever.to_dict() - with patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3"): - restored = S3VectorsEmbeddingRetriever.from_dict(d) - assert restored.top_k == 5 - - -class TestRun: - def test_run(self, doc_store, mock_client): - mock_client.query_vectors.return_value = { - "vectors": [ - { - "key": "1", - "distance": 0.05, - "metadata": {"_content": "Hello", "category": "news"}, - } - ], - "distanceMetric": "cosine", - } - retriever = S3VectorsEmbeddingRetriever(document_store=doc_store) - result = retriever.run(query_embedding=[0.1] * 768) - assert len(result["documents"]) == 1 - assert result["documents"][0].id == "1" - # cosine distance 0.05 -> similarity 0.95 - assert result["documents"][0].score == pytest.approx(0.95) - - def test_run_with_top_k_override(self, doc_store, mock_client): - mock_client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} - retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=10) - retriever.run(query_embedding=[0.1] * 768, top_k=3) - call_args = mock_client.query_vectors.call_args[1] - assert call_args["topK"] == 3 - - def test_run_with_filters(self, doc_store, mock_client): - mock_client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} - retriever = S3VectorsEmbeddingRetriever(document_store=doc_store) - filters = { - "operator": "AND", - "conditions": [ - {"field": "meta.category", "operator": "==", "value": "news"}, - ], - } - retriever.run(query_embedding=[0.1] * 768, filters=filters) - call_args = mock_client.query_vectors.call_args[1] - assert "filter" in call_args +def test_init_default(): + mock_store = Mock(spec=S3VectorsDocumentStore) + retriever = S3VectorsEmbeddingRetriever(document_store=mock_store) + assert retriever.document_store == mock_store + assert retriever.filters == {} + assert retriever.top_k == 10 + assert retriever.filter_policy == FilterPolicy.REPLACE + + retriever = S3VectorsEmbeddingRetriever(document_store=mock_store, filter_policy="replace") + assert retriever.filter_policy == FilterPolicy.REPLACE + + with pytest.raises(ValueError): + S3VectorsEmbeddingRetriever(document_store=mock_store, filter_policy="invalid") + + +def test_init_invalid_store(): + with pytest.raises(ValueError, match="must be an instance"): + S3VectorsEmbeddingRetriever(document_store="not a store") + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_to_dict(_mock_boto3): + store = S3VectorsDocumentStore( + vector_bucket_name="test-bucket", + index_name="test-index", + dimension=768, + region_name="us-east-1", + create_bucket_and_index=False, + ) + retriever = S3VectorsEmbeddingRetriever(document_store=store, top_k=5) + d = retriever.to_dict() + assert d == { + "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.S3VectorsEmbeddingRetriever", + "init_parameters": { + "document_store": { + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "init_parameters": { + "vector_bucket_name": "test-bucket", + "index_name": "test-index", + "dimension": 768, + "distance_metric": "cosine", + "region_name": "us-east-1", + "aws_access_key_id": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "create_bucket_and_index": False, + "non_filterable_metadata_keys": [], + }, + }, + "filters": {}, + "top_k": 5, + "filter_policy": "replace", + }, + } + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_from_dict(_mock_boto3): + data = { + "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.S3VectorsEmbeddingRetriever", + "init_parameters": { + "document_store": { + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "init_parameters": { + "vector_bucket_name": "test-bucket", + "index_name": "test-index", + "dimension": 768, + "distance_metric": "cosine", + "region_name": "us-east-1", + "aws_access_key_id": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "create_bucket_and_index": False, + "non_filterable_metadata_keys": [], + }, + }, + "filters": {}, + "top_k": 5, + "filter_policy": "replace", + }, + } + retriever = S3VectorsEmbeddingRetriever.from_dict(data) + assert retriever.top_k == 5 + assert retriever.filter_policy == FilterPolicy.REPLACE + assert retriever.document_store.vector_bucket_name == "test-bucket" + assert retriever.document_store.dimension == 768 + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_from_dict_no_filter_policy(_mock_boto3): + data = { + "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.S3VectorsEmbeddingRetriever", + "init_parameters": { + "document_store": { + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "init_parameters": { + "vector_bucket_name": "test-bucket", + "index_name": "test-index", + "dimension": 768, + "create_bucket_and_index": False, + }, + }, + "filters": {}, + "top_k": 10, + }, + } + retriever = S3VectorsEmbeddingRetriever.from_dict(data) + assert retriever.filter_policy == FilterPolicy.REPLACE # default + + +def test_run(): + mock_store = Mock(spec=S3VectorsDocumentStore) + mock_store._embedding_retrieval.return_value = [Document(content="Test doc", embedding=[0.1, 0.2])] + retriever = S3VectorsEmbeddingRetriever(document_store=mock_store) + res = retriever.run(query_embedding=[0.5, 0.7]) + mock_store._embedding_retrieval.assert_called_once_with( + query_embedding=[0.5, 0.7], + filters={}, + top_k=10, + ) + assert len(res["documents"]) == 1 + assert res["documents"][0].content == "Test doc" + + +def test_run_with_top_k_override(): + mock_store = Mock(spec=S3VectorsDocumentStore) + mock_store._embedding_retrieval.return_value = [] + retriever = S3VectorsEmbeddingRetriever(document_store=mock_store, top_k=10) + retriever.run(query_embedding=[0.1, 0.2], top_k=3) + mock_store._embedding_retrieval.assert_called_once_with( + query_embedding=[0.1, 0.2], + filters={}, + top_k=3, + ) + + +def test_run_with_filters(): + mock_store = Mock(spec=S3VectorsDocumentStore) + mock_store._embedding_retrieval.return_value = [] + retriever = S3VectorsEmbeddingRetriever(document_store=mock_store) + filters = {"operator": "AND", "conditions": [{"field": "meta.category", "operator": "==", "value": "news"}]} + retriever.run(query_embedding=[0.1, 0.2], filters=filters) + mock_store._embedding_retrieval.assert_called_once_with( + query_embedding=[0.1, 0.2], + filters=filters, + top_k=10, + ) From c9c6aed06d307cfc14e92e4ff1f233183e1f6242 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Mon, 13 Apr 2026 15:06:19 +0200 Subject: [PATCH 05/17] refactor: trim tests to high-signal only Remove tests that just verify mock plumbing (count, write, delete calling the mock client). Keep tests that verify our actual logic: - Serialization roundtrip (full dict structure) - Score conversion (cosine + euclidean) - Filter conversion (pure function with real logic) - Duplicate policy batch checks (SKIP/NONE) - Document <-> S3 vector conversion - Input validation Before: 49 unit tests (many testing mock behavior) After: 26 unit tests (all testing our code) + 12 integration tests --- .../tests/test_document_store.py | 155 +++++------------- .../tests/test_embedding_retriever.py | 31 +--- .../amazon_s3_vectors/tests/test_filters.py | 144 +++++++--------- 3 files changed, 102 insertions(+), 228 deletions(-) diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index d9ac3d836c..8df69ffc49 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -89,64 +89,9 @@ def test_from_dict(_mock_boto3): assert store.create_bucket_and_index is False -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_count_documents_empty(mock_boto3): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - client.list_vectors.return_value = {"vectors": []} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - assert store.count_documents() == 0 - - -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_count_documents_pagination(mock_boto3): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - client.list_vectors.side_effect = [ - {"vectors": [{"key": "1"}, {"key": "2"}], "nextToken": "tok"}, - {"vectors": [{"key": "3"}]}, - ] - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - assert store.count_documents() == 3 - - -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents(mock_boto3): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - client.put_vectors.return_value = {} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - docs = [ - Document(id="1", content="Hello", embedding=[0.1] * 4), - Document(id="2", content="World", embedding=[0.2] * 4), - ] - assert store.write_documents(docs) == 2 - client.put_vectors.assert_called_once() - - -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_empty(mock_boto3): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - assert store.write_documents([]) == 0 - client.put_vectors.assert_not_called() - - @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_write_documents_no_embedding_raises(mock_boto3): + """S3 Vectors requires embeddings — this tests our validation, not the store.""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} @@ -159,6 +104,7 @@ def test_write_documents_no_embedding_raises(mock_boto3): @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_write_documents_skip_existing(mock_boto3): + """Tests our batch existence check logic for SKIP policy.""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} @@ -166,12 +112,17 @@ def test_write_documents_skip_existing(mock_boto3): mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - docs = [Document(id="1", content="Hello", embedding=[0.1] * 4)] - assert store.write_documents(docs, policy=DuplicatePolicy.SKIP) == 0 + result = store.write_documents( + [Document(id="1", content="Hello", embedding=[0.1] * 4)], + policy=DuplicatePolicy.SKIP, + ) + assert result == 0 + client.put_vectors.assert_not_called() @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_write_documents_none_policy_raises(mock_boto3): + """Tests our batch existence check logic for NONE policy.""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} @@ -179,78 +130,56 @@ def test_write_documents_none_policy_raises(mock_boto3): mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - docs = [Document(id="1", content="Hello", embedding=[0.1] * 4)] with pytest.raises(DocumentStoreError, match="already exist"): - store.write_documents(docs, policy=DuplicatePolicy.NONE) + store.write_documents( + [Document(id="1", content="Hello", embedding=[0.1] * 4)], + policy=DuplicatePolicy.NONE, + ) @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_metadata(mock_boto3): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - client.put_vectors.return_value = {} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - docs = [Document(id="1", content="Hello", embedding=[0.1] * 4, meta={"category": "test", "year": 2024})] - store.write_documents(docs) - - vectors = client.put_vectors.call_args[1]["vectors"] - assert len(vectors) == 1 - assert vectors[0]["key"] == "1" - assert vectors[0]["metadata"]["_content"] == "Hello" - assert vectors[0]["metadata"]["category"] == "test" - assert vectors[0]["metadata"]["year"] == 2024 - - -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_delete_documents(mock_boto3): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - client.delete_vectors.return_value = {} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - store.delete_documents(["1", "2"]) - client.delete_vectors.assert_called_once_with(vectorBucketName="b", indexName="i", keys=["1", "2"]) - - -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_delete_documents_empty(mock_boto3): +def test_embedding_retrieval_score_conversion(mock_boto3): + """Tests our distance-to-score conversion logic — the only non-trivial transform in retrieval.""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} + client.query_vectors.return_value = { + "vectors": [{"key": "1", "distance": 0.05, "metadata": {"_content": "Hello", "category": "news"}}], + "distanceMetric": "cosine", + } mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - store.delete_documents([]) - client.delete_vectors.assert_not_called() + docs = store._embedding_retrieval(query_embedding=[0.1] * 4, top_k=5) + assert len(docs) == 1 + assert docs[0].id == "1" + assert docs[0].content == "Hello" + assert docs[0].score == pytest.approx(0.95) # cosine: 1.0 - 0.05 + assert docs[0].meta == {"category": "news"} @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_embedding_retrieval(mock_boto3): +def test_embedding_retrieval_euclidean_score(mock_boto3): + """Tests euclidean distance-to-score conversion (negated).""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} client.query_vectors.return_value = { - "vectors": [{"key": "1", "distance": 0.05, "metadata": {"_content": "Hello", "category": "news"}}], - "distanceMetric": "cosine", + "vectors": [{"key": "1", "distance": 1.5, "metadata": {}}], + "distanceMetric": "euclidean", } mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = S3VectorsDocumentStore( + vector_bucket_name="b", index_name="i", dimension=4, distance_metric="euclidean", region_name="us-east-1" + ) docs = store._embedding_retrieval(query_embedding=[0.1] * 4, top_k=5) - assert len(docs) == 1 - assert docs[0].id == "1" - assert docs[0].content == "Hello" - assert docs[0].score == pytest.approx(0.95) # 1 - 0.05 - assert docs[0].meta == {"category": "news"} + assert docs[0].score == pytest.approx(-1.5) # euclidean: negated @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_embedding_retrieval_with_filters(mock_boto3): +def test_embedding_retrieval_passes_filters(mock_boto3): + """Tests that Haystack filters are converted and passed to query_vectors.""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} @@ -265,19 +194,15 @@ def test_embedding_retrieval_with_filters(mock_boto3): assert call_args["filter"] == {"$and": [{"category": {"$eq": "news"}}]} -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_embedding_retrieval_empty_embedding_raises(mock_boto3): - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") +def test_embedding_retrieval_empty_embedding_raises(): + """Tests our input validation — no mocking needed.""" + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False) with pytest.raises(ValueError, match="non-empty"): store._embedding_retrieval(query_embedding=[]) def test_document_to_s3_vector(): + """Tests our Document → S3 vector conversion (pure function).""" doc = Document( id="test-1", content="Hello world", embedding=[0.1, 0.2, 0.3], meta={"category": "test", "year": 2024} ) @@ -290,6 +215,7 @@ def test_document_to_s3_vector(): def test_s3_vector_to_document(): + """Tests our S3 vector → Document conversion (pure function).""" vector = { "key": "test-1", "data": {"float32": [0.1, 0.2, 0.3]}, @@ -303,6 +229,7 @@ def test_s3_vector_to_document(): def test_document_roundtrip(): + """Tests Document → S3 vector → Document is lossless.""" doc = Document( id="test-1", content="Hello world", embedding=[0.1, 0.2, 0.3], meta={"category": "test", "year": 2024} ) diff --git a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py index a3566a5c1f..3f9e2e46da 100644 --- a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py +++ b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py @@ -27,11 +27,6 @@ def test_init_default(): S3VectorsEmbeddingRetriever(document_store=mock_store, filter_policy="invalid") -def test_init_invalid_store(): - with pytest.raises(ValueError, match="must be an instance"): - S3VectorsEmbeddingRetriever(document_store="not a store") - - @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_to_dict(_mock_boto3): store = S3VectorsDocumentStore( @@ -102,6 +97,7 @@ def test_from_dict(_mock_boto3): @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_from_dict_no_filter_policy(_mock_boto3): + """Pipelines serialized with older versions may not have filter_policy.""" data = { "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.S3VectorsEmbeddingRetriever", "init_parameters": { @@ -134,28 +130,3 @@ def test_run(): ) assert len(res["documents"]) == 1 assert res["documents"][0].content == "Test doc" - - -def test_run_with_top_k_override(): - mock_store = Mock(spec=S3VectorsDocumentStore) - mock_store._embedding_retrieval.return_value = [] - retriever = S3VectorsEmbeddingRetriever(document_store=mock_store, top_k=10) - retriever.run(query_embedding=[0.1, 0.2], top_k=3) - mock_store._embedding_retrieval.assert_called_once_with( - query_embedding=[0.1, 0.2], - filters={}, - top_k=3, - ) - - -def test_run_with_filters(): - mock_store = Mock(spec=S3VectorsDocumentStore) - mock_store._embedding_retrieval.return_value = [] - retriever = S3VectorsEmbeddingRetriever(document_store=mock_store) - filters = {"operator": "AND", "conditions": [{"field": "meta.category", "operator": "==", "value": "news"}]} - retriever.run(query_embedding=[0.1, 0.2], filters=filters) - mock_store._embedding_retrieval.assert_called_once_with( - query_embedding=[0.1, 0.2], - filters=filters, - top_k=10, - ) diff --git a/integrations/amazon_s3_vectors/tests/test_filters.py b/integrations/amazon_s3_vectors/tests/test_filters.py index 88739f8837..11d83a282d 100644 --- a/integrations/amazon_s3_vectors/tests/test_filters.py +++ b/integrations/amazon_s3_vectors/tests/test_filters.py @@ -2,112 +2,88 @@ # # SPDX-License-Identifier: Apache-2.0 +""" +Unit tests for the Haystack → S3 Vectors filter conversion. + +This is a pure function with real logic (operator mapping, type validation, +meta. prefix stripping, logical nesting) so unit testing is high signal. +""" + import pytest from haystack.errors import FilterError from haystack_integrations.document_stores.amazon_s3_vectors.filters import _normalize_filters, _validate_filters -class TestNormalizeFilters: - def test_simple_equality(self): - filters = {"field": "meta.category", "operator": "==", "value": "news"} - result = _normalize_filters(filters) - assert result == {"category": {"$eq": "news"}} - - def test_not_equal(self): - filters = {"field": "meta.status", "operator": "!=", "value": "draft"} - result = _normalize_filters(filters) - assert result == {"status": {"$ne": "draft"}} - - def test_greater_than(self): - filters = {"field": "meta.year", "operator": ">", "value": 2020} - result = _normalize_filters(filters) - assert result == {"year": {"$gt": 2020}} - - def test_greater_than_equal(self): - filters = {"field": "meta.year", "operator": ">=", "value": 2020} - result = _normalize_filters(filters) - assert result == {"year": {"$gte": 2020}} - - def test_less_than(self): - filters = {"field": "meta.price", "operator": "<", "value": 50.0} - result = _normalize_filters(filters) - assert result == {"price": {"$lt": 50.0}} - - def test_less_than_equal(self): - filters = {"field": "meta.price", "operator": "<=", "value": 50.0} - result = _normalize_filters(filters) - assert result == {"price": {"$lte": 50.0}} - - def test_in_operator(self): - filters = {"field": "meta.genre", "operator": "in", "value": ["comedy", "drama"]} - result = _normalize_filters(filters) - assert result == {"genre": {"$in": ["comedy", "drama"]}} - - def test_not_in_operator(self): - filters = {"field": "meta.genre", "operator": "not in", "value": ["horror"]} - result = _normalize_filters(filters) - assert result == {"genre": {"$nin": ["horror"]}} - - def test_and_logical(self): - filters = { +def test_comparison_operators(): + """All comparison operators produce the correct S3 Vectors filter.""" + cases = [ + ({"field": "meta.x", "operator": "==", "value": "a"}, {"x": {"$eq": "a"}}), + ({"field": "meta.x", "operator": "!=", "value": "a"}, {"x": {"$ne": "a"}}), + ({"field": "meta.x", "operator": ">", "value": 1}, {"x": {"$gt": 1}}), + ({"field": "meta.x", "operator": ">=", "value": 1}, {"x": {"$gte": 1}}), + ({"field": "meta.x", "operator": "<", "value": 1.5}, {"x": {"$lt": 1.5}}), + ({"field": "meta.x", "operator": "<=", "value": 1.5}, {"x": {"$lte": 1.5}}), + ({"field": "meta.x", "operator": "in", "value": [1, 2]}, {"x": {"$in": [1, 2]}}), + ({"field": "meta.x", "operator": "not in", "value": ["a"]}, {"x": {"$nin": ["a"]}}), + ] + for haystack_filter, expected in cases: + assert _normalize_filters(haystack_filter) == expected + + +def test_logical_operators(): + """AND/OR produce $and/$or with nested conditions.""" + assert _normalize_filters( + { "operator": "AND", "conditions": [ - {"field": "meta.genre", "operator": "==", "value": "drama"}, - {"field": "meta.year", "operator": ">=", "value": 2020}, + {"field": "meta.a", "operator": "==", "value": 1}, + {"field": "meta.b", "operator": ">", "value": 2}, ], } - result = _normalize_filters(filters) - assert result == {"$and": [{"genre": {"$eq": "drama"}}, {"year": {"$gte": 2020}}]} + ) == {"$and": [{"a": {"$eq": 1}}, {"b": {"$gt": 2}}]} - def test_or_logical(self): - filters = { + assert _normalize_filters( + { "operator": "OR", "conditions": [ - {"field": "meta.genre", "operator": "==", "value": "drama"}, - {"field": "meta.genre", "operator": "==", "value": "comedy"}, + {"field": "meta.a", "operator": "==", "value": "x"}, + {"field": "meta.a", "operator": "==", "value": "y"}, ], } - result = _normalize_filters(filters) - assert result == {"$or": [{"genre": {"$eq": "drama"}}, {"genre": {"$eq": "comedy"}}]} + ) == {"$or": [{"a": {"$eq": "x"}}, {"a": {"$eq": "y"}}]} - def test_field_without_meta_prefix(self): - filters = {"field": "category", "operator": "==", "value": "news"} - result = _normalize_filters(filters) - assert result == {"category": {"$eq": "news"}} - def test_unsupported_type_raises(self): - filters = {"field": "meta.data", "operator": "==", "value": [1, 2, 3]} - with pytest.raises(FilterError): - _normalize_filters(filters) +def test_meta_prefix_stripped(): + """The meta. prefix is stripped since S3 Vectors stores metadata flat.""" + assert _normalize_filters({"field": "meta.category", "operator": "==", "value": "news"}) == { + "category": {"$eq": "news"} + } + # Fields without meta. prefix also work + assert _normalize_filters({"field": "category", "operator": "==", "value": "news"}) == {"category": {"$eq": "news"}} - def test_in_with_non_list_raises(self): - filters = {"field": "meta.genre", "operator": "in", "value": "comedy"} - with pytest.raises(FilterError): - _normalize_filters(filters) - def test_missing_operator_raises(self): - filters = {"field": "meta.genre", "value": "comedy"} - with pytest.raises(FilterError): - _normalize_filters(filters) +def test_unsupported_type_raises(): + with pytest.raises(FilterError): + _normalize_filters({"field": "meta.x", "operator": "==", "value": [1, 2, 3]}) - def test_missing_value_raises(self): - filters = {"field": "meta.genre", "operator": "=="} - with pytest.raises(FilterError): - _normalize_filters(filters) - def test_invalid_dict_raises(self): - with pytest.raises(FilterError): - _normalize_filters("not a dict") +def test_in_requires_list(): + with pytest.raises(FilterError): + _normalize_filters({"field": "meta.x", "operator": "in", "value": "not-a-list"}) -class TestValidateFilters: - def test_valid_filters(self): - _validate_filters({"operator": "AND", "conditions": []}) +def test_missing_keys_raise(): + with pytest.raises(FilterError): + _normalize_filters({"field": "meta.x", "value": "a"}) # missing operator + with pytest.raises(FilterError): + _normalize_filters({"field": "meta.x", "operator": "=="}) # missing value + with pytest.raises(FilterError): + _normalize_filters("not a dict") - def test_none_is_valid(self): - _validate_filters(None) - def test_invalid_syntax_raises(self): - with pytest.raises(ValueError, match="Invalid filter syntax"): - _validate_filters({"field": "meta.x"}) +def test_validate_filters(): + _validate_filters(None) # None is valid + _validate_filters({"operator": "AND", "conditions": []}) # valid structure + with pytest.raises(ValueError, match="Invalid filter syntax"): + _validate_filters({"field": "meta.x"}) # missing operator/conditions From 037a7157e0f6dd27a89ddb069bb5932bf3590413 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Mon, 13 Apr 2026 15:12:00 +0200 Subject: [PATCH 06/17] docs: add service limits to component docstrings - Class docstring: top_k cap, dimension limit, metadata limits, float32 only - write_documents: embedding required, 40KB metadata limit - _embedding_retrieval: top_k=100 cap, no embeddings in response - Retriever run: top_k=100, server-side filters, no embeddings returned --- .../amazon_s3_vectors/embedding_retriever.py | 6 +++-- .../amazon_s3_vectors/document_store.py | 25 ++++++++++++++----- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py index e04ad20990..fbef8505d5 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py @@ -113,9 +113,11 @@ def run( :param query_embedding: Embedding of the query. :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on - the ``filter_policy`` chosen at retriever initialization. - :param top_k: Maximum number of Documents to return. + the ``filter_policy`` chosen at retriever initialization. Filters are applied server-side during + the vector search. + :param top_k: Maximum number of Documents to return. S3 Vectors caps this at 100. :returns: A dictionary with key ``"documents"`` containing the retrieved Documents. + Returned documents will not contain embeddings. """ filters = apply_filter_policy(self.filter_policy, self.filters, filters) top_k = top_k or self.top_k diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index 01c2c9ceb7..ba0241fb9d 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -48,10 +48,19 @@ class S3VectorsDocumentStore: """ A Document Store using [Amazon S3 Vectors](https://aws.amazon.com/s3/features/vectors/). - Amazon S3 Vectors provides native vector storage and similarity search within Amazon S3. + Amazon S3 Vectors provides serverless vector storage and similarity search within Amazon S3. This document store stores Haystack `Document` objects as vectors with associated metadata in an S3 vector bucket and index. + **Service limits:** + + - Maximum ``top_k``: 100 results per query + - Maximum vector dimension: 4,096 + - Metadata per vector: 40 KB total, 2 KB filterable + - All documents must have embeddings (``float32`` only) + - Distance metrics: ``cosine`` or ``euclidean`` (set at index creation, immutable) + - ``filter_documents()`` is client-side — prefer ``S3VectorsEmbeddingRetriever`` with filters + Usage example: ```python from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore @@ -229,11 +238,14 @@ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = D """ Write Documents to the S3 Vectors index. - S3 Vectors ``put_vectors`` is an upsert operation by default, so ``DuplicatePolicy.OVERWRITE`` is the - natural behavior. ``DuplicatePolicy.SKIP`` will check for existing documents first (slower). + All documents must have an embedding set. S3 Vectors ``put_vectors`` is an upsert operation + by default, so ``DuplicatePolicy.OVERWRITE`` is the natural behavior. + ``DuplicatePolicy.SKIP`` will check for existing documents first (slower). ``DuplicatePolicy.NONE`` will raise an error if a document already exists. - :param documents: A list of Documents to write. + Metadata per vector is limited to 40 KB total (2 KB filterable). + + :param documents: A list of Documents to write. Each document must have an embedding. :param policy: The duplicate policy. Defaults to ``DuplicatePolicy.OVERWRITE``. :returns: The number of documents written. """ @@ -372,8 +384,9 @@ def _embedding_retrieval( :param query_embedding: The query embedding vector. :param filters: Optional Haystack-format metadata filters. - :param top_k: Maximum number of results to return. - :returns: List of Documents sorted by similarity. + :param top_k: Maximum number of results to return. S3 Vectors caps this at 100. + :returns: List of Documents sorted by similarity. Returned documents will not contain + embeddings (S3 Vectors ``query_vectors`` does not return vector data). """ if not query_embedding: msg = "query_embedding must be a non-empty list of floats" From 0ecd075103820162e9323d9d677db156ef498535 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Mon, 13 Apr 2026 15:24:00 +0200 Subject: [PATCH 07/17] Refactor S3 Vectors document store: use Haystack built-in filter utility, deduplicate retrieval logic - Replace hand-rolled _apply_filters_in_memory/_document_matches/_compare with haystack.utils.filters.document_matches_filter (same utility used by InMemoryDocumentStore). Gains NOT operator, nested dotted field paths, and date comparison support for free. (-65 lines) - Deduplicate blob/content reconstruction in _embedding_retrieval() by reusing _s3_vector_to_document() + dataclasses.replace() (-20 lines) - Make filter_documents() warning conditional on filters actually being provided (no warning when listing all documents) --- .../amazon_s3_vectors/document_store.py | 111 +++--------------- 1 file changed, 15 insertions(+), 96 deletions(-) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index ba0241fb9d..3d6ddedf47 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -4,6 +4,7 @@ import base64 import json +from dataclasses import replace from typing import Any, Literal import boto3 @@ -13,6 +14,7 @@ from haystack.document_stores.errors import DocumentStoreError from haystack.document_stores.types import DuplicatePolicy from haystack.utils.auth import Secret, deserialize_secrets_inplace +from haystack.utils.filters import document_matches_filter from .filters import _normalize_filters, _validate_filters @@ -315,12 +317,13 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume :param filters: Haystack-format filters to apply. :returns: A list of matching Documents. """ - logger.warning( - "S3 Vectors does not support standalone filtered listing. " - "filter_documents() will fetch ALL vectors and apply filters client-side, " - "which can be very slow for large indexes. " - "Prefer using S3VectorsEmbeddingRetriever with filters for efficient filtered retrieval." - ) + if filters: + logger.warning( + "S3 Vectors does not support standalone filtered listing. " + "filter_documents() will fetch ALL vectors and apply filters client-side, " + "which can be very slow for large indexes. " + "Prefer using S3VectorsEmbeddingRetriever with filters for efficient filtered retrieval." + ) client = self._get_client() @@ -347,7 +350,7 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume if filters: _validate_filters(filters) - documents = _apply_filters_in_memory(documents, filters) + documents = [doc for doc in documents if document_matches_filter(filters=filters, document=doc)] return documents @@ -425,24 +428,11 @@ def _embedding_retrieval( documents = [] for v in result.get("vectors", []): - metadata = v.get("metadata", {}) - content = metadata.pop(_CONTENT_KEY, None) - - # Remove other internal keys from user-facing metadata - blob_data = metadata.pop(_BLOB_DATA_KEY, None) - blob_meta = metadata.pop(_BLOB_META_KEY, None) - blob_mime_type = metadata.pop(_BLOB_MIME_TYPE_KEY, None) - - blob = None - if blob_data is not None: - blob = ByteStream( - data=base64.b64decode(blob_data) if isinstance(blob_data, str) else blob_data, - meta=blob_meta or {}, - mime_type=blob_mime_type, - ) + doc = self._s3_vector_to_document(v) - raw_distance = v.get("distance") + # Compute score from distance score = None + raw_distance = v.get("distance") if raw_distance is not None: if distance_metric == "cosine": score = 1.0 - raw_distance @@ -450,14 +440,8 @@ def _embedding_retrieval( # euclidean: negate so higher = more similar score = -raw_distance - doc = Document( - id=v["key"], - content=content, - meta=metadata, - blob=blob, - score=score, - ) - documents.append(doc) + # query_vectors does not return vector data; attach score + documents.append(replace(doc, embedding=None, score=score)) return documents @@ -538,68 +522,3 @@ def _s3_vector_to_document(vector: dict[str, Any]) -> Document: embedding=embedding, blob=blob, ) - - -def _apply_filters_in_memory(documents: list[Document], filters: dict[str, Any]) -> list[Document]: - """ - Apply Haystack filters to a list of Documents in memory. - - This is used by ``filter_documents`` since S3 Vectors doesn't support - standalone metadata queries without a vector search. - """ - return [doc for doc in documents if _document_matches(doc, filters)] - - -def _document_matches(doc: Document, filters: dict[str, Any]) -> bool: - """Check if a single Document matches the given Haystack filter.""" - if "operator" in filters and "conditions" in filters: - operator = filters["operator"] - conditions = filters["conditions"] - if operator == "AND": - return all(_document_matches(doc, c) for c in conditions) - if operator == "OR": - return any(_document_matches(doc, c) for c in conditions) - msg = f"Unknown logical operator '{operator}'" - raise ValueError(msg) - - if "field" in filters: - field = filters["field"] - operator = filters["operator"] - value = filters["value"] - - # Resolve the field value from the Document - if field.startswith("meta."): - field_name = field[5:] - doc_value = doc.meta.get(field_name) if doc.meta else None - elif field == "content": - doc_value = doc.content - elif field == "id": - doc_value = doc.id - else: - doc_value = doc.meta.get(field) if doc.meta else None - - return _compare(doc_value, operator, value) - - return True - - -def _compare(doc_value: Any, operator: str, filter_value: Any) -> bool: - """Perform a comparison between a document field value and a filter value.""" - if operator == "==": - return doc_value == filter_value - if operator == "!=": - return doc_value != filter_value - if operator == ">": - return doc_value is not None and doc_value > filter_value - if operator == ">=": - return doc_value is not None and doc_value >= filter_value - if operator == "<": - return doc_value is not None and doc_value < filter_value - if operator == "<=": - return doc_value is not None and doc_value <= filter_value - if operator == "in": - return doc_value is not None and doc_value in filter_value - if operator == "not in": - return doc_value is not None and doc_value not in filter_value - msg = f"Unknown comparison operator '{operator}'" - raise ValueError(msg) From 06cd5885ae080db51db2eb99b53c263a3b1a4a1f Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Wed, 22 Apr 2026 21:02:17 +0200 Subject: [PATCH 08/17] ci: add AWS authentication step for integration tests Matches the pattern used by the amazon_bedrock workflow: - top-level id-token: write permission - AWS_REGION env var - configure-aws-credentials step (skipped on fork PRs and dependabot) - integration tests gated on successful auth --- .github/workflows/amazon_s3_vectors.yml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/.github/workflows/amazon_s3_vectors.yml b/.github/workflows/amazon_s3_vectors.yml index 1fa282e190..36787d359d 100644 --- a/.github/workflows/amazon_s3_vectors.yml +++ b/.github/workflows/amazon_s3_vectors.yml @@ -26,12 +26,18 @@ concurrency: group: amazon_s3_vectors-${{ github.head_ref || github.sha }} cancel-in-progress: true +permissions: + id-token: write + contents: read + env: PYTHONUNBUFFERED: "1" FORCE_COLOR: "1" TEST_MATRIX_OS: '["ubuntu-latest", "windows-latest", "macos-latest"]' TEST_MATRIX_PYTHON: '["3.10", "3.14"]' + AWS_REGION: "us-east-1" + jobs: compute-test-matrix: runs-on: ubuntu-slim @@ -51,6 +57,7 @@ jobs: name: Python ${{ matrix.python-version }} on ${{ startsWith(matrix.os, 'macos-') && 'macOS' || startsWith(matrix.os, 'windows-') && 'Windows' || 'Linux' }} needs: compute-test-matrix permissions: + id-token: write contents: write pull-requests: write runs-on: ${{ matrix.os }} @@ -101,7 +108,17 @@ jobs: name: coverage-comment-amazon_s3_vectors path: python-coverage-comment-action-amazon_s3_vectors.txt + # Do not authenticate on PRs from forks and on PRs created by dependabot + - name: AWS authentication + id: aws-auth + if: github.event_name == 'schedule' || (github.event.pull_request.head.repo.full_name == github.repository && !startsWith(github.event.pull_request.head.ref, 'dependabot/')) + uses: aws-actions/configure-aws-credentials@ec61189d14ec14c8efccab744f656cffd0e33f37 + with: + aws-region: ${{ env.AWS_REGION }} + role-to-assume: ${{ secrets.AWS_CI_ROLE_ARN }} + - name: Run integration tests + if: success() && steps.aws-auth.outcome == 'success' run: hatch run test:integration-cov-append-retry - name: Store combined coverage From 3f79ba9afa5ddc2f1cc1f48d795c0813e6562155 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Wed, 22 Apr 2026 21:02:20 +0200 Subject: [PATCH 09/17] docs: use single backticks for inline code in docstrings Matches the repo convention used across other integrations. --- .../amazon_s3_vectors/embedding_retriever.py | 10 +++---- .../amazon_s3_vectors/document_store.py | 28 +++++++++---------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py index fbef8505d5..f7f5bd68b2 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py @@ -15,7 +15,7 @@ @component class S3VectorsEmbeddingRetriever: """ - Retrieve documents from an ``S3VectorsDocumentStore`` based on their dense embeddings. + Retrieve documents from an `S3VectorsDocumentStore` based on their dense embeddings. Usage example: ```python @@ -64,11 +64,11 @@ def __init__( """ Initialize the S3VectorsEmbeddingRetriever. - :param document_store: An instance of ``S3VectorsDocumentStore``. + :param document_store: An instance of `S3VectorsDocumentStore`. :param filters: Filters applied to the retrieved Documents. :param top_k: Maximum number of Documents to return. :param filter_policy: Policy to determine how filters are applied. - :raises ValueError: If ``document_store`` is not an ``S3VectorsDocumentStore``. + :raises ValueError: If `document_store` is not an `S3VectorsDocumentStore`. """ if not isinstance(document_store, S3VectorsDocumentStore): msg = "document_store must be an instance of S3VectorsDocumentStore" @@ -113,10 +113,10 @@ def run( :param query_embedding: Embedding of the query. :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on - the ``filter_policy`` chosen at retriever initialization. Filters are applied server-side during + the `filter_policy` chosen at retriever initialization. Filters are applied server-side during the vector search. :param top_k: Maximum number of Documents to return. S3 Vectors caps this at 100. - :returns: A dictionary with key ``"documents"`` containing the retrieved Documents. + :returns: A dictionary with key `"documents"` containing the retrieved Documents. Returned documents will not contain embeddings. """ filters = apply_filter_policy(self.filter_policy, self.filters, filters) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index 3d6ddedf47..238436c787 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -56,12 +56,12 @@ class S3VectorsDocumentStore: **Service limits:** - - Maximum ``top_k``: 100 results per query + - Maximum `top_k`: 100 results per query - Maximum vector dimension: 4,096 - Metadata per vector: 40 KB total, 2 KB filterable - - All documents must have embeddings (``float32`` only) - - Distance metrics: ``cosine`` or ``euclidean`` (set at index creation, immutable) - - ``filter_documents()`` is client-side — prefer ``S3VectorsEmbeddingRetriever`` with filters + - All documents must have embeddings (`float32` only) + - Distance metrics: `cosine` or `euclidean` (set at index creation, immutable) + - `filter_documents()` is client-side — prefer `S3VectorsEmbeddingRetriever` with filters Usage example: ```python @@ -95,13 +95,13 @@ def __init__( :param vector_bucket_name: Name of the S3 vector bucket. :param index_name: Name of the vector index within the bucket. :param dimension: Dimensionality of the embeddings (e.g. 768, 1536). - :param distance_metric: Distance metric for similarity search: ``"cosine"`` or ``"euclidean"``. + :param distance_metric: Distance metric for similarity search: `"cosine"` or `"euclidean"`. :param region_name: AWS region. If not provided, uses the default from the environment/config. :param aws_access_key_id: AWS access key ID. If not provided, uses the default credential chain. :param aws_secret_access_key: AWS secret access key. :param aws_session_token: AWS session token for temporary credentials. :param create_bucket_and_index: Whether to automatically create the vector bucket and index - if they do not exist. Defaults to ``True``. + if they do not exist. Defaults to `True`. :param non_filterable_metadata_keys: Additional metadata keys to mark as non-filterable on the index (beyond the internal keys used for Document content/blob storage). """ @@ -240,15 +240,15 @@ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = D """ Write Documents to the S3 Vectors index. - All documents must have an embedding set. S3 Vectors ``put_vectors`` is an upsert operation - by default, so ``DuplicatePolicy.OVERWRITE`` is the natural behavior. - ``DuplicatePolicy.SKIP`` will check for existing documents first (slower). - ``DuplicatePolicy.NONE`` will raise an error if a document already exists. + All documents must have an embedding set. S3 Vectors `put_vectors` is an upsert operation + by default, so `DuplicatePolicy.OVERWRITE` is the natural behavior. + `DuplicatePolicy.SKIP` will check for existing documents first (slower). + `DuplicatePolicy.NONE` will raise an error if a document already exists. Metadata per vector is limited to 40 KB total (2 KB filterable). :param documents: A list of Documents to write. Each document must have an embedding. - :param policy: The duplicate policy. Defaults to ``DuplicatePolicy.OVERWRITE``. + :param policy: The duplicate policy. Defaults to `DuplicatePolicy.OVERWRITE`. :returns: The number of documents written. """ if len(documents) == 0: @@ -312,7 +312,7 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume S3 Vectors only supports metadata filtering during vector similarity queries, not as a standalone operation. This method lists all vectors and applies filters client-side, which can be very slow for large indexes. For filtered retrieval, prefer using - ``S3VectorsEmbeddingRetriever`` with filters instead. + `S3VectorsEmbeddingRetriever` with filters instead. :param filters: Haystack-format filters to apply. :returns: A list of matching Documents. @@ -383,13 +383,13 @@ def _embedding_retrieval( Retrieve documents most similar to the query embedding. This method is not part of the public interface. - Use ``S3VectorsEmbeddingRetriever`` instead. + Use `S3VectorsEmbeddingRetriever` instead. :param query_embedding: The query embedding vector. :param filters: Optional Haystack-format metadata filters. :param top_k: Maximum number of results to return. S3 Vectors caps this at 100. :returns: List of Documents sorted by similarity. Returned documents will not contain - embeddings (S3 Vectors ``query_vectors`` does not return vector data). + embeddings (S3 Vectors `query_vectors` does not return vector data). """ if not query_embedding: msg = "query_embedding must be a non-empty list of floats" From a701b9d50db8d00d37ba3de49919051658795d39 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Sat, 16 May 2026 12:38:29 +0200 Subject: [PATCH 10/17] docs(amazon-s3-vectors): minimize README to repo convention Trim the README down to badges + integration page / changelog links + the AWS auth note, matching the pattern used by amazon_bedrock, anthropic, qdrant, and other integrations in this repo. The rich usage / limitations content is preserved for the haystack.deepset.ai integration page. Addresses https://github.com/deepset-ai/haystack-core-integrations/pull/3149#discussion_r3108976852 --- integrations/amazon_s3_vectors/README.md | 206 ++--------------------- 1 file changed, 10 insertions(+), 196 deletions(-) diff --git a/integrations/amazon_s3_vectors/README.md b/integrations/amazon_s3_vectors/README.md index 8fd50ff2d1..7a291d6bd2 100644 --- a/integrations/amazon_s3_vectors/README.md +++ b/integrations/amazon_s3_vectors/README.md @@ -3,205 +3,19 @@ [![PyPI - Version](https://img.shields.io/pypi/v/amazon-s3-vectors-haystack.svg)](https://pypi.org/project/amazon-s3-vectors-haystack) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/amazon-s3-vectors-haystack.svg)](https://pypi.org/project/amazon-s3-vectors-haystack) +- [Integration page](https://haystack.deepset.ai/integrations/amazon-s3-vectors) +- [Changelog](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/amazon_s3_vectors/CHANGELOG.md) --- -A [Haystack](https://haystack.deepset.ai/) integration for [Amazon S3 Vectors](https://aws.amazon.com/s3/features/vectors/), providing a Document Store and Embedding Retriever backed by native vector storage in Amazon S3. +## Contributing -## Installation +Refer to the general [Contribution Guidelines](https://github.com/deepset-ai/haystack-core-integrations/blob/main/CONTRIBUTING.md). -```bash -pip install amazon-s3-vectors-haystack +To run integration tests locally, you need to authenticate with AWS. +For example, you can do that by exporting the following environment variables: ``` - -## Usage - -### Document Store - -```python -from haystack.dataclasses import Document -from haystack.document_stores.types import DuplicatePolicy -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore - -document_store = S3VectorsDocumentStore( - vector_bucket_name="my-vectors", - index_name="my-index", - dimension=768, - distance_metric="cosine", # or "euclidean" - region_name="us-east-1", -) - -# Write documents (embeddings are required) -docs = [ - Document(id="1", content="First document", embedding=[0.1] * 768, meta={"category": "news"}), - Document(id="2", content="Second document", embedding=[0.2] * 768, meta={"category": "sports"}), -] -document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) - -# Count documents -print(document_store.count_documents()) - -# Delete documents -document_store.delete_documents(["1", "2"]) -``` - -### Embedding Retriever in a Pipeline - -```python -from haystack import Pipeline -from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder -from haystack.document_stores.types import DuplicatePolicy -from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore - -document_store = S3VectorsDocumentStore( - vector_bucket_name="my-vectors", - index_name="my-index", - dimension=768, -) - -# Index documents -doc_embedder = SentenceTransformersDocumentEmbedder() -doc_embedder.warm_up() -# ... embed and write documents ... - -# Query pipeline -pipeline = Pipeline() -pipeline.add_component("embedder", SentenceTransformersTextEmbedder()) -pipeline.add_component("retriever", S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5)) -pipeline.connect("embedder.embedding", "retriever.query_embedding") - -result = pipeline.run({"embedder": {"text": "What is the latest news?"}}) -print(result["retriever"]["documents"]) -``` - -### Filtering - -The retriever supports Haystack metadata filters, which are converted to S3 Vectors filter syntax: - -```python -# Filter during retrieval -result = pipeline.run({ - "embedder": {"text": "sports news"}, - "retriever": { - "filters": { - "operator": "AND", - "conditions": [ - {"field": "meta.category", "operator": "==", "value": "sports"}, - {"field": "meta.year", "operator": ">=", "value": 2024}, - ], - } - }, -}) -``` - -**Supported filter operators:** `==`, `!=`, `>`, `>=`, `<`, `<=`, `in`, `not in`, `AND`, `OR` - -### AWS Authentication - -The document store uses the standard boto3 credential chain by default. You can also pass credentials explicitly: - -```python -from haystack.utils.auth import Secret - -document_store = S3VectorsDocumentStore( - vector_bucket_name="my-vectors", - index_name="my-index", - dimension=768, - region_name="us-east-1", - aws_access_key_id=Secret.from_env_var("AWS_ACCESS_KEY_ID"), - aws_secret_access_key=Secret.from_env_var("AWS_SECRET_ACCESS_KEY"), -) +export AWS_ACCESS_KEY_ID=... +export AWS_SECRET_ACCESS_KEY=... +export AWS_SESSION_TOKEN=... +export AWS_DEFAULT_REGION=... ``` - -## Configuration Reference - -| Parameter | Type | Default | Description | -|---|---|---|---| -| `vector_bucket_name` | `str` | *required* | Name of the S3 vector bucket | -| `index_name` | `str` | *required* | Name of the vector index within the bucket | -| `dimension` | `int` | *required* | Dimensionality of the embeddings (e.g. 768, 1536) | -| `distance_metric` | `str` | `"cosine"` | `"cosine"` or `"euclidean"` | -| `region_name` | `str \| None` | `None` | AWS region (uses default from env if not set) | -| `aws_access_key_id` | `Secret \| None` | `None` | AWS access key ID | -| `aws_secret_access_key` | `Secret \| None` | `None` | AWS secret access key | -| `aws_session_token` | `Secret \| None` | `None` | AWS session token (for temporary credentials) | -| `create_bucket_and_index` | `bool` | `True` | Auto-create bucket and index if they don't exist | -| `non_filterable_metadata_keys` | `list[str] \| None` | `None` | Additional metadata keys to mark as non-filterable | - -## Known Limitations & Considerations - -### No Native Document Count API -S3 Vectors does not provide a dedicated count endpoint. `count_documents()` paginates through all -vector keys, which can be slow for large indexes (millions of vectors). - -### `filter_documents()` Is Expensive -S3 Vectors only supports metadata filtering **during vector similarity queries** — there is no -standalone "list documents matching filter" API. As a result, `filter_documents()` must: -1. List all vectors with their data and metadata (paginated) -2. Apply filters client-side in memory - -**For filtered retrieval, always prefer `S3VectorsEmbeddingRetriever` with filters**, which uses -the native `query_vectors` API with server-side filtering. - -### Embedding Required -Every document written to the store **must** have an embedding. Documents without embeddings -will be rejected. This is a fundamental constraint of S3 Vectors as a pure vector store. - -### No Keyword / BM25 Retrieval -S3 Vectors only supports dense vector similarity search. There is no keyword or BM25 search -capability. If you need hybrid search, consider pairing this with Amazon OpenSearch. - -### Vector Data Type -Only `float32` vectors are supported. Higher-precision values are automatically downcast. - -### Metadata Size Limits -- **Total metadata per vector: 40 KB** (filterable + non-filterable combined) -- **Filterable metadata per vector: 2 KB** — user `meta` fields used in filters must fit in this budget -- **Non-filterable metadata keys per index: 10** — the integration reserves 4 internal keys - (`_content`, `_blob_data`, `_blob_meta`, `_blob_mime_type`), leaving 6 for user-defined keys -- Keys are set at index creation and **cannot be changed later** - -Large content (e.g. full document text) is stored as non-filterable metadata automatically. -If you store additional large metadata fields, declare them via `non_filterable_metadata_keys`. - -### Strict API Limits -- `put_vectors`: up to 500 vectors per call (handled automatically by the integration) -- `get_vectors`: up to 100 keys per call (handled automatically) -- `delete_vectors`: up to 500 keys per call (handled automatically) -- **`query_vectors`: maximum 100 results per query** — this is the hard cap on `top_k`. - If you need more than 100 results, you must implement pagination or use a different store. -- Combined PutVectors + DeleteVectors: up to 1,000 requests/second per index - -### Distance Metrics and Scoring -Only `cosine` and `euclidean` are supported. The metric is set at index creation time and cannot -be changed afterward. - -S3 Vectors returns raw **distances** (lower = more similar). The integration converts these to -Haystack-convention **scores** (higher = more similar): -- **Cosine:** `score = 1.0 - distance` (1.0 = identical, 0.0 = orthogonal) -- **Euclidean:** `score = -distance` (0.0 = identical, more negative = further) - -### No Embeddings Returned from Queries -The `query_vectors` API does not support returning the stored vector data alongside results. -Documents retrieved via `S3VectorsEmbeddingRetriever` will have `embedding=None`. If you need -the embedding vectors, use `filter_documents()` or fetch them separately via the boto3 client. - -### Eventual Consistency -Newly written vectors may not be immediately visible in query results. S3 Vectors provides -eventual consistency for write-then-read operations. - -## Running Tests - -```bash -cd integrations/amazon_s3_vectors - -# Unit tests (no AWS credentials required) -hatch run test:unit - -# Integration tests (requires AWS credentials and S3 Vectors access) -hatch run test:integration -``` - -## License - -`amazon-s3-vectors-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license. From fc846c2547073ebbe01dd65ac87c4c9a49bf2acd Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Sat, 16 May 2026 12:39:04 +0200 Subject: [PATCH 11/17] fix(amazon-s3-vectors): validate all documents up front in write_documents Move type / embedding validation out of the per-batch loop and into a single pass before any `put_vectors` call, so a bad input never leaves the store with a partial write. Also align with Haystack's duplicate-policy convention: `DuplicatePolicy.FAIL` is now supported and, together with `DuplicatePolicy.NONE`, raises `DuplicateDocumentError` (was `DocumentStoreError` for NONE) when any target id already exists. Unit tests updated accordingly. Addresses https://github.com/deepset-ai/haystack-core-integrations/pull/3149#discussion_r3128763028 --- .../amazon_s3_vectors/document_store.py | 36 +++++++++++----- .../tests/test_document_store.py | 41 ++++++++++++++++--- 2 files changed, 60 insertions(+), 17 deletions(-) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index 238436c787..caeb091800 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -11,7 +11,7 @@ from botocore.exceptions import ClientError from haystack import default_from_dict, default_to_dict, logging from haystack.dataclasses import ByteStream, Document -from haystack.document_stores.errors import DocumentStoreError +from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError from haystack.document_stores.types import DuplicatePolicy from haystack.utils.auth import Secret, deserialize_secrets_inplace from haystack.utils.filters import document_matches_filter @@ -243,32 +243,46 @@ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = D All documents must have an embedding set. S3 Vectors `put_vectors` is an upsert operation by default, so `DuplicatePolicy.OVERWRITE` is the natural behavior. `DuplicatePolicy.SKIP` will check for existing documents first (slower). - `DuplicatePolicy.NONE` will raise an error if a document already exists. + `DuplicatePolicy.FAIL` and `DuplicatePolicy.NONE` will raise `DuplicateDocumentError` + if any document already exists. Metadata per vector is limited to 40 KB total (2 KB filterable). :param documents: A list of Documents to write. Each document must have an embedding. :param policy: The duplicate policy. Defaults to `DuplicatePolicy.OVERWRITE`. :returns: The number of documents written. + :raises ValueError: If `documents` is not a list of `Document` instances. + :raises DocumentStoreError: If any document is missing an embedding. + :raises DuplicateDocumentError: If `policy` is `FAIL` or `NONE` and any document already + exists in the store. """ + # Validate input type up front, before any writes happen. + if not isinstance(documents, list) or any(not isinstance(d, Document) for d in documents): + msg = "Please provide a list of Documents." + raise ValueError(msg) + if len(documents) == 0: return 0 + # Validate ALL embeddings up front so we fail before any put_vectors call, + # never leaving the store with a partial write. + missing_embedding_ids = [doc.id for doc in documents if doc.embedding is None] + if missing_embedding_ids: + msg = ( + f"Document(s) {missing_embedding_ids} have no embedding. " + "S3VectorsDocumentStore requires every document to have an embedding." + ) + raise DocumentStoreError(msg) + client = self._get_client() written = 0 for i in range(0, len(documents), _WRITE_BATCH_SIZE): batch = documents[i : i + _WRITE_BATCH_SIZE] - # Validate embeddings upfront - for doc in batch: - if doc.embedding is None: - msg = f"Document '{doc.id}' has no embedding. S3VectorsDocumentStore requires embeddings." - raise DocumentStoreError(msg) - # Batch-check for existing documents when needed existing_ids: set[str] = set() - if policy in (DuplicatePolicy.SKIP, DuplicatePolicy.NONE): + if policy in (DuplicatePolicy.SKIP, DuplicatePolicy.NONE, DuplicatePolicy.FAIL): batch_ids = [doc.id for doc in batch] for j in range(0, len(batch_ids), _GET_BATCH_SIZE): id_chunk = batch_ids[j : j + _GET_BATCH_SIZE] @@ -280,12 +294,12 @@ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = D for v in response.get("vectors", []): existing_ids.add(v["key"]) - if policy == DuplicatePolicy.NONE and existing_ids: + if policy in (DuplicatePolicy.NONE, DuplicatePolicy.FAIL) and existing_ids: msg = ( f"Document(s) {sorted(existing_ids)} already exist in the document store. " "Use DuplicatePolicy.OVERWRITE or DuplicatePolicy.SKIP." ) - raise DocumentStoreError(msg) + raise DuplicateDocumentError(msg) vectors_to_write = [] for doc in batch: diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index 8df69ffc49..cfd3761242 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -6,7 +6,7 @@ import pytest from haystack.dataclasses import Document -from haystack.document_stores.errors import DocumentStoreError +from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError from haystack.document_stores.types import DuplicatePolicy from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore @@ -98,7 +98,7 @@ def test_write_documents_no_embedding_raises(mock_boto3): mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - with pytest.raises(DocumentStoreError, match="has no embedding"): + with pytest.raises(DocumentStoreError, match="no embedding"): store.write_documents([Document(id="1", content="Hello")]) @@ -120,9 +120,10 @@ def test_write_documents_skip_existing(mock_boto3): client.put_vectors.assert_not_called() +@pytest.mark.parametrize("policy", [DuplicatePolicy.FAIL, DuplicatePolicy.NONE]) @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_none_policy_raises(mock_boto3): - """Tests our batch existence check logic for NONE policy.""" +def test_write_documents_fail_policy_raises(mock_boto3, policy): + """Both FAIL and NONE raise DuplicateDocumentError on existing IDs.""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} @@ -130,13 +131,41 @@ def test_write_documents_none_policy_raises(mock_boto3): mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - with pytest.raises(DocumentStoreError, match="already exist"): + with pytest.raises(DuplicateDocumentError, match="already exist"): store.write_documents( [Document(id="1", content="Hello", embedding=[0.1] * 4)], - policy=DuplicatePolicy.NONE, + policy=policy, ) +def test_write_documents_invalid_input_raises(): + """write_documents must reject non-list / non-Document input before doing anything.""" + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False) + with pytest.raises(ValueError): + store.write_documents("not a list") # type: ignore[arg-type] + with pytest.raises(ValueError): + store.write_documents(["not a document"]) # type: ignore[list-item] + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents_missing_embedding_validated_upfront(mock_boto3): + """Missing embeddings are reported up front, before any put_vectors call.""" + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + docs = [ + Document(id="1", content="ok", embedding=[0.1] * 4), + Document(id="2", content="missing"), + Document(id="3", content="also-missing"), + ] + with pytest.raises(DocumentStoreError, match=r"\['2', '3'\].*no embedding"): + store.write_documents(docs) + client.put_vectors.assert_not_called() + + @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_embedding_retrieval_score_conversion(mock_boto3): """Tests our distance-to-score conversion logic — the only non-trivial transform in retrieval.""" From 4ae27eae698447e8bcdf827f58cc690024cc3170 Mon Sep 17 00:00:00 2001 From: Christian Kokott Date: Sat, 16 May 2026 12:39:21 +0200 Subject: [PATCH 12/17] test(amazon-s3-vectors): adopt Haystack base test mixins for integration tests Replace the handcrafted integration suite with Haystack's standard DocumentStore test contract, following the pgvector / pinecone pattern: * tests/conftest.py: shared `document_store` fixture. One vector bucket per session, one index per test for isolation. Wraps `write_documents` and `delete_documents` to (a) inject a default zero embedding for any Document missing one, since S3 Vectors requires embeddings, and (b) sleep briefly afterwards to absorb eventual consistency. * test_document_store.py: appends `TestDocumentStore` integration class inheriting CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest, FilterableDocsFixtureMixin. Overrides `assert_documents_are_equal` for float32 round-trip tolerance. * test_filters.py: appends `TestFilters` inheriting FilterDocumentsTest. `filter_documents` already routes matching through haystack's document_matches_filter, so no operators are skipped. * test_integration.py: trimmed to the retriever-specific tests not covered by the base mixins (embedding retrieval, retriever component, to_dict/from_dict roundtrip on a live store). Coverage goes from 12 handcrafted integration tests to 53 (10 store + 39 filter + 4 retriever). Addresses https://github.com/deepset-ai/haystack-core-integrations/pull/3149#discussion_r3128798133 --- .../amazon_s3_vectors/tests/conftest.py | 164 ++++++++++++ .../tests/test_document_store.py | 49 ++++ .../amazon_s3_vectors/tests/test_filters.py | 27 ++ .../tests/test_integration.py | 240 ++++-------------- 4 files changed, 288 insertions(+), 192 deletions(-) create mode 100644 integrations/amazon_s3_vectors/tests/conftest.py diff --git a/integrations/amazon_s3_vectors/tests/conftest.py b/integrations/amazon_s3_vectors/tests/conftest.py new file mode 100644 index 0000000000..b58f963c34 --- /dev/null +++ b/integrations/amazon_s3_vectors/tests/conftest.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Shared pytest fixtures for the Amazon S3 Vectors integration tests. + +The `document_store` fixture used by the Haystack base test mixins +(`CountDocumentsTest`, `WriteDocumentsTest`, `DeleteDocumentsTest`, +`FilterDocumentsTest`) is defined here. + +A single vector bucket is created per test session (creating buckets is the +slow part). Each test gets its own vector index inside that shared bucket so +state is isolated. + +Two quirks of S3 Vectors are smoothed over here so the generic base tests +can run unchanged: + +1. **Embeddings are required.** The base tests write `Document(content="...")` + without an embedding. We wrap `write_documents` so any document missing an + embedding gets a deterministic zero vector of the right dimension. +2. **Writes are eventually consistent.** We sleep briefly after `write_documents` + and `delete_documents` so the subsequent `filter_documents`/`count_documents` + reflects the new state. +""" + +from __future__ import annotations + +import os +import time +import uuid +import warnings +from collections.abc import Iterator + +import boto3 +import pytest +from botocore.exceptions import ClientError +from haystack.dataclasses import Document +from haystack.document_stores.types import DuplicatePolicy + +from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + +# Dimension used by Haystack's `FilterableDocsFixtureMixin` test corpus. +DIMENSION = 768 +REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") + +# Eventual-consistency budget for S3 Vectors write/delete propagation. +WRITE_SLEEP_SECONDS = 5 +DELETE_SLEEP_SECONDS = 5 + + +def _aws_credentials_available() -> bool: + """Return True if any boto3 credential source is configured.""" + if any(os.environ.get(k) for k in ("AWS_ACCESS_KEY_ID", "AWS_PROFILE", "AWS_ROLE_ARN")): + return True + try: + return boto3.Session().get_credentials() is not None + except Exception: + return False + + +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: # noqa: ARG001 + """Skip every integration test when AWS credentials are not configured.""" + if _aws_credentials_available(): + return + skip = pytest.mark.skip(reason="AWS credentials not configured") + for item in items: + if "integration" in item.keywords: + item.add_marker(skip) + + +@pytest.fixture(scope="session") +def s3_vectors_bucket() -> Iterator[str]: + """Create one vector bucket for the test session, tear it down at the end.""" + if not _aws_credentials_available(): + pytest.skip("AWS credentials not configured") + + bucket_name = f"haystack-test-{uuid.uuid4().hex[:8]}" + client = boto3.client("s3vectors", region_name=REGION) + client.create_vector_bucket(vectorBucketName=bucket_name) + + yield bucket_name + + # Tear down: delete every index in the bucket, then the bucket itself. + try: + next_token: str | None = None + while True: + kwargs: dict = {"vectorBucketName": bucket_name} + if next_token: + kwargs["nextToken"] = next_token + response = client.list_indexes(**kwargs) + for idx in response.get("indexes", []): + try: + client.delete_index(vectorBucketName=bucket_name, indexName=idx["indexName"]) + except ClientError: + pass + next_token = response.get("nextToken") + if not next_token: + break + except ClientError: + pass + + try: + client.delete_vector_bucket(vectorBucketName=bucket_name) + except ClientError: + pass + + +@pytest.fixture +def document_store(s3_vectors_bucket: str) -> Iterator[S3VectorsDocumentStore]: + """ + Provide a fresh S3VectorsDocumentStore (one per test) with `write_documents` + and `delete_documents` wrapped to: + + * inject a default embedding for any `Document` that doesn't have one, and + * sleep briefly afterwards to absorb S3 Vectors' eventual consistency. + """ + index_name = f"idx-{uuid.uuid4().hex[:10]}" + + store = S3VectorsDocumentStore( + vector_bucket_name=s3_vectors_bucket, + index_name=index_name, + dimension=DIMENSION, + distance_metric="cosine", + region_name=REGION, + create_bucket_and_index=True, + non_filterable_metadata_keys=[], + ) + + # Eagerly create the index so the first write doesn't race with index creation. + store._get_client() + + original_write = store.write_documents + + def write_with_defaults(documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.OVERWRITE) -> int: + # Only mutate well-formed input; bad input must surface as ValueError from + # the production code, not crash this wrapper. + if isinstance(documents, list): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + for d in documents: + if isinstance(d, Document) and d.embedding is None: + d.embedding = [0.0] * DIMENSION + result = original_write(documents, policy) + time.sleep(WRITE_SLEEP_SECONDS) + return result + + store.write_documents = write_with_defaults # type: ignore[method-assign] + + original_delete = store.delete_documents + + def delete_with_sleep(document_ids: list[str]) -> None: + original_delete(document_ids) + time.sleep(DELETE_SLEEP_SECONDS) + + store.delete_documents = delete_with_sleep # type: ignore[method-assign] + + yield store + + # Clean up: drop the per-test index. Best-effort. + try: + store._get_client().delete_index(vectorBucketName=s3_vectors_bucket, indexName=index_name) + except ClientError: + pass diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index cfd3761242..6fb0df0de0 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -2,12 +2,19 @@ # # SPDX-License-Identifier: Apache-2.0 +from dataclasses import replace from unittest.mock import MagicMock, patch import pytest from haystack.dataclasses import Document from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError from haystack.document_stores.types import DuplicatePolicy +from haystack.testing.document_store import ( + CountDocumentsTest, + DeleteDocumentsTest, + FilterableDocsFixtureMixin, + WriteDocumentsTest, +) from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore @@ -268,3 +275,45 @@ def test_document_roundtrip(): assert restored.content == doc.content assert restored.embedding == doc.embedding assert restored.meta == doc.meta + + +# --------------------------------------------------------------------------- +# Integration tests — exercise a real S3 Vectors bucket via the `document_store` +# fixture in conftest.py. The mixins below come from Haystack's test kit so we +# get its standard Document Store contract for free. +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestDocumentStore( + CountDocumentsTest, + WriteDocumentsTest, + DeleteDocumentsTest, + FilterableDocsFixtureMixin, +): + def assert_documents_are_equal(self, received: list[Document], expected: list[Document]) -> None: + """ + Compare documents while tolerating two S3 Vectors quirks: + + * embeddings round-trip through float32 storage, so we use `pytest.approx`; + * the `score` field is not set by `filter_documents`, only by retrieval, so + we ignore it for equality. + """ + assert len(received) == len(expected) + received.sort(key=lambda d: d.id) + expected.sort(key=lambda d: d.id) + for r, e in zip(received, expected, strict=True): + r_norm = replace(r, embedding=None, score=None) + e_norm = replace(e, embedding=None, score=None) + assert r_norm == e_norm + if r.embedding is not None and e.embedding is not None: + assert r.embedding == pytest.approx(e.embedding, abs=1e-5) + + def test_write_documents(self, document_store: S3VectorsDocumentStore) -> None: + """ + Default behaviour is OVERWRITE (S3 Vectors `put_vectors` is upsert), so + writing the same document twice succeeds and returns 1 each time. + """ + docs = [Document(id="1", content="hello", embedding=[0.1] * 768)] + assert document_store.write_documents(docs) == 1 + assert document_store.write_documents(docs) == 1 diff --git a/integrations/amazon_s3_vectors/tests/test_filters.py b/integrations/amazon_s3_vectors/tests/test_filters.py index 11d83a282d..6c744170ec 100644 --- a/integrations/amazon_s3_vectors/tests/test_filters.py +++ b/integrations/amazon_s3_vectors/tests/test_filters.py @@ -9,8 +9,12 @@ meta. prefix stripping, logical nesting) so unit testing is high signal. """ +from dataclasses import replace + import pytest +from haystack.dataclasses import Document from haystack.errors import FilterError +from haystack.testing.document_store import FilterDocumentsTest from haystack_integrations.document_stores.amazon_s3_vectors.filters import _normalize_filters, _validate_filters @@ -87,3 +91,26 @@ def test_validate_filters(): _validate_filters({"operator": "AND", "conditions": []}) # valid structure with pytest.raises(ValueError, match="Invalid filter syntax"): _validate_filters({"field": "meta.x"}) # missing operator/conditions + + +# --------------------------------------------------------------------------- +# Integration tests — run Haystack's full filter contract against a real S3 +# Vectors index. `S3VectorsDocumentStore.filter_documents` delegates the +# actual matching to `haystack.utils.filters.document_matches_filter`, so the +# only S3-specific quirk we have to absorb here is the float32 round-trip on +# embeddings. +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestFilters(FilterDocumentsTest): + def assert_documents_are_equal(self, received: list[Document], expected: list[Document]) -> None: + assert len(received) == len(expected) + received.sort(key=lambda d: d.id) + expected.sort(key=lambda d: d.id) + for r, e in zip(received, expected, strict=True): + r_norm = replace(r, embedding=None, score=None) + e_norm = replace(e, embedding=None, score=None) + assert r_norm == e_norm + if r.embedding is not None and e.embedding is not None: + assert r.embedding == pytest.approx(e.embedding, abs=1e-5) diff --git a/integrations/amazon_s3_vectors/tests/test_integration.py b/integrations/amazon_s3_vectors/tests/test_integration.py index 2ec71353f4..297c6c18aa 100644 --- a/integrations/amazon_s3_vectors/tests/test_integration.py +++ b/integrations/amazon_s3_vectors/tests/test_integration.py @@ -3,218 +3,74 @@ # SPDX-License-Identifier: Apache-2.0 """ -Integration tests for S3VectorsDocumentStore. +Retriever-focused integration tests. -Requires valid AWS credentials with s3vectors permissions. -Run with: hatch run test:integration +The Document Store contract is covered by `TestDocumentStore` / `TestFilters` +in `test_document_store.py` and `test_filters.py` (which inherit Haystack's +base test mixins). The tests here are the small set that don't fit those +mixins: the embedding-retrieval path and the `S3VectorsEmbeddingRetriever` +component wired up end-to-end. """ -import os -import time -import uuid +from __future__ import annotations import pytest from haystack.dataclasses import Document -from haystack.document_stores.errors import DocumentStoreError from haystack.document_stores.types import DuplicatePolicy from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore -def _aws_credentials_available() -> bool: - """Check if AWS credentials are available via any mechanism.""" - # Explicit env vars - if os.environ.get("AWS_ACCESS_KEY_ID") or os.environ.get("AWS_PROFILE") or os.environ.get("AWS_ROLE_ARN"): - return True - # Fallback: try to resolve credentials from the default chain - try: - import boto3 - - session = boto3.Session() - credentials = session.get_credentials() - return credentials is not None - except Exception: - return False - - -# Guard: skip all integration tests when AWS credentials are not available -pytestmark = pytest.mark.skipif( - not _aws_credentials_available(), - reason="AWS credentials not configured", -) - -# Use a small dimension to keep payloads light -DIMENSION = 4 -REGION = "us-east-1" - - -def _random_name(prefix: str = "haystack-test") -> str: - return f"{prefix}-{uuid.uuid4().hex[:8]}" - - -def _make_doc(doc_id: str, content: str, embedding: list[float] | None = None, meta: dict | None = None) -> Document: - return Document( - id=doc_id, - content=content, - embedding=embedding or [0.1] * DIMENSION, - meta=meta or {}, - ) - - -@pytest.fixture(scope="module") -def doc_store(): - """Create a real S3 Vectors document store with a unique bucket + index, tear down after.""" - bucket_name = _random_name("hs-integ") - index_name = _random_name("idx") - - store = S3VectorsDocumentStore( - vector_bucket_name=bucket_name, - index_name=index_name, - dimension=DIMENSION, - distance_metric="cosine", - region_name=REGION, - create_bucket_and_index=True, - non_filterable_metadata_keys=[], - ) - # Force initialization - store._get_client() - - yield store - - # Cleanup - client = store._client - try: - client.delete_index(vectorBucketName=bucket_name, indexName=index_name) - except Exception: - pass - try: - client.delete_vector_bucket(vectorBucketName=bucket_name) - except Exception: - pass - - -@pytest.mark.integration -class TestWriteAndCount: - def test_write_and_count(self, doc_store): - docs = [ - _make_doc("int-1", "First document"), - _make_doc("int-2", "Second document"), - _make_doc("int-3", "Third document"), - ] - written = doc_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) - assert written == 3 - - # S3 Vectors is eventually consistent — give it a moment - time.sleep(2) - - count = doc_store.count_documents() - assert count == 3 - - def test_overwrite(self, doc_store): - doc = _make_doc("int-1", "Updated first document", embedding=[0.9] * DIMENSION) - written = doc_store.write_documents([doc], policy=DuplicatePolicy.OVERWRITE) - assert written == 1 - - def test_skip_existing(self, doc_store): - doc = _make_doc("int-1", "Should be skipped") - written = doc_store.write_documents([doc], policy=DuplicatePolicy.SKIP) - assert written == 0 - - def test_none_policy_raises_on_existing(self, doc_store): - doc = _make_doc("int-1", "Should fail") - with pytest.raises(DocumentStoreError, match="already exist"): - doc_store.write_documents([doc], policy=DuplicatePolicy.NONE) - - def test_write_without_embedding_raises(self, doc_store): - doc = Document(id="no-emb", content="No embedding") - with pytest.raises(DocumentStoreError, match="has no embedding"): - doc_store.write_documents([doc]) - - @pytest.mark.integration -class TestQuery: - def test_embedding_retrieval(self, doc_store): - # Query with an embedding close to [0.1, 0.1, 0.1, 0.1] - docs = doc_store._embedding_retrieval( - query_embedding=[0.1] * DIMENSION, - top_k=10, - ) - assert len(docs) > 0 - # All returned docs should have a score (cosine similarity) - for doc in docs: - assert doc.score is not None - assert doc.content is not None - - def test_embedding_retrieval_with_metadata_filter(self, doc_store): - # Write a doc with distinctive metadata - tagged_doc = _make_doc( - "int-tagged", - "Tagged document", - embedding=[0.5] * DIMENSION, - meta={"category": "special"}, - ) - doc_store.write_documents([tagged_doc], policy=DuplicatePolicy.OVERWRITE) - time.sleep(2) - - # Query with filter - filters = {"field": "meta.category", "operator": "==", "value": "special"} - docs = doc_store._embedding_retrieval( - query_embedding=[0.5] * DIMENSION, - filters=filters, - top_k=10, - ) - assert len(docs) >= 1 - assert all(d.meta.get("category") == "special" for d in docs) - - def test_retriever_component(self, doc_store): - retriever = S3VectorsEmbeddingRetriever(document_store=doc_store, top_k=5) - result = retriever.run(query_embedding=[0.1] * DIMENSION) - assert "documents" in result - assert len(result["documents"]) > 0 +def test_embedding_retrieval(document_store: S3VectorsDocumentStore) -> None: + docs = [ + Document(id="r-1", content="First", embedding=[0.1] * 768), + Document(id="r-2", content="Second", embedding=[0.5] * 768), + Document(id="r-3", content="Third", embedding=[0.9] * 768), + ] + document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) + + results = document_store._embedding_retrieval(query_embedding=[0.1] * 768, top_k=10) + assert len(results) > 0 + for doc in results: + assert doc.score is not None + assert doc.content is not None + # query_vectors does not return vector data + assert doc.embedding is None @pytest.mark.integration -class TestFilterDocuments: - def test_filter_documents_no_filter(self, doc_store): - docs = doc_store.filter_documents() - assert len(docs) > 0 - for doc in docs: - assert doc.id is not None +def test_embedding_retrieval_with_metadata_filter(document_store: S3VectorsDocumentStore) -> None: + docs = [ + Document(id="f-1", content="Sports article", embedding=[0.1] * 768, meta={"category": "sports"}), + Document(id="f-2", content="News article", embedding=[0.1] * 768, meta={"category": "news"}), + ] + document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) - def test_filter_documents_with_filter(self, doc_store): - filters = {"field": "meta.category", "operator": "==", "value": "special"} - docs = doc_store.filter_documents(filters=filters) - assert all(d.meta.get("category") == "special" for d in docs) + filters = {"field": "meta.category", "operator": "==", "value": "sports"} + results = document_store._embedding_retrieval(query_embedding=[0.1] * 768, filters=filters, top_k=10) + assert len(results) >= 1 + assert all(d.meta.get("category") == "sports" for d in results) @pytest.mark.integration -class TestDelete: - def test_delete_documents(self, doc_store): - # Write a doc to delete - doc = _make_doc("int-to-delete", "Delete me") - doc_store.write_documents([doc], policy=DuplicatePolicy.OVERWRITE) - time.sleep(2) - - doc_store.delete_documents(["int-to-delete"]) - time.sleep(2) - - # Verify it's gone by trying to get it - client = doc_store._client - response = client.get_vectors( - vectorBucketName=doc_store.vector_bucket_name, - indexName=doc_store.index_name, - keys=["int-to-delete"], - ) - assert len(response.get("vectors", [])) == 0 +def test_retriever_component(document_store: S3VectorsDocumentStore) -> None: + docs = [Document(id="c-1", content="Hello", embedding=[0.1] * 768)] + document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) + + retriever = S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5) + result = retriever.run(query_embedding=[0.1] * 768) + assert "documents" in result + assert len(result["documents"]) > 0 @pytest.mark.integration -class TestSerialization: - def test_roundtrip(self, doc_store): - d = doc_store.to_dict() - restored = S3VectorsDocumentStore.from_dict(d) - assert restored.vector_bucket_name == doc_store.vector_bucket_name - assert restored.index_name == doc_store.index_name - assert restored.dimension == doc_store.dimension - assert restored.distance_metric == doc_store.distance_metric +def test_to_from_dict_roundtrip(document_store: S3VectorsDocumentStore) -> None: + """Serialization roundtrip against a real, fully-initialised store.""" + d = document_store.to_dict() + restored = S3VectorsDocumentStore.from_dict(d) + assert restored.vector_bucket_name == document_store.vector_bucket_name + assert restored.index_name == document_store.index_name + assert restored.dimension == document_store.dimension + assert restored.distance_metric == document_store.distance_metric From cdefc78adaec3fd2c601b54c3d00e81dc6c4f0e4 Mon Sep 17 00:00:00 2001 From: Oleksii Udovychenko Date: Fri, 4 Sep 2026 20:38:13 +0000 Subject: [PATCH 13/17] feat(amazon-s3-vectors): address review feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `write_documents` now only supports `DuplicatePolicy.OVERWRITE`. `put_vectors` is an upsert and S3 Vectors has no cheap existence check — `get_vectors` is capped at 100 keys and the only alternative is a full `ListVectors` scan — so the per-batch `get_vectors` pass is dropped rather than paying for a read on every write. Other policies are ignored with a warning. - Documents without an embedding are stored with a non-zero placeholder vector, as `PineconeDocumentStore` does, and reported in one warning before any write happens instead of raising. Haystack's `filterable_docs` fixture ships embedding-less documents on purpose, so refusing them makes `FilterDocumentsTest` unsatisfiable. The caller's `Document` is not mutated. - The test fixture no longer back-fills embeddings, which removes the `warnings.simplefilter("ignore")` that was hiding the store's own warnings. It now only absorbs S3 Vectors' eventual consistency. - Add `delete_all_documents()`, `delete_by_filter()` and `update_by_filter()`, and mix in `DeleteAllTest`, `DeleteByFilterTest` and `UpdateByFilterTest` so the integration suite covers everything in `DocumentStoreBaseExtendedTests`. `FilterDocumentsTest` stays in `test_filters.py`, as in pgvector. - Factor the `list_vectors` pagination into `_iter_vectors()` instead of writing a third copy of the loop. - Move the retriever integration tests into `test_embedding_retriever.py` and delete `test_integration.py`; its `to_dict`/`from_dict` roundtrip is already covered by the unit tests. - Drop the two decorative comment banners in the test files. Verified against a real S3 Vectors bucket in us-east-1: 59 passed, 4 skipped — the two unsupported duplicate policies, plus the two `recreate_index` variants that self-skip. Unit tests, `fmt-check` and mypy are green. --- .../amazon_s3_vectors/document_store.py | 226 +++++++++++------- .../amazon_s3_vectors/tests/conftest.py | 34 +-- .../tests/test_document_store.py | 119 ++++----- .../tests/test_embedding_retriever.py | 42 ++++ .../amazon_s3_vectors/tests/test_filters.py | 9 - .../tests/test_integration.py | 76 ------ 6 files changed, 253 insertions(+), 253 deletions(-) delete mode 100644 integrations/amazon_s3_vectors/tests/test_integration.py diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index caeb091800..cf1e3df233 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -4,6 +4,7 @@ import base64 import json +from collections.abc import Iterator from dataclasses import replace from typing import Any, Literal @@ -11,7 +12,6 @@ from botocore.exceptions import ClientError from haystack import default_from_dict, default_to_dict, logging from haystack.dataclasses import ByteStream, Document -from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError from haystack.document_stores.types import DuplicatePolicy from haystack.utils.auth import Secret, deserialize_secrets_inplace from haystack.utils.filters import document_matches_filter @@ -23,8 +23,7 @@ # S3 Vectors allows up to 500 vectors per put_vectors call _WRITE_BATCH_SIZE = 500 -# S3 Vectors allows up to 100 keys per get_vectors call, 500 per delete_vectors call -_GET_BATCH_SIZE = 100 +# S3 Vectors allows up to 500 keys per delete_vectors call _DELETE_BATCH_SIZE = 500 # Maximum number of results from query_vectors @@ -59,7 +58,10 @@ class S3VectorsDocumentStore: - Maximum `top_k`: 100 results per query - Maximum vector dimension: 4,096 - Metadata per vector: 40 KB total, 2 KB filterable - - All documents must have embeddings (`float32` only) + - Every record needs a vector (`float32` only), so documents written without an embedding are + stored with a placeholder vector (stripped again on read) and are only meaningfully + retrievable via `filter_documents()` + - Only `DuplicatePolicy.OVERWRITE` is supported, since `put_vectors` is an upsert - Distance metrics: `cosine` or `euclidean` (set at index creation, immutable) - `filter_documents()` is client-side — prefer `S3VectorsEmbeddingRetriever` with filters @@ -116,6 +118,10 @@ def __init__( self.create_bucket_and_index = create_bucket_and_index self.non_filterable_metadata_keys = non_filterable_metadata_keys or [] + # S3 Vectors requires a vector on every record, and a cosine index rejects zero-norm + # vectors, so documents written without an embedding get this non-zero placeholder. + self._dummy_vector = [-10.0] * dimension + self._client: Any = None def _get_client(self) -> Any: @@ -210,51 +216,62 @@ def from_dict(cls, data: dict[str, Any]) -> "S3VectorsDocumentStore": ) return default_from_dict(cls, data) - def count_documents(self) -> int: + def _iter_vectors(self, *, return_data: bool = False, return_metadata: bool = False) -> Iterator[dict[str, Any]]: """ - Return the number of documents in the document store. - - .. note:: + Yield every vector in the index, paginating through `list_vectors`. - S3 Vectors does not provide a dedicated count API. This method lists all vector keys - via pagination, which can be slow for large indexes. + :param return_data: Whether to ask S3 Vectors for the vector data. + :param return_metadata: Whether to ask S3 Vectors for the vector metadata. """ client = self._get_client() - count = 0 next_token = None while True: kwargs: dict[str, Any] = { "vectorBucketName": self.vector_bucket_name, "indexName": self.index_name, } + if return_data: + kwargs["returnData"] = True + if return_metadata: + kwargs["returnMetadata"] = True if next_token: kwargs["nextToken"] = next_token response = client.list_vectors(**kwargs) - count += len(response.get("vectors", [])) + yield from response.get("vectors", []) next_token = response.get("nextToken") if not next_token: break - return count + + def count_documents(self) -> int: + """ + Return the number of documents in the document store. + + .. note:: + + S3 Vectors does not provide a dedicated count API. This method lists all vector keys + via pagination, which can be slow for large indexes. + """ + return sum(1 for _ in self._iter_vectors()) def write_documents(self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.OVERWRITE) -> int: """ Write Documents to the S3 Vectors index. - All documents must have an embedding set. S3 Vectors `put_vectors` is an upsert operation - by default, so `DuplicatePolicy.OVERWRITE` is the natural behavior. - `DuplicatePolicy.SKIP` will check for existing documents first (slower). - `DuplicatePolicy.FAIL` and `DuplicatePolicy.NONE` will raise `DuplicateDocumentError` - if any document already exists. + `S3VectorsDocumentStore` only supports `DuplicatePolicy.OVERWRITE`: `put_vectors` is an + upsert, and S3 Vectors offers no cheap existence check (only `get_vectors` on batches of + 100 keys, or a full index scan), so enforcing any other policy would cost an extra read + pass on every write. Other policies are ignored with a warning. + + Documents without an embedding are stored with a placeholder vector, since S3 Vectors + requires a vector on every record. `filter_documents()` strips the placeholder, so they read + back with `embedding=None`, but they are meaningless as similarity-search results. Metadata per vector is limited to 40 KB total (2 KB filterable). - :param documents: A list of Documents to write. Each document must have an embedding. - :param policy: The duplicate policy. Defaults to `DuplicatePolicy.OVERWRITE`. + :param documents: A list of Documents to write. + :param policy: The duplicate policy. Only `DuplicatePolicy.OVERWRITE` is supported. :returns: The number of documents written. :raises ValueError: If `documents` is not a list of `Document` instances. - :raises DocumentStoreError: If any document is missing an embedding. - :raises DuplicateDocumentError: If `policy` is `FAIL` or `NONE` and any document already - exists in the store. """ # Validate input type up front, before any writes happen. if not isinstance(documents, list) or any(not isinstance(d, Document) for d in documents): @@ -264,56 +281,38 @@ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = D if len(documents) == 0: return 0 - # Validate ALL embeddings up front so we fail before any put_vectors call, - # never leaving the store with a partial write. + # DuplicatePolicy.NONE means "the store decides", so it needs no warning. + if policy not in (DuplicatePolicy.NONE, DuplicatePolicy.OVERWRITE): + logger.warning( + "S3VectorsDocumentStore only supports DuplicatePolicy.OVERWRITE because put_vectors is an " + "upsert. Ignoring policy={policy} and overwriting.", + policy=str(policy), + ) + + # Report every embedding-less document in one pass, before any put_vectors call, so the + # warning describes the whole write rather than arriving batch by batch. missing_embedding_ids = [doc.id for doc in documents if doc.embedding is None] if missing_embedding_ids: - msg = ( - f"Document(s) {missing_embedding_ids} have no embedding. " - "S3VectorsDocumentStore requires every document to have an embedding." + logger.warning( + "Document(s) {ids} have no embedding and will be stored with a placeholder vector. " + "They are retrievable with filter_documents() but will pollute similarity search results.", + ids=missing_embedding_ids, ) - raise DocumentStoreError(msg) client = self._get_client() written = 0 for i in range(0, len(documents), _WRITE_BATCH_SIZE): batch = documents[i : i + _WRITE_BATCH_SIZE] - - # Batch-check for existing documents when needed - existing_ids: set[str] = set() - if policy in (DuplicatePolicy.SKIP, DuplicatePolicy.NONE, DuplicatePolicy.FAIL): - batch_ids = [doc.id for doc in batch] - for j in range(0, len(batch_ids), _GET_BATCH_SIZE): - id_chunk = batch_ids[j : j + _GET_BATCH_SIZE] - response = client.get_vectors( - vectorBucketName=self.vector_bucket_name, - indexName=self.index_name, - keys=id_chunk, - ) - for v in response.get("vectors", []): - existing_ids.add(v["key"]) - - if policy in (DuplicatePolicy.NONE, DuplicatePolicy.FAIL) and existing_ids: - msg = ( - f"Document(s) {sorted(existing_ids)} already exist in the document store. " - "Use DuplicatePolicy.OVERWRITE or DuplicatePolicy.SKIP." - ) - raise DuplicateDocumentError(msg) - - vectors_to_write = [] - for doc in batch: - if policy == DuplicatePolicy.SKIP and doc.id in existing_ids: - continue - vectors_to_write.append(self._document_to_s3_vector(doc)) - - if vectors_to_write: - client.put_vectors( - vectorBucketName=self.vector_bucket_name, - indexName=self.index_name, - vectors=vectors_to_write, - ) - written += len(vectors_to_write) + vectors_to_write = [ + self._document_to_s3_vector(doc, fallback_embedding=self._dummy_vector) for doc in batch + ] + client.put_vectors( + vectorBucketName=self.vector_bucket_name, + indexName=self.index_name, + vectors=vectors_to_write, + ) + written += len(vectors_to_write) return written @@ -339,28 +338,13 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume "Prefer using S3VectorsEmbeddingRetriever with filters for efficient filtered retrieval." ) - client = self._get_client() - # list_vectors supports returnData and returnMetadata directly, # so we can read documents in a single paginated pass without # a separate get_vectors round-trip. - documents: list[Document] = [] - next_token = None - while True: - kwargs: dict[str, Any] = { - "vectorBucketName": self.vector_bucket_name, - "indexName": self.index_name, - "returnData": True, - "returnMetadata": True, - } - if next_token: - kwargs["nextToken"] = next_token - response = client.list_vectors(**kwargs) - for v in response.get("vectors", []): - documents.append(self._s3_vector_to_document(v)) - next_token = response.get("nextToken") - if not next_token: - break + documents = [ + self._s3_vector_to_document(v, placeholder_embedding=self._dummy_vector) + for v in self._iter_vectors(return_data=True, return_metadata=True) + ] if filters: _validate_filters(filters) @@ -386,6 +370,55 @@ def delete_documents(self, document_ids: list[str]) -> None: keys=batch, ) + def delete_all_documents(self) -> None: + """ + Delete all documents from the index. + + .. note:: + + S3 Vectors has no truncate operation, so every key is listed and then deleted in + batches, which can be slow for large indexes. + """ + self.delete_documents([v["key"] for v in self._iter_vectors()]) + + def delete_by_filter(self, filters: dict[str, Any]) -> int: + """ + Delete all documents matching the given filters. + + .. warning:: + + Filtering is client-side (see `filter_documents`), so this reads the whole index. + + :param filters: Haystack-format filters selecting the documents to delete. + :returns: The number of documents deleted. + """ + document_ids = [doc.id for doc in self.filter_documents(filters=filters)] + self.delete_documents(document_ids) + return len(document_ids) + + def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int: + """ + Merge `meta` into the metadata of every document matching the given filters. + + Keys present in `meta` are overwritten, the rest of each document's metadata is kept. + + .. warning:: + + S3 Vectors cannot update metadata in place, so matching documents are read and written + back in full. Filtering is client-side (see `filter_documents`), so this reads the + whole index. + + :param filters: Haystack-format filters selecting the documents to update. + :param meta: Metadata fields to set on the matching documents. + :returns: The number of documents updated. + """ + matching = self.filter_documents(filters=filters) + if not matching: + return 0 + + self.write_documents([replace(doc, meta={**doc.meta, **meta}) for doc in matching]) + return len(matching) + def _embedding_retrieval( self, query_embedding: list[float], @@ -460,8 +493,14 @@ def _embedding_retrieval( return documents @staticmethod - def _document_to_s3_vector(doc: Document) -> dict[str, Any]: - """Convert a Haystack Document to an S3 Vectors vector entry.""" + def _document_to_s3_vector(doc: Document, *, fallback_embedding: list[float] | None = None) -> dict[str, Any]: + """ + Convert a Haystack Document to an S3 Vectors vector entry. + + :param doc: The Document to convert. + :param fallback_embedding: Vector to send when `doc` has no embedding. S3 Vectors requires + a vector on every record. + """ metadata: dict[str, Any] = {} # Store content as non-filterable metadata @@ -503,13 +542,20 @@ def _document_to_s3_vector(doc: Document) -> dict[str, Any]: return { "key": doc.id, - "data": {"float32": doc.embedding}, + "data": {"float32": doc.embedding if doc.embedding is not None else fallback_embedding}, "metadata": metadata, } @staticmethod - def _s3_vector_to_document(vector: dict[str, Any]) -> Document: - """Convert an S3 Vectors vector response to a Haystack Document.""" + def _s3_vector_to_document(vector: dict[str, Any], *, placeholder_embedding: list[float] | None = None) -> Document: + """ + Convert an S3 Vectors vector response to a Haystack Document. + + :param vector: The vector as returned by `list_vectors` or `get_vectors`. + :param placeholder_embedding: Vector that stands for "no embedding". A stored vector equal to + it reads back as `embedding=None`, so a document written without an embedding round-trips + as one. See `write_documents`. + """ metadata = dict(vector.get("metadata", {})) content = metadata.pop(_CONTENT_KEY, None) blob_data = metadata.pop(_BLOB_DATA_KEY, None) @@ -528,6 +574,10 @@ def _s3_vector_to_document(vector: dict[str, Any]) -> Document: data = vector.get("data", {}) if "float32" in data: embedding = data["float32"] + # A placeholder vector means the document was written without an embedding; don't hand + # the caller a vector it never supplied. + if placeholder_embedding is not None and embedding == placeholder_embedding: + embedding = None return Document( id=vector["key"], diff --git a/integrations/amazon_s3_vectors/tests/conftest.py b/integrations/amazon_s3_vectors/tests/conftest.py index b58f963c34..a4f62e02c2 100644 --- a/integrations/amazon_s3_vectors/tests/conftest.py +++ b/integrations/amazon_s3_vectors/tests/conftest.py @@ -13,15 +13,10 @@ slow part). Each test gets its own vector index inside that shared bucket so state is isolated. -Two quirks of S3 Vectors are smoothed over here so the generic base tests -can run unchanged: - -1. **Embeddings are required.** The base tests write `Document(content="...")` - without an embedding. We wrap `write_documents` so any document missing an - embedding gets a deterministic zero vector of the right dimension. -2. **Writes are eventually consistent.** We sleep briefly after `write_documents` - and `delete_documents` so the subsequent `filter_documents`/`count_documents` - reflects the new state. +One quirk of S3 Vectors is smoothed over here so the generic base tests can run +unchanged: **writes are eventually consistent**. We sleep briefly after +`write_documents` and `delete_documents` so the subsequent +`filter_documents`/`count_documents` reflects the new state. """ from __future__ import annotations @@ -29,7 +24,6 @@ import os import time import uuid -import warnings from collections.abc import Iterator import boto3 @@ -109,11 +103,9 @@ def s3_vectors_bucket() -> Iterator[str]: @pytest.fixture def document_store(s3_vectors_bucket: str) -> Iterator[S3VectorsDocumentStore]: """ - Provide a fresh S3VectorsDocumentStore (one per test) with `write_documents` - and `delete_documents` wrapped to: - - * inject a default embedding for any `Document` that doesn't have one, and - * sleep briefly afterwards to absorb S3 Vectors' eventual consistency. + Provide a fresh S3VectorsDocumentStore (one per test), with `write_documents` and + `delete_documents` wrapped to sleep briefly afterwards so that S3 Vectors' eventual + consistency doesn't make the base tests flaky. """ index_name = f"idx-{uuid.uuid4().hex[:10]}" @@ -132,20 +124,12 @@ def document_store(s3_vectors_bucket: str) -> Iterator[S3VectorsDocumentStore]: original_write = store.write_documents - def write_with_defaults(documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.OVERWRITE) -> int: - # Only mutate well-formed input; bad input must surface as ValueError from - # the production code, not crash this wrapper. - if isinstance(documents, list): - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - for d in documents: - if isinstance(d, Document) and d.embedding is None: - d.embedding = [0.0] * DIMENSION + def write_with_sleep(documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.OVERWRITE) -> int: result = original_write(documents, policy) time.sleep(WRITE_SLEEP_SECONDS) return result - store.write_documents = write_with_defaults # type: ignore[method-assign] + store.write_documents = write_with_sleep # type: ignore[method-assign] original_delete = store.delete_documents diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index 6fb0df0de0..6a0a382801 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -7,12 +7,14 @@ import pytest from haystack.dataclasses import Document -from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError from haystack.document_stores.types import DuplicatePolicy from haystack.testing.document_store import ( CountDocumentsTest, + DeleteAllTest, + DeleteByFilterTest, DeleteDocumentsTest, FilterableDocsFixtureMixin, + UpdateByFilterTest, WriteDocumentsTest, ) @@ -97,52 +99,41 @@ def test_from_dict(_mock_boto3): @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_no_embedding_raises(mock_boto3): - """S3 Vectors requires embeddings — this tests our validation, not the store.""" +def test_write_documents_missing_embedding_uses_placeholder(mock_boto3): + """Every S3 Vectors record needs a vector, so embedding-less documents get the placeholder.""" client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - with pytest.raises(DocumentStoreError, match="no embedding"): - store.write_documents([Document(id="1", content="Hello")]) + docs = [Document(id="1", content="ok", embedding=[0.1] * 4), Document(id="2", content="missing")] + assert store.write_documents(docs) == 2 + written = client.put_vectors.call_args.kwargs["vectors"] + assert written[0]["data"]["float32"] == [0.1] * 4 + # Non-zero: a cosine index rejects zero-norm vectors. + assert written[1]["data"]["float32"] == [-10.0] * 4 + # The caller's Document is not mutated. + assert docs[1].embedding is None -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_skip_existing(mock_boto3): - """Tests our batch existence check logic for SKIP policy.""" - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - client.get_vectors.return_value = {"vectors": [{"key": "1"}]} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - result = store.write_documents( - [Document(id="1", content="Hello", embedding=[0.1] * 4)], - policy=DuplicatePolicy.SKIP, - ) - assert result == 0 - client.put_vectors.assert_not_called() - -@pytest.mark.parametrize("policy", [DuplicatePolicy.FAIL, DuplicatePolicy.NONE]) +@pytest.mark.parametrize("policy", [DuplicatePolicy.FAIL, DuplicatePolicy.SKIP]) @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_fail_policy_raises(mock_boto3, policy): - """Both FAIL and NONE raise DuplicateDocumentError on existing IDs.""" +def test_write_documents_unsupported_policy_overwrites(mock_boto3, policy): + """ + put_vectors is an upsert and S3 Vectors has no cheap existence check, so policies other than + OVERWRITE are ignored instead of paying for an extra read pass. + """ client = MagicMock() client.get_vector_bucket.return_value = {} client.get_index.return_value = {} - client.get_vectors.return_value = {"vectors": [{"key": "1"}]} mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - with pytest.raises(DuplicateDocumentError, match="already exist"): - store.write_documents( - [Document(id="1", content="Hello", embedding=[0.1] * 4)], - policy=policy, - ) + assert store.write_documents([Document(id="1", content="Hello", embedding=[0.1] * 4)], policy=policy) == 1 + client.get_vectors.assert_not_called() + client.put_vectors.assert_called_once() def test_write_documents_invalid_input_raises(): @@ -154,25 +145,6 @@ def test_write_documents_invalid_input_raises(): store.write_documents(["not a document"]) # type: ignore[list-item] -@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_missing_embedding_validated_upfront(mock_boto3): - """Missing embeddings are reported up front, before any put_vectors call.""" - client = MagicMock() - client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} - mock_boto3.client.return_value = client - - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") - docs = [ - Document(id="1", content="ok", embedding=[0.1] * 4), - Document(id="2", content="missing"), - Document(id="3", content="also-missing"), - ] - with pytest.raises(DocumentStoreError, match=r"\['2', '3'\].*no embedding"): - store.write_documents(docs) - client.put_vectors.assert_not_called() - - @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_embedding_retrieval_score_conversion(mock_boto3): """Tests our distance-to-score conversion logic — the only non-trivial transform in retrieval.""" @@ -277,11 +249,25 @@ def test_document_roundtrip(): assert restored.meta == doc.meta -# --------------------------------------------------------------------------- -# Integration tests — exercise a real S3 Vectors bucket via the `document_store` -# fixture in conftest.py. The mixins below come from Haystack's test kit so we -# get its standard Document Store contract for free. -# --------------------------------------------------------------------------- +def test_document_roundtrip_without_embedding(): + """A document written without an embedding must read back without one, not with the placeholder.""" + placeholder = [-10.0] * 3 + doc = Document(id="test-1", content="Hello world", meta={"category": "test"}) + + vector = S3VectorsDocumentStore._document_to_s3_vector(doc, fallback_embedding=placeholder) + assert vector["data"]["float32"] == placeholder + + restored = S3VectorsDocumentStore._s3_vector_to_document(vector, placeholder_embedding=placeholder) + assert restored.embedding is None + assert restored.content == doc.content + assert restored.meta == doc.meta + + +def test_s3_vector_to_document_keeps_embedding_matching_no_placeholder(): + """A real embedding is never mistaken for the placeholder.""" + vector = {"key": "test-1", "data": {"float32": [0.1, 0.2, 0.3]}, "metadata": {}} + doc = S3VectorsDocumentStore._s3_vector_to_document(vector, placeholder_embedding=[-10.0] * 3) + assert doc.embedding == [0.1, 0.2, 0.3] @pytest.mark.integration @@ -289,6 +275,9 @@ class TestDocumentStore( CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest, + DeleteAllTest, + DeleteByFilterTest, + UpdateByFilterTest, FilterableDocsFixtureMixin, ): def assert_documents_are_equal(self, received: list[Document], expected: list[Document]) -> None: @@ -317,3 +306,23 @@ def test_write_documents(self, document_store: S3VectorsDocumentStore) -> None: docs = [Document(id="1", content="hello", embedding=[0.1] * 768)] assert document_store.write_documents(docs) == 1 assert document_store.write_documents(docs) == 1 + + def test_write_documents_without_embedding_reads_back_as_none(self, document_store: S3VectorsDocumentStore) -> None: + """ + S3 Vectors stores a placeholder vector for embedding-less documents; `filter_documents` + must strip it again rather than hand back a vector the caller never supplied. + + `assert_documents_are_equal` ignores embeddings, so this needs its own assertion. + """ + document_store.write_documents([Document(id="no-emb", content="No embedding here")]) + + docs = document_store.filter_documents() + assert len(docs) == 1 + assert docs[0].content == "No embedding here" + assert docs[0].embedding is None + + @pytest.mark.skip(reason="S3 Vectors put_vectors is an upsert; only DuplicatePolicy.OVERWRITE is supported") + def test_write_documents_duplicate_fail(self, document_store: S3VectorsDocumentStore) -> None: ... + + @pytest.mark.skip(reason="S3 Vectors put_vectors is an upsert; only DuplicatePolicy.OVERWRITE is supported") + def test_write_documents_duplicate_skip(self, document_store: S3VectorsDocumentStore) -> None: ... diff --git a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py index 3f9e2e46da..3365f042a5 100644 --- a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py +++ b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py @@ -130,3 +130,45 @@ def test_run(): ) assert len(res["documents"]) == 1 assert res["documents"][0].content == "Test doc" + + +@pytest.mark.integration +def test_embedding_retrieval(document_store: S3VectorsDocumentStore) -> None: + docs = [ + Document(id="r-1", content="First", embedding=[0.1] * 768), + Document(id="r-2", content="Second", embedding=[0.5] * 768), + Document(id="r-3", content="Third", embedding=[0.9] * 768), + ] + document_store.write_documents(docs) + + results = document_store._embedding_retrieval(query_embedding=[0.1] * 768, top_k=10) + assert len(results) > 0 + for doc in results: + assert doc.score is not None + assert doc.content is not None + # query_vectors does not return vector data + assert doc.embedding is None + + +@pytest.mark.integration +def test_embedding_retrieval_with_metadata_filter(document_store: S3VectorsDocumentStore) -> None: + docs = [ + Document(id="f-1", content="Sports article", embedding=[0.1] * 768, meta={"category": "sports"}), + Document(id="f-2", content="News article", embedding=[0.1] * 768, meta={"category": "news"}), + ] + document_store.write_documents(docs) + + filters = {"field": "meta.category", "operator": "==", "value": "sports"} + results = document_store._embedding_retrieval(query_embedding=[0.1] * 768, filters=filters, top_k=10) + assert len(results) >= 1 + assert all(d.meta.get("category") == "sports" for d in results) + + +@pytest.mark.integration +def test_retriever_component(document_store: S3VectorsDocumentStore) -> None: + document_store.write_documents([Document(id="c-1", content="Hello", embedding=[0.1] * 768)]) + + retriever = S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5) + result = retriever.run(query_embedding=[0.1] * 768) + assert "documents" in result + assert len(result["documents"]) > 0 diff --git a/integrations/amazon_s3_vectors/tests/test_filters.py b/integrations/amazon_s3_vectors/tests/test_filters.py index 6c744170ec..f86ed17fb6 100644 --- a/integrations/amazon_s3_vectors/tests/test_filters.py +++ b/integrations/amazon_s3_vectors/tests/test_filters.py @@ -93,15 +93,6 @@ def test_validate_filters(): _validate_filters({"field": "meta.x"}) # missing operator/conditions -# --------------------------------------------------------------------------- -# Integration tests — run Haystack's full filter contract against a real S3 -# Vectors index. `S3VectorsDocumentStore.filter_documents` delegates the -# actual matching to `haystack.utils.filters.document_matches_filter`, so the -# only S3-specific quirk we have to absorb here is the float32 round-trip on -# embeddings. -# --------------------------------------------------------------------------- - - @pytest.mark.integration class TestFilters(FilterDocumentsTest): def assert_documents_are_equal(self, received: list[Document], expected: list[Document]) -> None: diff --git a/integrations/amazon_s3_vectors/tests/test_integration.py b/integrations/amazon_s3_vectors/tests/test_integration.py deleted file mode 100644 index 297c6c18aa..0000000000 --- a/integrations/amazon_s3_vectors/tests/test_integration.py +++ /dev/null @@ -1,76 +0,0 @@ -# SPDX-FileCopyrightText: 2022-present deepset GmbH -# -# SPDX-License-Identifier: Apache-2.0 - -""" -Retriever-focused integration tests. - -The Document Store contract is covered by `TestDocumentStore` / `TestFilters` -in `test_document_store.py` and `test_filters.py` (which inherit Haystack's -base test mixins). The tests here are the small set that don't fit those -mixins: the embedding-retrieval path and the `S3VectorsEmbeddingRetriever` -component wired up end-to-end. -""" - -from __future__ import annotations - -import pytest -from haystack.dataclasses import Document -from haystack.document_stores.types import DuplicatePolicy - -from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore - - -@pytest.mark.integration -def test_embedding_retrieval(document_store: S3VectorsDocumentStore) -> None: - docs = [ - Document(id="r-1", content="First", embedding=[0.1] * 768), - Document(id="r-2", content="Second", embedding=[0.5] * 768), - Document(id="r-3", content="Third", embedding=[0.9] * 768), - ] - document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) - - results = document_store._embedding_retrieval(query_embedding=[0.1] * 768, top_k=10) - assert len(results) > 0 - for doc in results: - assert doc.score is not None - assert doc.content is not None - # query_vectors does not return vector data - assert doc.embedding is None - - -@pytest.mark.integration -def test_embedding_retrieval_with_metadata_filter(document_store: S3VectorsDocumentStore) -> None: - docs = [ - Document(id="f-1", content="Sports article", embedding=[0.1] * 768, meta={"category": "sports"}), - Document(id="f-2", content="News article", embedding=[0.1] * 768, meta={"category": "news"}), - ] - document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) - - filters = {"field": "meta.category", "operator": "==", "value": "sports"} - results = document_store._embedding_retrieval(query_embedding=[0.1] * 768, filters=filters, top_k=10) - assert len(results) >= 1 - assert all(d.meta.get("category") == "sports" for d in results) - - -@pytest.mark.integration -def test_retriever_component(document_store: S3VectorsDocumentStore) -> None: - docs = [Document(id="c-1", content="Hello", embedding=[0.1] * 768)] - document_store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) - - retriever = S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5) - result = retriever.run(query_embedding=[0.1] * 768) - assert "documents" in result - assert len(result["documents"]) > 0 - - -@pytest.mark.integration -def test_to_from_dict_roundtrip(document_store: S3VectorsDocumentStore) -> None: - """Serialization roundtrip against a real, fully-initialised store.""" - d = document_store.to_dict() - restored = S3VectorsDocumentStore.from_dict(d) - assert restored.vector_bucket_name == document_store.vector_bucket_name - assert restored.index_name == document_store.index_name - assert restored.dimension == document_store.dimension - assert restored.distance_metric == document_store.distance_metric From 7c907b8e49c2a74f8df57933a41ad3c385a70a85 Mon Sep 17 00:00:00 2001 From: Oleksii Udovychenko Date: Fri, 4 Sep 2026 22:35:27 +0000 Subject: [PATCH 14/17] feat(amazon-s3-vectors): match PineconeDocumentStore's coverage Implement the whole optional DocumentStore extension surface Pinecone has, and mix in every matching test from haystack.testing.document_store: - count_documents_by_filter, count_unique_metadata_by_filter, get_metadata_fields_info, get_metadata_field_min_max and get_metadata_field_unique_values, with CountDocumentsByFilterTest, CountUniqueMetadataByFilterTest, GetMetadataFieldsInfoTest, GetMetadataFieldMinMaxTest and GetMetadataFieldUniqueValuesTest. - close(), which releases the boto3 client's connection pool. The store stays usable: the next call lazily builds a new client. - The store-specific Pinecone tests too (close-and-reopen, sentence-window retriever, metadata-introspection edge cases). Two deliberate divergences from Pinecone: no TOP_K_LIMIT truncation warnings, because list_vectors paginates over the whole index and these counts are exact; and get_metadata_field_unique_values keeps differently typed values distinct, so the base mixin's distinct-types test passes unmodified rather than needing Pinecone's override. Fix two metadata bugs that the new sentence-window test surfaced. S3 Vectors accepts a metadata value only if it is a string, a boolean, a finite number, or a non-empty array of either all strings or all non-boolean numbers; declaring a key non-filterable does not relax this. Anything else is a ValidationException from PutVectors, so: - Unstorable values are now dropped with a warning instead of failing the write, as PineconeDocumentStore does. This was not theoretical: DocumentSplitter with split_overlap adds a _split_overlap list of dicts to every chunk, so DocumentSplitter -> S3VectorsDocumentStore died with a ParamValidationError. Tuples are converted to lists, which the service accepts but botocore's document-type validator does not. - doc.blob.meta is an arbitrary dict, so writing any document with blob metadata failed. It is now JSON-encoded on write and decoded on read. --- .../amazon_s3_vectors/document_store.py | 289 +++++++++++++- .../tests/test_document_store.py | 352 +++++++++++++++++- 2 files changed, 637 insertions(+), 4 deletions(-) diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index cf1e3df233..8ba6592e9b 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -4,6 +4,7 @@ import base64 import json +import math from collections.abc import Iterator from dataclasses import replace from typing import Any, Literal @@ -58,6 +59,9 @@ class S3VectorsDocumentStore: - Maximum `top_k`: 100 results per query - Maximum vector dimension: 4,096 - Metadata per vector: 40 KB total, 2 KB filterable + - A metadata value must be a string, a boolean, a finite number, or a non-empty array of either + all strings or all non-boolean numbers. Anything else (`None`, nested objects, empty or mixed + arrays) is dropped with a warning rather than failing the write - Every record needs a vector (`float32` only), so documents written without an embedding are stored with a placeholder vector (stripped again on read) and are only meaningfully retrievable via `filter_documents()` @@ -419,6 +423,227 @@ def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int self.write_documents([replace(doc, meta={**doc.meta, **meta}) for doc in matching]) return len(matching) + def count_documents_by_filter(self, filters: dict[str, Any]) -> int: + """ + Return the number of documents matching the given filters. + + .. warning:: + + Filtering is client-side (see `filter_documents`), so this reads the whole index. + Unlike stores with a capped query API, the count is exact: `list_vectors` paginates + over every vector. + + :param filters: Haystack-format filters to apply. + :returns: The number of matching documents. + """ + return len(self.filter_documents(filters=filters)) + + def count_unique_metadata_by_filter(self, filters: dict[str, Any], metadata_fields: list[str]) -> dict[str, int]: + """ + Count the distinct values of each given metadata field, across documents matching `filters`. + + .. warning:: + + Filtering and aggregation are client-side (see `filter_documents`), so this reads the + whole index. + + :param filters: Haystack-format filters selecting the documents to consider. + :param metadata_fields: The metadata field names to count distinct values for. + :returns: A mapping of field name to its number of distinct values. + """ + documents = self.filter_documents(filters=filters) + return self._count_unique_metadata_impl(documents, metadata_fields) + + def get_metadata_fields_info(self) -> dict[str, dict[str, str]]: + """ + Infer the metadata fields present in the index and their types. + + S3 Vectors has no schema-introspection API, so types are inferred from the metadata of the + documents actually stored. Python types map to `boolean`, `long` (int/float), and `keyword` + (str); a field with mixed types across documents falls back to `keyword` with a warning. + + .. warning:: + + This reads the whole index (see `filter_documents`). + + :returns: A mapping of field name to `{"type": }`. + """ + return self._get_metadata_fields_info_impl(self.filter_documents()) + + def get_metadata_field_min_max(self, metadata_field: str) -> dict[str, Any]: + """ + Return the minimum and maximum value of a metadata field. + + Works for numeric, boolean and string values. A leading `meta.` on the field name is + stripped, since S3 Vectors stores metadata flat. + + .. warning:: + + This reads the whole index (see `filter_documents`). + + :param metadata_field: The metadata field to analyze. + :returns: `{"min": ..., "max": ...}`, both `None` when no document has a comparable value + for the field. + """ + return self._get_metadata_field_min_max_impl(self.filter_documents(), metadata_field) + + def get_metadata_field_unique_values( + self, + metadata_field: str, + search_term: str | None = None, + from_: int = 0, + size: int = 10, + filters: dict[str, Any] | None = None, + ) -> tuple[list[Any], int]: + """ + Return the distinct values of a metadata field, with optional search and pagination. + + Values keep their original Python type, and values that merely share a string form (the int + `1`, the float `1.0`, the str `"1"`, `True`) are kept distinct. A leading `meta.` on the + field name is stripped. + + .. warning:: + + This reads the whole index (see `filter_documents`). + + :param metadata_field: The metadata field to get distinct values for. + :param search_term: Optional case-insensitive substring the value must contain. + :param from_: Pagination offset into the sorted distinct values. + :param size: Maximum number of values to return. + :param filters: Optional Haystack-format filters restricting the documents considered. + :returns: A tuple of (page of distinct values, total number of matching distinct values). + """ + documents = self.filter_documents(filters=filters) + return self._get_metadata_field_unique_values_impl(documents, metadata_field, search_term, from_, size) + + @staticmethod + def _count_unique_metadata_impl(documents: list[Document], metadata_fields: list[str]) -> dict[str, int]: + """Count distinct metadata values per field over an already-filtered document list.""" + result = {} + for field in metadata_fields: + # Values of different types can compare equal in Python (`1 == True == 1.0`), so key on + # (type, value) to avoid silently merging them. + unique_values: set[tuple[type, Any]] = set() + for doc in documents: + if doc.meta and field in doc.meta: + value = doc.meta[field] + if isinstance(value, list): + unique_values.update((type(item), item) for item in value) + else: + unique_values.add((type(value), value)) + result[field] = len(unique_values) + return result + + @staticmethod + def _get_metadata_fields_info_impl(documents: list[Document]) -> dict[str, dict[str, str]]: + """Infer metadata field types from an already-filtered document list.""" + if not documents: + return {} + + field_types: dict[str, dict[str, str]] = {} + if any(doc.content is not None for doc in documents): + field_types["content"] = {"type": "text"} + + def _type_of(value: Any) -> str | None: + # bool must be checked before int/float: bool is a subclass of int in Python. + if isinstance(value, bool): + return "boolean" + if isinstance(value, (int, float)): + return "long" + if isinstance(value, str): + return "keyword" + return None + + field_samples: dict[str, set[str]] = {} + for doc in documents: + for field, value in (doc.meta or {}).items(): + # An empty list carries no type information, so it defaults to keyword. + inferred = (_type_of(value[0]) if value else "keyword") if isinstance(value, list) else _type_of(value) + if inferred is not None: + field_samples.setdefault(field, set()).add(inferred) + + for field, types_seen in field_samples.items(): + if len(types_seen) == 1: + field_types[field] = {"type": types_seen.pop()} + else: + logger.warning( + "Metadata field '{field}' has mixed types {types} across documents. Defaulting to 'keyword'.", + field=field, + types=sorted(types_seen), + ) + field_types[field] = {"type": "keyword"} + + return field_types + + @staticmethod + def _get_metadata_field_min_max_impl(documents: list[Document], metadata_field: str) -> dict[str, Any]: + """Compute min/max of a metadata field over an already-filtered document list.""" + field_name = metadata_field.removeprefix("meta.") + values: list[bool | int | float | str] = [] + for doc in documents: + if doc.meta and field_name in doc.meta: + value = doc.meta[field_name] + # bool is a subclass of int, so it is already covered by the numeric check; listing + # it explicitly documents that booleans are comparable and intentionally included. + if isinstance(value, (bool, int, float, str)): + values.append(value) + + if not values: + return {"min": None, "max": None} + + # min/max over mixed str and numeric values would raise, so only compare within one kind. + numeric = [v for v in values if not isinstance(v, str)] + comparable = numeric if numeric else values + return {"min": min(comparable), "max": max(comparable)} + + @staticmethod + def _get_metadata_field_unique_values_impl( + documents: list[Document], metadata_field: str, search_term: str | None, from_: int, size: int + ) -> tuple[list[Any], int]: + """Collect distinct values of a metadata field over an already-filtered document list.""" + field_name = metadata_field.removeprefix("meta.") + # Dedupe on (type, value): a plain set would merge `1`, `1.0` and `True`. + seen: set[tuple[type, Any]] = set() + unique_values: list[Any] = [] + + def _add(value: Any) -> None: + key = (type(value), value) + if key not in seen: + seen.add(key) + unique_values.append(value) + + for doc in documents: + if doc.meta and field_name in doc.meta: + value = doc.meta[field_name] + if isinstance(value, list): + for item in value: + _add(item) + else: + _add(value) + + # Sort by string form so pagination is stable across mixed types. + values_list = sorted(unique_values, key=str) + + if search_term: + needle = search_term.lower() + values_list = [v for v in values_list if needle in str(v).lower()] + + return values_list[from_ : from_ + size], len(values_list) + + def close(self) -> None: + """ + Close the underlying boto3 client and release its connection pool. + + The store can be used again afterwards: the next call lazily creates a new client. + """ + if self._client is not None: + try: + self._client.close() + except Exception: + # Closing is best-effort; a failure here must not mask the caller's own work. + logger.debug("Ignoring error while closing the S3 Vectors client.") + self._client = None + def _embedding_retrieval( self, query_embedding: list[float], @@ -507,15 +732,17 @@ def _document_to_s3_vector(doc: Document, *, fallback_embedding: list[float] | N if doc.content is not None: metadata[_CONTENT_KEY] = doc.content - # Store blob fields + # Store blob fields. `blob.meta` is an arbitrary dict, which S3 Vectors rejects, so it is + # JSON-encoded into a string (and decoded again in `_s3_vector_to_document`). if doc.blob is not None: metadata[_BLOB_DATA_KEY] = base64.b64encode(doc.blob.data).decode("ascii") if doc.blob.meta: - metadata[_BLOB_META_KEY] = doc.blob.meta + metadata[_BLOB_META_KEY] = json.dumps(doc.blob.meta) if doc.blob.mime_type: metadata[_BLOB_MIME_TYPE_KEY] = doc.blob.mime_type # Store user metadata + dropped: list[str] = [] if doc.meta: for key, value in doc.meta.items(): if key in _RESERVED_META_KEYS: @@ -524,7 +751,21 @@ def _document_to_s3_vector(doc: Document, *, fallback_embedding: list[float] | N key=key, ) continue - metadata[key] = value + if not S3VectorsDocumentStore._is_supported_metadata_value(value): + dropped.append(key) + continue + # A tuple is an array as far as S3 Vectors is concerned, but botocore's document-type + # validator only accepts `list`. + metadata[key] = list(value) if isinstance(value, tuple) else value + + if dropped: + logger.warning( + "Document '{doc_id}': metadata key(s) {keys} have values S3 Vectors cannot store and " + "were dropped. Metadata values must be a string, a boolean, a finite number, or a " + "non-empty array of either all strings or all non-boolean numbers.", + doc_id=doc.id, + keys=sorted(dropped), + ) # Warn if metadata is likely too large try: @@ -546,6 +787,37 @@ def _document_to_s3_vector(doc: Document, *, fallback_embedding: list[float] | N "metadata": metadata, } + @staticmethod + def _is_supported_metadata_value(value: Any) -> bool: + """ + Return whether S3 Vectors can store `value` as a metadata value. + + The accepted set was established against the live service: a string (empty is fine), a + boolean, a finite number (ints of any magnitude), or a non-empty array whose elements are + *either* all strings or all non-boolean numbers. Everything else is rejected with a + `ValidationException` — `None`, dicts, empty arrays, nested arrays, arrays mixing strings and + numbers, arrays containing booleans, and NaN/infinity. + + This matters in practice: `DocumentSplitter` with `split_overlap` writes a `_split_overlap` + list of dicts, so without this check a plain splitter-to-store pipeline would fail. + """ + # bool first: it is a subclass of int, and unlike numbers it is allowed only outside arrays. + if isinstance(value, (bool, str)): + return True + if isinstance(value, int): + return True + if isinstance(value, float): + return math.isfinite(value) + if isinstance(value, (list, tuple)): + if not value: + return False + if all(isinstance(item, str) for item in value): + return True + return all( + not isinstance(item, bool) and isinstance(item, (int, float)) and math.isfinite(item) for item in value + ) + return False + @staticmethod def _s3_vector_to_document(vector: dict[str, Any], *, placeholder_embedding: list[float] | None = None) -> Document: """ @@ -564,6 +836,17 @@ def _s3_vector_to_document(vector: dict[str, Any], *, placeholder_embedding: lis blob = None if blob_data is not None: + # `blob.meta` is stored JSON-encoded (see `_document_to_s3_vector`). + if isinstance(blob_meta, str): + try: + blob_meta = json.loads(blob_meta) + except json.JSONDecodeError: + logger.warning( + "Could not decode the stored blob metadata of document '{doc_id}'; " + "the blob is returned without it.", + doc_id=vector["key"], + ) + blob_meta = None blob = ByteStream( data=base64.b64decode(blob_data) if isinstance(blob_data, str) else blob_data, meta=blob_meta or {}, diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index 6a0a382801..0a4c5fe4a1 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -2,22 +2,32 @@ # # SPDX-License-Identifier: Apache-2.0 +import base64 from dataclasses import replace from unittest.mock import MagicMock, patch +import numpy as np import pytest -from haystack.dataclasses import Document +from haystack.components.preprocessors import DocumentSplitter +from haystack.components.retrievers import SentenceWindowRetriever +from haystack.dataclasses import ByteStream, Document from haystack.document_stores.types import DuplicatePolicy from haystack.testing.document_store import ( + CountDocumentsByFilterTest, CountDocumentsTest, + CountUniqueMetadataByFilterTest, DeleteAllTest, DeleteByFilterTest, DeleteDocumentsTest, FilterableDocsFixtureMixin, + GetMetadataFieldMinMaxTest, + GetMetadataFieldsInfoTest, + GetMetadataFieldUniqueValuesTest, UpdateByFilterTest, WriteDocumentsTest, ) +from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore @@ -270,6 +280,257 @@ def test_s3_vector_to_document_keeps_embedding_matching_no_placeholder(): assert doc.embedding == [0.1, 0.2, 0.3] +@pytest.mark.parametrize( + ("value", "supported"), + [ + # Every case below was verified against the live S3 Vectors API: the "False" rows all come + # back as `ValidationException` from PutVectors. + ("text", True), + ("", True), + (True, True), + (0, True), + (2**63, True), + (1.5, True), + (["x", "y"], True), + ([1, 2.5], True), + (None, False), + ({"a": 1}, False), + ([], False), + ([1, "x"], False), # an array may not mix strings and numbers + ([True, False], False), # booleans are allowed only outside arrays + ([{"doc_id": "a", "range": [0, 24]}], False), # DocumentSplitter's `_split_overlap` + ([[1, 2]], False), + (["x", None], False), + (float("nan"), False), + (float("inf"), False), + ], +) +def test_is_supported_metadata_value(value, supported): + assert S3VectorsDocumentStore._is_supported_metadata_value(value) is supported + + +def test_document_to_s3_vector_drops_unstorable_metadata(caplog): + """ + Values S3 Vectors rejects are dropped with one warning, instead of failing the whole write. + + The `_split_overlap` case is the one that bites in practice: `DocumentSplitter` with + `split_overlap` adds it to every chunk, so a plain splitter-to-store pipeline would otherwise + fail with a `ParamValidationError`. + """ + doc = Document( + id="test-1", + content="Hello world", + meta={ + "keep_str": "x", + "keep_list": ["a", "b"], + "_split_overlap": [{"doc_id": "a", "range": (0, 24)}], + "nothing": None, + "nested": {"a": 1}, + "empty": [], + }, + ) + result = S3VectorsDocumentStore._document_to_s3_vector(doc) + + assert result["metadata"] == {"_content": "Hello world", "keep_str": "x", "keep_list": ["a", "b"]} + assert "['_split_overlap', 'empty', 'nested', 'nothing']" in caplog.text + + +def test_document_to_s3_vector_converts_tuples_to_lists(): + """A tuple is a valid array for S3 Vectors, but botocore's validator only accepts `list`.""" + doc = Document(id="test-1", meta={"tags": ("a", "b")}, embedding=[0.1]) + result = S3VectorsDocumentStore._document_to_s3_vector(doc) + assert result["metadata"]["tags"] == ["a", "b"] + + +def test_blob_roundtrip(): + """A blob and its (dict) metadata survive the round-trip; the dict is stored JSON-encoded.""" + doc = Document( + id="test-1", + content="Hello world", + blob=ByteStream(data=b"\x00\x01binary", meta={"name": "f.pdf", "pages": 3}, mime_type="application/pdf"), + embedding=[0.1, 0.2, 0.3], + ) + vector = S3VectorsDocumentStore._document_to_s3_vector(doc) + # S3 Vectors rejects dicts, so the blob meta must go over the wire as a string. + assert vector["metadata"]["_blob_meta"] == '{"name": "f.pdf", "pages": 3}' + + restored = S3VectorsDocumentStore._s3_vector_to_document(vector) + assert restored.blob is not None + assert restored.blob.data == b"\x00\x01binary" + assert restored.blob.meta == {"name": "f.pdf", "pages": 3} + assert restored.blob.mime_type == "application/pdf" + assert restored.meta == {} + + +def test_blob_roundtrip_with_undecodable_meta(caplog): + """Corrupt stored blob metadata degrades to an empty dict rather than raising.""" + vector = { + "key": "test-1", + "data": {"float32": [0.1]}, + "metadata": {"_blob_data": base64.b64encode(b"x").decode("ascii"), "_blob_meta": "not json"}, + } + doc = S3VectorsDocumentStore._s3_vector_to_document(vector) + assert doc.blob is not None + assert doc.blob.meta == {} + assert "Could not decode the stored blob metadata" in caplog.text + + +@pytest.mark.parametrize( + ("documents", "expected", "warns"), + [ + ([], {}, False), + ([Document(content="a", meta={"n": 1})], {"content": {"type": "text"}, "n": {"type": "long"}}, False), + ([Document(content="a", meta={"f": 1.5})], {"content": {"type": "text"}, "f": {"type": "long"}}, False), + ([Document(content="a", meta={"s": "x"})], {"content": {"type": "text"}, "s": {"type": "keyword"}}, False), + # bool must win over int: bool is a subclass of int in Python. + ([Document(content="a", meta={"b": True})], {"content": {"type": "text"}, "b": {"type": "boolean"}}, False), + # A list is typed by its first element; an empty list carries no type info. + ( + [Document(content="a", meta={"tags": ["x"]})], + {"content": {"type": "text"}, "tags": {"type": "keyword"}}, + False, + ), + ([Document(content="a", meta={"tags": []})], {"content": {"type": "text"}, "tags": {"type": "keyword"}}, False), + # Mixed types across documents fall back to keyword, with a warning. + ( + [Document(content="a", meta={"m": 1}), Document(content="b", meta={"m": "x"})], + {"content": {"type": "text"}, "m": {"type": "keyword"}}, + True, + ), + # No content anywhere means no synthetic "content" field. + ([Document(meta={"n": 1}, embedding=[0.1])], {"n": {"type": "long"}}, False), + ], +) +def test_get_metadata_fields_info_impl(documents, expected, warns, caplog): + """Type inference is a pure function over already-fetched documents.""" + assert S3VectorsDocumentStore._get_metadata_fields_info_impl(documents) == expected + assert ("mixed types" in caplog.text) is warns + + +def test_get_metadata_field_min_max_impl(): + """min/max strips the meta. prefix, compares numerically, and reports None when absent.""" + docs = [ + Document(content="a", meta={"priority": 5}), + Document(content="b", meta={"priority": 1}), + # 10 vs 5 catches an implementation comparing numbers as strings ("10" < "5"). + Document(content="c", meta={"priority": 10}), + Document(content="d", meta={}), + ] + assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(docs, "priority") == {"min": 1, "max": 10} + assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(docs, "meta.priority") == {"min": 1, "max": 10} + assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(docs, "absent") == {"min": None, "max": None} + assert S3VectorsDocumentStore._get_metadata_field_min_max_impl([], "priority") == {"min": None, "max": None} + + strings = [Document(content="a", meta={"k": "beta"}), Document(content="b", meta={"k": "alpha"})] + assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(strings, "k") == {"min": "alpha", "max": "beta"} + + # Mixed str/numeric would raise on a naive min(); numerics win instead of blowing up. + mixed = [Document(content="a", meta={"k": 3}), Document(content="b", meta={"k": "zzz"})] + assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(mixed, "k") == {"min": 3, "max": 3} + + +def test_get_metadata_field_unique_values_impl_search_and_pagination(): + """Distinct values sort stably, then search_term and from_/size are applied.""" + docs = [Document(content=f"d{i}", meta={"category": f"category_{i}"}) for i in range(5)] + docs.append(Document(content="dup", meta={"category": "category_0"})) + + values, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "category", None, 0, 10) + assert values == [f"category_{i}" for i in range(5)] + assert total == 5 + + first, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "category", None, 0, 2) + second, _ = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "category", None, 2, 2) + assert first == ["category_0", "category_1"] + assert second == ["category_2", "category_3"] + assert total == 5 + + # search_term matches the value, case-insensitively — not the document content. + hits, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "meta.category", "GORY_3", 0, 10) + assert hits == ["category_3"] + assert total == 1 + + +def test_get_metadata_field_unique_values_impl_keeps_types_distinct(): + """`1`, `1.0`, `"1"` and `True` all compare equal in Python but must stay distinct.""" + docs = [ + Document(content="a", meta={"p": 1}), + Document(content="b", meta={"p": "1"}), + Document(content="c", meta={"p": 1.0}), + Document(content="d", meta={"p": True}), + Document(content="e", meta={"p": 1}), + ] + values, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "p", None, 0, 10) + assert total == 4 + assert {type(v) for v in values} == {int, str, float, bool} + + # List-valued metadata contributes each item, not the list. + listed = [Document(content="a", meta={"tags": ["x", "y"]}), Document(content="b", meta={"tags": ["y"]})] + values, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(listed, "tags", None, 0, 10) + assert values == ["x", "y"] + assert total == 2 + + +def test_count_unique_metadata_impl(): + """Distinct-value counts are per field and keep differently-typed values apart.""" + docs = [ + Document(content="a", meta={"category": "A", "status": "active"}), + Document(content="b", meta={"category": "B", "status": "active"}), + Document(content="c", meta={"category": "A", "status": "inactive"}), + ] + assert S3VectorsDocumentStore._count_unique_metadata_impl(docs, ["category", "status"]) == { + "category": 2, + "status": 2, + } + assert S3VectorsDocumentStore._count_unique_metadata_impl(docs, ["absent"]) == {"absent": 0} + assert S3VectorsDocumentStore._count_unique_metadata_impl([], ["category"]) == {"category": 0} + + listed = [Document(content="a", meta={"tags": ["x", "y"]}), Document(content="b", meta={"tags": ["y", "z"]})] + assert S3VectorsDocumentStore._count_unique_metadata_impl(listed, ["tags"]) == {"tags": 3} + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_close_releases_the_client_and_allows_reopen(mock_boto3): + """close() closes the boto3 client; the store stays usable and lazily builds a new one.""" + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store._get_client() + assert store._client is not None + + store.close() + client.close.assert_called_once() + assert store._client is None + + store._get_client() + assert store._client is not None + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_close_is_a_noop_when_never_connected(mock_boto3): + """close() before any call must not create a client just to close it.""" + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store.close() + mock_boto3.client.assert_not_called() + + +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_close_is_exception_safe(mock_boto3): + """A failure while closing must not propagate to the caller.""" + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.close.side_effect = RuntimeError("boom") + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store._get_client() + store.close() # must not raise + assert store._client is None + + @pytest.mark.integration class TestDocumentStore( CountDocumentsTest, @@ -278,6 +539,11 @@ class TestDocumentStore( DeleteAllTest, DeleteByFilterTest, UpdateByFilterTest, + CountDocumentsByFilterTest, + CountUniqueMetadataByFilterTest, + GetMetadataFieldsInfoTest, + GetMetadataFieldMinMaxTest, + GetMetadataFieldUniqueValuesTest, FilterableDocsFixtureMixin, ): def assert_documents_are_equal(self, received: list[Document], expected: list[Document]) -> None: @@ -307,6 +573,90 @@ def test_write_documents(self, document_store: S3VectorsDocumentStore) -> None: assert document_store.write_documents(docs) == 1 assert document_store.write_documents(docs) == 1 + def test_close_and_reopen(self, document_store: S3VectorsDocumentStore) -> None: + """close() drops the boto3 client; the store keeps working against the same index.""" + document_store._get_client() + assert document_store._client is not None + + document_store.close() + assert document_store._client is None + + assert document_store.count_documents() == 0 + assert document_store._client is not None + + def test_sentence_window_retriever(self, document_store: S3VectorsDocumentStore) -> None: + """The store composes with SentenceWindowRetriever, which pulls neighbours via filter_documents.""" + splitter = DocumentSplitter(split_length=10, split_overlap=5, split_by="word") + splitter.warm_up() + text = ( + "Whose woods these are I think I know. His house is in the village though; He will not see me stopping " + "here To watch his woods fill up with snow." + ) + docs = splitter.run(documents=[Document(content=text)])["documents"] + + rng = np.random.default_rng(seed=42) + # One chunk gets the query vector itself, so retrieval deterministically lands on it. + docs = [ + replace(doc, embedding=[0.1] * 768 if idx == 2 else rng.random(768).tolist()) + for idx, doc in enumerate(docs) + ] + document_store.write_documents(docs) + + retrieved = S3VectorsEmbeddingRetriever(document_store=document_store).run(query_embedding=[0.1] * 768, top_k=1) + result = SentenceWindowRetriever(document_store=document_store, window_size=2).run( + retrieved_documents=[retrieved["documents"][0]] + ) + assert len(result["context_windows"]) == 1 + + def test_get_metadata_fields_info_consistent_types(self, document_store: S3VectorsDocumentStore) -> None: + """A field with the same type in every document keeps that type, with no mixed-type fallback.""" + docs = [ + Document(content="Doc 1", meta={"score": 85}), + Document(content="Doc 2", meta={"score": 90}), + Document(content="Doc 3", meta={"score": 78}), + ] + document_store.write_documents(docs) + + fields_info = document_store.get_metadata_fields_info() + assert fields_info["score"] == {"type": "long"} + + def test_get_metadata_field_min_max_boolean_and_string(self, document_store: S3VectorsDocumentStore) -> None: + """min/max works on booleans (False < True) and on strings (alphabetical), not just numbers.""" + docs = [ + Document(content="Doc 1", meta={"active": True, "category": "Zebra"}), + Document(content="Doc 2", meta={"active": False, "category": "Alpha"}), + Document(content="Doc 3", meta={"active": True, "category": "Beta"}), + Document(content="Doc 4", meta={"active": False, "category": "Gamma"}), + ] + document_store.write_documents(docs) + + assert document_store.get_metadata_field_min_max("active") == {"min": False, "max": True} + assert document_store.get_metadata_field_min_max("category") == {"min": "Alpha", "max": "Zebra"} + + def test_get_metadata_field_min_max_no_values(self, document_store: S3VectorsDocumentStore) -> None: + """A list-valued field has no comparable scalar, so min/max is None — as for a missing field.""" + docs = [ + Document(content="Doc 1", meta={"tags": ["tag1", "tag2"]}), + Document(content="Doc 2", meta={"tags": ["tag3", "tag4"]}), + ] + document_store.write_documents(docs) + + assert document_store.get_metadata_field_min_max("tags") == {"min": None, "max": None} + assert document_store.get_metadata_field_min_max("nonexistent") == {"min": None, "max": None} + + def test_get_metadata_field_unique_values_with_lists(self, document_store: S3VectorsDocumentStore) -> None: + """A list-valued field contributes each item as a distinct value, not the list itself.""" + docs = [ + Document(content="Doc 1", meta={"tags": ["python", "java"]}), + Document(content="Doc 2", meta={"tags": ["python", "rust"]}), + Document(content="Doc 3", meta={"tags": ["java", "go"]}), + ] + document_store.write_documents(docs) + + values, total = document_store.get_metadata_field_unique_values("tags", size=10) + assert total == 4 + assert set(values) == {"go", "java", "python", "rust"} + def test_write_documents_without_embedding_reads_back_as_none(self, document_store: S3VectorsDocumentStore) -> None: """ S3 Vectors stores a placeholder vector for embedding-less documents; `filter_documents` From 51d92c3ab46c287a801cbf952d7ca0de81b07070 Mon Sep 17 00:00:00 2001 From: Oleksii Udovychenko Date: Thu, 15 Oct 2026 15:38:01 +0000 Subject: [PATCH 15/17] feat(amazon-s3-vectors): align with current S3 Vectors and repo style - Test on Python 3.11 and 3.15, require Python >= 3.11 and haystack-ai >= 3.0.0 - Raise the top_k cap to 10,000 and follow QueryVectors pagination, which needs boto3 >= 1.43.31 - Stop storing Document.blob - Validate filters before listing in filter_documents and note that local filtering can differ from S3 Vectors filtering on arrays - Remove the optional methods built on filter_documents - Add S3VectorsEmbeddingRetriever.close - Remove the example file and the eventual-consistency sleeps in the test fixtures - Fix the retriever usage example for Haystack 3, where the Sentence Transformers embedders live in sentence-transformers-haystack - Note in the class docstring that search is approximate, so a top_k close to the number of stored or matching documents returns fewer --- .github/workflows/amazon_s3_vectors.yml | 29 +- .../amazon_s3_vectors/examples/example.py | 108 ------ integrations/amazon_s3_vectors/pyproject.toml | 9 +- .../amazon_s3_vectors/embedding_retriever.py | 12 +- .../amazon_s3_vectors/document_store.py | 347 ++---------------- .../amazon_s3_vectors/tests/conftest.py | 35 +- .../tests/test_document_store.py | 230 ++---------- 7 files changed, 92 insertions(+), 678 deletions(-) delete mode 100644 integrations/amazon_s3_vectors/examples/example.py diff --git a/.github/workflows/amazon_s3_vectors.yml b/.github/workflows/amazon_s3_vectors.yml index 36787d359d..041ba4a153 100644 --- a/.github/workflows/amazon_s3_vectors.yml +++ b/.github/workflows/amazon_s3_vectors.yml @@ -34,7 +34,7 @@ env: PYTHONUNBUFFERED: "1" FORCE_COLOR: "1" TEST_MATRIX_OS: '["ubuntu-latest", "windows-latest", "macos-latest"]' - TEST_MATRIX_PYTHON: '["3.10", "3.14"]' + TEST_MATRIX_PYTHON: '["3.11", "3.15"]' AWS_REGION: "us-east-1" @@ -51,7 +51,7 @@ jobs: - id: set run: | echo 'os=${{ github.event_name == 'push' && '["ubuntu-latest"]' || env.TEST_MATRIX_OS }}' >> "$GITHUB_OUTPUT" - echo 'python-version=${{ github.event_name == 'push' && '["3.10"]' || env.TEST_MATRIX_PYTHON }}' >> "$GITHUB_OUTPUT" + echo 'python-version=${{ github.event_name == 'push' && '["3.11"]' || env.TEST_MATRIX_PYTHON }}' >> "$GITHUB_OUTPUT" run: name: Python ${{ matrix.python-version }} on ${{ startsWith(matrix.os, 'macos-') && 'macOS' || startsWith(matrix.os, 'windows-') && 'Windows' || 'Linux' }} @@ -73,17 +73,19 @@ jobs: working-directory: . run: git config --system core.longpaths true - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: python-version: ${{ matrix.python-version }} - name: Install Hatch - run: pip install --upgrade hatch + run: | + python -m pip install --upgrade pip + pip install hatch --uploaded-prior-to=P1D - name: Lint - if: matrix.python-version == '3.10' && runner.os == 'Linux' + if: matrix.python-version == '3.11' && runner.os == 'Linux' run: hatch run fmt-check && hatch run test:types - name: Run unit tests @@ -92,8 +94,8 @@ jobs: # On PR: posts coverage comment (directly on same-repo PRs; via artifact for fork PRs). On push to main: stores coverage baseline on data branch. - name: Store unit tests coverage id: coverage_comment - if: matrix.python-version == '3.10' && runner.os == 'Linux' && github.event_name != 'schedule' - uses: py-cov-action/python-coverage-comment-action@7188638f871f721a365d644f505d1ff3df20d683 # v3.40 + if: matrix.python-version == '3.11' && runner.os == 'Linux' && github.event_name != 'schedule' + uses: py-cov-action/python-coverage-comment-action@e5fb2218463c3377f3504b0c5b5646ed53da7a2f # v4.5 with: GITHUB_TOKEN: ${{ github.token }} COVERAGE_PATH: integrations/amazon_s3_vectors @@ -102,8 +104,8 @@ jobs: MINIMUM_ORANGE: 60 - name: Upload coverage comment to be posted - if: matrix.python-version == '3.10' && runner.os == 'Linux' && github.event_name == 'pull_request' && steps.coverage_comment.outputs.COMMENT_FILE_WRITTEN == 'true' - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + if: matrix.python-version == '3.11' && runner.os == 'Linux' && github.event_name == 'pull_request' && steps.coverage_comment.outputs.COMMENT_FILE_WRITTEN == 'true' + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 with: name: coverage-comment-amazon_s3_vectors path: python-coverage-comment-action-amazon_s3_vectors.txt @@ -112,7 +114,7 @@ jobs: - name: AWS authentication id: aws-auth if: github.event_name == 'schedule' || (github.event.pull_request.head.repo.full_name == github.repository && !startsWith(github.event.pull_request.head.ref, 'dependabot/')) - uses: aws-actions/configure-aws-credentials@ec61189d14ec14c8efccab744f656cffd0e33f37 + uses: aws-actions/configure-aws-credentials@e1253824e5c10ff9df46874f81ed3ec929e19cfd with: aws-region: ${{ env.AWS_REGION }} role-to-assume: ${{ secrets.AWS_CI_ROLE_ARN }} @@ -123,7 +125,7 @@ jobs: - name: Store combined coverage if: github.event_name == 'push' - uses: py-cov-action/python-coverage-comment-action@7188638f871f721a365d644f505d1ff3df20d683 # v3.40 + uses: py-cov-action/python-coverage-comment-action@e5fb2218463c3377f3504b0c5b5646ed53da7a2f # v4.5 with: GITHUB_TOKEN: ${{ github.token }} COVERAGE_PATH: integrations/amazon_s3_vectors @@ -145,12 +147,11 @@ jobs: hatch -e test env run -- uv pip install git+https://github.com/deepset-ai/haystack.git@main hatch run test:unit - notify-slack-on-failure: needs: run if: failure() && github.event_name == 'schedule' runs-on: ubuntu-slim steps: - - uses: deepset-ai/notify-slack-action@3cda73b77a148f16f703274198e7771340cf862b # v1 + - uses: deepset-ai/notify-slack-action@a65def0c8bf91d6520286ab34280151c76a5a008 # v1.1.0 with: slack-webhook-url: ${{ secrets.SLACK_WEBHOOK_URL_NOTIFICATIONS }} diff --git a/integrations/amazon_s3_vectors/examples/example.py b/integrations/amazon_s3_vectors/examples/example.py deleted file mode 100644 index 67d3a67c20..0000000000 --- a/integrations/amazon_s3_vectors/examples/example.py +++ /dev/null @@ -1,108 +0,0 @@ -# Amazon S3 Vectors — Haystack Integration Example -# -# A serverless vector store on AWS: no database to provision, no cluster to manage. -# Just a bucket name and AWS credentials. -# -# This example indexes documents, runs unfiltered and filtered queries, and cleans up. -# -# Prerequisites: -# pip install amazon-s3-vectors-haystack "sentence-transformers>=2.2.0" -# AWS credentials configured (env vars, ~/.aws/credentials, or IAM role) - -import time - -from haystack import Document, Pipeline -from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder -from haystack.components.writers import DocumentWriter -from haystack.document_stores.types import DuplicatePolicy - -from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore - -MODEL = "sentence-transformers/all-MiniLM-L6-v2" - -# --- Create the document store (bucket + index are created automatically) --- -document_store = S3VectorsDocumentStore( - vector_bucket_name="haystack-example", - index_name="products", - dimension=384, # all-MiniLM-L6-v2 produces 384-dim embeddings - distance_metric="cosine", - region_name="us-east-1", -) - -# --- Index documents with metadata --- -products = [ - Document( - content="Lightweight running shoes with breathable mesh", - meta={"category": "shoes", "price": 89.99, "in_stock": True}, - ), - Document( - content="Waterproof hiking boots with ankle support", - meta={"category": "shoes", "price": 149.99, "in_stock": True}, - ), - Document( - content="Classic leather dress shoes", - meta={"category": "shoes", "price": 199.99, "in_stock": False}, - ), - Document( - content="Insulated winter jacket with down filling", - meta={"category": "jackets", "price": 249.99, "in_stock": True}, - ), - Document( - content="Lightweight windbreaker for running", - meta={"category": "jackets", "price": 79.99, "in_stock": True}, - ), -] - -indexing = Pipeline() -indexing.add_component("embedder", SentenceTransformersDocumentEmbedder(model=MODEL)) -indexing.add_component("writer", DocumentWriter(document_store=document_store, policy=DuplicatePolicy.OVERWRITE)) -indexing.connect("embedder", "writer") - -print("Indexing products...") -indexing.run({"embedder": {"documents": products}}) -print(f"Indexed {document_store.count_documents()} products.\n") - -# --- Build the query pipeline --- -querying = Pipeline() -querying.add_component("embedder", SentenceTransformersTextEmbedder(model=MODEL)) -querying.add_component("retriever", S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5)) -querying.connect("embedder", "retriever") - - -def search(query: str, filters: dict | None = None) -> None: - print(f"Query: '{query}'" + (f" Filter: {filters}" if filters else "")) - start = time.time() - result = querying.run({"embedder": {"text": query}, "retriever": {"filters": filters}}) - elapsed = (time.time() - start) * 1000 - for i, doc in enumerate(result["retriever"]["documents"], 1): - meta = doc.meta - print(f" {i}. [{doc.score:.3f}] {doc.content}") - print(f" category={meta['category']} price=${meta['price']} in_stock={meta['in_stock']}") - print(f" ({elapsed:.0f}ms)\n") - - -# Unfiltered: returns shoes AND jackets -search("something for running") - -# Filtered to shoes only — server-side, inside the vector search -search("something for running", filters={"field": "meta.category", "operator": "==", "value": "shoes"}) - -# Combined filter: in-stock AND under $100 -search( - "something for running", - filters={ - "operator": "AND", - "conditions": [ - {"field": "meta.in_stock", "operator": "==", "value": True}, - {"field": "meta.price", "operator": "<", "value": 100.0}, - ], - }, -) - -# --- Cleanup --- -document_store.delete_documents([doc.id for doc in products]) -client = document_store._get_client() -client.delete_index(vectorBucketName="haystack-example", indexName="products") -client.delete_vector_bucket(vectorBucketName="haystack-example") -print("Cleaned up all AWS resources.") diff --git a/integrations/amazon_s3_vectors/pyproject.toml b/integrations/amazon_s3_vectors/pyproject.toml index c44b588f4b..940ea46036 100644 --- a/integrations/amazon_s3_vectors/pyproject.toml +++ b/integrations/amazon_s3_vectors/pyproject.toml @@ -7,7 +7,7 @@ name = "amazon-s3-vectors-haystack" dynamic = ["version"] description = "Haystack integration for amazon_s3_vectors" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" license = "Apache-2.0" keywords = ["aws", "s3", "vectors", "document-store", "haystack", "vector-search"] authors = [{ name = "deepset GmbH", email = "info@deepset.ai" }] @@ -15,15 +15,15 @@ classifiers = [ "License :: OSI Approved :: Apache Software License", "Development Status :: 4 - Beta", "Programming Language :: Python", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: 3.15", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["haystack-ai>=2.26.1", "boto3>=1.42.0"] +dependencies = ["haystack-ai>=3.0.0", "boto3>=1.43.31"] [project.urls] Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_s3_vectors#readme" @@ -135,6 +135,7 @@ ignore = [ "PLR0912", "PLR0913", "PLR0915", + "PLR0917", # Allow `Any` type - used legitimately for dynamic types and SDK boundaries "ANN401", ] @@ -148,8 +149,6 @@ ban-relative-imports = "parents" [tool.ruff.lint.per-file-ignores] # Tests can use magic values, assertions, relative imports, and don't need type annotations "tests/**/*" = ["E501", "PLC0415", "PLR2004", "S101", "S110", "TID252", "D", "ANN"] -# Examples can contain print statements and long lines -"examples/**/*" = ["D", "T201", "E501"] [tool.coverage.run] source = ["haystack_integrations"] diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py index f7f5bd68b2..a9fa55566b 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py @@ -20,8 +20,10 @@ class S3VectorsEmbeddingRetriever: Usage example: ```python from haystack import Document, Pipeline - from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder from haystack.document_stores.types import DuplicatePolicy + # Requires: pip install sentence-transformers-haystack + from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersTextEmbedder + from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersDocumentEmbedder from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore @@ -115,7 +117,7 @@ def run( :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on the `filter_policy` chosen at retriever initialization. Filters are applied server-side during the vector search. - :param top_k: Maximum number of Documents to return. S3 Vectors caps this at 100. + :param top_k: Maximum number of Documents to return. S3 Vectors caps this at 10,000. :returns: A dictionary with key `"documents"` containing the retrieved Documents. Returned documents will not contain embeddings. """ @@ -128,3 +130,9 @@ def run( top_k=top_k, ) return {"documents": docs} + + def close(self) -> None: + """ + Release the resources of the underlying Document Store. + """ + self.document_store.close() diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index 8ba6592e9b..e0fad768d9 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -2,7 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -import base64 import json import math from collections.abc import Iterator @@ -12,7 +11,7 @@ import boto3 from botocore.exceptions import ClientError from haystack import default_from_dict, default_to_dict, logging -from haystack.dataclasses import ByteStream, Document +from haystack.dataclasses import Document from haystack.document_stores.types import DuplicatePolicy from haystack.utils.auth import Secret, deserialize_secrets_inplace from haystack.utils.filters import document_matches_filter @@ -27,23 +26,20 @@ # S3 Vectors allows up to 500 keys per delete_vectors call _DELETE_BATCH_SIZE = 500 -# Maximum number of results from query_vectors -_MAX_TOP_K = 100 +# Maximum top_k of a query_vectors request. Results come back in pages of up to 100. +_MAX_TOP_K = 10_000 # S3 Vectors metadata limits _MAX_TOTAL_METADATA_BYTES = 40 * 1024 # 40 KB total -# Reserved metadata keys used to store Haystack Document fields +# Reserved metadata key used to store the Document content _CONTENT_KEY = "_content" -_BLOB_DATA_KEY = "_blob_data" -_BLOB_META_KEY = "_blob_meta" -_BLOB_MIME_TYPE_KEY = "_blob_mime_type" -_RESERVED_META_KEYS = {_CONTENT_KEY, _BLOB_DATA_KEY, _BLOB_META_KEY, _BLOB_MIME_TYPE_KEY} +_RESERVED_META_KEYS = {_CONTENT_KEY} # These keys are stored but should not be used in query filters. # They are configured as nonFilterableMetadataKeys on the index. -_NON_FILTERABLE_KEYS = [_CONTENT_KEY, _BLOB_DATA_KEY, _BLOB_META_KEY, _BLOB_MIME_TYPE_KEY] +_NON_FILTERABLE_KEYS = [_CONTENT_KEY] class S3VectorsDocumentStore: @@ -54,9 +50,10 @@ class S3VectorsDocumentStore: This document store stores Haystack `Document` objects as vectors with associated metadata in an S3 vector bucket and index. - **Service limits:** + Service limits: - - Maximum `top_k`: 100 results per query + - Maximum `top_k`: 10,000 results per query. Search is approximate, so a `top_k` close to the number of + stored or matching documents can return fewer results - Maximum vector dimension: 4,096 - Metadata per vector: 40 KB total, 2 KB filterable - A metadata value must be a string, a boolean, a finite number, or a non-empty array of either @@ -67,7 +64,8 @@ class S3VectorsDocumentStore: retrievable via `filter_documents()` - Only `DuplicatePolicy.OVERWRITE` is supported, since `put_vectors` is an upsert - Distance metrics: `cosine` or `euclidean` (set at index creation, immutable) - - `filter_documents()` is client-side — prefer `S3VectorsEmbeddingRetriever` with filters + - `filter_documents()` filters locally, so prefer `S3VectorsEmbeddingRetriever` with filters + - `Document.blob` is not stored Usage example: ```python @@ -109,7 +107,7 @@ def __init__( :param create_bucket_and_index: Whether to automatically create the vector bucket and index if they do not exist. Defaults to `True`. :param non_filterable_metadata_keys: Additional metadata keys to mark as non-filterable - on the index (beyond the internal keys used for Document content/blob storage). + on the index (beyond the internal key used for Document content). """ self.vector_bucket_name = vector_bucket_name self.index_name = index_name @@ -250,10 +248,8 @@ def count_documents(self) -> int: """ Return the number of documents in the document store. - .. note:: - - S3 Vectors does not provide a dedicated count API. This method lists all vector keys - via pagination, which can be slow for large indexes. + Note: S3 Vectors does not provide a dedicated count API. This method lists all vector keys + via pagination, which can be slow for large indexes. """ return sum(1 for _ in self._iter_vectors()) @@ -324,16 +320,18 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume """ Return documents matching the provided filters. - .. warning:: - - S3 Vectors only supports metadata filtering during vector similarity queries, not as a - standalone operation. This method lists all vectors and applies filters client-side, - which can be very slow for large indexes. For filtered retrieval, prefer using - `S3VectorsEmbeddingRetriever` with filters instead. + Note: S3 Vectors only supports metadata filtering during vector similarity queries, not as a + standalone operation. This method lists all vectors and filters them locally with Haystack, + which can be slow for large indexes. `S3VectorsEmbeddingRetriever` filters through S3 + Vectors instead, and the two can match different documents for the same filter, especially + on array metadata: S3 Vectors matches an array field when any element matches, while + Haystack compares the whole array. :param filters: Haystack-format filters to apply. :returns: A list of matching Documents. """ + _validate_filters(filters) + if filters: logger.warning( "S3 Vectors does not support standalone filtered listing. " @@ -351,7 +349,6 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume ] if filters: - _validate_filters(filters) documents = [doc for doc in documents if document_matches_filter(filters=filters, document=doc)] return documents @@ -378,258 +375,11 @@ def delete_all_documents(self) -> None: """ Delete all documents from the index. - .. note:: - - S3 Vectors has no truncate operation, so every key is listed and then deleted in - batches, which can be slow for large indexes. + Note: S3 Vectors has no truncate operation, so every key is listed and then deleted in + batches, which can be slow for large indexes. """ self.delete_documents([v["key"] for v in self._iter_vectors()]) - def delete_by_filter(self, filters: dict[str, Any]) -> int: - """ - Delete all documents matching the given filters. - - .. warning:: - - Filtering is client-side (see `filter_documents`), so this reads the whole index. - - :param filters: Haystack-format filters selecting the documents to delete. - :returns: The number of documents deleted. - """ - document_ids = [doc.id for doc in self.filter_documents(filters=filters)] - self.delete_documents(document_ids) - return len(document_ids) - - def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int: - """ - Merge `meta` into the metadata of every document matching the given filters. - - Keys present in `meta` are overwritten, the rest of each document's metadata is kept. - - .. warning:: - - S3 Vectors cannot update metadata in place, so matching documents are read and written - back in full. Filtering is client-side (see `filter_documents`), so this reads the - whole index. - - :param filters: Haystack-format filters selecting the documents to update. - :param meta: Metadata fields to set on the matching documents. - :returns: The number of documents updated. - """ - matching = self.filter_documents(filters=filters) - if not matching: - return 0 - - self.write_documents([replace(doc, meta={**doc.meta, **meta}) for doc in matching]) - return len(matching) - - def count_documents_by_filter(self, filters: dict[str, Any]) -> int: - """ - Return the number of documents matching the given filters. - - .. warning:: - - Filtering is client-side (see `filter_documents`), so this reads the whole index. - Unlike stores with a capped query API, the count is exact: `list_vectors` paginates - over every vector. - - :param filters: Haystack-format filters to apply. - :returns: The number of matching documents. - """ - return len(self.filter_documents(filters=filters)) - - def count_unique_metadata_by_filter(self, filters: dict[str, Any], metadata_fields: list[str]) -> dict[str, int]: - """ - Count the distinct values of each given metadata field, across documents matching `filters`. - - .. warning:: - - Filtering and aggregation are client-side (see `filter_documents`), so this reads the - whole index. - - :param filters: Haystack-format filters selecting the documents to consider. - :param metadata_fields: The metadata field names to count distinct values for. - :returns: A mapping of field name to its number of distinct values. - """ - documents = self.filter_documents(filters=filters) - return self._count_unique_metadata_impl(documents, metadata_fields) - - def get_metadata_fields_info(self) -> dict[str, dict[str, str]]: - """ - Infer the metadata fields present in the index and their types. - - S3 Vectors has no schema-introspection API, so types are inferred from the metadata of the - documents actually stored. Python types map to `boolean`, `long` (int/float), and `keyword` - (str); a field with mixed types across documents falls back to `keyword` with a warning. - - .. warning:: - - This reads the whole index (see `filter_documents`). - - :returns: A mapping of field name to `{"type": }`. - """ - return self._get_metadata_fields_info_impl(self.filter_documents()) - - def get_metadata_field_min_max(self, metadata_field: str) -> dict[str, Any]: - """ - Return the minimum and maximum value of a metadata field. - - Works for numeric, boolean and string values. A leading `meta.` on the field name is - stripped, since S3 Vectors stores metadata flat. - - .. warning:: - - This reads the whole index (see `filter_documents`). - - :param metadata_field: The metadata field to analyze. - :returns: `{"min": ..., "max": ...}`, both `None` when no document has a comparable value - for the field. - """ - return self._get_metadata_field_min_max_impl(self.filter_documents(), metadata_field) - - def get_metadata_field_unique_values( - self, - metadata_field: str, - search_term: str | None = None, - from_: int = 0, - size: int = 10, - filters: dict[str, Any] | None = None, - ) -> tuple[list[Any], int]: - """ - Return the distinct values of a metadata field, with optional search and pagination. - - Values keep their original Python type, and values that merely share a string form (the int - `1`, the float `1.0`, the str `"1"`, `True`) are kept distinct. A leading `meta.` on the - field name is stripped. - - .. warning:: - - This reads the whole index (see `filter_documents`). - - :param metadata_field: The metadata field to get distinct values for. - :param search_term: Optional case-insensitive substring the value must contain. - :param from_: Pagination offset into the sorted distinct values. - :param size: Maximum number of values to return. - :param filters: Optional Haystack-format filters restricting the documents considered. - :returns: A tuple of (page of distinct values, total number of matching distinct values). - """ - documents = self.filter_documents(filters=filters) - return self._get_metadata_field_unique_values_impl(documents, metadata_field, search_term, from_, size) - - @staticmethod - def _count_unique_metadata_impl(documents: list[Document], metadata_fields: list[str]) -> dict[str, int]: - """Count distinct metadata values per field over an already-filtered document list.""" - result = {} - for field in metadata_fields: - # Values of different types can compare equal in Python (`1 == True == 1.0`), so key on - # (type, value) to avoid silently merging them. - unique_values: set[tuple[type, Any]] = set() - for doc in documents: - if doc.meta and field in doc.meta: - value = doc.meta[field] - if isinstance(value, list): - unique_values.update((type(item), item) for item in value) - else: - unique_values.add((type(value), value)) - result[field] = len(unique_values) - return result - - @staticmethod - def _get_metadata_fields_info_impl(documents: list[Document]) -> dict[str, dict[str, str]]: - """Infer metadata field types from an already-filtered document list.""" - if not documents: - return {} - - field_types: dict[str, dict[str, str]] = {} - if any(doc.content is not None for doc in documents): - field_types["content"] = {"type": "text"} - - def _type_of(value: Any) -> str | None: - # bool must be checked before int/float: bool is a subclass of int in Python. - if isinstance(value, bool): - return "boolean" - if isinstance(value, (int, float)): - return "long" - if isinstance(value, str): - return "keyword" - return None - - field_samples: dict[str, set[str]] = {} - for doc in documents: - for field, value in (doc.meta or {}).items(): - # An empty list carries no type information, so it defaults to keyword. - inferred = (_type_of(value[0]) if value else "keyword") if isinstance(value, list) else _type_of(value) - if inferred is not None: - field_samples.setdefault(field, set()).add(inferred) - - for field, types_seen in field_samples.items(): - if len(types_seen) == 1: - field_types[field] = {"type": types_seen.pop()} - else: - logger.warning( - "Metadata field '{field}' has mixed types {types} across documents. Defaulting to 'keyword'.", - field=field, - types=sorted(types_seen), - ) - field_types[field] = {"type": "keyword"} - - return field_types - - @staticmethod - def _get_metadata_field_min_max_impl(documents: list[Document], metadata_field: str) -> dict[str, Any]: - """Compute min/max of a metadata field over an already-filtered document list.""" - field_name = metadata_field.removeprefix("meta.") - values: list[bool | int | float | str] = [] - for doc in documents: - if doc.meta and field_name in doc.meta: - value = doc.meta[field_name] - # bool is a subclass of int, so it is already covered by the numeric check; listing - # it explicitly documents that booleans are comparable and intentionally included. - if isinstance(value, (bool, int, float, str)): - values.append(value) - - if not values: - return {"min": None, "max": None} - - # min/max over mixed str and numeric values would raise, so only compare within one kind. - numeric = [v for v in values if not isinstance(v, str)] - comparable = numeric if numeric else values - return {"min": min(comparable), "max": max(comparable)} - - @staticmethod - def _get_metadata_field_unique_values_impl( - documents: list[Document], metadata_field: str, search_term: str | None, from_: int, size: int - ) -> tuple[list[Any], int]: - """Collect distinct values of a metadata field over an already-filtered document list.""" - field_name = metadata_field.removeprefix("meta.") - # Dedupe on (type, value): a plain set would merge `1`, `1.0` and `True`. - seen: set[tuple[type, Any]] = set() - unique_values: list[Any] = [] - - def _add(value: Any) -> None: - key = (type(value), value) - if key not in seen: - seen.add(key) - unique_values.append(value) - - for doc in documents: - if doc.meta and field_name in doc.meta: - value = doc.meta[field_name] - if isinstance(value, list): - for item in value: - _add(item) - else: - _add(value) - - # Sort by string form so pagination is stable across mixed types. - values_list = sorted(unique_values, key=str) - - if search_term: - needle = search_term.lower() - values_list = [v for v in values_list if needle in str(v).lower()] - - return values_list[from_ : from_ + size], len(values_list) - def close(self) -> None: """ Close the underlying boto3 client and release its connection pool. @@ -659,7 +409,7 @@ def _embedding_retrieval( :param query_embedding: The query embedding vector. :param filters: Optional Haystack-format metadata filters. - :param top_k: Maximum number of results to return. S3 Vectors caps this at 100. + :param top_k: Maximum number of results to return. S3 Vectors caps this at 10,000. :returns: List of Documents sorted by similarity. Returned documents will not contain embeddings (S3 Vectors `query_vectors` does not return vector data). """ @@ -682,7 +432,7 @@ def _embedding_retrieval( query_kwargs: dict[str, Any] = { "vectorBucketName": self.vector_bucket_name, "indexName": self.index_name, - "topK": min(top_k, _MAX_TOP_K), # S3 Vectors caps at 100 + "topK": min(top_k, _MAX_TOP_K), "queryVector": {"float32": query_embedding}, "returnMetadata": True, "returnDistance": True, @@ -690,7 +440,15 @@ def _embedding_retrieval( if s3_filter: query_kwargs["filter"] = s3_filter - result = client.query_vectors(**query_kwargs) + # query_vectors returns up to 100 results per page. + vectors: list[dict[str, Any]] = [] + while True: + result = client.query_vectors(**query_kwargs) + vectors.extend(result.get("vectors", [])) + next_token = result.get("nextToken") + if not next_token: + break + query_kwargs["nextToken"] = next_token # Convert distance to a score. # For cosine: S3 Vectors returns cosine *distance* (0 = identical, 2 = opposite). @@ -699,7 +457,7 @@ def _embedding_retrieval( distance_metric = result.get("distanceMetric", self.distance_metric) documents = [] - for v in result.get("vectors", []): + for v in vectors: doc = self._s3_vector_to_document(v) # Compute score from distance @@ -732,14 +490,10 @@ def _document_to_s3_vector(doc: Document, *, fallback_embedding: list[float] | N if doc.content is not None: metadata[_CONTENT_KEY] = doc.content - # Store blob fields. `blob.meta` is an arbitrary dict, which S3 Vectors rejects, so it is - # JSON-encoded into a string (and decoded again in `_s3_vector_to_document`). if doc.blob is not None: - metadata[_BLOB_DATA_KEY] = base64.b64encode(doc.blob.data).decode("ascii") - if doc.blob.meta: - metadata[_BLOB_META_KEY] = json.dumps(doc.blob.meta) - if doc.blob.mime_type: - metadata[_BLOB_MIME_TYPE_KEY] = doc.blob.mime_type + logger.warning( + "Document '{doc_id}' has a blob, which S3VectorsDocumentStore does not store.", doc_id=doc.id + ) # Store user metadata dropped: list[str] = [] @@ -830,28 +584,6 @@ def _s3_vector_to_document(vector: dict[str, Any], *, placeholder_embedding: lis """ metadata = dict(vector.get("metadata", {})) content = metadata.pop(_CONTENT_KEY, None) - blob_data = metadata.pop(_BLOB_DATA_KEY, None) - blob_meta = metadata.pop(_BLOB_META_KEY, None) - blob_mime_type = metadata.pop(_BLOB_MIME_TYPE_KEY, None) - - blob = None - if blob_data is not None: - # `blob.meta` is stored JSON-encoded (see `_document_to_s3_vector`). - if isinstance(blob_meta, str): - try: - blob_meta = json.loads(blob_meta) - except json.JSONDecodeError: - logger.warning( - "Could not decode the stored blob metadata of document '{doc_id}'; " - "the blob is returned without it.", - doc_id=vector["key"], - ) - blob_meta = None - blob = ByteStream( - data=base64.b64decode(blob_data) if isinstance(blob_data, str) else blob_data, - meta=blob_meta or {}, - mime_type=blob_mime_type, - ) embedding = None data = vector.get("data", {}) @@ -867,5 +599,4 @@ def _s3_vector_to_document(vector: dict[str, Any], *, placeholder_embedding: lis content=content, meta=metadata, embedding=embedding, - blob=blob, ) diff --git a/integrations/amazon_s3_vectors/tests/conftest.py b/integrations/amazon_s3_vectors/tests/conftest.py index a4f62e02c2..2621026168 100644 --- a/integrations/amazon_s3_vectors/tests/conftest.py +++ b/integrations/amazon_s3_vectors/tests/conftest.py @@ -12,25 +12,17 @@ A single vector bucket is created per test session (creating buckets is the slow part). Each test gets its own vector index inside that shared bucket so state is isolated. - -One quirk of S3 Vectors is smoothed over here so the generic base tests can run -unchanged: **writes are eventually consistent**. We sleep briefly after -`write_documents` and `delete_documents` so the subsequent -`filter_documents`/`count_documents` reflects the new state. """ from __future__ import annotations import os -import time import uuid from collections.abc import Iterator import boto3 import pytest from botocore.exceptions import ClientError -from haystack.dataclasses import Document -from haystack.document_stores.types import DuplicatePolicy from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore @@ -38,10 +30,6 @@ DIMENSION = 768 REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1") -# Eventual-consistency budget for S3 Vectors write/delete propagation. -WRITE_SLEEP_SECONDS = 5 -DELETE_SLEEP_SECONDS = 5 - def _aws_credentials_available() -> bool: """Return True if any boto3 credential source is configured.""" @@ -102,11 +90,7 @@ def s3_vectors_bucket() -> Iterator[str]: @pytest.fixture def document_store(s3_vectors_bucket: str) -> Iterator[S3VectorsDocumentStore]: - """ - Provide a fresh S3VectorsDocumentStore (one per test), with `write_documents` and - `delete_documents` wrapped to sleep briefly afterwards so that S3 Vectors' eventual - consistency doesn't make the base tests flaky. - """ + """Provide a fresh S3VectorsDocumentStore (one per test).""" index_name = f"idx-{uuid.uuid4().hex[:10]}" store = S3VectorsDocumentStore( @@ -122,23 +106,6 @@ def document_store(s3_vectors_bucket: str) -> Iterator[S3VectorsDocumentStore]: # Eagerly create the index so the first write doesn't race with index creation. store._get_client() - original_write = store.write_documents - - def write_with_sleep(documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.OVERWRITE) -> int: - result = original_write(documents, policy) - time.sleep(WRITE_SLEEP_SECONDS) - return result - - store.write_documents = write_with_sleep # type: ignore[method-assign] - - original_delete = store.delete_documents - - def delete_with_sleep(document_ids: list[str]) -> None: - original_delete(document_ids) - time.sleep(DELETE_SLEEP_SECONDS) - - store.delete_documents = delete_with_sleep # type: ignore[method-assign] - yield store # Clean up: drop the per-test index. Best-effort. diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index 0a4c5fe4a1..09a0f6d30f 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -2,7 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -import base64 from dataclasses import replace from unittest.mock import MagicMock, patch @@ -13,17 +12,10 @@ from haystack.dataclasses import ByteStream, Document from haystack.document_stores.types import DuplicatePolicy from haystack.testing.document_store import ( - CountDocumentsByFilterTest, CountDocumentsTest, - CountUniqueMetadataByFilterTest, DeleteAllTest, - DeleteByFilterTest, DeleteDocumentsTest, FilterableDocsFixtureMixin, - GetMetadataFieldMinMaxTest, - GetMetadataFieldsInfoTest, - GetMetadataFieldUniqueValuesTest, - UpdateByFilterTest, WriteDocumentsTest, ) @@ -212,6 +204,25 @@ def test_embedding_retrieval_passes_filters(mock_boto3): assert call_args["filter"] == {"$and": [{"category": {"$eq": "news"}}]} +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_embedding_retrieval_follows_pagination(mock_boto3): + client = MagicMock() + client.get_vector_bucket.return_value = {} + client.get_index.return_value = {} + client.query_vectors.side_effect = [ + {"vectors": [{"key": "1", "distance": 0.1, "metadata": {}}], "distanceMetric": "cosine", "nextToken": "t"}, + {"vectors": [{"key": "2", "distance": 0.2, "metadata": {}}], "distanceMetric": "cosine"}, + ] + mock_boto3.client.return_value = client + + store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + docs = store._embedding_retrieval(query_embedding=[0.1] * 4, top_k=200) + + assert [d.id for d in docs] == ["1", "2"] + assert client.query_vectors.call_args_list[0][1]["topK"] == 200 + assert client.query_vectors.call_args_list[1][1]["nextToken"] == "t" + + def test_embedding_retrieval_empty_embedding_raises(): """Tests our input validation — no mocking needed.""" store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False) @@ -342,150 +353,11 @@ def test_document_to_s3_vector_converts_tuples_to_lists(): assert result["metadata"]["tags"] == ["a", "b"] -def test_blob_roundtrip(): - """A blob and its (dict) metadata survive the round-trip; the dict is stored JSON-encoded.""" - doc = Document( - id="test-1", - content="Hello world", - blob=ByteStream(data=b"\x00\x01binary", meta={"name": "f.pdf", "pages": 3}, mime_type="application/pdf"), - embedding=[0.1, 0.2, 0.3], - ) +def test_document_to_s3_vector_skips_blob(caplog): + doc = Document(id="1", content="x", embedding=[0.1], blob=ByteStream(data=b"binary", mime_type="application/pdf")) vector = S3VectorsDocumentStore._document_to_s3_vector(doc) - # S3 Vectors rejects dicts, so the blob meta must go over the wire as a string. - assert vector["metadata"]["_blob_meta"] == '{"name": "f.pdf", "pages": 3}' - - restored = S3VectorsDocumentStore._s3_vector_to_document(vector) - assert restored.blob is not None - assert restored.blob.data == b"\x00\x01binary" - assert restored.blob.meta == {"name": "f.pdf", "pages": 3} - assert restored.blob.mime_type == "application/pdf" - assert restored.meta == {} - - -def test_blob_roundtrip_with_undecodable_meta(caplog): - """Corrupt stored blob metadata degrades to an empty dict rather than raising.""" - vector = { - "key": "test-1", - "data": {"float32": [0.1]}, - "metadata": {"_blob_data": base64.b64encode(b"x").decode("ascii"), "_blob_meta": "not json"}, - } - doc = S3VectorsDocumentStore._s3_vector_to_document(vector) - assert doc.blob is not None - assert doc.blob.meta == {} - assert "Could not decode the stored blob metadata" in caplog.text - - -@pytest.mark.parametrize( - ("documents", "expected", "warns"), - [ - ([], {}, False), - ([Document(content="a", meta={"n": 1})], {"content": {"type": "text"}, "n": {"type": "long"}}, False), - ([Document(content="a", meta={"f": 1.5})], {"content": {"type": "text"}, "f": {"type": "long"}}, False), - ([Document(content="a", meta={"s": "x"})], {"content": {"type": "text"}, "s": {"type": "keyword"}}, False), - # bool must win over int: bool is a subclass of int in Python. - ([Document(content="a", meta={"b": True})], {"content": {"type": "text"}, "b": {"type": "boolean"}}, False), - # A list is typed by its first element; an empty list carries no type info. - ( - [Document(content="a", meta={"tags": ["x"]})], - {"content": {"type": "text"}, "tags": {"type": "keyword"}}, - False, - ), - ([Document(content="a", meta={"tags": []})], {"content": {"type": "text"}, "tags": {"type": "keyword"}}, False), - # Mixed types across documents fall back to keyword, with a warning. - ( - [Document(content="a", meta={"m": 1}), Document(content="b", meta={"m": "x"})], - {"content": {"type": "text"}, "m": {"type": "keyword"}}, - True, - ), - # No content anywhere means no synthetic "content" field. - ([Document(meta={"n": 1}, embedding=[0.1])], {"n": {"type": "long"}}, False), - ], -) -def test_get_metadata_fields_info_impl(documents, expected, warns, caplog): - """Type inference is a pure function over already-fetched documents.""" - assert S3VectorsDocumentStore._get_metadata_fields_info_impl(documents) == expected - assert ("mixed types" in caplog.text) is warns - - -def test_get_metadata_field_min_max_impl(): - """min/max strips the meta. prefix, compares numerically, and reports None when absent.""" - docs = [ - Document(content="a", meta={"priority": 5}), - Document(content="b", meta={"priority": 1}), - # 10 vs 5 catches an implementation comparing numbers as strings ("10" < "5"). - Document(content="c", meta={"priority": 10}), - Document(content="d", meta={}), - ] - assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(docs, "priority") == {"min": 1, "max": 10} - assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(docs, "meta.priority") == {"min": 1, "max": 10} - assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(docs, "absent") == {"min": None, "max": None} - assert S3VectorsDocumentStore._get_metadata_field_min_max_impl([], "priority") == {"min": None, "max": None} - - strings = [Document(content="a", meta={"k": "beta"}), Document(content="b", meta={"k": "alpha"})] - assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(strings, "k") == {"min": "alpha", "max": "beta"} - - # Mixed str/numeric would raise on a naive min(); numerics win instead of blowing up. - mixed = [Document(content="a", meta={"k": 3}), Document(content="b", meta={"k": "zzz"})] - assert S3VectorsDocumentStore._get_metadata_field_min_max_impl(mixed, "k") == {"min": 3, "max": 3} - - -def test_get_metadata_field_unique_values_impl_search_and_pagination(): - """Distinct values sort stably, then search_term and from_/size are applied.""" - docs = [Document(content=f"d{i}", meta={"category": f"category_{i}"}) for i in range(5)] - docs.append(Document(content="dup", meta={"category": "category_0"})) - - values, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "category", None, 0, 10) - assert values == [f"category_{i}" for i in range(5)] - assert total == 5 - - first, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "category", None, 0, 2) - second, _ = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "category", None, 2, 2) - assert first == ["category_0", "category_1"] - assert second == ["category_2", "category_3"] - assert total == 5 - - # search_term matches the value, case-insensitively — not the document content. - hits, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "meta.category", "GORY_3", 0, 10) - assert hits == ["category_3"] - assert total == 1 - - -def test_get_metadata_field_unique_values_impl_keeps_types_distinct(): - """`1`, `1.0`, `"1"` and `True` all compare equal in Python but must stay distinct.""" - docs = [ - Document(content="a", meta={"p": 1}), - Document(content="b", meta={"p": "1"}), - Document(content="c", meta={"p": 1.0}), - Document(content="d", meta={"p": True}), - Document(content="e", meta={"p": 1}), - ] - values, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(docs, "p", None, 0, 10) - assert total == 4 - assert {type(v) for v in values} == {int, str, float, bool} - - # List-valued metadata contributes each item, not the list. - listed = [Document(content="a", meta={"tags": ["x", "y"]}), Document(content="b", meta={"tags": ["y"]})] - values, total = S3VectorsDocumentStore._get_metadata_field_unique_values_impl(listed, "tags", None, 0, 10) - assert values == ["x", "y"] - assert total == 2 - - -def test_count_unique_metadata_impl(): - """Distinct-value counts are per field and keep differently-typed values apart.""" - docs = [ - Document(content="a", meta={"category": "A", "status": "active"}), - Document(content="b", meta={"category": "B", "status": "active"}), - Document(content="c", meta={"category": "A", "status": "inactive"}), - ] - assert S3VectorsDocumentStore._count_unique_metadata_impl(docs, ["category", "status"]) == { - "category": 2, - "status": 2, - } - assert S3VectorsDocumentStore._count_unique_metadata_impl(docs, ["absent"]) == {"absent": 0} - assert S3VectorsDocumentStore._count_unique_metadata_impl([], ["category"]) == {"category": 0} - - listed = [Document(content="a", meta={"tags": ["x", "y"]}), Document(content="b", meta={"tags": ["y", "z"]})] - assert S3VectorsDocumentStore._count_unique_metadata_impl(listed, ["tags"]) == {"tags": 3} + assert vector["metadata"] == {"_content": "x"} + assert "does not store" in caplog.text @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") @@ -537,13 +409,6 @@ class TestDocumentStore( WriteDocumentsTest, DeleteDocumentsTest, DeleteAllTest, - DeleteByFilterTest, - UpdateByFilterTest, - CountDocumentsByFilterTest, - CountUniqueMetadataByFilterTest, - GetMetadataFieldsInfoTest, - GetMetadataFieldMinMaxTest, - GetMetadataFieldUniqueValuesTest, FilterableDocsFixtureMixin, ): def assert_documents_are_equal(self, received: list[Document], expected: list[Document]) -> None: @@ -608,55 +473,6 @@ def test_sentence_window_retriever(self, document_store: S3VectorsDocumentStore) ) assert len(result["context_windows"]) == 1 - def test_get_metadata_fields_info_consistent_types(self, document_store: S3VectorsDocumentStore) -> None: - """A field with the same type in every document keeps that type, with no mixed-type fallback.""" - docs = [ - Document(content="Doc 1", meta={"score": 85}), - Document(content="Doc 2", meta={"score": 90}), - Document(content="Doc 3", meta={"score": 78}), - ] - document_store.write_documents(docs) - - fields_info = document_store.get_metadata_fields_info() - assert fields_info["score"] == {"type": "long"} - - def test_get_metadata_field_min_max_boolean_and_string(self, document_store: S3VectorsDocumentStore) -> None: - """min/max works on booleans (False < True) and on strings (alphabetical), not just numbers.""" - docs = [ - Document(content="Doc 1", meta={"active": True, "category": "Zebra"}), - Document(content="Doc 2", meta={"active": False, "category": "Alpha"}), - Document(content="Doc 3", meta={"active": True, "category": "Beta"}), - Document(content="Doc 4", meta={"active": False, "category": "Gamma"}), - ] - document_store.write_documents(docs) - - assert document_store.get_metadata_field_min_max("active") == {"min": False, "max": True} - assert document_store.get_metadata_field_min_max("category") == {"min": "Alpha", "max": "Zebra"} - - def test_get_metadata_field_min_max_no_values(self, document_store: S3VectorsDocumentStore) -> None: - """A list-valued field has no comparable scalar, so min/max is None — as for a missing field.""" - docs = [ - Document(content="Doc 1", meta={"tags": ["tag1", "tag2"]}), - Document(content="Doc 2", meta={"tags": ["tag3", "tag4"]}), - ] - document_store.write_documents(docs) - - assert document_store.get_metadata_field_min_max("tags") == {"min": None, "max": None} - assert document_store.get_metadata_field_min_max("nonexistent") == {"min": None, "max": None} - - def test_get_metadata_field_unique_values_with_lists(self, document_store: S3VectorsDocumentStore) -> None: - """A list-valued field contributes each item as a distinct value, not the list itself.""" - docs = [ - Document(content="Doc 1", meta={"tags": ["python", "java"]}), - Document(content="Doc 2", meta={"tags": ["python", "rust"]}), - Document(content="Doc 3", meta={"tags": ["java", "go"]}), - ] - document_store.write_documents(docs) - - values, total = document_store.get_metadata_field_unique_values("tags", size=10) - assert total == 4 - assert set(values) == {"go", "java", "python", "rust"} - def test_write_documents_without_embedding_reads_back_as_none(self, document_store: S3VectorsDocumentStore) -> None: """ S3 Vectors stores a placeholder vector for embedding-less documents; `filter_documents` From bf42a0f59bf341a32b1a461bad8b7ca732334d1c Mon Sep 17 00:00:00 2001 From: anakin87 Date: Fri, 6 Nov 2026 17:36:51 +0100 Subject: [PATCH 16/17] improvements --- .github/workflows/amazon_bedrock.yml | 2 +- .github/workflows/amazon_s3_vectors.yml | 2 +- .../amazon_s3_vectors/document_store.py | 20 ++++++--- .../tests/test_document_store.py | 44 +++++++++++-------- .../tests/test_embedding_retriever.py | 19 ++++---- .../amazon_s3_vectors/tests/test_filters.py | 11 ++--- 6 files changed, 54 insertions(+), 44 deletions(-) diff --git a/.github/workflows/amazon_bedrock.yml b/.github/workflows/amazon_bedrock.yml index 103eab2784..ede0ef024f 100644 --- a/.github/workflows/amazon_bedrock.yml +++ b/.github/workflows/amazon_bedrock.yml @@ -117,7 +117,7 @@ jobs: # Do not authenticate on PRs from forks and on PRs created by dependabot - name: AWS authentication id: aws-auth - if: github.event_name == 'schedule' || (github.event.pull_request.head.repo.full_name == github.repository && !startsWith(github.event.pull_request.head.ref, 'dependabot/')) + if: github.event_name == 'push' || github.event_name == 'schedule' || (github.event.pull_request.head.repo.full_name == github.repository && !startsWith(github.event.pull_request.head.ref, 'dependabot/')) uses: aws-actions/configure-aws-credentials@e1253824e5c10ff9df46874f81ed3ec929e19cfd with: aws-region: ${{ env.AWS_REGION }} diff --git a/.github/workflows/amazon_s3_vectors.yml b/.github/workflows/amazon_s3_vectors.yml index 041ba4a153..e44ab1f4fe 100644 --- a/.github/workflows/amazon_s3_vectors.yml +++ b/.github/workflows/amazon_s3_vectors.yml @@ -113,7 +113,7 @@ jobs: # Do not authenticate on PRs from forks and on PRs created by dependabot - name: AWS authentication id: aws-auth - if: github.event_name == 'schedule' || (github.event.pull_request.head.repo.full_name == github.repository && !startsWith(github.event.pull_request.head.ref, 'dependabot/')) + if: github.event_name == 'push' || github.event_name == 'schedule' || (github.event.pull_request.head.repo.full_name == github.repository && !startsWith(github.event.pull_request.head.ref, 'dependabot/')) uses: aws-actions/configure-aws-credentials@e1253824e5c10ff9df46874f81ed3ec929e19cfd with: aws-region: ${{ env.AWS_REGION }} diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index e0fad768d9..1efbdd4744 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -98,7 +98,8 @@ def __init__( :param vector_bucket_name: Name of the S3 vector bucket. :param index_name: Name of the vector index within the bucket. - :param dimension: Dimensionality of the embeddings (e.g. 768, 1536). + :param dimension: Dimensionality of the embeddings when creating an index (e.g. 768, 1536). + When connecting to an existing index, its actual dimension is used. :param distance_metric: Distance metric for similarity search: `"cosine"` or `"euclidean"`. :param region_name: AWS region. If not provided, uses the default from the environment/config. :param aws_access_key_id: AWS access key ID. If not provided, uses the default credential chain. @@ -120,9 +121,7 @@ def __init__( self.create_bucket_and_index = create_bucket_and_index self.non_filterable_metadata_keys = non_filterable_metadata_keys or [] - # S3 Vectors requires a vector on every record, and a cosine index rejects zero-norm - # vectors, so documents written without an embedding get this non-zero placeholder. - self._dummy_vector = [-10.0] * dimension + self._dummy_vector: list[float] self._client: Any = None @@ -143,8 +142,17 @@ def _get_client(self) -> Any: self._client = boto3.client(**kwargs) - if self.create_bucket_and_index: - self._ensure_bucket_and_index() + try: + if self.create_bucket_and_index: + self._ensure_bucket_and_index() + + response = self._client.get_index(vectorBucketName=self.vector_bucket_name, indexName=self.index_name) + self.dimension = response["index"]["dimension"] + # Cosine indexes reject zero-norm vectors, even for documents without embeddings. + self._dummy_vector = [-10.0] * self.dimension + except Exception: + self.close() + raise return self._client diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index 09a0f6d30f..d4246013fe 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -100,16 +100,29 @@ def test_from_dict(_mock_boto3): assert store.create_bucket_and_index is False +@pytest.mark.parametrize("create_bucket_and_index", [True, False]) @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") -def test_write_documents_missing_embedding_uses_placeholder(mock_boto3): +def test_write_documents_missing_embedding_uses_placeholder(mock_boto3, create_bucket_and_index): """Every S3 Vectors record needs a vector, so embedding-less documents get the placeholder.""" client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = S3VectorsDocumentStore( + vector_bucket_name="b", + index_name="i", + dimension=3, + region_name="us-east-1", + create_bucket_and_index=create_bucket_and_index, + ) docs = [Document(id="1", content="ok", embedding=[0.1] * 4), Document(id="2", content="missing")] + client.list_vectors.return_value = { + "vectors": [{"key": "2", "data": {"float32": [-10.0] * 4}, "metadata": {"_content": "missing"}}] + } + # Reading first initializes the placeholder using the index dimension, not the constructor's. + assert store.filter_documents() == [docs[1]] + assert store.dimension == 4 assert store.write_documents(docs) == 2 written = client.put_vectors.call_args.kwargs["vectors"] @@ -129,7 +142,7 @@ def test_write_documents_unsupported_policy_overwrites(mock_boto3, policy): """ client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") @@ -152,7 +165,7 @@ def test_embedding_retrieval_score_conversion(mock_boto3): """Tests our distance-to-score conversion logic — the only non-trivial transform in retrieval.""" client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} client.query_vectors.return_value = { "vectors": [{"key": "1", "distance": 0.05, "metadata": {"_content": "Hello", "category": "news"}}], "distanceMetric": "cosine", @@ -173,7 +186,7 @@ def test_embedding_retrieval_euclidean_score(mock_boto3): """Tests euclidean distance-to-score conversion (negated).""" client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} client.query_vectors.return_value = { "vectors": [{"key": "1", "distance": 1.5, "metadata": {}}], "distanceMetric": "euclidean", @@ -192,7 +205,7 @@ def test_embedding_retrieval_passes_filters(mock_boto3): """Tests that Haystack filters are converted and passed to query_vectors.""" client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} mock_boto3.client.return_value = client @@ -208,7 +221,7 @@ def test_embedding_retrieval_passes_filters(mock_boto3): def test_embedding_retrieval_follows_pagination(mock_boto3): client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} client.query_vectors.side_effect = [ {"vectors": [{"key": "1", "distance": 0.1, "metadata": {}}], "distanceMetric": "cosine", "nextToken": "t"}, {"vectors": [{"key": "2", "distance": 0.2, "metadata": {}}], "distanceMetric": "cosine"}, @@ -365,7 +378,7 @@ def test_close_releases_the_client_and_allows_reopen(mock_boto3): """close() closes the boto3 client; the store stays usable and lazily builds a new one.""" client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} mock_boto3.client.return_value = client store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") @@ -393,7 +406,7 @@ def test_close_is_exception_safe(mock_boto3): """A failure while closing must not propagate to the caller.""" client = MagicMock() client.get_vector_bucket.return_value = {} - client.get_index.return_value = {} + client.get_index.return_value = {"index": {"dimension": 4}} client.close.side_effect = RuntimeError("boom") mock_boto3.client.return_value = client @@ -420,14 +433,9 @@ def assert_documents_are_equal(self, received: list[Document], expected: list[Do we ignore it for equality. """ assert len(received) == len(expected) - received.sort(key=lambda d: d.id) - expected.sort(key=lambda d: d.id) - for r, e in zip(received, expected, strict=True): - r_norm = replace(r, embedding=None, score=None) - e_norm = replace(e, embedding=None, score=None) - assert r_norm == e_norm - if r.embedding is not None and e.embedding is not None: - assert r.embedding == pytest.approx(e.embedding, abs=1e-5) + for r, e in zip(sorted(received, key=lambda d: d.id), sorted(expected, key=lambda d: d.id), strict=True): + assert replace(r, embedding=None, score=None) == replace(e, embedding=None, score=None) + assert r.embedding == pytest.approx(e.embedding, abs=1e-5) def test_write_documents(self, document_store: S3VectorsDocumentStore) -> None: """ diff --git a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py index 3365f042a5..0c56623d93 100644 --- a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py +++ b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py @@ -135,19 +135,18 @@ def test_run(): @pytest.mark.integration def test_embedding_retrieval(document_store: S3VectorsDocumentStore) -> None: docs = [ - Document(id="r-1", content="First", embedding=[0.1] * 768), - Document(id="r-2", content="Second", embedding=[0.5] * 768), - Document(id="r-3", content="Third", embedding=[0.9] * 768), + Document(id="r-1", content="First", embedding=[1.0, 0.0] + [0.0] * 766), + Document(id="r-2", content="Second", embedding=[0.8, 0.6] + [0.0] * 766), + Document(id="r-3", content="Third", embedding=[0.0, 1.0] + [0.0] * 766), ] document_store.write_documents(docs) - results = document_store._embedding_retrieval(query_embedding=[0.1] * 768, top_k=10) - assert len(results) > 0 - for doc in results: - assert doc.score is not None - assert doc.content is not None - # query_vectors does not return vector data - assert doc.embedding is None + results = document_store._embedding_retrieval(query_embedding=[1.0, 0.0] + [0.0] * 766, top_k=2) + assert [doc.id for doc in results] == ["r-1", "r-2"] + assert [doc.score for doc in results] == pytest.approx([1.0, 0.8], abs=1e-5) + assert [doc.content for doc in results] == ["First", "Second"] + # query_vectors does not return vector data + assert all(doc.embedding is None for doc in results) @pytest.mark.integration diff --git a/integrations/amazon_s3_vectors/tests/test_filters.py b/integrations/amazon_s3_vectors/tests/test_filters.py index f86ed17fb6..16db36f691 100644 --- a/integrations/amazon_s3_vectors/tests/test_filters.py +++ b/integrations/amazon_s3_vectors/tests/test_filters.py @@ -97,11 +97,6 @@ def test_validate_filters(): class TestFilters(FilterDocumentsTest): def assert_documents_are_equal(self, received: list[Document], expected: list[Document]) -> None: assert len(received) == len(expected) - received.sort(key=lambda d: d.id) - expected.sort(key=lambda d: d.id) - for r, e in zip(received, expected, strict=True): - r_norm = replace(r, embedding=None, score=None) - e_norm = replace(e, embedding=None, score=None) - assert r_norm == e_norm - if r.embedding is not None and e.embedding is not None: - assert r.embedding == pytest.approx(e.embedding, abs=1e-5) + for r, e in zip(sorted(received, key=lambda d: d.id), sorted(expected, key=lambda d: d.id), strict=True): + assert replace(r, embedding=None, score=None) == replace(e, embedding=None, score=None) + assert r.embedding == pytest.approx(e.embedding, abs=1e-5) From bed150c8ee8c96105697fc7a8cbe382d4363ff49 Mon Sep 17 00:00:00 2001 From: anakin87 Date: Fri, 6 Nov 2026 18:02:47 +0100 Subject: [PATCH 17/17] rename and align --- .../pydoc/config_docusaurus.yml | 2 +- .../retrievers/amazon_s3_vectors/__init__.py | 4 +- .../amazon_s3_vectors/embedding_retriever.py | 32 ++-- .../components/{ => retrievers}/py.typed | 0 .../amazon_s3_vectors/__init__.py | 4 +- .../amazon_s3_vectors/document_store.py | 26 +-- .../amazon_s3_vectors/tests/conftest.py | 8 +- .../tests/test_document_store.py | 174 ++++++++++++++---- .../tests/test_embedding_retriever.py | 65 ++++--- .../amazon_s3_vectors/tests/test_filters.py | 76 ++++---- 10 files changed, 250 insertions(+), 141 deletions(-) rename integrations/amazon_s3_vectors/src/haystack_integrations/components/{ => retrievers}/py.typed (100%) diff --git a/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml b/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml index 4320722107..8b32659296 100644 --- a/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml +++ b/integrations/amazon_s3_vectors/pydoc/config_docusaurus.yml @@ -9,6 +9,6 @@ processors: skip_empty_modules: true renderer: description: Amazon S3 Vectors integration for Haystack - id: integrations-amazon_s3_vectors + id: integrations-amazon-s3-vectors filename: amazon_s3_vectors.md title: Amazon S3 Vectors diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py index 39ebcd2f58..b7c6130e08 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/__init__.py @@ -2,6 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -from .embedding_retriever import S3VectorsEmbeddingRetriever +from .embedding_retriever import AmazonS3VectorsEmbeddingRetriever -__all__ = ["S3VectorsEmbeddingRetriever"] +__all__ = ["AmazonS3VectorsEmbeddingRetriever"] diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py index a9fa55566b..8b80b3e227 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/amazon_s3_vectors/embedding_retriever.py @@ -9,13 +9,13 @@ from haystack.document_stores.types import FilterPolicy from haystack.document_stores.types.filter_policy import apply_filter_policy -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore +from haystack_integrations.document_stores.amazon_s3_vectors import AmazonS3VectorsDocumentStore @component -class S3VectorsEmbeddingRetriever: +class AmazonS3VectorsEmbeddingRetriever: """ - Retrieve documents from an `S3VectorsDocumentStore` based on their dense embeddings. + Retrieve documents from an `AmazonS3VectorsDocumentStore` based on their dense embeddings. Usage example: ```python @@ -24,10 +24,10 @@ class S3VectorsEmbeddingRetriever: # Requires: pip install sentence-transformers-haystack from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersTextEmbedder from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersDocumentEmbedder - from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever - from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + from haystack_integrations.components.retrievers.amazon_s3_vectors import AmazonS3VectorsEmbeddingRetriever + from haystack_integrations.document_stores.amazon_s3_vectors import AmazonS3VectorsDocumentStore - document_store = S3VectorsDocumentStore( + document_store = AmazonS3VectorsDocumentStore( vector_bucket_name="my-vectors", index_name="my-index", dimension=768, @@ -47,7 +47,7 @@ class S3VectorsEmbeddingRetriever: query_pipeline = Pipeline() query_pipeline.add_component("text_embedder", SentenceTransformersTextEmbedder()) - query_pipeline.add_component("retriever", S3VectorsEmbeddingRetriever(document_store=document_store)) + query_pipeline.add_component("retriever", AmazonS3VectorsEmbeddingRetriever(document_store=document_store)) query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding") query = "How many languages are there?" @@ -58,22 +58,22 @@ class S3VectorsEmbeddingRetriever: def __init__( self, *, - document_store: S3VectorsDocumentStore, + document_store: AmazonS3VectorsDocumentStore, filters: dict[str, Any] | None = None, top_k: int = 10, filter_policy: str | FilterPolicy = FilterPolicy.REPLACE, ) -> None: """ - Initialize the S3VectorsEmbeddingRetriever. + Initialize the AmazonS3VectorsEmbeddingRetriever. - :param document_store: An instance of `S3VectorsDocumentStore`. + :param document_store: An instance of `AmazonS3VectorsDocumentStore`. :param filters: Filters applied to the retrieved Documents. :param top_k: Maximum number of Documents to return. :param filter_policy: Policy to determine how filters are applied. - :raises ValueError: If `document_store` is not an `S3VectorsDocumentStore`. + :raises ValueError: If `document_store` is not an `AmazonS3VectorsDocumentStore`. """ - if not isinstance(document_store, S3VectorsDocumentStore): - msg = "document_store must be an instance of S3VectorsDocumentStore" + if not isinstance(document_store, AmazonS3VectorsDocumentStore): + msg = "document_store must be an instance of AmazonS3VectorsDocumentStore" raise ValueError(msg) self.document_store = document_store @@ -94,9 +94,9 @@ def to_dict(self) -> dict[str, Any]: ) @classmethod - def from_dict(cls, data: dict[str, Any]) -> "S3VectorsEmbeddingRetriever": + def from_dict(cls, data: dict[str, Any]) -> "AmazonS3VectorsEmbeddingRetriever": """Deserialize the component from a dictionary.""" - data["init_parameters"]["document_store"] = S3VectorsDocumentStore.from_dict( + data["init_parameters"]["document_store"] = AmazonS3VectorsDocumentStore.from_dict( data["init_parameters"]["document_store"] ) if filter_policy := data["init_parameters"].get("filter_policy"): @@ -111,7 +111,7 @@ def run( top_k: int | None = None, ) -> dict[str, list[Document]]: """ - Retrieve documents from the S3VectorsDocumentStore based on dense embeddings. + Retrieve documents from the AmazonS3VectorsDocumentStore based on dense embeddings. :param query_embedding: Embedding of the query. :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/components/py.typed b/integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/py.typed similarity index 100% rename from integrations/amazon_s3_vectors/src/haystack_integrations/components/py.typed rename to integrations/amazon_s3_vectors/src/haystack_integrations/components/retrievers/py.typed diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py index 4e6beddd97..9c3c6414c1 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/__init__.py @@ -2,6 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -from .document_store import S3VectorsDocumentStore +from .document_store import AmazonS3VectorsDocumentStore -__all__ = ["S3VectorsDocumentStore"] +__all__ = ["AmazonS3VectorsDocumentStore"] diff --git a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py index 1efbdd4744..6cd9f49dbe 100644 --- a/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py +++ b/integrations/amazon_s3_vectors/src/haystack_integrations/document_stores/amazon_s3_vectors/document_store.py @@ -42,7 +42,7 @@ _NON_FILTERABLE_KEYS = [_CONTENT_KEY] -class S3VectorsDocumentStore: +class AmazonS3VectorsDocumentStore: """ A Document Store using [Amazon S3 Vectors](https://aws.amazon.com/s3/features/vectors/). @@ -64,14 +64,14 @@ class S3VectorsDocumentStore: retrievable via `filter_documents()` - Only `DuplicatePolicy.OVERWRITE` is supported, since `put_vectors` is an upsert - Distance metrics: `cosine` or `euclidean` (set at index creation, immutable) - - `filter_documents()` filters locally, so prefer `S3VectorsEmbeddingRetriever` with filters + - `filter_documents()` filters locally, so prefer `AmazonS3VectorsEmbeddingRetriever` with filters - `Document.blob` is not stored Usage example: ```python - from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore + from haystack_integrations.document_stores.amazon_s3_vectors import AmazonS3VectorsDocumentStore - document_store = S3VectorsDocumentStore( + document_store = AmazonS3VectorsDocumentStore( vector_bucket_name="my-vectors", index_name="my-index", dimension=768, @@ -94,7 +94,7 @@ def __init__( non_filterable_metadata_keys: list[str] | None = None, ) -> None: """ - Create an S3VectorsDocumentStore instance. + Create an AmazonS3VectorsDocumentStore instance. :param vector_bucket_name: Name of the S3 vector bucket. :param index_name: Name of the vector index within the bucket. @@ -218,7 +218,7 @@ def to_dict(self) -> dict[str, Any]: ) @classmethod - def from_dict(cls, data: dict[str, Any]) -> "S3VectorsDocumentStore": + def from_dict(cls, data: dict[str, Any]) -> "AmazonS3VectorsDocumentStore": """Deserialize a document store from a dictionary.""" deserialize_secrets_inplace( data["init_parameters"], @@ -265,7 +265,7 @@ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = D """ Write Documents to the S3 Vectors index. - `S3VectorsDocumentStore` only supports `DuplicatePolicy.OVERWRITE`: `put_vectors` is an + `AmazonS3VectorsDocumentStore` only supports `DuplicatePolicy.OVERWRITE`: `put_vectors` is an upsert, and S3 Vectors offers no cheap existence check (only `get_vectors` on batches of 100 keys, or a full index scan), so enforcing any other policy would cost an extra read pass on every write. Other policies are ignored with a warning. @@ -292,7 +292,7 @@ def write_documents(self, documents: list[Document], policy: DuplicatePolicy = D # DuplicatePolicy.NONE means "the store decides", so it needs no warning. if policy not in (DuplicatePolicy.NONE, DuplicatePolicy.OVERWRITE): logger.warning( - "S3VectorsDocumentStore only supports DuplicatePolicy.OVERWRITE because put_vectors is an " + "AmazonS3VectorsDocumentStore only supports DuplicatePolicy.OVERWRITE because put_vectors is an " "upsert. Ignoring policy={policy} and overwriting.", policy=str(policy), ) @@ -330,7 +330,7 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume Note: S3 Vectors only supports metadata filtering during vector similarity queries, not as a standalone operation. This method lists all vectors and filters them locally with Haystack, - which can be slow for large indexes. `S3VectorsEmbeddingRetriever` filters through S3 + which can be slow for large indexes. `AmazonS3VectorsEmbeddingRetriever` filters through S3 Vectors instead, and the two can match different documents for the same filter, especially on array metadata: S3 Vectors matches an array field when any element matches, while Haystack compares the whole array. @@ -345,7 +345,7 @@ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Docume "S3 Vectors does not support standalone filtered listing. " "filter_documents() will fetch ALL vectors and apply filters client-side, " "which can be very slow for large indexes. " - "Prefer using S3VectorsEmbeddingRetriever with filters for efficient filtered retrieval." + "Prefer using AmazonS3VectorsEmbeddingRetriever with filters for efficient filtered retrieval." ) # list_vectors supports returnData and returnMetadata directly, @@ -413,7 +413,7 @@ def _embedding_retrieval( Retrieve documents most similar to the query embedding. This method is not part of the public interface. - Use `S3VectorsEmbeddingRetriever` instead. + Use `AmazonS3VectorsEmbeddingRetriever` instead. :param query_embedding: The query embedding vector. :param filters: Optional Haystack-format metadata filters. @@ -500,7 +500,7 @@ def _document_to_s3_vector(doc: Document, *, fallback_embedding: list[float] | N if doc.blob is not None: logger.warning( - "Document '{doc_id}' has a blob, which S3VectorsDocumentStore does not store.", doc_id=doc.id + "Document '{doc_id}' has a blob, which AmazonS3VectorsDocumentStore does not store.", doc_id=doc.id ) # Store user metadata @@ -513,7 +513,7 @@ def _document_to_s3_vector(doc: Document, *, fallback_embedding: list[float] | N key=key, ) continue - if not S3VectorsDocumentStore._is_supported_metadata_value(value): + if not AmazonS3VectorsDocumentStore._is_supported_metadata_value(value): dropped.append(key) continue # A tuple is an array as far as S3 Vectors is concerned, but botocore's document-type diff --git a/integrations/amazon_s3_vectors/tests/conftest.py b/integrations/amazon_s3_vectors/tests/conftest.py index 2621026168..cbdc60ba63 100644 --- a/integrations/amazon_s3_vectors/tests/conftest.py +++ b/integrations/amazon_s3_vectors/tests/conftest.py @@ -24,7 +24,7 @@ import pytest from botocore.exceptions import ClientError -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore +from haystack_integrations.document_stores.amazon_s3_vectors import AmazonS3VectorsDocumentStore # Dimension used by Haystack's `FilterableDocsFixtureMixin` test corpus. DIMENSION = 768 @@ -89,11 +89,11 @@ def s3_vectors_bucket() -> Iterator[str]: @pytest.fixture -def document_store(s3_vectors_bucket: str) -> Iterator[S3VectorsDocumentStore]: - """Provide a fresh S3VectorsDocumentStore (one per test).""" +def document_store(s3_vectors_bucket: str) -> Iterator[AmazonS3VectorsDocumentStore]: + """Provide a fresh AmazonS3VectorsDocumentStore (one per test).""" index_name = f"idx-{uuid.uuid4().hex[:10]}" - store = S3VectorsDocumentStore( + store = AmazonS3VectorsDocumentStore( vector_bucket_name=s3_vectors_bucket, index_name=index_name, dimension=DIMENSION, diff --git a/integrations/amazon_s3_vectors/tests/test_document_store.py b/integrations/amazon_s3_vectors/tests/test_document_store.py index d4246013fe..0a59a9a1a5 100644 --- a/integrations/amazon_s3_vectors/tests/test_document_store.py +++ b/integrations/amazon_s3_vectors/tests/test_document_store.py @@ -2,11 +2,13 @@ # # SPDX-License-Identifier: Apache-2.0 +from contextlib import nullcontext from dataclasses import replace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import numpy as np import pytest +from botocore.exceptions import ClientError from haystack.components.preprocessors import DocumentSplitter from haystack.components.retrievers import SentenceWindowRetriever from haystack.dataclasses import ByteStream, Document @@ -18,13 +20,14 @@ FilterableDocsFixtureMixin, WriteDocumentsTest, ) +from haystack.utils import Secret -from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore +from haystack_integrations.components.retrievers.amazon_s3_vectors import AmazonS3VectorsEmbeddingRetriever +from haystack_integrations.document_stores.amazon_s3_vectors import AmazonS3VectorsDocumentStore def test_init_is_lazy(): - store = S3VectorsDocumentStore( + store = AmazonS3VectorsDocumentStore( vector_bucket_name="my-bucket", index_name="my-index", dimension=768, @@ -33,7 +36,7 @@ def test_init_is_lazy(): def test_init_default_params(): - store = S3VectorsDocumentStore( + store = AmazonS3VectorsDocumentStore( vector_bucket_name="my-bucket", index_name="my-index", dimension=768, @@ -48,7 +51,7 @@ def test_init_default_params(): @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_to_dict(_mock_boto3): - store = S3VectorsDocumentStore( + store = AmazonS3VectorsDocumentStore( vector_bucket_name="my-bucket", index_name="my-index", dimension=768, @@ -58,7 +61,7 @@ def test_to_dict(_mock_boto3): ) d = store.to_dict() assert d == { - "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.AmazonS3VectorsDocumentStore", "init_parameters": { "vector_bucket_name": "my-bucket", "index_name": "my-index", @@ -74,10 +77,38 @@ def test_to_dict(_mock_boto3): } +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_to_dict_with_credentials(mock_boto3, monkeypatch): + credentials = {} + for name in ("aws_access_key_id", "aws_secret_access_key", "aws_session_token"): + monkeypatch.setenv(name.upper(), "test-" + name) + credentials[name] = Secret.from_env_var(name.upper()) + store = AmazonS3VectorsDocumentStore( + vector_bucket_name="my-bucket", + index_name="my-index", + dimension=768, + region_name="us-west-2", + create_bucket_and_index=False, + **credentials, + ) + data = store.to_dict() + for name, secret in credentials.items(): + assert data["init_parameters"][name] == secret.to_dict() + restored = AmazonS3VectorsDocumentStore.from_dict(data) + mock_boto3.client.assert_not_called() + mock_boto3.client.return_value.get_index.return_value = {"index": {"dimension": 768}} + restored._get_client() + mock_boto3.client.assert_called_once_with( + service_name="s3vectors", + region_name="us-west-2", + **{name: "test-" + name for name in credentials}, + ) + + @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_from_dict(_mock_boto3): data = { - "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.AmazonS3VectorsDocumentStore", "init_parameters": { "vector_bucket_name": "my-bucket", "index_name": "my-index", @@ -91,7 +122,7 @@ def test_from_dict(_mock_boto3): "non_filterable_metadata_keys": [], }, } - store = S3VectorsDocumentStore.from_dict(data) + store = AmazonS3VectorsDocumentStore.from_dict(data) assert store.vector_bucket_name == "my-bucket" assert store.index_name == "my-index" assert store.dimension == 768 @@ -109,7 +140,7 @@ def test_write_documents_missing_embedding_uses_placeholder(mock_boto3, create_b client.get_index.return_value = {"index": {"dimension": 4}} mock_boto3.client.return_value = client - store = S3VectorsDocumentStore( + store = AmazonS3VectorsDocumentStore( vector_bucket_name="b", index_name="i", dimension=3, @@ -133,6 +164,22 @@ def test_write_documents_missing_embedding_uses_placeholder(mock_boto3, create_b assert docs[1].embedding is None +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_filter_documents_follows_pagination(mock_boto3): + client = mock_boto3.client.return_value + client.get_index.return_value = {"index": {"dimension": 4}} + client.list_vectors.side_effect = [ + {"vectors": [{"key": "1", "metadata": {"_content": "First"}}], "nextToken": "page-2"}, + {"vectors": [{"key": "2", "metadata": {"_content": "Second"}}]}, + ] + store = AmazonS3VectorsDocumentStore( + vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False + ) + assert store.filter_documents() == [Document(id="1", content="First"), Document(id="2", content="Second")] + kwargs = {"vectorBucketName": "b", "indexName": "i", "returnData": True, "returnMetadata": True} + assert client.list_vectors.call_args_list == [call(**kwargs), call(**kwargs, nextToken="page-2")] + + @pytest.mark.parametrize("policy", [DuplicatePolicy.FAIL, DuplicatePolicy.SKIP]) @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_write_documents_unsupported_policy_overwrites(mock_boto3, policy): @@ -145,15 +192,60 @@ def test_write_documents_unsupported_policy_overwrites(mock_boto3, policy): client.get_index.return_value = {"index": {"dimension": 4}} mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = AmazonS3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") assert store.write_documents([Document(id="1", content="Hello", embedding=[0.1] * 4)], policy=policy) == 1 client.get_vectors.assert_not_called() client.put_vectors.assert_called_once() +@pytest.mark.parametrize("fail_second_batch", [False, True]) +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_write_documents_batches(mock_boto3, fail_second_batch): + client = mock_boto3.client.return_value + client.get_index.return_value = {"index": {"dimension": 4}} + if fail_second_batch: + client.put_vectors.side_effect = [ + {}, + ClientError({"Error": {"Code": "InternalServerException", "Message": "second batch"}}, "PutVectors"), + ] + store = AmazonS3VectorsDocumentStore( + vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False + ) + docs = [Document(id=str(i), embedding=[0.1] * 4) for i in range(501)] + with pytest.raises(ClientError, match="second batch") if fail_second_batch else nullcontext(): + assert store.write_documents(docs) == 501 + batches = [c.kwargs["vectors"] for c in client.put_vectors.call_args_list] + assert [len(batch) for batch in batches] == [500, 1] + assert [v["key"] for batch in batches for v in batch] == [doc.id for doc in docs] + + +@pytest.mark.parametrize("fail_second_batch", [False, True]) +@patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") +def test_delete_documents_batches(mock_boto3, fail_second_batch): + client = mock_boto3.client.return_value + client.get_index.return_value = {"index": {"dimension": 4}} + if fail_second_batch: + client.delete_vectors.side_effect = [ + {}, + ClientError({"Error": {"Code": "InternalServerException", "Message": "second batch"}}, "DeleteVectors"), + ] + store = AmazonS3VectorsDocumentStore( + vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False + ) + ids = [str(i) for i in range(501)] + with pytest.raises(ClientError, match="second batch") if fail_second_batch else nullcontext(): + store.delete_documents(ids) + assert client.delete_vectors.call_args_list == [ + call(vectorBucketName="b", indexName="i", keys=ids[:500]), + call(vectorBucketName="b", indexName="i", keys=ids[500:]), + ] + + def test_write_documents_invalid_input_raises(): """write_documents must reject non-list / non-Document input before doing anything.""" - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False) + store = AmazonS3VectorsDocumentStore( + vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False + ) with pytest.raises(ValueError): store.write_documents("not a list") # type: ignore[arg-type] with pytest.raises(ValueError): @@ -172,7 +264,7 @@ def test_embedding_retrieval_score_conversion(mock_boto3): } mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = AmazonS3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") docs = store._embedding_retrieval(query_embedding=[0.1] * 4, top_k=5) assert len(docs) == 1 assert docs[0].id == "1" @@ -193,7 +285,7 @@ def test_embedding_retrieval_euclidean_score(mock_boto3): } mock_boto3.client.return_value = client - store = S3VectorsDocumentStore( + store = AmazonS3VectorsDocumentStore( vector_bucket_name="b", index_name="i", dimension=4, distance_metric="euclidean", region_name="us-east-1" ) docs = store._embedding_retrieval(query_embedding=[0.1] * 4, top_k=5) @@ -209,7 +301,7 @@ def test_embedding_retrieval_passes_filters(mock_boto3): client.query_vectors.return_value = {"vectors": [], "distanceMetric": "cosine"} mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = AmazonS3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") filters = {"operator": "AND", "conditions": [{"field": "meta.category", "operator": "==", "value": "news"}]} store._embedding_retrieval(query_embedding=[0.1] * 4, filters=filters, top_k=5) @@ -228,7 +320,7 @@ def test_embedding_retrieval_follows_pagination(mock_boto3): ] mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = AmazonS3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") docs = store._embedding_retrieval(query_embedding=[0.1] * 4, top_k=200) assert [d.id for d in docs] == ["1", "2"] @@ -238,7 +330,9 @@ def test_embedding_retrieval_follows_pagination(mock_boto3): def test_embedding_retrieval_empty_embedding_raises(): """Tests our input validation — no mocking needed.""" - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False) + store = AmazonS3VectorsDocumentStore( + vector_bucket_name="b", index_name="i", dimension=4, create_bucket_and_index=False + ) with pytest.raises(ValueError, match="non-empty"): store._embedding_retrieval(query_embedding=[]) @@ -248,7 +342,7 @@ def test_document_to_s3_vector(): doc = Document( id="test-1", content="Hello world", embedding=[0.1, 0.2, 0.3], meta={"category": "test", "year": 2024} ) - result = S3VectorsDocumentStore._document_to_s3_vector(doc) + result = AmazonS3VectorsDocumentStore._document_to_s3_vector(doc) assert result["key"] == "test-1" assert result["data"] == {"float32": [0.1, 0.2, 0.3]} assert result["metadata"]["_content"] == "Hello world" @@ -263,7 +357,7 @@ def test_s3_vector_to_document(): "data": {"float32": [0.1, 0.2, 0.3]}, "metadata": {"_content": "Hello world", "category": "test"}, } - doc = S3VectorsDocumentStore._s3_vector_to_document(vector) + doc = AmazonS3VectorsDocumentStore._s3_vector_to_document(vector) assert doc.id == "test-1" assert doc.content == "Hello world" assert doc.embedding == [0.1, 0.2, 0.3] @@ -275,8 +369,8 @@ def test_document_roundtrip(): doc = Document( id="test-1", content="Hello world", embedding=[0.1, 0.2, 0.3], meta={"category": "test", "year": 2024} ) - vector = S3VectorsDocumentStore._document_to_s3_vector(doc) - restored = S3VectorsDocumentStore._s3_vector_to_document(vector) + vector = AmazonS3VectorsDocumentStore._document_to_s3_vector(doc) + restored = AmazonS3VectorsDocumentStore._s3_vector_to_document(vector) assert restored.id == doc.id assert restored.content == doc.content assert restored.embedding == doc.embedding @@ -288,10 +382,10 @@ def test_document_roundtrip_without_embedding(): placeholder = [-10.0] * 3 doc = Document(id="test-1", content="Hello world", meta={"category": "test"}) - vector = S3VectorsDocumentStore._document_to_s3_vector(doc, fallback_embedding=placeholder) + vector = AmazonS3VectorsDocumentStore._document_to_s3_vector(doc, fallback_embedding=placeholder) assert vector["data"]["float32"] == placeholder - restored = S3VectorsDocumentStore._s3_vector_to_document(vector, placeholder_embedding=placeholder) + restored = AmazonS3VectorsDocumentStore._s3_vector_to_document(vector, placeholder_embedding=placeholder) assert restored.embedding is None assert restored.content == doc.content assert restored.meta == doc.meta @@ -300,7 +394,7 @@ def test_document_roundtrip_without_embedding(): def test_s3_vector_to_document_keeps_embedding_matching_no_placeholder(): """A real embedding is never mistaken for the placeholder.""" vector = {"key": "test-1", "data": {"float32": [0.1, 0.2, 0.3]}, "metadata": {}} - doc = S3VectorsDocumentStore._s3_vector_to_document(vector, placeholder_embedding=[-10.0] * 3) + doc = AmazonS3VectorsDocumentStore._s3_vector_to_document(vector, placeholder_embedding=[-10.0] * 3) assert doc.embedding == [0.1, 0.2, 0.3] @@ -330,7 +424,7 @@ def test_s3_vector_to_document_keeps_embedding_matching_no_placeholder(): ], ) def test_is_supported_metadata_value(value, supported): - assert S3VectorsDocumentStore._is_supported_metadata_value(value) is supported + assert AmazonS3VectorsDocumentStore._is_supported_metadata_value(value) is supported def test_document_to_s3_vector_drops_unstorable_metadata(caplog): @@ -353,7 +447,7 @@ def test_document_to_s3_vector_drops_unstorable_metadata(caplog): "empty": [], }, ) - result = S3VectorsDocumentStore._document_to_s3_vector(doc) + result = AmazonS3VectorsDocumentStore._document_to_s3_vector(doc) assert result["metadata"] == {"_content": "Hello world", "keep_str": "x", "keep_list": ["a", "b"]} assert "['_split_overlap', 'empty', 'nested', 'nothing']" in caplog.text @@ -362,13 +456,13 @@ def test_document_to_s3_vector_drops_unstorable_metadata(caplog): def test_document_to_s3_vector_converts_tuples_to_lists(): """A tuple is a valid array for S3 Vectors, but botocore's validator only accepts `list`.""" doc = Document(id="test-1", meta={"tags": ("a", "b")}, embedding=[0.1]) - result = S3VectorsDocumentStore._document_to_s3_vector(doc) + result = AmazonS3VectorsDocumentStore._document_to_s3_vector(doc) assert result["metadata"]["tags"] == ["a", "b"] def test_document_to_s3_vector_skips_blob(caplog): doc = Document(id="1", content="x", embedding=[0.1], blob=ByteStream(data=b"binary", mime_type="application/pdf")) - vector = S3VectorsDocumentStore._document_to_s3_vector(doc) + vector = AmazonS3VectorsDocumentStore._document_to_s3_vector(doc) assert vector["metadata"] == {"_content": "x"} assert "does not store" in caplog.text @@ -381,7 +475,7 @@ def test_close_releases_the_client_and_allows_reopen(mock_boto3): client.get_index.return_value = {"index": {"dimension": 4}} mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = AmazonS3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") store._get_client() assert store._client is not None @@ -396,7 +490,7 @@ def test_close_releases_the_client_and_allows_reopen(mock_boto3): @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_close_is_a_noop_when_never_connected(mock_boto3): """close() before any call must not create a client just to close it.""" - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = AmazonS3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") store.close() mock_boto3.client.assert_not_called() @@ -410,7 +504,7 @@ def test_close_is_exception_safe(mock_boto3): client.close.side_effect = RuntimeError("boom") mock_boto3.client.return_value = client - store = S3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") + store = AmazonS3VectorsDocumentStore(vector_bucket_name="b", index_name="i", dimension=4, region_name="us-east-1") store._get_client() store.close() # must not raise assert store._client is None @@ -437,7 +531,7 @@ def assert_documents_are_equal(self, received: list[Document], expected: list[Do assert replace(r, embedding=None, score=None) == replace(e, embedding=None, score=None) assert r.embedding == pytest.approx(e.embedding, abs=1e-5) - def test_write_documents(self, document_store: S3VectorsDocumentStore) -> None: + def test_write_documents(self, document_store: AmazonS3VectorsDocumentStore) -> None: """ Default behaviour is OVERWRITE (S3 Vectors `put_vectors` is upsert), so writing the same document twice succeeds and returns 1 each time. @@ -446,7 +540,7 @@ def test_write_documents(self, document_store: S3VectorsDocumentStore) -> None: assert document_store.write_documents(docs) == 1 assert document_store.write_documents(docs) == 1 - def test_close_and_reopen(self, document_store: S3VectorsDocumentStore) -> None: + def test_close_and_reopen(self, document_store: AmazonS3VectorsDocumentStore) -> None: """close() drops the boto3 client; the store keeps working against the same index.""" document_store._get_client() assert document_store._client is not None @@ -457,7 +551,7 @@ def test_close_and_reopen(self, document_store: S3VectorsDocumentStore) -> None: assert document_store.count_documents() == 0 assert document_store._client is not None - def test_sentence_window_retriever(self, document_store: S3VectorsDocumentStore) -> None: + def test_sentence_window_retriever(self, document_store: AmazonS3VectorsDocumentStore) -> None: """The store composes with SentenceWindowRetriever, which pulls neighbours via filter_documents.""" splitter = DocumentSplitter(split_length=10, split_overlap=5, split_by="word") splitter.warm_up() @@ -475,13 +569,17 @@ def test_sentence_window_retriever(self, document_store: S3VectorsDocumentStore) ] document_store.write_documents(docs) - retrieved = S3VectorsEmbeddingRetriever(document_store=document_store).run(query_embedding=[0.1] * 768, top_k=1) + retrieved = AmazonS3VectorsEmbeddingRetriever(document_store=document_store).run( + query_embedding=[0.1] * 768, top_k=1 + ) result = SentenceWindowRetriever(document_store=document_store, window_size=2).run( retrieved_documents=[retrieved["documents"][0]] ) assert len(result["context_windows"]) == 1 - def test_write_documents_without_embedding_reads_back_as_none(self, document_store: S3VectorsDocumentStore) -> None: + def test_write_documents_without_embedding_reads_back_as_none( + self, document_store: AmazonS3VectorsDocumentStore + ) -> None: """ S3 Vectors stores a placeholder vector for embedding-less documents; `filter_documents` must strip it again rather than hand back a vector the caller never supplied. @@ -496,7 +594,7 @@ def test_write_documents_without_embedding_reads_back_as_none(self, document_sto assert docs[0].embedding is None @pytest.mark.skip(reason="S3 Vectors put_vectors is an upsert; only DuplicatePolicy.OVERWRITE is supported") - def test_write_documents_duplicate_fail(self, document_store: S3VectorsDocumentStore) -> None: ... + def test_write_documents_duplicate_fail(self, document_store: AmazonS3VectorsDocumentStore) -> None: ... @pytest.mark.skip(reason="S3 Vectors put_vectors is an upsert; only DuplicatePolicy.OVERWRITE is supported") - def test_write_documents_duplicate_skip(self, document_store: S3VectorsDocumentStore) -> None: ... + def test_write_documents_duplicate_skip(self, document_store: AmazonS3VectorsDocumentStore) -> None: ... diff --git a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py index 0c56623d93..8ffc8dd0a2 100644 --- a/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py +++ b/integrations/amazon_s3_vectors/tests/test_embedding_retriever.py @@ -8,41 +8,41 @@ from haystack.dataclasses import Document from haystack.document_stores.types import FilterPolicy -from haystack_integrations.components.retrievers.amazon_s3_vectors import S3VectorsEmbeddingRetriever -from haystack_integrations.document_stores.amazon_s3_vectors import S3VectorsDocumentStore +from haystack_integrations.components.retrievers.amazon_s3_vectors import AmazonS3VectorsEmbeddingRetriever +from haystack_integrations.document_stores.amazon_s3_vectors import AmazonS3VectorsDocumentStore def test_init_default(): - mock_store = Mock(spec=S3VectorsDocumentStore) - retriever = S3VectorsEmbeddingRetriever(document_store=mock_store) + mock_store = Mock(spec=AmazonS3VectorsDocumentStore) + retriever = AmazonS3VectorsEmbeddingRetriever(document_store=mock_store) assert retriever.document_store == mock_store assert retriever.filters == {} assert retriever.top_k == 10 assert retriever.filter_policy == FilterPolicy.REPLACE - retriever = S3VectorsEmbeddingRetriever(document_store=mock_store, filter_policy="replace") + retriever = AmazonS3VectorsEmbeddingRetriever(document_store=mock_store, filter_policy="replace") assert retriever.filter_policy == FilterPolicy.REPLACE with pytest.raises(ValueError): - S3VectorsEmbeddingRetriever(document_store=mock_store, filter_policy="invalid") + AmazonS3VectorsEmbeddingRetriever(document_store=mock_store, filter_policy="invalid") @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_to_dict(_mock_boto3): - store = S3VectorsDocumentStore( + store = AmazonS3VectorsDocumentStore( vector_bucket_name="test-bucket", index_name="test-index", dimension=768, region_name="us-east-1", create_bucket_and_index=False, ) - retriever = S3VectorsEmbeddingRetriever(document_store=store, top_k=5) + retriever = AmazonS3VectorsEmbeddingRetriever(document_store=store, top_k=5) d = retriever.to_dict() assert d == { - "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.S3VectorsEmbeddingRetriever", + "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.AmazonS3VectorsEmbeddingRetriever", "init_parameters": { "document_store": { - "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.AmazonS3VectorsDocumentStore", "init_parameters": { "vector_bucket_name": "test-bucket", "index_name": "test-index", @@ -66,10 +66,10 @@ def test_to_dict(_mock_boto3): @patch("haystack_integrations.document_stores.amazon_s3_vectors.document_store.boto3") def test_from_dict(_mock_boto3): data = { - "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.S3VectorsEmbeddingRetriever", + "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.AmazonS3VectorsEmbeddingRetriever", "init_parameters": { "document_store": { - "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.AmazonS3VectorsDocumentStore", "init_parameters": { "vector_bucket_name": "test-bucket", "index_name": "test-index", @@ -88,7 +88,7 @@ def test_from_dict(_mock_boto3): "filter_policy": "replace", }, } - retriever = S3VectorsEmbeddingRetriever.from_dict(data) + retriever = AmazonS3VectorsEmbeddingRetriever.from_dict(data) assert retriever.top_k == 5 assert retriever.filter_policy == FilterPolicy.REPLACE assert retriever.document_store.vector_bucket_name == "test-bucket" @@ -99,10 +99,10 @@ def test_from_dict(_mock_boto3): def test_from_dict_no_filter_policy(_mock_boto3): """Pipelines serialized with older versions may not have filter_policy.""" data = { - "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.S3VectorsEmbeddingRetriever", + "type": "haystack_integrations.components.retrievers.amazon_s3_vectors.embedding_retriever.AmazonS3VectorsEmbeddingRetriever", "init_parameters": { "document_store": { - "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.S3VectorsDocumentStore", + "type": "haystack_integrations.document_stores.amazon_s3_vectors.document_store.AmazonS3VectorsDocumentStore", "init_parameters": { "vector_bucket_name": "test-bucket", "index_name": "test-index", @@ -114,26 +114,37 @@ def test_from_dict_no_filter_policy(_mock_boto3): "top_k": 10, }, } - retriever = S3VectorsEmbeddingRetriever.from_dict(data) + retriever = AmazonS3VectorsEmbeddingRetriever.from_dict(data) assert retriever.filter_policy == FilterPolicy.REPLACE # default def test_run(): - mock_store = Mock(spec=S3VectorsDocumentStore) - mock_store._embedding_retrieval.return_value = [Document(content="Test doc", embedding=[0.1, 0.2])] - retriever = S3VectorsEmbeddingRetriever(document_store=mock_store) - res = retriever.run(query_embedding=[0.5, 0.7]) + init_filters = {"field": "meta.category", "operator": "==", "value": "news"} + runtime_filters = {"field": "meta.year", "operator": ">", "value": 2020} + mock_store = Mock(spec=AmazonS3VectorsDocumentStore) + mock_store._embedding_retrieval.return_value = [Document(content="Test doc")] + retriever = AmazonS3VectorsEmbeddingRetriever( + document_store=mock_store, + filter_policy=FilterPolicy.MERGE, + filters=init_filters, + ) + res = retriever.run(query_embedding=[0.5, 0.7], filters=runtime_filters) mock_store._embedding_retrieval.assert_called_once_with( query_embedding=[0.5, 0.7], - filters={}, + filters={"operator": "AND", "conditions": [init_filters, runtime_filters]}, top_k=10, ) - assert len(res["documents"]) == 1 - assert res["documents"][0].content == "Test doc" + assert res["documents"] == [Document(content="Test doc")] + + +def test_close(): + mock_store = Mock(spec=AmazonS3VectorsDocumentStore) + AmazonS3VectorsEmbeddingRetriever(document_store=mock_store).close() + mock_store.close.assert_called_once_with() @pytest.mark.integration -def test_embedding_retrieval(document_store: S3VectorsDocumentStore) -> None: +def test_embedding_retrieval(document_store: AmazonS3VectorsDocumentStore) -> None: docs = [ Document(id="r-1", content="First", embedding=[1.0, 0.0] + [0.0] * 766), Document(id="r-2", content="Second", embedding=[0.8, 0.6] + [0.0] * 766), @@ -150,7 +161,7 @@ def test_embedding_retrieval(document_store: S3VectorsDocumentStore) -> None: @pytest.mark.integration -def test_embedding_retrieval_with_metadata_filter(document_store: S3VectorsDocumentStore) -> None: +def test_embedding_retrieval_with_metadata_filter(document_store: AmazonS3VectorsDocumentStore) -> None: docs = [ Document(id="f-1", content="Sports article", embedding=[0.1] * 768, meta={"category": "sports"}), Document(id="f-2", content="News article", embedding=[0.1] * 768, meta={"category": "news"}), @@ -164,10 +175,10 @@ def test_embedding_retrieval_with_metadata_filter(document_store: S3VectorsDocum @pytest.mark.integration -def test_retriever_component(document_store: S3VectorsDocumentStore) -> None: +def test_retriever_component(document_store: AmazonS3VectorsDocumentStore) -> None: document_store.write_documents([Document(id="c-1", content="Hello", embedding=[0.1] * 768)]) - retriever = S3VectorsEmbeddingRetriever(document_store=document_store, top_k=5) + retriever = AmazonS3VectorsEmbeddingRetriever(document_store=document_store, top_k=5) result = retriever.run(query_embedding=[0.1] * 768) assert "documents" in result assert len(result["documents"]) > 0 diff --git a/integrations/amazon_s3_vectors/tests/test_filters.py b/integrations/amazon_s3_vectors/tests/test_filters.py index 16db36f691..08f9a94eee 100644 --- a/integrations/amazon_s3_vectors/tests/test_filters.py +++ b/integrations/amazon_s3_vectors/tests/test_filters.py @@ -35,27 +35,24 @@ def test_comparison_operators(): assert _normalize_filters(haystack_filter) == expected -def test_logical_operators(): - """AND/OR produce $and/$or with nested conditions.""" - assert _normalize_filters( - { - "operator": "AND", - "conditions": [ - {"field": "meta.a", "operator": "==", "value": 1}, - {"field": "meta.b", "operator": ">", "value": 2}, - ], - } - ) == {"$and": [{"a": {"$eq": 1}}, {"b": {"$gt": 2}}]} - - assert _normalize_filters( - { - "operator": "OR", - "conditions": [ - {"field": "meta.a", "operator": "==", "value": "x"}, - {"field": "meta.a", "operator": "==", "value": "y"}, - ], - } - ) == {"$or": [{"a": {"$eq": "x"}}, {"a": {"$eq": "y"}}]} +@pytest.mark.parametrize(("operator", "expected_operator"), [("AND", "$and"), ("OR", "$or")]) +def test_logical_operators(operator, expected_operator): + filters = { + "operator": operator, + "conditions": [ + {"field": "meta.a", "operator": "==", "value": 1}, + { + "operator": "OR", + "conditions": [ + {"field": "meta.b", "operator": ">", "value": 2}, + {"field": "meta.c", "operator": "==", "value": "x"}, + ], + }, + ], + } + assert _normalize_filters(filters) == { + expected_operator: [{"a": {"$eq": 1}}, {"$or": [{"b": {"$gt": 2}}, {"c": {"$eq": "x"}}]}] + } def test_meta_prefix_stripped(): @@ -67,23 +64,26 @@ def test_meta_prefix_stripped(): assert _normalize_filters({"field": "category", "operator": "==", "value": "news"}) == {"category": {"$eq": "news"}} -def test_unsupported_type_raises(): - with pytest.raises(FilterError): - _normalize_filters({"field": "meta.x", "operator": "==", "value": [1, 2, 3]}) - - -def test_in_requires_list(): - with pytest.raises(FilterError): - _normalize_filters({"field": "meta.x", "operator": "in", "value": "not-a-list"}) - - -def test_missing_keys_raise(): - with pytest.raises(FilterError): - _normalize_filters({"field": "meta.x", "value": "a"}) # missing operator - with pytest.raises(FilterError): - _normalize_filters({"field": "meta.x", "operator": "=="}) # missing value - with pytest.raises(FilterError): - _normalize_filters("not a dict") +@pytest.mark.parametrize( + ("filters", "error"), + [ + ("not a dict", "Filters must be a dictionary"), + ({"field": "meta.x", "value": "a"}, "'operator' key missing"), + ({"field": "meta.x", "operator": "=="}, "'value' key missing"), + ({"conditions": []}, "'operator' key missing"), + ({"operator": "AND"}, "'conditions' key missing"), + ({"operator": "XOR", "conditions": []}, "Unknown logical operator"), + ({"field": "meta.x", "operator": "unknown", "value": 1}, "Unknown comparison operator"), + ({"field": "meta.x", "operator": "==", "value": [1]}, "Unsupported type"), + ({"field": "meta.x", "operator": ">", "value": "text"}, "Unsupported type"), + ({"field": "meta.x", "operator": "in", "value": "text"}, "must be a list"), + ({"field": "meta.x", "operator": "not in", "value": "text"}, "must be a list"), + ({"field": "meta.x", "operator": "in", "value": [{}]}, "Unsupported type"), + ], +) +def test_invalid_filters(filters, error): + with pytest.raises(FilterError, match=error): + _normalize_filters(filters) def test_validate_filters():