diff --git a/.github/workflows/CI_license_compliance.yml b/.github/workflows/CI_license_compliance.yml index 54e79f2c87..c71935a661 100644 --- a/.github/workflows/CI_license_compliance.yml +++ b/.github/workflows/CI_license_compliance.yml @@ -12,7 +12,7 @@ on: env: PYTHON_VERSION: "3.10" - EXCLUDE_PACKAGES: "(?i)^(azure-identity|azure-search-documents|fastembed|llama-cpp-python|tqdm|psycopg|mistralai|pgvector|twelvelabs).*" + EXCLUDE_PACKAGES: "(?i)^(azure-identity|azure-search-documents|fastembed|llama-cpp-python|tqdm|psycopg|mistralai|pgvector|twelvelabs|mariadb).*" # Exclusions must be explicitly motivated # @@ -26,6 +26,7 @@ env: # - tqdm is MLP but there are no better alternatives # - psycopg is LGPL-3.0 but FOSSA is fine with it + # - mariadb (MariaDB Connector/Python) is LGPL-2.1 but FOSSA is fine with it jobs: license_check_direct: diff --git a/.github/workflows/mariadb.yml b/.github/workflows/mariadb.yml new file mode 100644 index 0000000000..7b48479885 --- /dev/null +++ b/.github/workflows/mariadb.yml @@ -0,0 +1,149 @@ +name: Test / mariadb + +on: + schedule: + - cron: "0 0 * * *" + pull_request: + paths: + - "integrations/mariadb/**" + - "!integrations/mariadb/*.md" + - ".github/workflows/mariadb.yml" + push: + branches: + - main + paths: + - "integrations/mariadb/**" + - "!integrations/mariadb/*.md" + - ".github/workflows/mariadb.yml" + +concurrency: + group: mariadb-${{ github.head_ref || github.sha }} + cancel-in-progress: true + +env: + PYTHONUNBUFFERED: "1" + FORCE_COLOR: "1" + TEST_MATRIX_OS: '["ubuntu-latest"]' + TEST_MATRIX_PYTHON: '["3.10", "3.14"]' + MARIADB_HOST: "127.0.0.1" + MARIADB_PORT: "3306" + MARIADB_DATABASE: "haystack" + MARIADB_USER: "root" + MARIADB_PASSWORD: "password" + +defaults: + run: + working-directory: integrations/mariadb + +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) }} + services: + mariadb: + image: mariadb:11.7 + env: + MARIADB_ROOT_PASSWORD: password + MARIADB_DATABASE: haystack + ports: + - 3306:3306 + + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: ${{ matrix.python-version }} + + - name: Install MariaDB Connector/C + run: sudo apt-get install -y libmariadb-dev + + - name: Install 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' + run: hatch run fmt-check && hatch run test:types + + - name: Run unit tests + run: hatch run test:unit-cov-retry + + - 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@5d8df5979747514c914e1c5a12335a7cf9a2745f # v4.1 + with: + GITHUB_TOKEN: ${{ github.token }} + COVERAGE_PATH: integrations/mariadb + SUBPROJECT_ID: mariadb + 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@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: coverage-comment-mariadb + path: python-coverage-comment-action-mariadb.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@5d8df5979747514c914e1c5a12335a7cf9a2745f # v4.1 + with: + GITHUB_TOKEN: ${{ github.token }} + COVERAGE_PATH: integrations/mariadb + SUBPROJECT_ID: mariadb-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@a65def0c8bf91d6520286ab34280151c76a5a008 # v1.1.0 + with: + slack-webhook-url: ${{ secrets.SLACK_WEBHOOK_URL_NOTIFICATIONS }} diff --git a/integrations/mariadb/CHANGELOG.md b/integrations/mariadb/CHANGELOG.md new file mode 100644 index 0000000000..076c57fa33 --- /dev/null +++ b/integrations/mariadb/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +## [unreleased] + +### Added + +- Add `MariaDBDocumentStore` with native VECTOR support (MariaDB 11.7+), MHNSW indexing, and full-text keyword search +- Add `MariaDBEmbeddingRetriever` for approximate nearest-neighbour vector search +- Add `MariaDBKeywordRetriever` for full-text search using `MATCH ... AGAINST` +- Support for Haystack metadata filtering via JSON path expressions diff --git a/integrations/mariadb/LICENSE.txt b/integrations/mariadb/LICENSE.txt new file mode 100644 index 0000000000..e866204bf1 --- /dev/null +++ b/integrations/mariadb/LICENSE.txt @@ -0,0 +1,73 @@ +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 passes 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 [yyyy] [name of copyright owner] + +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/mariadb/README.md b/integrations/mariadb/README.md new file mode 100644 index 0000000000..81404a924f --- /dev/null +++ b/integrations/mariadb/README.md @@ -0,0 +1,77 @@ +# MariaDB Document Store for Haystack + +[![PyPI - Version](https://img.shields.io/pypi/v/mariadb-haystack.svg)](https://pypi.org/project/mariadb-haystack) +[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mariadb-haystack.svg)](https://pypi.org/project/mariadb-haystack) + +An integration of [MariaDB 11.7+](https://mariadb.org/) with [Haystack](https://github.com/deepset-ai/haystack) that uses MariaDB's native `VECTOR` datatype and `MHNSW` indexing for vector search, and `MATCH ... AGAINST` for full-text keyword search. + +## Requirements + +- MariaDB 11.7 or later +- Python 3.10+ + +## Installation + +```bash +pip install mariadb-haystack +``` + +## Usage + +### Document Store + +```python +import os +from haystack.dataclasses import Document +from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore + +os.environ["MARIADB_USER"] = "root" +os.environ["MARIADB_PASSWORD"] = "password" + +store = MariaDBDocumentStore( + host="localhost", + port=3306, + database="haystack", + embedding_dimension=768, +) + +docs = [ + Document(content="Haystack is an open-source framework for building NLP pipelines."), + Document(content="MariaDB is a community-developed fork of MySQL."), +] +store.write_documents(docs) +print(store.count_documents()) # 2 +``` + +### Embedding Retriever + +```python +from haystack_integrations.components.retrievers.mariadb import MariaDBEmbeddingRetriever + +retriever = MariaDBEmbeddingRetriever(document_store=store, top_k=5) +results = retriever.run(query_embedding=[0.1] * 768) +``` + +### Keyword Retriever + +```python +from haystack_integrations.components.retrievers.mariadb import MariaDBKeywordRetriever + +retriever = MariaDBKeywordRetriever(document_store=store, top_k=5) +results = retriever.run(query="NLP pipelines") +``` + +## Running MariaDB locally with Docker + +```bash +docker run -d \ + --name mariadb-haystack \ + -e MARIADB_ROOT_PASSWORD=password \ + -e MARIADB_DATABASE=haystack \ + -p 3306:3306 \ + mariadb:11.7 +``` + +## License + +`mariadb-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license. diff --git a/integrations/mariadb/pydoc/config_docusaurus.yml b/integrations/mariadb/pydoc/config_docusaurus.yml new file mode 100644 index 0000000000..0605a1bf0c --- /dev/null +++ b/integrations/mariadb/pydoc/config_docusaurus.yml @@ -0,0 +1,15 @@ +loaders: + - modules: + - haystack_integrations.components.retrievers.mariadb.embedding_retriever + - haystack_integrations.components.retrievers.mariadb.keyword_retriever + - haystack_integrations.document_stores.mariadb.document_store + search_path: [../src] +processors: + - type: filter + documented_only: true + skip_empty_modules: true +renderer: + description: MariaDB integration for Haystack + id: integrations-mariadb + filename: mariadb.md + title: MariaDB diff --git a/integrations/mariadb/pyproject.toml b/integrations/mariadb/pyproject.toml new file mode 100644 index 0000000000..888c83f1f4 --- /dev/null +++ b/integrations/mariadb/pyproject.toml @@ -0,0 +1,169 @@ +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" + +[project] +name = "mariadb-haystack" +dynamic = ["version"] +description = "An integration of MariaDB 11.7+ (native VECTOR support) with Haystack" +readme = "README.md" +requires-python = ">=3.10" +license = "Apache-2.0" +keywords = [] +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>=2.28.0", + "mariadb>=1.1.0", +] + +[project.urls] +Source = "https://github.com/deepset-ai/haystack-core-integrations" +Documentation = "https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/mariadb/README.md" +Issues = "https://github.com/deepset-ai/haystack-core-integrations/issues" + +[tool.hatch.build.targets.wheel] +packages = ["src/haystack_integrations"] + +[tool.hatch.version] +source = "vcs" +tag-pattern = 'integrations\/mariadb-v(?P.*)' + +[tool.hatch.version.raw-options] +root = "../.." +git_describe_command = 'git describe --tags --match="integrations/mariadb-v[0-9]*"' + +[tool.hatch.envs.default] +installer = "uv" +# skip-install prevents hatch from installing the project (and thus the mariadb C extension) +# in the docs environment. The pydoc search_path points to src/ so modules are found directly. +skip-install = true +dependencies = ["haystack-pydoc-tools", "ruff", "haystack-ai>=2.28.0"] + +[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] +skip-install = false +dependencies = [ + "pytest", + "pytest-asyncio", + "pytest-cov", + "pytest-rerunfailures", + "mypy", + "pip", + "setuptools", +] + +[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.mariadb -p haystack_integrations.components.retrievers.mariadb {args}" + +[tool.mypy] +install_types = true +non_interactive = true +check_untyped_defs = true +disallow_incomplete_defs = true + +[[tool.mypy.overrides]] +module = [ + "mariadb.*", +] +ignore_missing_imports = true + +[tool.ruff] +line-length = 120 + +[tool.ruff.lint] +select = [ + "A", + "ANN", + "ARG", + "B", + "C", + "D102", + "D103", + "D205", + "D209", + "D213", + "D417", + "D419", + "DTZ", + "E", + "EM", + "F", + "FBT", + "I", + "ICN", + "ISC", + "N", + "PLC", + "PLE", + "PLR", + "PLW", + "Q", + "RUF", + "S", + "T", + "TID", + "UP", + "W", + "YTT", +] +ignore = [ + "B027", + "ANN401", + "FBT003", + "S105", + "S106", + "S107", + "C901", + "PLR0911", + "PLR0912", + "PLR0913", + "PLR0915", + "B008", + "S101", +] + +[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/**/*" = ["D", "PLR2004", "S101", "TID252", "ANN"] +"examples/**/*" = ["D", "T201"] + +[tool.coverage.run] +source = ["haystack_integrations"] +branch = true +relative_files = true +parallel = false + +[tool.coverage.report] +omit = ["*/tests/*", "*/__init__.py"] +show_missing = true +exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"] + +[tool.pytest.ini_options] +markers = ["integration: integration tests"] diff --git a/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/__init__.py b/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/__init__.py new file mode 100644 index 0000000000..8a60137804 --- /dev/null +++ b/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from haystack_integrations.components.retrievers.mariadb.embedding_retriever import MariaDBEmbeddingRetriever +from haystack_integrations.components.retrievers.mariadb.keyword_retriever import MariaDBKeywordRetriever + +__all__ = ["MariaDBEmbeddingRetriever", "MariaDBKeywordRetriever"] diff --git a/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/embedding_retriever.py b/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/embedding_retriever.py new file mode 100644 index 0000000000..1c60afeccd --- /dev/null +++ b/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/embedding_retriever.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Literal + +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.mariadb import MariaDBDocumentStore +from haystack_integrations.document_stores.mariadb.document_store import VALID_VECTOR_FUNCTIONS + + +@component +class MariaDBEmbeddingRetriever: + """ + Retrieves documents from `MariaDBDocumentStore` using vector similarity search. + + Uses MariaDB's native `VEC_DISTANCE_COSINE` or `VEC_DISTANCE_EUCLIDEAN` functions + with MHNSW indexing for efficient approximate nearest-neighbour search. + + ### Usage example + + ```python + from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore + from haystack_integrations.components.retrievers.mariadb import MariaDBEmbeddingRetriever + + store = MariaDBDocumentStore(host="localhost", database="haystack", embedding_dimension=768) + retriever = MariaDBEmbeddingRetriever(document_store=store, top_k=5) + result = retriever.run(query_embedding=[0.1] * 768) + documents = result["documents"] + ``` + """ + + def __init__( + self, + *, + document_store: MariaDBDocumentStore, + filters: dict[str, Any] | None = None, + top_k: int = 10, + vector_function: Literal["cosine_similarity", "l2_distance"] | None = None, + filter_policy: str | FilterPolicy = FilterPolicy.REPLACE, + ) -> None: + """ + Initialize the MariaDBEmbeddingRetriever. + + :param document_store: A `MariaDBDocumentStore` instance. + :param filters: Default Haystack metadata filters applied to every query. + :param top_k: Maximum number of documents to return. + :param vector_function: Override the store's default vector function for this retriever. + :param filter_policy: How runtime filters interact with init-time filters. + :raises ValueError: If `document_store` is not a `MariaDBDocumentStore` or + `vector_function` is invalid. + """ + if not isinstance(document_store, MariaDBDocumentStore): + msg = "document_store must be an instance of MariaDBDocumentStore" + raise ValueError(msg) + if vector_function and vector_function not in VALID_VECTOR_FUNCTIONS: + msg = f"vector_function must be one of {VALID_VECTOR_FUNCTIONS}" + raise ValueError(msg) + + self.document_store = document_store + self.filters = filters or {} + self.top_k = top_k + self.vector_function = vector_function or document_store.vector_function + self.filter_policy = ( + filter_policy if isinstance(filter_policy, FilterPolicy) else FilterPolicy.from_str(filter_policy) + ) + + def to_dict(self) -> dict[str, Any]: + """Serialize the component to a dictionary.""" + return default_to_dict( + self, + filters=self.filters, + top_k=self.top_k, + vector_function=self.vector_function, + filter_policy=self.filter_policy.value, + document_store=self.document_store.to_dict(), + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "MariaDBEmbeddingRetriever": + """Deserialize the component from a dictionary.""" + doc_store_params = data["init_parameters"]["document_store"] + data["init_parameters"]["document_store"] = MariaDBDocumentStore.from_dict(doc_store_params) + 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, + vector_function: Literal["cosine_similarity", "l2_distance"] | None = None, + ) -> dict[str, list[Document]]: + """ + Retrieve documents similar to the query embedding. + + :param query_embedding: The query vector. + :param filters: Runtime filters merged with init-time filters per `filter_policy`. + :param top_k: Override the retriever's `top_k`. + :param vector_function: Override the vector function for this call. + :returns: Dictionary with `"documents"` key containing the ranked results. + """ + filters = apply_filter_policy(self.filter_policy, self.filters, filters) + return { + "documents": self.document_store._embedding_retrieval( + query_embedding=query_embedding, + filters=filters, + top_k=top_k or self.top_k, + vector_function=vector_function or self.vector_function, + ) + } diff --git a/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/keyword_retriever.py b/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/keyword_retriever.py new file mode 100644 index 0000000000..16f04755e4 --- /dev/null +++ b/integrations/mariadb/src/haystack_integrations/components/retrievers/mariadb/keyword_retriever.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: 2023-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.mariadb import MariaDBDocumentStore + + +@component +class MariaDBKeywordRetriever: + """ + Retrieves documents from `MariaDBDocumentStore` using full-text keyword search. + + Uses MariaDB's `MATCH ... AGAINST` full-text search in natural language mode, + backed by a FULLTEXT index on the `content` column. + + ### Usage example + + ```python + from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore + from haystack_integrations.components.retrievers.mariadb import MariaDBKeywordRetriever + + store = MariaDBDocumentStore(host="localhost", database="haystack", embedding_dimension=768) + retriever = MariaDBKeywordRetriever(document_store=store, top_k=5) + result = retriever.run(query="climate change") + documents = result["documents"] + ``` + """ + + def __init__( + self, + *, + document_store: MariaDBDocumentStore, + filters: dict[str, Any] | None = None, + top_k: int = 10, + filter_policy: str | FilterPolicy = FilterPolicy.REPLACE, + ) -> None: + """ + Initialize the MariaDBKeywordRetriever. + + :param document_store: A `MariaDBDocumentStore` instance. + :param filters: Default Haystack metadata filters. + :param top_k: Maximum number of documents to return. + :param filter_policy: How runtime filters interact with init-time filters. + :raises ValueError: If `document_store` is not a `MariaDBDocumentStore`. + """ + if not isinstance(document_store, MariaDBDocumentStore): + msg = "document_store must be an instance of MariaDBDocumentStore" + 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 the 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]) -> "MariaDBKeywordRetriever": + """Deserialize the component from a dictionary.""" + doc_store_params = data["init_parameters"]["document_store"] + data["init_parameters"]["document_store"] = MariaDBDocumentStore.from_dict(doc_store_params) + 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: str, + filters: dict[str, Any] | None = None, + top_k: int | None = None, + ) -> dict[str, list[Document]]: + """ + Retrieve documents matching the query via full-text search. + + :param query: The keyword query string. + :param filters: Runtime filters merged with init-time filters per `filter_policy`. + :param top_k: Override the retriever's `top_k`. + :returns: Dictionary with `"documents"` key containing results ranked by relevance. + """ + filters = apply_filter_policy(self.filter_policy, self.filters, filters) + return { + "documents": self.document_store._keyword_retrieval( + query=query, + filters=filters, + top_k=top_k or self.top_k, + ) + } diff --git a/integrations/mariadb/src/haystack_integrations/components/retrievers/py.typed b/integrations/mariadb/src/haystack_integrations/components/retrievers/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/__init__.py b/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/__init__.py new file mode 100644 index 0000000000..3ddcc9d29f --- /dev/null +++ b/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from haystack_integrations.document_stores.mariadb.document_store import MariaDBDocumentStore + +__all__ = ["MariaDBDocumentStore"] diff --git a/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/document_store.py b/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/document_store.py new file mode 100644 index 0000000000..193f8d6862 --- /dev/null +++ b/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/document_store.py @@ -0,0 +1,482 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import struct +from dataclasses import replace +from typing import Any, Literal + +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 DocumentStore, DuplicatePolicy +from haystack.utils.auth import Secret, deserialize_secrets_inplace + +from .filters import _convert_filters_to_where_clause_and_params, _validate_filters + +logger = logging.getLogger(__name__) + +VALID_VECTOR_FUNCTIONS = ["cosine_similarity", "l2_distance"] + +VECTOR_FUNCTION_TO_SQL = { + "cosine_similarity": "VEC_DISTANCE_COSINE", + "l2_distance": "VEC_DISTANCE_EUCLIDEAN", +} + +CREATE_TABLE_STATEMENT = """ +CREATE TABLE IF NOT EXISTS `{table_name}` ( + id VARCHAR(128) PRIMARY KEY, + embedding VECTOR({embedding_dimension}) COMMENT 'MHNSW(M=16)', + content LONGTEXT, + blob_data LONGBLOB, + blob_meta JSON, + blob_mime_type VARCHAR(255), + meta JSON, + FULLTEXT KEY content_ft_idx (content) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 +""" + +INSERT_STATEMENT = """ +INSERT INTO `{table_name}` +(id, embedding, content, blob_data, blob_meta, blob_mime_type, meta) +VALUES (?, ?, ?, ?, ?, ?, ?) +""" + +INSERT_IGNORE_STATEMENT = """ +INSERT IGNORE INTO `{table_name}` +(id, embedding, content, blob_data, blob_meta, blob_mime_type, meta) +VALUES (?, ?, ?, ?, ?, ?, ?) +""" + +UPSERT_STATEMENT = """ +INSERT INTO `{table_name}` +(id, embedding, content, blob_data, blob_meta, blob_mime_type, meta) +VALUES (?, ?, ?, ?, ?, ?, ?) +ON DUPLICATE KEY UPDATE + embedding = VALUES(embedding), + content = VALUES(content), + blob_data = VALUES(blob_data), + blob_meta = VALUES(blob_meta), + blob_mime_type = VALUES(blob_mime_type), + meta = VALUES(meta) +""" + +KEYWORD_QUERY = """ +SELECT *, MATCH(content) AGAINST(? IN NATURAL LANGUAGE MODE) AS score +FROM `{table_name}` +{where_clause} +ORDER BY score DESC +LIMIT ? +""" + +EMBEDDING_QUERY = """ +SELECT *, {vec_func}(embedding, ?) AS score +FROM `{table_name}` +WHERE embedding IS NOT NULL +{extra_where} +ORDER BY score ASC +LIMIT ? +""" + + +class MariaDBDocumentStore(DocumentStore): + """ + A Document Store backed by MariaDB 11.7+ using native VECTOR support. + + Uses MariaDB's `VECTOR` datatype with `MHNSW` indexing for approximate nearest-neighbour + vector search, and `MATCH ... AGAINST` for full-text keyword search. + + Requires MariaDB 11.7 or later. Connect string parameters are passed individually; + credentials are managed via Haystack `Secret` (defaults to `MARIADB_USER` and + `MARIADB_PASSWORD` environment variables). + + ### Usage example + + ```python + from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore + + store = MariaDBDocumentStore( + host="localhost", + port=3306, + database="haystack", + embedding_dimension=768, + ) + store.write_documents(documents) + ``` + """ + + def __init__( + self, + *, + host: str = "localhost", + port: int = 3306, + database: str = "haystack", + user: Secret | str = Secret.from_env_var("MARIADB_USER"), + password: Secret | str = Secret.from_env_var("MARIADB_PASSWORD"), + table_name: str = "haystack_documents", + embedding_dimension: int = 768, + vector_function: Literal["cosine_similarity", "l2_distance"] = "cosine_similarity", + recreate_table: bool = False, + ) -> None: + """ + Initialize the MariaDBDocumentStore. + + :param host: MariaDB host. Defaults to `"localhost"`. + :param port: MariaDB port. Defaults to `3306`. + :param database: Database name. + :param user: Database user. Reads `MARIADB_USER` env var by default. + :param password: Database password. Reads `MARIADB_PASSWORD` env var by default. + :param table_name: Table used to store documents. + :param embedding_dimension: Dimension of embedding vectors. + :param vector_function: Similarity function — `"cosine_similarity"` or `"l2_distance"`. + :param recreate_table: Drop and recreate the table on init. **Deletes all data.** + """ + self._connection: Any = None + self._cursor: Any = None + self._table_initialized = False + + if vector_function not in VALID_VECTOR_FUNCTIONS: + msg = f"vector_function must be one of {VALID_VECTOR_FUNCTIONS}, got '{vector_function}'" + raise ValueError(msg) + + self.host = host + self.port = port + self.database = database + self.user = Secret.from_token(user) if isinstance(user, str) else user + self.password = Secret.from_token(password) if isinstance(password, str) else password + self.table_name = table_name + self.embedding_dimension = embedding_dimension + self.vector_function = vector_function + self.recreate_table = recreate_table + + def to_dict(self) -> dict[str, Any]: + """Serialize the component to a dictionary.""" + return default_to_dict( + self, + host=self.host, + port=self.port, + database=self.database, + user=self.user.to_dict(), + password=self.password.to_dict(), + table_name=self.table_name, + embedding_dimension=self.embedding_dimension, + vector_function=self.vector_function, + recreate_table=self.recreate_table, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> MariaDBDocumentStore: + """Deserialize the component from a dictionary.""" + deserialize_secrets_inplace(data["init_parameters"], ["user", "password"]) + return default_from_dict(cls, data) + + def _ensure_connection(self) -> None: + """Lazily establish the DB connection and initialize the table.""" + if self._connection is not None and self._cursor is not None: + try: + self._connection.ping() + return + except Exception: + self._close_connection() + + import mariadb # noqa: PLC0415 + + try: + self._connection = mariadb.connect( + host=self.host, + port=self.port, + database=self.database, + user=self.user.resolve_value() or "", + password=self.password.resolve_value() or "", + autocommit=True, + ) + except mariadb.Error as e: + msg = ( + f"Failed to connect to MariaDB at {self.host}:{self.port}/{self.database}. " + "Ensure MariaDB 11.7+ is running and credentials are correct." + ) + raise DocumentStoreError(msg) from e + + self._cursor = self._connection.cursor(dictionary=True) + + if not self._table_initialized: + self._initialize_table() + + def _initialize_table(self) -> None: + import mariadb # noqa: PLC0415 + + if self.recreate_table: + self._drop_table() + + sql = CREATE_TABLE_STATEMENT.format( + table_name=self.table_name, + embedding_dimension=self.embedding_dimension, + ) + try: + self._cursor.execute(sql) + self._table_initialized = True + except mariadb.Error as e: + msg = f"Could not create table '{self.table_name}'" + raise DocumentStoreError(msg) from e + + def _drop_table(self) -> None: + import mariadb # noqa: PLC0415 + + try: + self._cursor.execute(f"DROP TABLE IF EXISTS `{self.table_name}`") + self._table_initialized = False + except mariadb.Error as e: + msg = f"Could not drop table '{self.table_name}'" + raise DocumentStoreError(msg) from e + + def _close_connection(self) -> None: + if self._cursor is not None: + try: + self._cursor.close() + except Exception: # noqa: S110 + pass + self._cursor = None + if self._connection is not None: + try: + self._connection.close() + except Exception: # noqa: S110 + pass + self._connection = None + self._table_initialized = False + + def close(self) -> None: + """Close the database connection.""" + self._close_connection() + + def __del__(self) -> None: + self._close_connection() + + def count_documents(self) -> int: + """Return the number of documents in the store.""" + self._ensure_connection() + self._cursor.execute(f"SELECT COUNT(*) AS cnt FROM `{self.table_name}`") # noqa: S608 + row = self._cursor.fetchone() + return row["cnt"] if row else 0 + + def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]: + """ + Return documents matching the given Haystack filters. + + :param filters: Optional Haystack metadata filters. + :returns: List of matching Documents. + """ + _validate_filters(filters) + self._ensure_connection() + + sql = f"SELECT * FROM `{self.table_name}`" # noqa: S608 + params: list[Any] = [] + + if filters: + where_clause, params = _convert_filters_to_where_clause_and_params(filters) + sql += where_clause + + self._cursor.execute(sql, params) + records = self._cursor.fetchall() + return _rows_to_documents(records) + + def write_documents(self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.FAIL) -> int: + """ + Write documents to the store. + + :param documents: Documents to write. + :param policy: `FAIL`, `OVERWRITE`, or `SKIP`. + :returns: Number of documents written. + """ + if not documents: + return 0 + if not isinstance(documents[0], Document): + msg = "param 'documents' must be a list of Document objects" + raise ValueError(msg) + + self._ensure_connection() + + if policy == DuplicatePolicy.OVERWRITE: + sql = UPSERT_STATEMENT.format(table_name=self.table_name) + elif policy == DuplicatePolicy.SKIP: + sql = INSERT_IGNORE_STATEMENT.format(table_name=self.table_name) + else: + sql = INSERT_STATEMENT.format(table_name=self.table_name) + + import mariadb # noqa: PLC0415 + + written = 0 + for doc in documents: + row = _document_to_row(doc) + try: + self._cursor.execute(sql, row) + written += 1 + except mariadb.IntegrityError as e: + if policy == DuplicatePolicy.FAIL: + msg = f"Document with id '{doc.id}' already exists" + raise DuplicateDocumentError(msg) from e + + return written + + def delete_documents(self, document_ids: list[str]) -> None: + """ + Delete documents by ID. + + :param document_ids: IDs to delete. + """ + import mariadb # noqa: PLC0415 + + if not document_ids: + return + self._ensure_connection() + placeholders = ", ".join(["?"] * len(document_ids)) + try: + self._cursor.execute( + f"DELETE FROM `{self.table_name}` WHERE id IN ({placeholders})", # noqa: S608 + tuple(document_ids), + ) + except mariadb.Error as e: + msg = "Failed to delete documents" + raise DocumentStoreError(msg) from e + + def _embedding_retrieval( + self, + query_embedding: list[float], + *, + filters: dict[str, Any] | None = None, + top_k: int = 10, + vector_function: str | None = None, + ) -> list[Document]: + """ + Retrieve documents by vector similarity. + + :param query_embedding: Query vector. + :param filters: Optional Haystack filters. + :param top_k: Maximum results. + :param vector_function: Override the store's vector function for this query. + :returns: List of Documents ordered by similarity (most similar first). + """ + _validate_filters(filters) + self._ensure_connection() + + vec_func = VECTOR_FUNCTION_TO_SQL[vector_function or self.vector_function] + embedding_bytes = _embedding_to_bytes(query_embedding) + + extra_where = "" + params: list[Any] = [embedding_bytes] + + if filters: + where_clause, filter_params = _convert_filters_to_where_clause_and_params(filters, operator="AND") + extra_where = where_clause + params.extend(filter_params) + + params.append(top_k) + + sql = EMBEDDING_QUERY.format( + vec_func=vec_func, + table_name=self.table_name, + extra_where=extra_where, + ) + + self._cursor.execute(sql, params) + records = self._cursor.fetchall() + + docs = _rows_to_documents(records) + # VEC_DISTANCE_* returns distance (lower = more similar); convert to a positive score + docs = [ + replace(doc, score=float(1.0 - record["score"]) if "COSINE" in vec_func else float(-record["score"])) + if record.get("score") is not None + else doc + for doc, record in zip(docs, records, strict=True) + ] + return docs + + def _keyword_retrieval( + self, + query: str, + *, + filters: dict[str, Any] | None = None, + top_k: int = 10, + ) -> list[Document]: + """ + Retrieve documents by full-text keyword search. + + :param query: Search query string. + :param filters: Optional Haystack filters. + :param top_k: Maximum results. + :returns: List of Documents ordered by relevance. + """ + _validate_filters(filters) + self._ensure_connection() + + params: list[Any] = [query] + where_clause = "" + + if filters: + where_clause, filter_params = _convert_filters_to_where_clause_and_params(filters) + params = [query, *filter_params] + + params.append(top_k) + + sql = KEYWORD_QUERY.format(table_name=self.table_name, where_clause=where_clause) + self._cursor.execute(sql, params) + records = self._cursor.fetchall() + + docs = _rows_to_documents(records) + docs = [ + replace(doc, score=float(record.get("score") or 0.0)) for doc, record in zip(docs, records, strict=True) + ] + return docs + + +def _embedding_to_bytes(embedding: list[float]) -> bytes: + return struct.pack(f"{len(embedding)}f", *embedding) + + +def _bytes_to_embedding(data: bytes | bytearray) -> list[float]: + n = len(data) // 4 + return list(struct.unpack(f"{n}f", data)) + + +def _document_to_row(doc: Document) -> tuple: + embedding_bytes = _embedding_to_bytes(doc.embedding) if doc.embedding is not None else None + blob_data = doc.blob.data if doc.blob else None + blob_meta = json.dumps(doc.blob.meta) if doc.blob and doc.blob.meta else None + blob_mime_type = doc.blob.mime_type if doc.blob else None + meta = json.dumps(doc.meta) if doc.meta else json.dumps({}) + return (doc.id, embedding_bytes, doc.content, blob_data, blob_meta, blob_mime_type, meta) + + +def _rows_to_documents(records: list[dict[str, Any]]) -> list[Document]: + docs = [] + for record in records: + row = dict(record) + row.pop("score", None) + + blob_data = row.pop("blob_data", None) + blob_meta_raw = row.pop("blob_meta", None) + blob_mime_type = row.pop("blob_mime_type", None) + + if isinstance(row.get("meta"), str): + row["meta"] = json.loads(row["meta"]) + elif row.get("meta") is None: + row.pop("meta", None) + + emb = row.get("embedding") + if emb is not None: + if hasattr(emb, "tolist"): + row["embedding"] = emb.tolist() + elif isinstance(emb, (bytes, bytearray)): + row["embedding"] = _bytes_to_embedding(emb) + + doc = Document.from_dict(row) + + if blob_data: + if isinstance(blob_meta_raw, str): + blob_meta_raw = json.loads(blob_meta_raw) + doc = replace(doc, blob=ByteStream(data=blob_data, meta=blob_meta_raw or {}, mime_type=blob_mime_type)) + + docs.append(doc) + return docs diff --git a/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/filters.py b/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/filters.py new file mode 100644 index 0000000000..8248e00e44 --- /dev/null +++ b/integrations/mariadb/src/haystack_integrations/document_stores/mariadb/filters.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from datetime import datetime +from typing import Any + +from haystack.errors import FilterError + +NO_VALUE = "no_value" + + +def _validate_filters(filters: dict[str, Any] | None = None) -> None: + if filters: + if not isinstance(filters, dict): + msg = "Filters must be a dictionary" + raise TypeError(msg) + if "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) + + +def _convert_filters_to_where_clause_and_params( + filters: dict[str, Any], + operator: str = "WHERE", +) -> tuple[str, list[Any]]: + """Convert Haystack filters to a MariaDB WHERE clause string and parameter list.""" + if "field" in filters: + clause, values = _parse_comparison_condition(filters) + else: + clause, values = _parse_logical_condition(filters) + + params = [v for v in values if v != NO_VALUE] + return f" {operator} {clause}", params + + +def _parse_logical_condition(condition: dict[str, Any]) -> tuple[str, list[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"] + if operator not in ("AND", "OR"): + msg = f"Unknown logical operator '{operator}'. Valid operators are: 'AND', 'OR'" + raise FilterError(msg) + + parts, all_params = [], [] + for c in condition["conditions"]: + if "field" in c: + clause, vals = _parse_comparison_condition(c) + else: + clause, vals = _parse_logical_condition(c) + parts.append(clause) + all_params.extend(vals) + + joined = f" {operator} ".join(parts) + return f"({joined})", all_params + + +def _parse_comparison_condition(condition: dict[str, Any]) -> tuple[str, list[Any]]: + 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"] + if operator not in COMPARISON_OPERATORS: + msg = f"Unknown comparison operator '{operator}'. Valid operators are: {list(COMPARISON_OPERATORS.keys())}" + raise FilterError(msg) + + value: Any = condition["value"] + sql_field = _build_field_expr(field, value) + clause, val = COMPARISON_OPERATORS[operator](sql_field, value) + # _in/_not_in return the value list directly; flatten it so params stay flat + if isinstance(val, list): + return clause, val + return clause, [val] + + +def _build_field_expr(field: str, value: Any) -> str: + """Return a safe SQL expression for a document field or meta JSON path.""" + if field.startswith("meta."): + field_name = field.split(".", 1)[1] + # JSON_UNQUOTE(JSON_EXTRACT(meta, '$.field')) returns the value as a string. + # Cast to numeric types when the comparison value is numeric. + base = f"JSON_UNQUOTE(JSON_EXTRACT(meta, '$.{field_name}'))" + if isinstance(value, bool): + return base + if isinstance(value, int): + return f"CAST({base} AS SIGNED)" + if isinstance(value, float): + return f"CAST({base} AS DECIMAL(65,30))" + if isinstance(value, list) and value: + first = value[0] + if isinstance(first, int) and not isinstance(first, bool): + return f"CAST({base} AS SIGNED)" + if isinstance(first, float): + return f"CAST({base} AS DECIMAL(65,30))" + return base + # Top-level document field (id, content, etc.) — backtick-quoted + return f"`{field}`" + + +def _equal(field: str, value: Any) -> tuple[str, Any]: + if value is None: + return f"{field} IS NULL", NO_VALUE + return f"{field} = ?", value + + +def _not_equal(field: str, value: Any) -> tuple[str, Any]: + return f"({field} IS NULL OR {field} != ?)", value + + +def _greater_than(field: str, value: Any) -> tuple[str, Any]: + _check_comparable(value, ">") + return f"{field} > ?", value + + +def _greater_than_equal(field: str, value: Any) -> tuple[str, Any]: + _check_comparable(value, ">=") + return f"{field} >= ?", value + + +def _less_than(field: str, value: Any) -> tuple[str, Any]: + _check_comparable(value, "<") + return f"{field} < ?", value + + +def _less_than_equal(field: str, value: Any) -> tuple[str, Any]: + _check_comparable(value, "<=") + return f"{field} <= ?", value + + +def _in(field: str, value: Any) -> tuple[str, list]: + if not isinstance(value, list): + msg = f"'in' operator requires a list value, got {type(value)}" + raise FilterError(msg) + placeholders = ", ".join(["?"] * len(value)) + return f"{field} IN ({placeholders})", value + + +def _not_in(field: str, value: Any) -> tuple[str, list]: + if not isinstance(value, list): + msg = f"'not in' operator requires a list value, got {type(value)}" + raise FilterError(msg) + placeholders = ", ".join(["?"] * len(value)) + return f"({field} IS NULL OR {field} NOT IN ({placeholders}))", value + + +def _like(field: str, value: Any) -> tuple[str, Any]: + if not isinstance(value, str): + msg = f"'like' operator requires a string value, got {type(value)}" + raise FilterError(msg) + return f"{field} LIKE ?", value + + +def _not_like(field: str, value: Any) -> tuple[str, Any]: + if not isinstance(value, str): + msg = f"'not like' operator requires a string value, got {type(value)}" + raise FilterError(msg) + return f"{field} NOT LIKE ?", value + + +def _check_comparable(value: Any, op: str) -> None: + if isinstance(value, str): + try: + datetime.fromisoformat(value) + except (ValueError, TypeError) as exc: + msg = f"Cannot compare strings with '{op}'. Strings are only comparable if they are ISO formatted dates." + raise FilterError(msg) from exc + if isinstance(value, list): + msg = f"Filter value cannot be a list when using '{op}'" + raise FilterError(msg) + + +COMPARISON_OPERATORS = { + "==": _equal, + "!=": _not_equal, + ">": _greater_than, + ">=": _greater_than_equal, + "<": _less_than, + "<=": _less_than_equal, + "in": _in, + "not in": _not_in, + "like": _like, + "not like": _not_like, +} diff --git a/integrations/mariadb/src/haystack_integrations/document_stores/py.typed b/integrations/mariadb/src/haystack_integrations/document_stores/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/integrations/mariadb/tests/__init__.py b/integrations/mariadb/tests/__init__.py new file mode 100644 index 0000000000..e873bc3327 --- /dev/null +++ b/integrations/mariadb/tests/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/mariadb/tests/test_document_store.py b/integrations/mariadb/tests/test_document_store.py new file mode 100644 index 0000000000..3ca54115c7 --- /dev/null +++ b/integrations/mariadb/tests/test_document_store.py @@ -0,0 +1,401 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import os +from unittest.mock import MagicMock + +import pytest +from haystack.dataclasses import ByteStream, Document +from haystack.document_stores.errors import DuplicateDocumentError +from haystack.document_stores.types import DuplicatePolicy + +from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore +from haystack_integrations.document_stores.mariadb.document_store import ( + _bytes_to_embedding, + _document_to_row, + _embedding_to_bytes, + _rows_to_documents, +) + +# --------------------------------------------------------------------------- +# Helper: build a store with a mocked DB connection +# --------------------------------------------------------------------------- + + +def _mock_store(**kwargs) -> MariaDBDocumentStore: + store = MariaDBDocumentStore( + host="localhost", + database="test_db", + user="root", + password="secret", + **kwargs, + ) + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_cursor.dictionary = True + mock_conn.cursor.return_value = mock_cursor + store._connection = mock_conn + store._cursor = mock_cursor + store._table_initialized = True + return store + + +# --------------------------------------------------------------------------- +# Serialization +# --------------------------------------------------------------------------- + + +class TestSerialization: + def test_to_dict(self, monkeypatch): + monkeypatch.setenv("MARIADB_USER", "test_user") + monkeypatch.setenv("MARIADB_PASSWORD", "test_pass") + store = MariaDBDocumentStore() + d = store.to_dict() + assert d["type"] == "haystack_integrations.document_stores.mariadb.document_store.MariaDBDocumentStore" + params = d["init_parameters"] + assert params["host"] == "localhost" + assert params["embedding_dimension"] == 768 + assert params["vector_function"] == "cosine_similarity" + + def test_from_dict_roundtrip(self, monkeypatch): + monkeypatch.setenv("MARIADB_USER", "test_user") + monkeypatch.setenv("MARIADB_PASSWORD", "test_pass") + store = MariaDBDocumentStore(embedding_dimension=512) + d = store.to_dict() + restored = MariaDBDocumentStore.from_dict(d) + assert restored.embedding_dimension == 512 + assert restored.host == "localhost" + + def test_invalid_vector_function_raises(self): + with pytest.raises(ValueError, match="vector_function must be one of"): + MariaDBDocumentStore(user="u", password="p", vector_function="bad_func") + + +# --------------------------------------------------------------------------- +# Embedding helpers +# --------------------------------------------------------------------------- + + +class TestEmbeddingHelpers: + def test_round_trip_embedding(self): + original = [0.1, 0.2, 0.3, -0.5] + raw = _embedding_to_bytes(original) + recovered = _bytes_to_embedding(raw) + assert len(recovered) == len(original) + for a, b in zip(recovered, original, strict=True): + assert abs(a - b) < 1e-6 + + def test_embedding_bytes_length(self): + emb = [1.0] * 8 + raw = _embedding_to_bytes(emb) + assert len(raw) == 8 * 4 # 4 bytes per float32 + + +# --------------------------------------------------------------------------- +# Document conversion helpers +# --------------------------------------------------------------------------- + + +class TestDocumentHelpers: + def test_document_to_row_simple(self): + doc = Document(id="abc", content="hello", meta={"x": 1}) + row = _document_to_row(doc) + assert row[0] == "abc" + assert row[1] is None # no embedding + assert row[2] == "hello" + assert row[3] is None # no blob + assert row[6] == '{"x": 1}' + + def test_document_to_row_with_embedding(self): + doc = Document(id="d1", content="text", embedding=[0.5, 0.5]) + row = _document_to_row(doc) + assert isinstance(row[1], bytes) + + def test_document_to_row_with_blob(self): + blob = ByteStream(data=b"raw bytes", meta={"k": "v"}, mime_type="application/octet-stream") + doc = Document(id="d2", content=None, blob=blob) + row = _document_to_row(doc) + assert row[3] == b"raw bytes" + assert row[4] == '{"k": "v"}' + assert row[5] == "application/octet-stream" + + def test_rows_to_documents_simple(self): + records = [ + { + "id": "x1", + "content": "text", + "embedding": None, + "blob_data": None, + "blob_meta": None, + "blob_mime_type": None, + "meta": '{"a": 1}', + } + ] + docs = _rows_to_documents(records) + assert len(docs) == 1 + assert docs[0].id == "x1" + assert docs[0].meta == {"a": 1} + + def test_rows_to_documents_bytes_embedding(self): + emb = [0.1, 0.9] + raw = _embedding_to_bytes(emb) + records = [ + { + "id": "e1", + "content": "hi", + "embedding": raw, + "blob_data": None, + "blob_meta": None, + "blob_mime_type": None, + "meta": "{}", + } + ] + docs = _rows_to_documents(records) + assert docs[0].embedding is not None + assert len(docs[0].embedding) == 2 + + def test_rows_to_documents_strips_score(self): + records = [ + { + "id": "s1", + "content": "text", + "embedding": None, + "blob_data": None, + "blob_meta": None, + "blob_mime_type": None, + "meta": "{}", + "score": 0.99, + } + ] + docs = _rows_to_documents(records) + # Document.score should not be set from the row (no score field in Document) + assert docs[0].id == "s1" + + def test_rows_to_documents_with_blob(self): + records = [ + { + "id": "b1", + "content": None, + "embedding": None, + "blob_data": b"data", + "blob_meta": '{"k": "v"}', + "blob_mime_type": "image/png", + "meta": "{}", + } + ] + docs = _rows_to_documents(records) + assert docs[0].blob is not None + assert docs[0].blob.data == b"data" + assert docs[0].blob.mime_type == "image/png" + + +# --------------------------------------------------------------------------- +# count_documents +# --------------------------------------------------------------------------- + + +class TestCountDocuments: + def test_count_documents(self): + store = _mock_store() + store._cursor.fetchone.return_value = {"cnt": 5} + assert store.count_documents() == 5 + + def test_count_empty(self): + store = _mock_store() + store._cursor.fetchone.return_value = {"cnt": 0} + assert store.count_documents() == 0 + + +# --------------------------------------------------------------------------- +# write_documents +# --------------------------------------------------------------------------- + + +class TestWriteDocuments: + def test_write_empty_returns_zero(self): + store = _mock_store() + assert store.write_documents([]) == 0 + + def test_write_single_doc(self): + store = _mock_store() + doc = Document(content="test") + result = store.write_documents([doc]) + assert result == 1 + store._cursor.execute.assert_called_once() + + def test_write_overwrite_uses_upsert(self): + store = _mock_store() + doc = Document(content="test") + store.write_documents([doc], policy=DuplicatePolicy.OVERWRITE) + call_args = store._cursor.execute.call_args[0] + assert "ON DUPLICATE KEY UPDATE" in call_args[0] + + def test_write_skip_uses_insert_ignore(self): + store = _mock_store() + doc = Document(content="test") + store.write_documents([doc], policy=DuplicatePolicy.SKIP) + call_args = store._cursor.execute.call_args[0] + assert "INSERT IGNORE" in call_args[0] + + def test_write_fail_raises_on_duplicate(self): + import mariadb # noqa: PLC0415 + + store = _mock_store() + store._cursor.execute.side_effect = mariadb.IntegrityError("Duplicate entry") + doc = Document(id="dup", content="test") + with pytest.raises(DuplicateDocumentError): + store.write_documents([doc], policy=DuplicatePolicy.FAIL) + + def test_write_invalid_type_raises(self): + store = _mock_store() + with pytest.raises(ValueError, match="list of Document objects"): + store.write_documents(["not a doc"]) # type: ignore[list-item] + + +# --------------------------------------------------------------------------- +# delete_documents +# --------------------------------------------------------------------------- + + +class TestDeleteDocuments: + def test_delete_empty_noop(self): + store = _mock_store() + store.delete_documents([]) + store._cursor.execute.assert_not_called() + + def test_delete_calls_execute(self): + store = _mock_store() + store.delete_documents(["id1", "id2"]) + call_args = store._cursor.execute.call_args[0] + assert "DELETE FROM" in call_args[0] + assert "IN (?, ?)" in call_args[0] + + +# --------------------------------------------------------------------------- +# filter_documents +# --------------------------------------------------------------------------- + + +class TestFilterDocuments: + def test_filter_no_filters(self): + store = _mock_store() + store._cursor.fetchall.return_value = [] + result = store.filter_documents() + assert result == [] + call_args = store._cursor.execute.call_args[0] + assert "WHERE" not in call_args[0] + + def test_filter_with_filters(self): + store = _mock_store() + store._cursor.fetchall.return_value = [] + store.filter_documents(filters={"field": "meta.x", "operator": "==", "value": 1}) + call_args = store._cursor.execute.call_args[0] + assert "WHERE" in call_args[0] + + +# --------------------------------------------------------------------------- +# Integration tests (require real MariaDB) +# --------------------------------------------------------------------------- + +MARIADB_HOST = os.environ.get("MARIADB_HOST", "localhost") +MARIADB_PORT = int(os.environ.get("MARIADB_PORT", "3306")) +MARIADB_DB = os.environ.get("MARIADB_DATABASE", "haystack") +MARIADB_USER = os.environ.get("MARIADB_USER", "root") +MARIADB_PASSWORD = os.environ.get("MARIADB_PASSWORD", "password") + + +@pytest.fixture +def integration_store(): + store = MariaDBDocumentStore( + host=MARIADB_HOST, + port=MARIADB_PORT, + database=MARIADB_DB, + user=MARIADB_USER, + password=MARIADB_PASSWORD, + table_name="test_haystack_docs", + embedding_dimension=4, + recreate_table=True, + ) + yield store + store.delete_documents(store.filter_documents() and [d.id for d in store.filter_documents()]) + try: + store._cursor.execute("DROP TABLE IF EXISTS `test_haystack_docs`") + except Exception: # noqa: S110 + pass + store.close() + + +@pytest.mark.integration +class TestIntegrationDocumentStore: + def test_write_and_count(self, integration_store): + docs = [Document(content="Hello"), Document(content="World")] + written = integration_store.write_documents(docs) + assert written == 2 + assert integration_store.count_documents() == 2 + + def test_filter_by_content_id(self, integration_store): + doc = Document(id="fixed-id", content="unique content") + integration_store.write_documents([doc]) + results = integration_store.filter_documents({"field": "id", "operator": "==", "value": "fixed-id"}) + assert len(results) == 1 + assert results[0].id == "fixed-id" + + def test_filter_by_meta(self, integration_store): + docs = [ + Document(content="A", meta={"category": "cat"}), + Document(content="B", meta={"category": "dog"}), + ] + integration_store.write_documents(docs) + results = integration_store.filter_documents({"field": "meta.category", "operator": "==", "value": "cat"}) + assert all(d.meta.get("category") == "cat" for d in results) + + def test_delete_documents(self, integration_store): + doc = Document(id="to-delete", content="bye") + integration_store.write_documents([doc]) + assert integration_store.count_documents() >= 1 + integration_store.delete_documents(["to-delete"]) + results = integration_store.filter_documents({"field": "id", "operator": "==", "value": "to-delete"}) + assert results == [] + + def test_duplicate_fail(self, integration_store): + doc = Document(id="dup-test", content="once") + integration_store.write_documents([doc]) + with pytest.raises(DuplicateDocumentError): + integration_store.write_documents([doc], policy=DuplicatePolicy.FAIL) + + def test_duplicate_overwrite(self, integration_store): + doc = Document(id="ow-test", content="original") + integration_store.write_documents([doc]) + doc2 = Document(id="ow-test", content="updated") + integration_store.write_documents([doc2], policy=DuplicatePolicy.OVERWRITE) + results = integration_store.filter_documents({"field": "id", "operator": "==", "value": "ow-test"}) + assert results[0].content == "updated" + + def test_embedding_retrieval(self, integration_store): + docs = [ + Document(content="red apple", embedding=[1.0, 0.0, 0.0, 0.0]), + Document(content="green tea", embedding=[0.0, 1.0, 0.0, 0.0]), + ] + integration_store.write_documents(docs) + results = integration_store._embedding_retrieval(query_embedding=[1.0, 0.0, 0.0, 0.0], top_k=1) + assert len(results) == 1 + assert results[0].content == "red apple" + + def test_keyword_retrieval(self, integration_store): + docs = [ + Document(content="machine learning algorithms"), + Document(content="database management systems"), + ] + integration_store.write_documents(docs) + results = integration_store._keyword_retrieval(query="machine learning", top_k=5) + assert len(results) >= 1 + assert any("machine" in d.content for d in results) + + def test_write_with_blob(self, integration_store): + blob = ByteStream(data=b"binary content", meta={"type": "test"}, mime_type="text/plain") + doc = Document(id="blob-doc", blob=blob) + integration_store.write_documents([doc]) + results = integration_store.filter_documents({"field": "id", "operator": "==", "value": "blob-doc"}) + assert results[0].blob is not None + assert results[0].blob.data == b"binary content" diff --git a/integrations/mariadb/tests/test_filters.py b/integrations/mariadb/tests/test_filters.py new file mode 100644 index 0000000000..e941f3d64a --- /dev/null +++ b/integrations/mariadb/tests/test_filters.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from haystack.errors import FilterError + +from haystack_integrations.document_stores.mariadb.filters import ( + _convert_filters_to_where_clause_and_params, + _validate_filters, +) + + +class TestValidateFilters: + def test_none_filters_pass(self): + _validate_filters(None) + + def test_dict_without_operator_raises(self): + with pytest.raises(ValueError, match="Invalid filter syntax"): + _validate_filters({"unknown_key": "value"}) + + def test_invalid_type_raises(self): + with pytest.raises(TypeError): + _validate_filters("not a dict") # type: ignore[arg-type] + + def test_valid_filter_passes(self): + _validate_filters({"field": "meta.x", "operator": "==", "value": 1}) + + +class TestFilterConversion: + def test_equal_meta_string(self): + filters = {"field": "meta.name", "operator": "==", "value": "Alice"} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "JSON_UNQUOTE" in clause + assert "= ?" in clause + assert params == ["Alice"] + + def test_equal_meta_int(self): + filters = {"field": "meta.age", "operator": "==", "value": 30} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "CAST" in clause + assert "SIGNED" in clause + assert params == [30] + + def test_equal_meta_float(self): + filters = {"field": "meta.score", "operator": "==", "value": 0.9} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "DECIMAL" in clause + assert params == [0.9] + + def test_equal_none(self): + filters = {"field": "meta.x", "operator": "==", "value": None} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "IS NULL" in clause + assert params == [] + + def test_not_equal(self): + filters = {"field": "meta.x", "operator": "!=", "value": "foo"} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "!= ?" in clause + assert params == ["foo"] + + def test_greater_than(self): + filters = {"field": "meta.count", "operator": ">", "value": 5} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "> ?" in clause + assert params == [5] + + def test_less_than(self): + filters = {"field": "meta.count", "operator": "<", "value": 100} + clause, _ = _convert_filters_to_where_clause_and_params(filters) + assert "< ?" in clause + + def test_in_operator(self): + filters = {"field": "meta.tag", "operator": "in", "value": ["a", "b", "c"]} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "IN (?, ?, ?)" in clause + assert params == ["a", "b", "c"] + + def test_not_in_operator(self): + filters = {"field": "meta.tag", "operator": "not in", "value": ["x", "y"]} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "NOT IN" in clause + assert params == ["x", "y"] + + def test_like_operator(self): + filters = {"field": "meta.title", "operator": "like", "value": "%python%"} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "LIKE ?" in clause + assert params == ["%python%"] + + def test_not_like_operator(self): + filters = {"field": "meta.title", "operator": "not like", "value": "%java%"} + clause, _ = _convert_filters_to_where_clause_and_params(filters) + assert "NOT LIKE ?" in clause + + def test_top_level_field(self): + filters = {"field": "id", "operator": "==", "value": "doc-1"} + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "`id`" in clause + assert params == ["doc-1"] + + def test_and_logical(self): + filters = { + "operator": "AND", + "conditions": [ + {"field": "meta.age", "operator": ">=", "value": 18}, + {"field": "meta.active", "operator": "==", "value": "true"}, + ], + } + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "AND" in clause + assert len(params) == 2 + + def test_or_logical(self): + filters = { + "operator": "OR", + "conditions": [ + {"field": "meta.tag", "operator": "==", "value": "ai"}, + {"field": "meta.tag", "operator": "==", "value": "ml"}, + ], + } + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "OR" in clause + assert len(params) == 2 + + def test_nested_logical(self): + filters = { + "operator": "AND", + "conditions": [ + {"field": "meta.x", "operator": ">", "value": 0}, + { + "operator": "OR", + "conditions": [ + {"field": "meta.tag", "operator": "==", "value": "a"}, + {"field": "meta.tag", "operator": "==", "value": "b"}, + ], + }, + ], + } + clause, params = _convert_filters_to_where_clause_and_params(filters) + assert "AND" in clause + assert "OR" in clause + assert len(params) == 3 + + def test_where_operator(self): + filters = {"field": "meta.x", "operator": "==", "value": 1} + clause, _ = _convert_filters_to_where_clause_and_params(filters, operator="WHERE") + assert clause.startswith(" WHERE ") + + def test_and_operator(self): + filters = {"field": "meta.x", "operator": "==", "value": 1} + clause, _ = _convert_filters_to_where_clause_and_params(filters, operator="AND") + assert clause.startswith(" AND ") + + def test_invalid_operator_raises(self): + filters = {"field": "meta.x", "operator": "INVALID", "value": 1} + with pytest.raises(FilterError, match="Unknown comparison operator"): + _convert_filters_to_where_clause_and_params(filters) + + def test_in_non_list_raises(self): + filters = {"field": "meta.x", "operator": "in", "value": "not-a-list"} + with pytest.raises(FilterError): + _convert_filters_to_where_clause_and_params(filters) + + def test_gt_with_string_non_date_raises(self): + filters = {"field": "meta.x", "operator": ">", "value": "notadate"} + with pytest.raises(FilterError): + _convert_filters_to_where_clause_and_params(filters) + + def test_gt_with_iso_date_passes(self): + filters = {"field": "meta.created", "operator": ">", "value": "2024-01-01T00:00:00"} + clause, _ = _convert_filters_to_where_clause_and_params(filters) + assert "> ?" in clause diff --git a/integrations/mariadb/tests/test_retrievers.py b/integrations/mariadb/tests/test_retrievers.py new file mode 100644 index 0000000000..f1685341a8 --- /dev/null +++ b/integrations/mariadb/tests/test_retrievers.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import os +from unittest.mock import MagicMock + +import pytest +from haystack.dataclasses import Document +from haystack.document_stores.types import FilterPolicy + +from haystack_integrations.components.retrievers.mariadb import ( + MariaDBEmbeddingRetriever, + MariaDBKeywordRetriever, +) +from haystack_integrations.document_stores.mariadb import MariaDBDocumentStore + + +def _make_store(**kwargs) -> MariaDBDocumentStore: + os.environ.setdefault("MARIADB_USER", "test_user") + os.environ.setdefault("MARIADB_PASSWORD", "test_pass") + store = MariaDBDocumentStore(**kwargs) + return store + + +# --------------------------------------------------------------------------- +# MariaDBEmbeddingRetriever unit tests +# --------------------------------------------------------------------------- + + +class TestEmbeddingRetrieverInit: + def test_requires_mariadb_store(self): + with pytest.raises(ValueError, match="MariaDBDocumentStore"): + MariaDBEmbeddingRetriever(document_store=MagicMock()) + + def test_invalid_vector_function_raises(self): + store = _make_store() + with pytest.raises(ValueError, match="vector_function"): + MariaDBEmbeddingRetriever(document_store=store, vector_function="bad") + + def test_defaults(self): + store = _make_store() + r = MariaDBEmbeddingRetriever(document_store=store) + assert r.top_k == 10 + assert r.filter_policy == FilterPolicy.REPLACE + assert r.vector_function == "cosine_similarity" + + def test_custom_vector_function(self): + store = _make_store() + r = MariaDBEmbeddingRetriever(document_store=store, vector_function="l2_distance") + assert r.vector_function == "l2_distance" + + +class TestEmbeddingRetrieverSerialization: + def test_to_dict(self): + store = _make_store() + r = MariaDBEmbeddingRetriever(document_store=store, top_k=3) + d = r.to_dict() + assert d["init_parameters"]["top_k"] == 3 + assert "document_store" in d["init_parameters"] + + def test_from_dict_roundtrip(self): + store = _make_store() + r = MariaDBEmbeddingRetriever(document_store=store, top_k=5) + d = r.to_dict() + restored = MariaDBEmbeddingRetriever.from_dict(d) + assert restored.top_k == 5 + + def test_from_dict_with_filter_policy(self): + store = _make_store() + r = MariaDBEmbeddingRetriever(document_store=store, filter_policy=FilterPolicy.MERGE) + d = r.to_dict() + restored = MariaDBEmbeddingRetriever.from_dict(d) + assert restored.filter_policy == FilterPolicy.MERGE + + +class TestEmbeddingRetrieverRun: + def test_run_calls_embedding_retrieval(self): + store = _make_store() + store._embedding_retrieval = MagicMock(return_value=[Document(content="hit")]) + r = MariaDBEmbeddingRetriever(document_store=store, top_k=3) + result = r.run(query_embedding=[0.1, 0.2]) + assert "documents" in result + assert len(result["documents"]) == 1 + store._embedding_retrieval.assert_called_once_with( + query_embedding=[0.1, 0.2], + filters={}, + top_k=3, + vector_function="cosine_similarity", + ) + + def test_run_with_runtime_top_k(self): + store = _make_store() + store._embedding_retrieval = MagicMock(return_value=[]) + r = MariaDBEmbeddingRetriever(document_store=store, top_k=5) + r.run(query_embedding=[0.1], top_k=2) + call_kwargs = store._embedding_retrieval.call_args[1] + assert call_kwargs["top_k"] == 2 + + def test_run_with_runtime_vector_function(self): + store = _make_store() + store._embedding_retrieval = MagicMock(return_value=[]) + r = MariaDBEmbeddingRetriever(document_store=store) + r.run(query_embedding=[0.1], vector_function="l2_distance") + call_kwargs = store._embedding_retrieval.call_args[1] + assert call_kwargs["vector_function"] == "l2_distance" + + def test_filter_policy_replace(self): + store = _make_store() + store._embedding_retrieval = MagicMock(return_value=[]) + r = MariaDBEmbeddingRetriever( + document_store=store, + filters={"field": "meta.x", "operator": "==", "value": 1}, + filter_policy=FilterPolicy.REPLACE, + ) + runtime_filter = {"field": "meta.y", "operator": "==", "value": 2} + r.run(query_embedding=[0.1], filters=runtime_filter) + call_kwargs = store._embedding_retrieval.call_args[1] + assert call_kwargs["filters"] == runtime_filter + + def test_filter_policy_merge(self): + store = _make_store() + store._embedding_retrieval = MagicMock(return_value=[]) + init_filter = {"field": "meta.x", "operator": "==", "value": 1} + r = MariaDBEmbeddingRetriever( + document_store=store, + filters=init_filter, + filter_policy=FilterPolicy.MERGE, + ) + runtime_filter = {"field": "meta.y", "operator": "==", "value": 2} + r.run(query_embedding=[0.1], filters=runtime_filter) + call_kwargs = store._embedding_retrieval.call_args[1] + # MERGE combines both filters into an AND condition + assert call_kwargs["filters"]["operator"] == "AND" + + +# --------------------------------------------------------------------------- +# MariaDBKeywordRetriever unit tests +# --------------------------------------------------------------------------- + + +class TestKeywordRetrieverInit: + def test_requires_mariadb_store(self): + with pytest.raises(ValueError, match="MariaDBDocumentStore"): + MariaDBKeywordRetriever(document_store=MagicMock()) + + def test_defaults(self): + store = _make_store() + r = MariaDBKeywordRetriever(document_store=store) + assert r.top_k == 10 + assert r.filter_policy == FilterPolicy.REPLACE + + +class TestKeywordRetrieverSerialization: + def test_to_dict(self): + store = _make_store() + r = MariaDBKeywordRetriever(document_store=store, top_k=7) + d = r.to_dict() + assert d["init_parameters"]["top_k"] == 7 + + def test_from_dict_roundtrip(self): + store = _make_store() + r = MariaDBKeywordRetriever(document_store=store, top_k=4) + d = r.to_dict() + restored = MariaDBKeywordRetriever.from_dict(d) + assert restored.top_k == 4 + + +class TestKeywordRetrieverRun: + def test_run_calls_keyword_retrieval(self): + store = _make_store() + store._keyword_retrieval = MagicMock(return_value=[Document(content="result")]) + r = MariaDBKeywordRetriever(document_store=store, top_k=5) + result = r.run(query="haystack") + assert "documents" in result + assert len(result["documents"]) == 1 + store._keyword_retrieval.assert_called_once_with( + query="haystack", + filters={}, + top_k=5, + ) + + def test_run_with_runtime_top_k(self): + store = _make_store() + store._keyword_retrieval = MagicMock(return_value=[]) + r = MariaDBKeywordRetriever(document_store=store, top_k=10) + r.run(query="test", top_k=3) + call_kwargs = store._keyword_retrieval.call_args[1] + assert call_kwargs["top_k"] == 3 + + def test_run_with_runtime_filters(self): + store = _make_store() + store._keyword_retrieval = MagicMock(return_value=[]) + r = MariaDBKeywordRetriever(document_store=store) + runtime_filter = {"field": "meta.lang", "operator": "==", "value": "en"} + r.run(query="test", filters=runtime_filter) + call_kwargs = store._keyword_retrieval.call_args[1] + assert call_kwargs["filters"] == runtime_filter + + +# --------------------------------------------------------------------------- +# Integration tests (require real MariaDB) +# --------------------------------------------------------------------------- + +MARIADB_HOST = os.environ.get("MARIADB_HOST", "localhost") +MARIADB_PORT = int(os.environ.get("MARIADB_PORT", "3306")) +MARIADB_DB = os.environ.get("MARIADB_DATABASE", "haystack") +MARIADB_USER = os.environ.get("MARIADB_USER", "root") +MARIADB_PASSWORD = os.environ.get("MARIADB_PASSWORD", "password") + + +@pytest.fixture +def integration_store(): + store = MariaDBDocumentStore( + host=MARIADB_HOST, + port=MARIADB_PORT, + database=MARIADB_DB, + user=MARIADB_USER, + password=MARIADB_PASSWORD, + table_name="test_retrievers_docs", + embedding_dimension=4, + recreate_table=True, + ) + yield store + try: + store._cursor.execute("DROP TABLE IF EXISTS `test_retrievers_docs`") + except Exception: # noqa: S110 + pass + store.close() + + +@pytest.mark.integration +class TestIntegrationRetrievers: + def test_embedding_retriever(self, integration_store): + docs = [ + Document(content="cats are great", embedding=[1.0, 0.0, 0.0, 0.0]), + Document(content="dogs are loyal", embedding=[0.0, 1.0, 0.0, 0.0]), + ] + integration_store.write_documents(docs) + retriever = MariaDBEmbeddingRetriever(document_store=integration_store, top_k=1) + result = retriever.run(query_embedding=[1.0, 0.0, 0.0, 0.0]) + docs_out = result["documents"] + assert len(docs_out) == 1 + assert docs_out[0].content == "cats are great" + + def test_keyword_retriever(self, integration_store): + docs = [ + Document(content="natural language processing"), + Document(content="computer vision research"), + ] + integration_store.write_documents(docs) + retriever = MariaDBKeywordRetriever(document_store=integration_store, top_k=5) + result = retriever.run(query="language") + assert any("language" in d.content for d in result["documents"]) + + def test_embedding_retriever_l2(self, integration_store): + docs = [ + Document(content="near origin", embedding=[0.1, 0.1, 0.1, 0.1]), + Document(content="far origin", embedding=[10.0, 10.0, 10.0, 10.0]), + ] + integration_store.write_documents(docs) + retriever = MariaDBEmbeddingRetriever(document_store=integration_store, top_k=1, vector_function="l2_distance") + result = retriever.run(query_embedding=[0.0, 0.0, 0.0, 0.0]) + assert result["documents"][0].content == "near origin"