diff --git a/CHANGES.md b/CHANGES.md index 537c3a32447c..57ebb5b747f5 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -65,6 +65,7 @@ ## I/Os * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* Qdrant enrichment handler added (Python) ([#40628](https://github.com/apache/beam/issues/40628)). `QdrantSearchEnrichmentHandler` supports dense, sparse and hybrid (RRF/DBSF fusion) search with payload filters, and reads collections written by `QdrantWriteConfig`. * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)). diff --git a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py new file mode 100644 index 000000000000..85b510c105e3 --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py @@ -0,0 +1,458 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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. +# + +from dataclasses import dataclass +from dataclasses import field +from dataclasses import replace +from typing import Any +from typing import Optional +from typing import Union + +import grpc +import httpx +from qdrant_client import QdrantClient +from qdrant_client import models +from qdrant_client.common.client_exceptions import ResourceExhaustedResponse +from qdrant_client.http.exceptions import ResponseHandlingException +from qdrant_client.http.exceptions import UnexpectedResponse + +from apache_beam.io.requestresponse import UserCodeQuotaException +from apache_beam.io.requestresponse import UserCodeTimeoutException +from apache_beam.ml.rag.ingestion.qdrant import QdrantConnectionParameters +from apache_beam.ml.rag.types import EmbeddableItem +from apache_beam.transforms.enrichment import EnrichmentSourceHandler + +__all__ = [ + 'QdrantConnectionParameters', + 'QdrantDenseSearchParameters', + 'QdrantSparseSearchParameters', + 'QdrantHybridSearchParameters', + 'QdrantSearchParameters', + 'QdrantSearchResult', + 'QdrantSearchEnrichmentHandler', +] + + +@dataclass +class QdrantDenseSearchParameters: + """Parameters for dense vector similarity search. + + Args: + vector_name: Name of the dense vector in the Qdrant collection. Defaults to + "dense", matching the default `dense_embedding_key` of + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + limit: Maximum number of results to return per query. Must be positive. + Defaults to 3 search results. + filter: Optional payload filter applied to the search, e.g. + `models.Filter(must=[models.FieldCondition(key="source", + match=models.MatchValue(value="docs"))])`. + score_threshold: Optional minimal score; results scoring worse are dropped. + search_params: Optional `models.SearchParams` such as `hnsw_ef` or `exact`. + kwargs: Optional keyword arguments for additional `models.QueryRequest` + fields. Enables forward compatibility. + """ + vector_name: str = "dense" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + search_params: Optional[models.SearchParams] = None + kwargs: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not self.vector_name: + raise ValueError("Vector name must be provided") + + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +@dataclass +class QdrantSparseSearchParameters: + """Parameters for sparse vector (keyword) search. + + Args: + vector_name: Name of the sparse vector in the Qdrant collection. Defaults to + "sparse", matching the default `sparse_embedding_key` of + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + limit: Maximum number of results to return per query. Must be positive. + Defaults to 3 search results. + filter: Optional payload filter applied to the search. + score_threshold: Optional minimal score; results scoring worse are dropped. + search_params: Optional `models.SearchParams`. + kwargs: Optional keyword arguments for additional `models.QueryRequest` + fields. Enables forward compatibility. + """ + vector_name: str = "sparse" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + search_params: Optional[models.SearchParams] = None + kwargs: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not self.vector_name: + raise ValueError("Vector name must be provided") + + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +@dataclass +class QdrantHybridSearchParameters: + """Parameters for hybrid (dense + sparse) search. + + The dense and sparse searches run as prefetches and their results are fused + server side with the Qdrant Query API. Configure `search_params` on `dense` + and `sparse` to control their respective prefetches; fusion does not run a + vector search. + + Args: + dense: Parameters for the dense search component. Its `limit` is the + number of candidates prefetched before fusion. + sparse: Parameters for the sparse search component. Its `limit` is the + number of candidates prefetched before fusion. + fusion: Fusion method used to combine both result lists. Defaults to + reciprocal rank fusion (`models.Fusion.RRF`). + limit: Maximum number of fused results to return per query. Defaults to 3 + search results. + filter: Optional common payload filter on the outer query. Qdrant also + applies it to the prefetches, together with their component filters. + score_threshold: Optional minimum fused score. Component thresholds apply + to the individual dense or sparse scores before fusion instead. + kwargs: Optional keyword arguments for additional `models.QueryRequest` + fields. Explicit non-None `filter` and `score_threshold` fields take + precedence over the corresponding entries in `kwargs`. + """ + dense: QdrantDenseSearchParameters + sparse: QdrantSparseSearchParameters + fusion: models.Fusion = models.Fusion.RRF + limit: int = 3 + kwargs: dict[str, Any] = field(default_factory=dict) + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + + def __post_init__(self): + if not self.dense or not self.sparse: + raise ValueError( + "Dense and sparse search parameters must be provided for " + "hybrid search") + + if not self.fusion: + raise ValueError("Fusion must be provided for hybrid search") + + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +QdrantSearchStrategyType = Union[QdrantDenseSearchParameters, + QdrantSparseSearchParameters, + QdrantHybridSearchParameters] + + +@dataclass +class QdrantSearchParameters: + """Parameters configuring Qdrant search operations. + + Args: + collection_name: Name of the collection to search in. + search_strategy: Dense, sparse or hybrid search parameters. + with_payload: Whether to return the payload of each point. Either a bool + or a list of payload keys to return. Defaults to True. + consistency: Optional read consistency for the search, e.g. + `models.ReadConsistencyType.MAJORITY`. + timeout: Optional search timeout in seconds. If not specified, the client's + default timeout is used. + """ + collection_name: str + search_strategy: QdrantSearchStrategyType + with_payload: Union[bool, list[str]] = True + consistency: Optional[models.ReadConsistency] = None + timeout: Optional[int] = None + + def __post_init__(self): + if not self.collection_name: + raise ValueError("Collection name must be provided") + + if not self.search_strategy: + raise ValueError("Search strategy must be provided") + + +@dataclass +class QdrantSearchResult: + """Search result from Qdrant per embeddable_item. + + Args: + id: List of point IDs returned from the search. Can be either string (UUID) + or integer IDs. + score: List of scores for each returned point. Higher is more similar for + cosine and dot product, lower is more similar for Euclidean and Manhattan + distances. + payload: List of payload dictionaries, one for each returned point. Empty + if the point has no payload or `with_payload` is False. + """ + id: list[Union[str, int]] = field(default_factory=list) + score: list[float] = field(default_factory=list) + payload: list[dict[str, Any]] = field(default_factory=list) + + +InputT, OutputT = (Union[EmbeddableItem, list[EmbeddableItem]], + list[tuple[EmbeddableItem, dict[str, Any]]]) + + +class QdrantSearchEnrichmentHandler(EnrichmentSourceHandler[InputT, OutputT]): + """Enrichment handler for Qdrant vector database searches. + + This handler is designed to work with the + :class:`apache_beam.transforms.enrichment.EnrichmentSourceHandler` transform. + It enables enriching data through dense, sparse, or hybrid searches against + Qdrant collections, such as ones written by + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + + The handler supports different search strategies: + * Dense search - For finding similar embeddings based on vector similarity + * Sparse search - For keyword-style retrieval with sparse embeddings + * Hybrid search - For fusing dense and sparse search results server side + + This handler queries Qdrant per element by default. To enable batching for + improved performance, set the `min_batch_size` and `max_batch_size` + parameters. These control the batching behavior in the + :class:`apache_beam.transforms.utils.BatchElements` transform. Each batch is + sent to Qdrant as a single batch query request. + + REST and gRPC quota and timeout errors are translated to Beam exceptions + so that `Enrichment` can retry them using its configured repeater. Other + client errors propagate unchanged. + """ + def __init__( + self, + connection_parameters: QdrantConnectionParameters, + search_parameters: QdrantSearchParameters, + *, + min_batch_size: int = 1, + max_batch_size: int = 1000, + **kwargs): + """ + Example Usage: + connection_parameters = QdrantConnectionParameters( + url="http://localhost:6333") + search_parameters = QdrantSearchParameters( + collection_name="my_collection", + search_strategy=QdrantDenseSearchParameters(limit=5)) + qdrant_handler = QdrantSearchEnrichmentHandler( + connection_parameters, + search_parameters, + min_batch_size=10, + max_batch_size=100) + + Args: + connection_parameters (QdrantConnectionParameters): Configuration for + connecting to Qdrant, shared with + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + search_parameters (QdrantSearchParameters): Configuration for search + operations, including collection name and search strategy. + min_batch_size (int): Minimum number of elements to batch together when + querying Qdrant. Default is 1 (no batching when max_batch_size is 1). + max_batch_size (int): Maximum number of elements to batch together. + Default is 1000. Higher values may improve throughput but increase + memory usage. + **kwargs: Additional keyword arguments for Qdrant Enrichment Handler. + """ + self._connection_parameters = connection_parameters + self._search_parameters = search_parameters + self._batching_kwargs = { + 'min_batch_size': min_batch_size, 'max_batch_size': max_batch_size + } + self.kwargs = kwargs + self._client = None + self.join_fn = join_fn + self.use_custom_types = True + + def __enter__(self): + """Enters the context manager and creates the Qdrant client. + + Returns: + Self, enabling use in 'with' statements. + """ + if not self._client: + params = self._connection_parameters + self._client = QdrantClient( + location=params.location, + url=params.url, + port=params.port, + grpc_port=params.grpc_port, + prefer_grpc=params.prefer_grpc, + https=params.https, + api_key=params.api_key, + prefix=params.prefix, + timeout=params.timeout, + host=params.host, + path=params.path, + check_compatibility=False, + **params.kwargs, + ) + return self + + def __call__( + self, + request: Union[EmbeddableItem, list[EmbeddableItem]], + *args, + **kwargs) -> list[tuple[EmbeddableItem, dict[str, Any]]]: + reqs = request if isinstance(request, list) else [request] + # Early return for empty requests to avoid unnecessary round trips. + if not reqs: + return [] + requests = [self._get_query_request(item) for item in reqs] + try: + responses = self._client.query_batch_points( + collection_name=self.collection_name, + requests=requests, + consistency=self._search_parameters.consistency, + timeout=self._search_parameters.timeout) + except ResourceExhaustedResponse as exc: + raise UserCodeQuotaException(str(exc)) from exc + except UnexpectedResponse as exc: + if exc.status_code == 429: + raise UserCodeQuotaException(str(exc)) from exc + if exc.status_code in (408, 504): + raise UserCodeTimeoutException(str(exc)) from exc + raise + except ResponseHandlingException as exc: + if isinstance(exc.source, httpx.TimeoutException): + raise UserCodeTimeoutException(str(exc)) from exc + raise + except grpc.RpcError as exc: + if exc.code() == grpc.StatusCode.RESOURCE_EXHAUSTED: + raise UserCodeQuotaException(str(exc)) from exc + if exc.code() == grpc.StatusCode.DEADLINE_EXCEEDED: + raise UserCodeTimeoutException(str(exc)) from exc + raise + return self._get_call_response(reqs, responses) + + def _get_query_request( + self, embeddable_item: EmbeddableItem) -> models.QueryRequest: + strategy = self.search_strategy + with_payload = self._search_parameters.with_payload + if isinstance(strategy, QdrantHybridSearchParameters): + query_kwargs = dict(strategy.kwargs) + if strategy.filter is not None: + query_kwargs['filter'] = strategy.filter + if strategy.score_threshold is not None: + query_kwargs['score_threshold'] = strategy.score_threshold + return models.QueryRequest( + prefetch=[ + self._get_prefetch( + strategy.dense, self._get_dense_query(embeddable_item)), + self._get_prefetch( + strategy.sparse, self._get_sparse_query(embeddable_item)), + ], + query=models.FusionQuery(fusion=strategy.fusion), + limit=strategy.limit, + with_payload=with_payload, + **query_kwargs) + elif isinstance(strategy, QdrantDenseSearchParameters): + return self._get_single_query_request( + strategy, self._get_dense_query(embeddable_item), with_payload) + elif isinstance(strategy, QdrantSparseSearchParameters): + return self._get_single_query_request( + strategy, self._get_sparse_query(embeddable_item), with_payload) + else: + raise ValueError(f"Not supported search strategy yet: {strategy}") + + def _get_single_query_request( + self, + strategy: Union[QdrantDenseSearchParameters, + QdrantSparseSearchParameters], + query: Union[list[float], models.SparseVector], + with_payload: Union[bool, list[str]]) -> models.QueryRequest: + return models.QueryRequest( + query=query, + using=strategy.vector_name, + filter=strategy.filter, + score_threshold=strategy.score_threshold, + params=strategy.search_params, + limit=strategy.limit, + with_payload=with_payload, + **strategy.kwargs) + + def _get_prefetch( + self, + strategy: Union[QdrantDenseSearchParameters, + QdrantSparseSearchParameters], + query: Union[list[float], models.SparseVector]) -> models.Prefetch: + return models.Prefetch( + query=query, + using=strategy.vector_name, + filter=strategy.filter, + score_threshold=strategy.score_threshold, + params=strategy.search_params, + limit=strategy.limit, + **strategy.kwargs) + + def _get_dense_query(self, embeddable_item: EmbeddableItem) -> list[float]: + if not embeddable_item.dense_embedding: + raise ValueError( + f"Item {embeddable_item.id} missing dense embedding required for" + " dense search") + return embeddable_item.dense_embedding + + def _get_sparse_query( + self, embeddable_item: EmbeddableItem) -> models.SparseVector: + if not embeddable_item.sparse_embedding: + raise ValueError( + f"Item {embeddable_item.id} missing sparse embedding required for" + " sparse search") + indices, values = embeddable_item.sparse_embedding + return models.SparseVector(indices=indices, values=values) + + def _get_call_response( + self, + embeddable_items: list[EmbeddableItem], + responses: list[models.QueryResponse]): + response = [] + for embeddable_item, query_response in zip(embeddable_items, responses): + result = QdrantSearchResult() + for point in query_response.points: + result.id.append(point.id) + result.score.append(point.score) + result.payload.append(point.payload or {}) + response.append((embeddable_item, result.__dict__)) + return response + + @property + def collection_name(self): + """Getter method for collection_name property""" + return self._search_parameters.collection_name + + @property + def search_strategy(self): + """Getter method for search_strategy property""" + return self._search_parameters.search_strategy + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._client: + try: + self._client.close() + finally: + self._client = None + + def batch_elements_kwargs(self) -> dict[str, int]: + """Returns kwargs for beam.BatchElements.""" + return self._batching_kwargs + + +def join_fn(left: EmbeddableItem, right: dict[str, Any]) -> EmbeddableItem: + return replace(left, metadata={**left.metadata, 'enrichment_data': right}) diff --git a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py new file mode 100644 index 000000000000..2d423ec2c5ab --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py @@ -0,0 +1,210 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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. +# + +import unittest + +import pytest + +import apache_beam as beam +from apache_beam.ml.rag.types import Content +from apache_beam.ml.rag.types import EmbeddableItem +from apache_beam.ml.rag.types import Embedding +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to + +# pylint: disable=ungrouped-imports +try: + from qdrant_client import models + + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantConnectionParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantDenseSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantHybridSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchEnrichmentHandler + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSparseSearchParameters + from apache_beam.ml.rag.ingestion.qdrant import QdrantWriteConfig + from apache_beam.transforms.enrichment import Enrichment + + QDRANT_AVAILABLE = True +except ImportError: + QDRANT_AVAILABLE = False + +try: + from testcontainers.qdrant import QdrantContainer +except ImportError: + QdrantContainer = None +# pylint: enable=ungrouped-imports + +TEST_CORPUS = [ + EmbeddableItem( + id="1", + content=Content(text="Beam runs batch and streaming pipelines."), + metadata={"source": "beam"}, + embedding=Embedding( + dense_embedding=[1.0, 0.0], sparse_embedding=([0, 1], [1.0, 0.5])), + ), + EmbeddableItem( + id="2", + content=Content(text="Qdrant is a vector database."), + metadata={"source": "qdrant"}, + embedding=Embedding( + dense_embedding=[0.0, 1.0], sparse_embedding=([2, 3], [1.0, 0.5])), + ), + EmbeddableItem( + id="3", + content=Content(text="Enrichment joins external data."), + metadata={"source": "beam"}, + embedding=Embedding( + dense_embedding=[-1.0, 0.0], sparse_embedding=([4, 5], [1.0, 0.5])), + ), +] + + +def _query(id, dense=None, sparse=None): + return EmbeddableItem( + id=id, + content=Content(text=f"query {id}"), + embedding=Embedding(dense_embedding=dense, sparse_embedding=sparse)) + + +def _retrieved(item): + """Extracts (query id, retrieved ids, retrieved contents) for assertions.""" + data = item.metadata["enrichment_data"] + return ( + item.id, + data["id"], [payload.get("content") for payload in data["payload"]]) + + +@unittest.skipIf(not QDRANT_AVAILABLE, "qdrant dependencies not installed.") +@unittest.skipIf(not QdrantContainer, "qdrant test_container not installed.") +@pytest.mark.require_docker_in_docker +class TestQdrantSearchEnrichment(unittest.TestCase): + """Writes TEST_CORPUS with QdrantWriteConfig, then retrieves it.""" + _collection_name = "test_qdrant_search" + + @classmethod + def setUpClass(cls): + cls._container = QdrantContainer() + cls._container.start() + cls.addClassCleanup(cls._container.stop) + cls._connection_params = QdrantConnectionParameters( + host=cls._container.get_container_host_ip(), + port=int(cls._container.get_exposed_port(6333))) + + client = cls._container.get_client() + client.create_collection( + collection_name=cls._collection_name, + vectors_config={ + "dense": models.VectorParams( + size=2, distance=models.Distance.COSINE) + }, + sparse_vectors_config={"sparse": models.SparseVectorParams()}) + client.close() + + write_config = QdrantWriteConfig( + connection_params=cls._connection_params, + collection_name=cls._collection_name, + content_payload_key="content") + with TestPipeline() as p: + _ = (p | beam.Create(TEST_CORPUS) | write_config.create_write_transform()) + + def _run_enrichment(self, search_strategy, queries, expected, **kwargs): + handler = QdrantSearchEnrichmentHandler( + self._connection_params, + QdrantSearchParameters( + collection_name=self._collection_name, + search_strategy=search_strategy), + **kwargs) + with TestPipeline() as p: + result = ( + p + | beam.Create(queries) + | Enrichment(handler) + | beam.Map(_retrieved)) + assert_that(result, equal_to(expected)) + + def test_dense_search(self): + self._run_enrichment( + QdrantDenseSearchParameters(limit=1), + [_query("a", dense=[0.9, 0.1]), _query("b", dense=[0.1, 0.9])], + [ + ("a", [1], ["Beam runs batch and streaming pipelines."]), + ("b", [2], ["Qdrant is a vector database."]), + ]) + + def test_dense_search_with_filter(self): + self._run_enrichment( + QdrantDenseSearchParameters( + limit=3, + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="beam")) + ])), [_query("a", dense=[0.1, 0.9])], + [( + "a", [1, 3], + [ + "Beam runs batch and streaming pipelines.", + "Enrichment joins external data." + ])]) + + def test_sparse_search(self): + self._run_enrichment( + QdrantSparseSearchParameters(limit=3), + [_query("a", sparse=([4], [1.0]))], + [("a", [3], ["Enrichment joins external data."])]) + + def test_hybrid_search(self): + self._run_enrichment( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + limit=1), + [_query("a", dense=[0.0, 1.0], sparse=([2, 3], [1.0, 1.0]))], + [("a", [2], ["Qdrant is a vector database."])]) + + def test_batched_search(self): + queries = [_query(str(i), dense=[1.0, 0.0]) for i in range(10)] + self._run_enrichment( + QdrantDenseSearchParameters(limit=1), + queries, + [(q.id, [1], ["Beam runs batch and streaming pipelines."]) + for q in queries], + min_batch_size=5, + max_batch_size=5) + + def test_hybrid_search_with_filter_and_fused_score_threshold(self): + self._run_enrichment( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="beam")) + ]), + score_threshold=0.7), + [_query("a", dense=[0.1, 0.9], sparse=([4], [1.0]))], + [("a", [3], ["Enrichment joins external data."])]) + + def test_empty_input(self): + self._run_enrichment(QdrantDenseSearchParameters(), [], []) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py new file mode 100644 index 000000000000..0c3a6a706c88 --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py @@ -0,0 +1,388 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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. +# + +import unittest +from unittest.mock import patch + +import grpc +from parameterized import parameterized + +try: + import httpx + from qdrant_client import models + from qdrant_client.common.client_exceptions import ResourceExhaustedResponse + from qdrant_client.http.exceptions import ResponseHandlingException + from qdrant_client.http.exceptions import UnexpectedResponse + from qdrant_client.http.models import QueryResponse + + from apache_beam.io.requestresponse import ExponentialBackOffRepeater + from apache_beam.io.requestresponse import UserCodeQuotaException + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantConnectionParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantDenseSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantHybridSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchEnrichmentHandler + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSparseSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import join_fn + from apache_beam.ml.rag.types import Content + from apache_beam.ml.rag.types import EmbeddableItem + from apache_beam.ml.rag.types import Embedding +except ImportError as e: + raise unittest.SkipTest(f'Qdrant dependencies not installed: {str(e)}') + +_COLLECTION = "test_collection" + +# Points stored in the in-memory collection. Dense vectors use cosine +# distance, so a query of [1.0, 0.0] ranks point 1 first, then 2, then 3. +_POINTS = [ + models.PointStruct( + id=1, + vector={ + "dense": [1.0, 0.0], + "sparse": models.SparseVector(indices=[0, 1], values=[1.0, 0.5]), + }, + payload={ + "content": "doc one", "source": "a" + }, + ), + models.PointStruct( + id=2, + vector={ + "dense": [0.8, 0.6], + "sparse": models.SparseVector(indices=[1, 2], values=[1.0, 0.5]), + }, + payload={ + "content": "doc two", "source": "b" + }, + ), + models.PointStruct( + id=3, + vector={ + "dense": [-1.0, 0.0], + "sparse": models.SparseVector(indices=[2, 3], values=[1.0, 0.5]), + }, + payload={ + "content": "doc three", "source": "a" + }, + ), +] + + +def _item(id, dense=None, sparse=None): + return EmbeddableItem( + id=id, + content=Content(text=f"query {id}"), + embedding=Embedding(dense_embedding=dense, sparse_embedding=sparse)) + + +class _RpcError(grpc.RpcError): + def __init__(self, status_code): + self._status_code = status_code + + def code(self): + return self._status_code + + +class TestQdrantSearchParameters(unittest.TestCase): + """Unit tests for validation of Qdrant search parameters.""" + @parameterized.expand([ + (lambda: QdrantDenseSearchParameters(limit=0), "positive"), + (lambda: QdrantDenseSearchParameters(vector_name=""), "Vector name"), + (lambda: QdrantSparseSearchParameters(limit=-1), "positive"), + (lambda: QdrantSparseSearchParameters(vector_name=""), "Vector name"), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(), sparse= + QdrantSparseSearchParameters(), limit=0), + "positive"), + ( + lambda: QdrantHybridSearchParameters( + dense=None, sparse=QdrantSparseSearchParameters()), + "Dense and sparse"), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(), sparse= + QdrantSparseSearchParameters(), fusion=None), + "Fusion"), + ( + lambda: QdrantSearchParameters( + collection_name="", search_strategy=QdrantDenseSearchParameters() + ), + "Collection name"), + ( + lambda: QdrantSearchParameters( + collection_name=_COLLECTION, search_strategy=None), + "Search strategy"), + ]) + def test_invalid_parameters(self, create, expected_message): + with self.assertRaises(ValueError) as context: + create() + self.assertIn(expected_message, str(context.exception)) + + def test_batch_elements_kwargs(self): + handler = QdrantSearchEnrichmentHandler( + QdrantConnectionParameters.in_memory(), + QdrantSearchParameters( + collection_name=_COLLECTION, + search_strategy=QdrantDenseSearchParameters()), + min_batch_size=5, + max_batch_size=50) + self.assertEqual( + handler.batch_elements_kwargs(), { + 'min_batch_size': 5, 'max_batch_size': 50 + }) + + +class TestQdrantSearchEnrichmentHandler(unittest.TestCase): + """Runs the handler against an in-memory Qdrant collection.""" + def _create_handler(self, search_strategy, **search_kwargs): + handler = QdrantSearchEnrichmentHandler( + QdrantConnectionParameters.in_memory(), + QdrantSearchParameters( + collection_name=_COLLECTION, + search_strategy=search_strategy, + **search_kwargs)) + handler.__enter__() + self.addCleanup(handler.__exit__, None, None, None) + # The in-memory client is created on __enter__, so populate it here. + client = handler._client + client.create_collection( + collection_name=_COLLECTION, + vectors_config={ + "dense": models.VectorParams( + size=2, distance=models.Distance.COSINE) + }, + sparse_vectors_config={"sparse": models.SparseVectorParams()}) + client.upsert(collection_name=_COLLECTION, points=_POINTS) + return handler + + def test_dense_search(self): + handler = self._create_handler(QdrantDenseSearchParameters(limit=2)) + item = _item("q", dense=[1.0, 0.0]) + + [(returned_item, result)] = handler(item) + + self.assertIs(returned_item, item) + self.assertEqual(result["id"], [1, 2]) + self.assertAlmostEqual(result["score"][0], 1.0, places=5) + self.assertAlmostEqual(result["score"][1], 0.8, places=5) + self.assertEqual( + result["payload"][0], { + "content": "doc one", "source": "a" + }) + + def test_dense_search_with_filter_and_score_threshold(self): + handler = self._create_handler( + QdrantDenseSearchParameters( + limit=3, + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="a")) + ]), + score_threshold=0.0)) + + [(_, result)] = handler(_item("q", dense=[1.0, 0.0])) + + # Point 2 is filtered out by payload and point 3 by score (-1.0). + self.assertEqual(result["id"], [1]) + + def test_sparse_search(self): + handler = self._create_handler(QdrantSparseSearchParameters(limit=3)) + + [(_, result)] = handler(_item("q", sparse=([2, 3], [1.0, 1.0]))) + + # Dot products: point 3 = 1.5, point 2 = 0.5, point 1 shares no index. + self.assertEqual(result["id"], [3, 2]) + self.assertEqual(result["score"], [1.5, 0.5]) + + @parameterized.expand([(models.Fusion.RRF, ), (models.Fusion.DBSF, )]) + def test_hybrid_search(self, fusion): + handler = self._create_handler( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + fusion=fusion, + limit=2)) + + [(_, result) + ] = handler(_item("q", dense=[0.8, 0.6], sparse=([1, 2], [1.0, 1.0]))) + + # Point 2 ranks first in both prefetches, so it wins after fusion. + self.assertEqual(len(result["id"]), 2) + self.assertEqual(result["id"][0], 2) + self.assertEqual(result["payload"][0]["content"], "doc two") + + def test_hybrid_common_filter_applies_before_prefetch_limit(self): + handler = self._create_handler( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=1), + sparse=QdrantSparseSearchParameters(limit=1), + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="a")) + ]))) + + [(_, result)] = handler(_item("q", dense=[0.8, 0.6], sparse=([1], [1.0]))) + + # Point 2 would win both prefetches without the common filter. + self.assertEqual(result["id"], [1]) + + def test_hybrid_fused_score_threshold(self): + handler = self._create_handler( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + score_threshold=0.8)) + + [(_, result) + ] = handler(_item("q", dense=[0.8, 0.6], sparse=([1, 2], [1.0, 1.0]))) + + self.assertEqual(result["id"], [2]) + self.assertEqual(result["score"], [1.0]) + + @parameterized.expand([ + (ResourceExhaustedResponse("rate limited", 1), ), + (UnexpectedResponse(429, "rate limited", b"", httpx.Headers()), ), + (UnexpectedResponse(408, "timeout", b"", httpx.Headers()), ), + (UnexpectedResponse(504, "timeout", b"", httpx.Headers()), ), + (ResponseHandlingException(httpx.ReadTimeout("timeout")), ), + (_RpcError(grpc.StatusCode.RESOURCE_EXHAUSTED), ), + (_RpcError(grpc.StatusCode.DEADLINE_EXCEEDED), ), + ]) + def test_default_repeater_recovers_from_quota_and_timeout(self, error): + handler = self._create_handler(QdrantDenseSearchParameters(limit=1)) + item = _item("q", dense=[1.0, 0.0]) + response = QueryResponse(points=[]) + with patch.object( + handler._client, 'query_batch_points', + side_effect=[error, [response]]) as query, patch( + 'apache_beam.utils.retry.time.sleep'): + result = ExponentialBackOffRepeater().repeat(handler, [item], 30) + + self.assertEqual(result, [(item, {"id": [], "score": [], "payload": []})]) + self.assertEqual(query.call_count, 2) + + @parameterized.expand([ + (UnexpectedResponse(400, "bad request", b"", httpx.Headers()), ), + (UnexpectedResponse(401, "unauthorized", b"", httpx.Headers()), ), + (UnexpectedResponse(404, "not found", b"", httpx.Headers()), ), + (ResponseHandlingException(ValueError("invalid response")), ), + (_RpcError(grpc.StatusCode.INVALID_ARGUMENT), ), + (ValueError("invalid query"), ), + ]) + def test_default_repeater_does_not_retry_permanent_errors(self, error): + handler = self._create_handler(QdrantDenseSearchParameters()) + with patch.object(handler._client, 'query_batch_points', + side_effect=error) as query: + with self.assertRaises(type(error)) as context: + ExponentialBackOffRepeater().repeat( + handler, [_item("q", dense=[1.0, 0.0])], 30) + + self.assertIs(context.exception, error) + self.assertEqual(query.call_count, 1) + + def test_default_repeater_exhausts_retries(self): + handler = self._create_handler(QdrantDenseSearchParameters()) + error = ResourceExhaustedResponse("rate limited", 1) + with patch.object(handler._client, 'query_batch_points', + side_effect=error) as query, \ + patch('apache_beam.utils.retry.time.sleep'): + with self.assertRaises(UserCodeQuotaException) as context: + ExponentialBackOffRepeater().repeat( + handler, [_item("q", dense=[1.0, 0.0])], 30) + + self.assertIs(context.exception.__cause__, error) + self.assertEqual(query.call_count, 3) + + def test_batched_request_returns_result_per_item(self): + handler = self._create_handler(QdrantDenseSearchParameters(limit=1)) + items = [ + _item("q1", dense=[1.0, 0.0]), + _item("q2", dense=[-1.0, 0.0]), + ] + + response = handler(items) + + self.assertEqual([item for item, _ in response], items) + self.assertEqual([result["id"] for _, result in response], [[1], [3]]) + + def test_with_payload_keys(self): + handler = self._create_handler( + QdrantDenseSearchParameters(limit=1), with_payload=["content"]) + + [(_, result)] = handler(_item("q", dense=[1.0, 0.0])) + + self.assertEqual(result["payload"], [{"content": "doc one"}]) + + def test_without_payload(self): + handler = self._create_handler( + QdrantDenseSearchParameters(limit=1), with_payload=False) + + [(_, result)] = handler(_item("q", dense=[1.0, 0.0])) + + self.assertEqual(result["id"], [1]) + self.assertEqual(result["payload"], [{}]) + + def test_empty_request(self): + handler = self._create_handler(QdrantDenseSearchParameters()) + self.assertEqual(handler([]), []) + + @parameterized.expand([ + (QdrantDenseSearchParameters(), _item("q", sparse=([0], [1.0])), "dense"), + (QdrantSparseSearchParameters(), _item("q", dense=[1.0, 0.0]), "sparse"), + ( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(), + sparse=QdrantSparseSearchParameters()), + _item("q", dense=[1.0, 0.0]), + "sparse"), + ]) + def test_missing_embedding(self, search_strategy, item, embedding_type): + handler = self._create_handler(search_strategy) + with self.assertRaises(ValueError) as context: + handler(item) + self.assertIn( + f"Item q missing {embedding_type} embedding", str(context.exception)) + + def test_unsupported_search_strategy(self): + handler = self._create_handler({"vector_name": "dense"}) + with self.assertRaises(ValueError) as context: + handler(_item("q", dense=[1.0, 0.0])) + self.assertIn("Not supported search strategy", str(context.exception)) + + def test_join_fn(self): + item = _item("q", dense=[1.0, 0.0]) + item.metadata = {"source": "query", "enrichment_data": {"previous": True}} + original_metadata = dict(item.metadata) + result = {"id": [1], "score": [1.0], "payload": [{"content": "doc one"}]} + + joined = join_fn(item, result) + + self.assertEqual(joined.metadata["enrichment_data"], result) + self.assertEqual(joined.metadata["source"], "query") + self.assertEqual(item.metadata, original_metadata) + self.assertIsNot(joined, item) + self.assertIsNot(joined.metadata, item.metadata) + self.assertEqual(joined.id, item.id) + self.assertEqual(joined.content, item.content) + self.assertEqual(joined.embedding, item.embedding) + + +if __name__ == '__main__': + unittest.main()