From 859c74eb0563e54e56f140c78504141a1415e848 Mon Sep 17 00:00:00 2001 From: Anush008 Date: Fri, 25 Jul 2025 17:16:05 +0530 Subject: [PATCH 1/4] feat: Qdrant connector Signed-off-by: Anush008 --- .../ml/rag/enrichment/qdrant_search.py | 266 ++++++++++++++++++ .../ml/rag/enrichment/qdrant_search_test.py | 194 +++++++++++++ sdks/python/setup.py | 3 +- 3 files changed, 462 insertions(+), 1 deletion(-) create mode 100644 sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py create mode 100644 sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py 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..b4819143633b --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py @@ -0,0 +1,266 @@ +# +# 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, field +from typing import Any, Dict, List, Optional, Union, Tuple + +from qdrant_client import QdrantClient, models + +from apache_beam.ml.rag.types import Chunk, Embedding +from apache_beam.transforms.enrichment import EnrichmentSourceHandler + + +@dataclass +class QdrantConnectionParameters: + location: Optional[str] = None + url: Optional[str] = None + port: Optional[int] = 6333 + grpc_port: int = 6334 + prefer_grpc: bool = False + https: Optional[bool] = None + api_key: Optional[str] = None + prefix: Optional[str] = None + timeout: Optional[int] = None + host: Optional[str] = None + path: Optional[str] = None + kwargs: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not (self.location or self.url or self.host or self.path): + raise ValueError( + "At least one of location, url, host, or path must be provided for Qdrant connection" + ) + + +@dataclass +class QdrantDenseSearchParameters: + vector_name: str = "dense" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + kwargs: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +@dataclass +class QdrantSparseSearchParameters: + vector_name: str = "sparse" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + kwargs: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +@dataclass +class QdrantHybridSearchParameters: + dense: QdrantDenseSearchParameters + sparse: QdrantSparseSearchParameters + limit: int = 3 + fusion: models.Fusion = models.Fusion.RRF + kwargs: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not self.dense or not self.sparse: + raise ValueError( + "Both dense and sparse search parameters 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: + collection_name: str + search_strategy: QdrantSearchStrategyType + kwargs: Dict[str, Any] = field(default_factory=dict) + + 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") + + +InputT = Union[Chunk, List[Chunk]] +OutputT = List[Tuple[Chunk, Dict[str, Any]]] + + +class QdrantSearchEnrichmentHandler(EnrichmentSourceHandler[InputT, OutputT]): + def __init__( + self, + connection_parameters: QdrantConnectionParameters, + search_parameters: QdrantSearchParameters, + **kwargs, + ): + self._connection_parameters = connection_parameters + self._search_parameters = search_parameters + self.kwargs = kwargs + self.join_fn = join_fn + self.use_custom_types = True + + def __enter__(self): + 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, + ) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._client: + self._client.close() + + def __call__(self, request: Union[Chunk, List[Chunk]], *args, **kwargs) -> OutputT: + reqs = request if isinstance(request, list) else [request] + search_result = self._search_documents(reqs) + return self._get_call_response(reqs, search_result) + + def _search_documents(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + if isinstance(strategy, QdrantHybridSearchParameters): + return self._hybrid_search(chunks) + elif isinstance(strategy, QdrantDenseSearchParameters): + return self._dense_search(chunks) + elif isinstance(strategy, QdrantSparseSearchParameters): + return self._sparse_search(chunks) + else: + raise ValueError("Unknown Qdrant search strategy") + + def _dense_search(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + results = [] + for chunk in chunks: + if not chunk.embedding or chunk.embedding.dense_embedding is None: + raise ValueError( + f"Chunk {getattr(chunk, 'id', None)} missing dense embedding required for dense search" + ) + res = self._client.query_points( + collection_name=self._search_parameters.collection_name, + query=chunk.embedding.dense_embedding, + using=strategy.vector_name, + limit=strategy.limit, + query_filter=strategy.filter, + with_payload=True, + with_vectors=False, + score_threshold=strategy.score_threshold, + **strategy.kwargs, + ).points + results.append(res) + return results + + def _sparse_search(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + results = [] + for chunk in chunks: + if not chunk.embedding or chunk.embedding.sparse_embedding is None: + raise ValueError( + f"Chunk {getattr(chunk, 'id', None)} missing sparse embedding required for sparse search" + ) + indices, values = chunk.embedding.sparse_embedding + sparse_vec = models.SparseVector(indices=indices, values=values) + res = self._client.query_points( + collection_name=self._search_parameters.collection_name, + query=sparse_vec, + using=strategy.vector_name, + limit=strategy.limit, + query_filter=strategy.filter, + with_payload=True, + with_vectors=False, + score_threshold=strategy.score_threshold, + **strategy.kwargs, + ).points + results.append(res) + return results + + def _hybrid_search(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + results = [] + for chunk in chunks: + dense_vec = chunk.embedding.dense_embedding if chunk.embedding else None + sparse_emb = chunk.embedding.sparse_embedding if chunk.embedding else None + if dense_vec is None or sparse_emb is None: + raise ValueError( + f"Chunk {getattr(chunk, 'id', None)} missing dense or sparse embedding required for hybrid search" + ) + indices, values = sparse_emb + sparse_vec = models.SparseVector(indices=indices, values=values) + res = self._client.query_points( + collection_name=self._search_parameters.collection_name, + prefetch=[ + models.Prefetch( + using=strategy.dense.vector_name, + query=dense_vec, + filter=strategy.dense.filter, + limit=strategy.dense.limit, + params=None, + ), + models.Prefetch( + using=strategy.sparse.vector_name, + query=sparse_vec, + filter=strategy.sparse.filter, + limit=strategy.sparse.limit, + params=None, + ), + ], + query=models.FusionQuery(fusion=strategy.fusion), + limit=strategy.limit, + with_payload=True, + with_vectors=False, + **strategy.kwargs, + ).points + results.append(res) + return results + + def _get_call_response( + self, chunks: List[Chunk], search_result: List[List[models.ScoredPoint]] + ): + response = [] + for chunk, result in zip(chunks, search_result): + hits = [] + for point in result: + hit = {"id": point.id, "score": point.score, "payload": point.payload} + hits.append(hit) + response.append((chunk, {"hits": hits})) + return response + +def join_fn(left: Embedding, right: Dict[str, Any]) -> Embedding: + left.metadata['enrichment_data'] = right + return left 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..9521f7664cde --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py @@ -0,0 +1,194 @@ +# +# 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 parameterized import parameterized + +try: + from apache_beam.ml.rag.types import Chunk, Embedding, Content + from apache_beam.ml.rag.enrichment.qdrant_search import ( + QdrantSearchEnrichmentHandler, + QdrantConnectionParameters, + QdrantSearchParameters, + QdrantDenseSearchParameters, + QdrantSparseSearchParameters, + QdrantHybridSearchParameters, + ) +except ImportError as e: + raise unittest.SkipTest(f"Qdrant dependencies not installed: {str(e)}") + + +class TestQdrantSearchEnrichment(unittest.TestCase): + def test_invalid_connection_parameters(self): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters(url="") + search_params = QdrantSearchParameters( + collection_name="test_collection", + search_strategy=QdrantDenseSearchParameters(vector_name="dense"), + ) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, search_parameters=search_params + ) + self.assertIn("At least one of location, url, host, or path must be provided for Qdrant connection", str(context.exception)) + + @parameterized.expand( + [ + ( + lambda: QdrantSearchParameters( + collection_name="", + search_strategy=QdrantDenseSearchParameters(vector_name="dense"), + ), + "Collection name must be provided", + ), + ( + lambda: QdrantSearchParameters( + collection_name="test_collection", search_strategy=None + ), + "Search strategy must be provided", + ), + ] + ) + def test_invalid_search_parameters(self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters(url="http://localhost:6333") + search_params = create_params() + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, search_parameters=search_params + ) + self.assertIn(expected_error_msg, str(context.exception)) + + +class TestQdrantDenseSearchEnrichment(unittest.TestCase): + @parameterized.expand( + [ + ( + lambda: QdrantDenseSearchParameters(vector_name="dense", limit=-1), + "Search limit must be positive, got -1", + ), + ] + ) + def test_invalid_dense_search_parameters(self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters(url="http://localhost:6333") + dense_search_params = create_params() + search_params = QdrantSearchParameters( + collection_name="test_collection", search_strategy=dense_search_params + ) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, search_parameters=search_params + ) + self.assertIn(expected_error_msg, str(context.exception)) + + def test_missing_dense_embedding(self): + with self.assertRaises(ValueError) as context: + chunk = Chunk(id=1, content=None, embedding=Embedding(dense_embedding=None)) + connection_params = QdrantConnectionParameters(url="http://localhost:6333") + dense_search_params = QdrantDenseSearchParameters(vector_name="dense") + search_params = QdrantSearchParameters( + collection_name="test_collection", search_strategy=dense_search_params + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, search_parameters=search_params + ) + _ = handler._dense_search([chunk]) + self.assertIn( + "missing dense embedding required for dense search", str(context.exception) + ) + + +class TestQdrantSparseSearchEnrichment(unittest.TestCase): + @parameterized.expand( + [ + ( + lambda: QdrantSparseSearchParameters(vector_name="sparse", limit=-1), + "Search limit must be positive, got -1", + ), + ] + ) + def test_invalid_sparse_search_parameters(self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters(url="http://localhost:6333") + sparse_search_params = create_params() + search_params = QdrantSearchParameters( + collection_name="test_collection", search_strategy=sparse_search_params + ) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, search_parameters=search_params + ) + self.assertIn(expected_error_msg, str(context.exception)) + + def test_missing_sparse_embedding(self): + with self.assertRaises(ValueError) as context: + chunk = Chunk( + id=1, content=None, embedding=Embedding(sparse_embedding=None) + ) + connection_params = QdrantConnectionParameters(url="http://localhost:6333") + sparse_search_params = QdrantSparseSearchParameters(vector_name="sparse") + search_params = QdrantSearchParameters( + collection_name="test_collection", search_strategy=sparse_search_params + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, search_parameters=search_params + ) + _ = handler._sparse_search([chunk]) + self.assertIn( + "missing sparse embedding required for sparse search", + str(context.exception), + ) + + +class TestQdrantHybridSearchEnrichment(unittest.TestCase): + @parameterized.expand( + [ + ( + lambda: QdrantHybridSearchParameters( + dense=None, + sparse=QdrantSparseSearchParameters(vector_name="sparse"), + ), + "Both dense and sparse search parameters must be provided for hybrid search", + ), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(vector_name="dense"), sparse=None + ), + "Both dense and sparse search parameters must be provided for hybrid search", + ), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(vector_name="dense"), + sparse=QdrantSparseSearchParameters(vector_name="sparse"), + limit=-1, + ), + "Search limit must be positive, got -1", + ), + ] + ) + def test_invalid_hybrid_search_parameters(self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters(url="http://localhost:6333") + hybrid_search_params = create_params() + search_params = QdrantSearchParameters( + collection_name="test_collection", search_strategy=hybrid_search_params + ) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, search_parameters=search_params + ) + self.assertIn(expected_error_msg, str(context.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/sdks/python/setup.py b/sdks/python/setup.py index f60fe902bb7a..04fb2e370ae5 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -402,6 +402,7 @@ def get_portability_package_data(): 'zstandard>=0.18.0,<1', 'pyyaml>=3.12,<7.0.0', 'pymilvus>=2.5.10,<3.0.0', + 'qdrant-client>=1.15.0,<1.16.0' # Dynamic dependencies must be specified in a separate list, otherwise # Dependabot won't be able to parse the main list. Any dynamic # dependencies will not receive updates from Dependabot. @@ -438,7 +439,7 @@ def get_portability_package_data(): 'sqlalchemy>=1.3,<3.0', 'psycopg2-binary>=2.8.5,<2.9.10; python_version <= "3.9"', 'psycopg2-binary>=2.8.5,<3.0; python_version >= "3.10"', - 'testcontainers[mysql,kafka,milvus]>=4.0.0,<5.0.0', + 'testcontainers[mysql,kafka,milvus,qdrant]>=4.0.0,<5.0.0', 'cryptography>=41.0.2', 'hypothesis>5.0.0,<7.0.0', 'virtualenv-clone>=0.5,<1.0', From 6c7579ec86a79e744d55b832f4ec25f161c7d541 Mon Sep 17 00:00:00 2001 From: Anush008 Date: Mon, 28 Jul 2025 14:35:15 +0530 Subject: [PATCH 2/4] test: Qdrant IT test Signed-off-by: Anush008 --- .../ml/rag/enrichment/qdrant_search.py | 419 ++++++------- .../rag/enrichment/qdrant_search_it_test.py | 557 ++++++++++++++++++ .../ml/rag/enrichment/qdrant_search_test.py | 309 +++++----- 3 files changed, 932 insertions(+), 353 deletions(-) create mode 100644 sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py diff --git a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py index b4819143633b..05bd64bbbbcb 100644 --- a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py @@ -26,67 +26,64 @@ @dataclass class QdrantConnectionParameters: - location: Optional[str] = None - url: Optional[str] = None - port: Optional[int] = 6333 - grpc_port: int = 6334 - prefer_grpc: bool = False - https: Optional[bool] = None - api_key: Optional[str] = None - prefix: Optional[str] = None - timeout: Optional[int] = None - host: Optional[str] = None - path: Optional[str] = None - kwargs: Dict[str, Any] = field(default_factory=dict) - - def __post_init__(self): - if not (self.location or self.url or self.host or self.path): - raise ValueError( - "At least one of location, url, host, or path must be provided for Qdrant connection" - ) + location: Optional[str] = None + url: Optional[str] = None + port: Optional[int] = 6333 + grpc_port: int = 6334 + prefer_grpc: bool = False + https: Optional[bool] = None + api_key: Optional[str] = None + prefix: Optional[str] = None + timeout: Optional[int] = None + host: Optional[str] = None + path: Optional[str] = None + kwargs: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not (self.location or self.url or self.host or self.path): + raise ValueError( + "One of location, url, host, or path must be provided for Qdrant") @dataclass class QdrantDenseSearchParameters: - vector_name: str = "dense" - limit: int = 3 - filter: Optional[models.Filter] = None - score_threshold: Optional[float] = None - kwargs: Dict[str, Any] = field(default_factory=dict) + vector_name: str = "dense" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + kwargs: Dict[str, Any] = field(default_factory=dict) - def __post_init__(self): - if self.limit <= 0: - raise ValueError(f"Search limit must be positive, got {self.limit}") + def __post_init__(self): + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") @dataclass class QdrantSparseSearchParameters: - vector_name: str = "sparse" - limit: int = 3 - filter: Optional[models.Filter] = None - score_threshold: Optional[float] = None - kwargs: Dict[str, Any] = field(default_factory=dict) + vector_name: str = "sparse" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + kwargs: Dict[str, Any] = field(default_factory=dict) - def __post_init__(self): - if self.limit <= 0: - raise ValueError(f"Search limit must be positive, got {self.limit}") + def __post_init__(self): + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") @dataclass class QdrantHybridSearchParameters: - dense: QdrantDenseSearchParameters - sparse: QdrantSparseSearchParameters - limit: int = 3 - fusion: models.Fusion = models.Fusion.RRF - kwargs: Dict[str, Any] = field(default_factory=dict) + dense: QdrantDenseSearchParameters + sparse: QdrantSparseSearchParameters + limit: int = 3 + fusion: models.Fusion = models.Fusion.RRF + kwargs: Dict[str, Any] = field(default_factory=dict) - def __post_init__(self): - if not self.dense or not self.sparse: - raise ValueError( - "Both dense and sparse search parameters must be provided for hybrid search" - ) - if self.limit <= 0: - raise ValueError(f"Search limit must be positive, got {self.limit}") + def __post_init__(self): + if not self.dense or not self.sparse: + raise ValueError("Both dense and sparse search params are required") + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") QdrantSearchStrategyType = Union[ @@ -98,15 +95,31 @@ def __post_init__(self): @dataclass class QdrantSearchParameters: - collection_name: str - search_strategy: QdrantSearchStrategyType - kwargs: Dict[str, Any] = field(default_factory=dict) + collection_name: str + search_strategy: QdrantSearchStrategyType + kwargs: Dict[str, Any] = field(default_factory=dict) - 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") + 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 chunk. + + Args: + id: List of entity IDs returned from the search. + score: List of scores (similarity or distance) for each returned entity. + payload: List of dictionaries containing additional values for entities. + Each dictionary corresponds to one returned entity. + """ + + 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 = Union[Chunk, List[Chunk]] @@ -114,153 +127,161 @@ def __post_init__(self): class QdrantSearchEnrichmentHandler(EnrichmentSourceHandler[InputT, OutputT]): - def __init__( - self, - connection_parameters: QdrantConnectionParameters, - search_parameters: QdrantSearchParameters, - **kwargs, - ): - self._connection_parameters = connection_parameters - self._search_parameters = search_parameters - self.kwargs = kwargs - self.join_fn = join_fn - self.use_custom_types = True - - def __enter__(self): - 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, + def __init__( + self, + connection_parameters: QdrantConnectionParameters, + search_parameters: QdrantSearchParameters, + **kwargs, + ): + self._connection_parameters = connection_parameters + self._search_parameters = search_parameters + self.kwargs = kwargs + self.join_fn = join_fn + self.use_custom_types = True + + def __enter__(self): + 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, + ) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._client: + self._client.close() + + def __call__( + self, request: Union[Chunk, List[Chunk]], *args, **kwargs) -> OutputT: + reqs = request if isinstance(request, list) else [request] + search_result = self._search_documents(reqs) + + results = self._get_call_response(reqs, search_result) + + if isinstance(request, list): + return results + else: + return results[0] if results else (request, {}) + + def _search_documents(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + if isinstance(strategy, QdrantHybridSearchParameters): + return self._hybrid_search(chunks) + elif isinstance(strategy, QdrantDenseSearchParameters): + return self._dense_search(chunks) + elif isinstance(strategy, QdrantSparseSearchParameters): + return self._sparse_search(chunks) + else: + raise ValueError("Unknown Qdrant search strategy") + + def _dense_search(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + results = [] + for chunk in chunks: + if not chunk.embedding or chunk.embedding.dense_embedding is None: + raise ValueError( + f"Chunk {getattr(chunk, 'id', None)} missing dense embedding") + res = self._client.query_points( + collection_name=self._search_parameters.collection_name, + query=chunk.embedding.dense_embedding, + using=strategy.vector_name, + limit=strategy.limit, + query_filter=strategy.filter, + with_payload=True, + with_vectors=False, + score_threshold=strategy.score_threshold, + **strategy.kwargs, + ).points + results.append(res) + return results + + def _sparse_search(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + results = [] + for chunk in chunks: + if not chunk.embedding or chunk.embedding.sparse_embedding is None: + raise ValueError( + f"Chunk {getattr(chunk, 'id', None)} missing sparse embedding") + indices, values = chunk.embedding.sparse_embedding + sparse_vec = models.SparseVector(indices=indices, values=values) + res = self._client.query_points( + collection_name=self._search_parameters.collection_name, + query=sparse_vec, + using=strategy.vector_name, + limit=strategy.limit, + query_filter=strategy.filter, + with_payload=True, + with_vectors=False, + score_threshold=strategy.score_threshold, + **strategy.kwargs, + ).points + results.append(res) + return results + + def _hybrid_search(self, chunks: List[Chunk]): + strategy = self._search_parameters.search_strategy + results = [] + for chunk in chunks: + dense_vec = chunk.embedding.dense_embedding if chunk.embedding else None + sparse_emb = chunk.embedding.sparse_embedding if chunk.embedding else None + if dense_vec is None or sparse_emb is None: + raise ValueError( + f"Chunk {getattr(chunk, 'id', None)} missing dense/sparse embedding" ) + indices, values = sparse_emb + sparse_vec = models.SparseVector(indices=indices, values=values) + res = self._client.query_points( + collection_name=self._search_parameters.collection_name, + prefetch=[ + models.Prefetch( + using=strategy.dense.vector_name, + query=dense_vec, + filter=strategy.dense.filter, + limit=strategy.dense.limit, + params=None, + ), + models.Prefetch( + using=strategy.sparse.vector_name, + query=sparse_vec, + filter=strategy.sparse.filter, + limit=strategy.sparse.limit, + params=None, + ), + ], + query=models.FusionQuery(fusion=strategy.fusion), + limit=strategy.limit, + with_payload=True, + with_vectors=False, + **strategy.kwargs, + ).points + results.append(res) + return results + + def _get_call_response( + self, chunks: List[Chunk], search_result: List[List[models.ScoredPoint]]): + response = [] + for i in range(len(chunks)): + chunk = chunks[i] + hits = search_result[i] + result = QdrantSearchResult() + for hit in hits: + result.id.append(hit.id) + result.score.append(hit.score) + result.payload.append(hit.payload) + response.append((chunk, result.__dict__)) + return response - def __exit__(self, exc_type, exc_val, exc_tb): - if self._client: - self._client.close() - - def __call__(self, request: Union[Chunk, List[Chunk]], *args, **kwargs) -> OutputT: - reqs = request if isinstance(request, list) else [request] - search_result = self._search_documents(reqs) - return self._get_call_response(reqs, search_result) - - def _search_documents(self, chunks: List[Chunk]): - strategy = self._search_parameters.search_strategy - if isinstance(strategy, QdrantHybridSearchParameters): - return self._hybrid_search(chunks) - elif isinstance(strategy, QdrantDenseSearchParameters): - return self._dense_search(chunks) - elif isinstance(strategy, QdrantSparseSearchParameters): - return self._sparse_search(chunks) - else: - raise ValueError("Unknown Qdrant search strategy") - - def _dense_search(self, chunks: List[Chunk]): - strategy = self._search_parameters.search_strategy - results = [] - for chunk in chunks: - if not chunk.embedding or chunk.embedding.dense_embedding is None: - raise ValueError( - f"Chunk {getattr(chunk, 'id', None)} missing dense embedding required for dense search" - ) - res = self._client.query_points( - collection_name=self._search_parameters.collection_name, - query=chunk.embedding.dense_embedding, - using=strategy.vector_name, - limit=strategy.limit, - query_filter=strategy.filter, - with_payload=True, - with_vectors=False, - score_threshold=strategy.score_threshold, - **strategy.kwargs, - ).points - results.append(res) - return results - - def _sparse_search(self, chunks: List[Chunk]): - strategy = self._search_parameters.search_strategy - results = [] - for chunk in chunks: - if not chunk.embedding or chunk.embedding.sparse_embedding is None: - raise ValueError( - f"Chunk {getattr(chunk, 'id', None)} missing sparse embedding required for sparse search" - ) - indices, values = chunk.embedding.sparse_embedding - sparse_vec = models.SparseVector(indices=indices, values=values) - res = self._client.query_points( - collection_name=self._search_parameters.collection_name, - query=sparse_vec, - using=strategy.vector_name, - limit=strategy.limit, - query_filter=strategy.filter, - with_payload=True, - with_vectors=False, - score_threshold=strategy.score_threshold, - **strategy.kwargs, - ).points - results.append(res) - return results - - def _hybrid_search(self, chunks: List[Chunk]): - strategy = self._search_parameters.search_strategy - results = [] - for chunk in chunks: - dense_vec = chunk.embedding.dense_embedding if chunk.embedding else None - sparse_emb = chunk.embedding.sparse_embedding if chunk.embedding else None - if dense_vec is None or sparse_emb is None: - raise ValueError( - f"Chunk {getattr(chunk, 'id', None)} missing dense or sparse embedding required for hybrid search" - ) - indices, values = sparse_emb - sparse_vec = models.SparseVector(indices=indices, values=values) - res = self._client.query_points( - collection_name=self._search_parameters.collection_name, - prefetch=[ - models.Prefetch( - using=strategy.dense.vector_name, - query=dense_vec, - filter=strategy.dense.filter, - limit=strategy.dense.limit, - params=None, - ), - models.Prefetch( - using=strategy.sparse.vector_name, - query=sparse_vec, - filter=strategy.sparse.filter, - limit=strategy.sparse.limit, - params=None, - ), - ], - query=models.FusionQuery(fusion=strategy.fusion), - limit=strategy.limit, - with_payload=True, - with_vectors=False, - **strategy.kwargs, - ).points - results.append(res) - return results - - def _get_call_response( - self, chunks: List[Chunk], search_result: List[List[models.ScoredPoint]] - ): - response = [] - for chunk, result in zip(chunks, search_result): - hits = [] - for point in result: - hit = {"id": point.id, "score": point.score, "payload": point.payload} - hits.append(hit) - response.append((chunk, {"hits": hits})) - return response def join_fn(left: Embedding, right: Dict[str, Any]) -> Embedding: - left.metadata['enrichment_data'] = right + left.metadata["enrichment_data"] = right return left 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..d239a9e16ecf --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py @@ -0,0 +1,557 @@ +# +# 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 dataclasses import dataclass + +import pytest +from testcontainers.qdrant import QdrantContainer + +import apache_beam as beam +from apache_beam.ml.rag.types import Chunk, Content, Embedding +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that, equal_to + +try: + from apache_beam.transforms.enrichment import Enrichment + from apache_beam.ml.rag.enrichment.qdrant_search import ( + QdrantSearchEnrichmentHandler, + QdrantConnectionParameters, + QdrantSearchParameters, + QdrantDenseSearchParameters, + ) +except ImportError as e: + raise unittest.SkipTest(f"Qdrant dependencies not installed: {str(e)}") + + +@dataclass +class QdrantITDataConstruct: + id: int + content: str + metadata: dict + dense_embedding: list + sparse_embedding: dict + + def __getitem__(self, key): + return getattr(self, key) + + +QDRANT_IT_CONFIG = { + "collection_name": "docs_catalog", + "corpus": [ + QdrantITDataConstruct( + id=1, + content="This is a test document", + metadata={"language": "en"}, + dense_embedding=[0.1, 0.2, 0.3], + sparse_embedding={ + 1: 0.05, 2: 0.41, 3: 0.05, 4: 0.41 + }, + ), + QdrantITDataConstruct( + id=2, + content="Another test document", + metadata={"language": "en"}, + dense_embedding=[0.2, 0.3, 0.4], + sparse_embedding={ + 1: 0.07, 3: 3.07, 0: 0.53 + }, + ), + QdrantITDataConstruct( + id=3, + content="وثيقة اختبار", + metadata={"language": "ar"}, + dense_embedding=[0.3, 0.4, 0.5], + sparse_embedding={ + 6: 0.62, 5: 0.62 + }, + ), + ], +} + + +def _to_sparse_dict(sparse_embedding): + if isinstance(sparse_embedding, dict): + return sparse_embedding + if sparse_embedding is None: + return None + indices, values = sparse_embedding + return dict(zip(indices, values)) + + +def assert_chunks_equivalent( + actual_chunks, + expected_chunks, + check_dense=True, + check_sparse=True, + check_id_order=True, +): + actual_sorted = sorted(actual_chunks, key=lambda c: c.id) + expected_sorted = sorted(expected_chunks, key=lambda c: c.id) + assert len(actual_sorted) == len(expected_sorted) + for actual, expected in zip(actual_sorted, expected_sorted): + assert actual.id == expected.id + if check_dense: + assert ( + actual.embedding.dense_embedding == expected.embedding.dense_embedding + ) + if check_sparse: + assert _to_sparse_dict( + actual.embedding.sparse_embedding) == _to_sparse_dict( + expected.embedding.sparse_embedding) + assert "enrichment_data" in actual.metadata + actual_data = actual.metadata["enrichment_data"] + expected_data = expected.metadata["enrichment_data"] + if check_id_order: + assert actual_data["id"] == expected_data["id"] + else: + assert set(actual_data["id"]) == set(expected_data["id"]) + for a, e in zip(actual_data["payload"], expected_data["payload"]): + assert a["content"] == e["content"] + assert a["metadata"] == e["metadata"] + + +@pytest.mark.uses_testcontainer +class TestQdrantSearchEnrichmentIT(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls._container = QdrantContainer() + cls._container.start() + cls._host = cls._container.get_container_host_ip() + cls._port = int(cls._container.get_exposed_port(6333)) + cls._connection_params = QdrantConnectionParameters( + host=cls._host, port=cls._port) + from qdrant_client import QdrantClient, models + + client = QdrantClient( + host=cls._host, port=cls._port, check_compatibility=False) + points = [] + for doc in QDRANT_IT_CONFIG["corpus"]: + points.append( + models.PointStruct( + id=doc.id, + vector={ + "dense": doc.dense_embedding, + "sparse": models.SparseVector( + indices=list(doc.sparse_embedding.keys()), + values=list(doc.sparse_embedding.values()), + ), + }, + payload={ + "content": doc.content, "metadata": doc.metadata + }, + )) + client.create_collection( + collection_name=QDRANT_IT_CONFIG["collection_name"], + vectors_config={ + "dense": models.VectorParams( + size=3, distance=models.Distance.COSINE), + }, + sparse_vectors_config={ + "sparse": models.SparseVectorParams(), + }, + ) + client.upsert( + collection_name=QDRANT_IT_CONFIG["collection_name"], points=points) + + @classmethod + def tearDownClass(cls): + cls._container.stop() + + def test_dense_search(self): + test_chunks = [ + Chunk( + id="query1", + embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), + content=Content(), + ), + Chunk( + id="query2", + embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), + content=Content(), + ), + ] + search_parameters = QdrantSearchParameters( + collection_name=QDRANT_IT_CONFIG["collection_name"], + search_strategy=QdrantDenseSearchParameters( + vector_name="dense", limit=2), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + with TestPipeline(is_integration_test=True) as p: + result = p | beam.Create(test_chunks) | Enrichment(handler) + + assert_that(result, lambda actual: len(actual) == 2) + + def test_dense_search_empty(self): + test_chunks = [] + search_parameters = QdrantSearchParameters( + collection_name=QDRANT_IT_CONFIG["collection_name"], + search_strategy=QdrantDenseSearchParameters( + vector_name="dense", limit=2), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + with TestPipeline(is_integration_test=True) as p: + result = p | beam.Create(test_chunks) | Enrichment(handler) + assert_that(result, equal_to([])) + + def test_invalid_collection_name(self): + search_parameters = QdrantSearchParameters( + collection_name="nonexistent_collection", + search_strategy=QdrantDenseSearchParameters( + vector_name="dense", limit=2), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + test_chunks = [ + Chunk( + id="query1", + embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), + content=Content(), + ) + ] + with self.assertRaises(Exception) as context: + with TestPipeline() as p: + _ = p | beam.Create(test_chunks) | Enrichment(handler) + self.assertIn("not found", str(context.exception).lower()) + + def test_missing_embedding(self): + search_parameters = QdrantSearchParameters( + collection_name=QDRANT_IT_CONFIG["collection_name"], + search_strategy=QdrantDenseSearchParameters( + vector_name="dense", limit=2), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + test_chunks = [ + Chunk( + id="query1", + embedding=Embedding(dense_embedding=None), + content=Content(), + ) + ] + with self.assertRaises(ValueError) as context: + with TestPipeline() as p: + _ = p | beam.Create(test_chunks) | Enrichment(handler) + self.assertIn("missing dense embedding", str(context.exception).lower()) + + def test_batched_dense_search(self): + test_chunks = [ + Chunk( + id="query1", + embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), + content=Content(), + ), + Chunk( + id="query2", + embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), + content=Content(), + ), + Chunk( + id="query3", + embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5]), + content=Content(), + ), + ] + search_parameters = QdrantSearchParameters( + collection_name=QDRANT_IT_CONFIG["collection_name"], + search_strategy=QdrantDenseSearchParameters( + vector_name="dense", limit=2), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + with TestPipeline(is_integration_test=True) as p: + result = p | beam.Create(test_chunks) | Enrichment(handler) + assert_that(result, lambda actual: len(actual) == 3) + + def test_dense_search_values(self): + test_chunks = [ + Chunk( + id="query1", + embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), + content=Content(), + ), + Chunk( + id="query2", + embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), + content=Content(), + ), + ] + search_parameters = QdrantSearchParameters( + collection_name=QDRANT_IT_CONFIG["collection_name"], + search_strategy=QdrantDenseSearchParameters( + vector_name="dense", limit=2), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + expected_chunks = [ + Chunk( + id="query1", + content=Content(), + embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), + metadata={ + "enrichment_data": { + "id": [1, 2], + "score": None, + "payload": [ + { + "content": "This is a test document", + "metadata": { + "language": "en" + }, + }, + { + "content": "Another test document", + "metadata": { + "language": "en" + }, + }, + ], + } + }, + ), + Chunk( + id="query2", + content=Content(), + embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), + metadata={ + "enrichment_data": { + "id": [2, 3], + "score": None, + "payload": [ + { + "content": "Another test document", + "metadata": { + "language": "en" + }, + }, + { + "content": "وثيقة اختبار", + "metadata": { + "language": "ar" + } + }, + ], + } + }, + ), + ] + with TestPipeline(is_integration_test=True) as p: + result = p | beam.Create(test_chunks) | Enrichment(handler) + assert_that( + result, + lambda actual: assert_chunks_equivalent( + actual, + expected_chunks, + check_dense=True, + check_sparse=False, + check_id_order=True, ), + ) + + def test_sparse_search_values(self): + from apache_beam.ml.rag.enrichment.qdrant_search import ( + QdrantSparseSearchParameters, ) + + test_chunks = [ + Chunk( + id="query1", + embedding=Embedding( + dense_embedding=None, + sparse_embedding=([1, 2, 3, 4], [0.05, 0.41, 0.05, 0.41]), + ), + content=Content(), + ), + Chunk( + id="query2", + embedding=Embedding( + dense_embedding=None, + sparse_embedding=([1, 3, 0], [0.07, 3.07, 0.53]), + ), + content=Content(), + ), + ] + search_parameters = QdrantSearchParameters( + collection_name=QDRANT_IT_CONFIG["collection_name"], + search_strategy=QdrantSparseSearchParameters( + vector_name="sparse", limit=2), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + expected_chunks = [ + Chunk( + id="query1", + content=Content(), + embedding=Embedding( + dense_embedding=None, + sparse_embedding={ + 1: 0.05, 2: 0.41, 3: 0.05, 4: 0.41 + }, + ), + metadata={ + "enrichment_data": { + "id": [1, 2], + "score": None, + "payload": [ + { + "content": "This is a test document", + "metadata": { + "language": "en" + }, + }, + { + "content": "Another test document", + "metadata": { + "language": "en" + }, + }, + ], + } + }, + ), + Chunk( + id="query2", + content=Content(), + embedding=Embedding( + dense_embedding=None, + sparse_embedding={ + 1: 0.07, 3: 3.07, 0: 0.53 + }), + metadata={ + "enrichment_data": { + "id": [2, 1], + "score": None, + "payload": [ + { + "content": "Another test document", + "metadata": { + "language": "en" + }, + }, + { + "content": "This is a test document", + "metadata": { + "language": "en" + }, + }, + ], + } + }, + ), + ] + with TestPipeline(is_integration_test=True) as p: + result = p | beam.Create(test_chunks) | Enrichment(handler) + assert_that( + result, + lambda actual: assert_chunks_equivalent( + actual, + expected_chunks, + check_dense=False, + check_sparse=True, + check_id_order=True, ), + ) + + def test_hybrid_search_values(self): + from apache_beam.ml.rag.enrichment.qdrant_search import ( + QdrantHybridSearchParameters, + QdrantDenseSearchParameters, + QdrantSparseSearchParameters, + ) + + test_chunks = [ + Chunk( + id="query1", + embedding=Embedding( + dense_embedding=[0.1, 0.2, 0.3], + sparse_embedding=([1, 2, 3, 4], [0.05, 0.41, 0.05, 0.41]), + ), + content=Content(), + ) + ] + search_parameters = QdrantSearchParameters( + collection_name=QDRANT_IT_CONFIG["collection_name"], + search_strategy=QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(vector_name="dense", limit=2), + sparse=QdrantSparseSearchParameters(vector_name="sparse", limit=2), + limit=2, + ), + ) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=self._connection_params, + search_parameters=search_parameters, + ) + expected_chunks = [ + Chunk( + id="query1", + content=Content(), + embedding=Embedding( + dense_embedding=[0.1, 0.2, 0.3], + sparse_embedding={ + 1: 0.05, 2: 0.41, 3: 0.05, 4: 0.41 + }, + ), + metadata={ + "enrichment_data": { + "id": [1, 2], + "score": None, + "payload": [ + { + "content": "This is a test document", + "metadata": { + "language": "en" + }, + }, + { + "content": "Another test document", + "metadata": { + "language": "en" + }, + }, + ], + } + }, + ) + ] + with TestPipeline(is_integration_test=True) as p: + result = p | beam.Create(test_chunks) | Enrichment(handler) + assert_that( + result, + lambda actual: assert_chunks_equivalent( + actual, + expected_chunks, + check_dense=True, + check_sparse=True, + check_id_order=False, ), + ) + + +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 index 9521f7664cde..acb90bc69490 100644 --- a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py @@ -19,176 +19,177 @@ from parameterized import parameterized try: - from apache_beam.ml.rag.types import Chunk, Embedding, Content - from apache_beam.ml.rag.enrichment.qdrant_search import ( - QdrantSearchEnrichmentHandler, - QdrantConnectionParameters, - QdrantSearchParameters, - QdrantDenseSearchParameters, - QdrantSparseSearchParameters, - QdrantHybridSearchParameters, - ) + from apache_beam.ml.rag.types import Chunk, Embedding + from apache_beam.ml.rag.enrichment.qdrant_search import ( + QdrantSearchEnrichmentHandler, + QdrantConnectionParameters, + QdrantSearchParameters, + QdrantDenseSearchParameters, + QdrantSparseSearchParameters, + QdrantHybridSearchParameters, + ) except ImportError as e: - raise unittest.SkipTest(f"Qdrant dependencies not installed: {str(e)}") + raise unittest.SkipTest(f"Qdrant dependencies not installed: {str(e)}") class TestQdrantSearchEnrichment(unittest.TestCase): - def test_invalid_connection_parameters(self): - with self.assertRaises(ValueError) as context: - connection_params = QdrantConnectionParameters(url="") - search_params = QdrantSearchParameters( - collection_name="test_collection", - search_strategy=QdrantDenseSearchParameters(vector_name="dense"), - ) - _ = QdrantSearchEnrichmentHandler( - connection_parameters=connection_params, search_parameters=search_params - ) - self.assertIn("At least one of location, url, host, or path must be provided for Qdrant connection", str(context.exception)) - - @parameterized.expand( - [ - ( - lambda: QdrantSearchParameters( - collection_name="", - search_strategy=QdrantDenseSearchParameters(vector_name="dense"), - ), - "Collection name must be provided", - ), - ( - lambda: QdrantSearchParameters( - collection_name="test_collection", search_strategy=None - ), - "Search strategy must be provided", - ), - ] + def test_invalid_connection_parameters(self): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters(url="") + search_params = QdrantSearchParameters( + collection_name="test_collection", + search_strategy=QdrantDenseSearchParameters(vector_name="dense"), + ) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, + search_parameters=search_params) + self.assertIn( + "One of location, url, host, or path must be provided for Qdrant", + str(context.exception), ) - def test_invalid_search_parameters(self, create_params, expected_error_msg): - with self.assertRaises(ValueError) as context: - connection_params = QdrantConnectionParameters(url="http://localhost:6333") - search_params = create_params() - _ = QdrantSearchEnrichmentHandler( - connection_parameters=connection_params, search_parameters=search_params - ) - self.assertIn(expected_error_msg, str(context.exception)) + + @parameterized.expand([ + ( + lambda: QdrantSearchParameters( + collection_name="", + search_strategy=QdrantDenseSearchParameters(vector_name="dense"), + ), + "Collection name must be provided", + ), + ( + lambda: QdrantSearchParameters( + collection_name="test_collection", search_strategy=None), + "Search strategy must be provided", + ), + ]) + def test_invalid_search_parameters(self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters( + url="http://localhost:6333") + search_params = create_params() + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, + search_parameters=search_params) + self.assertIn(expected_error_msg, str(context.exception)) class TestQdrantDenseSearchEnrichment(unittest.TestCase): - @parameterized.expand( - [ - ( - lambda: QdrantDenseSearchParameters(vector_name="dense", limit=-1), - "Search limit must be positive, got -1", - ), - ] - ) - def test_invalid_dense_search_parameters(self, create_params, expected_error_msg): - with self.assertRaises(ValueError) as context: - connection_params = QdrantConnectionParameters(url="http://localhost:6333") - dense_search_params = create_params() - search_params = QdrantSearchParameters( - collection_name="test_collection", search_strategy=dense_search_params - ) - _ = QdrantSearchEnrichmentHandler( - connection_parameters=connection_params, search_parameters=search_params - ) - self.assertIn(expected_error_msg, str(context.exception)) + @parameterized.expand([ + ( + lambda: QdrantDenseSearchParameters(vector_name="dense", limit=-1), + "Search limit must be positive, got -1", + ), + ]) + def test_invalid_dense_search_parameters( + self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters( + url="http://localhost:6333") + dense_search_params = create_params() + search_params = QdrantSearchParameters( + collection_name="test_collection", + search_strategy=dense_search_params) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, + search_parameters=search_params) + self.assertIn(expected_error_msg, str(context.exception)) - def test_missing_dense_embedding(self): - with self.assertRaises(ValueError) as context: - chunk = Chunk(id=1, content=None, embedding=Embedding(dense_embedding=None)) - connection_params = QdrantConnectionParameters(url="http://localhost:6333") - dense_search_params = QdrantDenseSearchParameters(vector_name="dense") - search_params = QdrantSearchParameters( - collection_name="test_collection", search_strategy=dense_search_params - ) - handler = QdrantSearchEnrichmentHandler( - connection_parameters=connection_params, search_parameters=search_params - ) - _ = handler._dense_search([chunk]) - self.assertIn( - "missing dense embedding required for dense search", str(context.exception) - ) + def test_missing_dense_embedding(self): + with self.assertRaises(ValueError) as context: + chunk = Chunk( + id=1, content=None, embedding=Embedding(dense_embedding=None)) + connection_params = QdrantConnectionParameters( + url="http://localhost:6333") + dense_search_params = QdrantDenseSearchParameters(vector_name="dense") + search_params = QdrantSearchParameters( + collection_name="test_collection", + search_strategy=dense_search_params) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, + search_parameters=search_params) + _ = handler._dense_search([chunk]) + self.assertIn( + "missing dense embedding required for dense search", + str(context.exception)) class TestQdrantSparseSearchEnrichment(unittest.TestCase): - @parameterized.expand( - [ - ( - lambda: QdrantSparseSearchParameters(vector_name="sparse", limit=-1), - "Search limit must be positive, got -1", - ), - ] - ) - def test_invalid_sparse_search_parameters(self, create_params, expected_error_msg): - with self.assertRaises(ValueError) as context: - connection_params = QdrantConnectionParameters(url="http://localhost:6333") - sparse_search_params = create_params() - search_params = QdrantSearchParameters( - collection_name="test_collection", search_strategy=sparse_search_params - ) - _ = QdrantSearchEnrichmentHandler( - connection_parameters=connection_params, search_parameters=search_params - ) - self.assertIn(expected_error_msg, str(context.exception)) + @parameterized.expand([ + ( + lambda: QdrantSparseSearchParameters(vector_name="sparse", limit=-1), + "Search limit must be positive, got -1", + ), + ]) + def test_invalid_sparse_search_parameters( + self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters( + url="http://localhost:6333") + sparse_search_params = create_params() + search_params = QdrantSearchParameters( + collection_name="test_collection", + search_strategy=sparse_search_params) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, + search_parameters=search_params) + self.assertIn(expected_error_msg, str(context.exception)) - def test_missing_sparse_embedding(self): - with self.assertRaises(ValueError) as context: - chunk = Chunk( - id=1, content=None, embedding=Embedding(sparse_embedding=None) - ) - connection_params = QdrantConnectionParameters(url="http://localhost:6333") - sparse_search_params = QdrantSparseSearchParameters(vector_name="sparse") - search_params = QdrantSearchParameters( - collection_name="test_collection", search_strategy=sparse_search_params - ) - handler = QdrantSearchEnrichmentHandler( - connection_parameters=connection_params, search_parameters=search_params - ) - _ = handler._sparse_search([chunk]) - self.assertIn( - "missing sparse embedding required for sparse search", - str(context.exception), - ) + def test_missing_sparse_embedding(self): + with self.assertRaises(ValueError) as context: + chunk = Chunk( + id=1, content=None, embedding=Embedding(sparse_embedding=None)) + connection_params = QdrantConnectionParameters( + url="http://localhost:6333") + sparse_search_params = QdrantSparseSearchParameters(vector_name="sparse") + search_params = QdrantSearchParameters( + collection_name="test_collection", + search_strategy=sparse_search_params) + handler = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, + search_parameters=search_params) + _ = handler._sparse_search([chunk]) + self.assertIn( + "missing sparse embedding", + str(context.exception), + ) class TestQdrantHybridSearchEnrichment(unittest.TestCase): - @parameterized.expand( - [ - ( - lambda: QdrantHybridSearchParameters( - dense=None, - sparse=QdrantSparseSearchParameters(vector_name="sparse"), - ), - "Both dense and sparse search parameters must be provided for hybrid search", - ), - ( - lambda: QdrantHybridSearchParameters( - dense=QdrantDenseSearchParameters(vector_name="dense"), sparse=None - ), - "Both dense and sparse search parameters must be provided for hybrid search", - ), - ( - lambda: QdrantHybridSearchParameters( - dense=QdrantDenseSearchParameters(vector_name="dense"), - sparse=QdrantSparseSearchParameters(vector_name="sparse"), - limit=-1, - ), - "Search limit must be positive, got -1", - ), - ] - ) - def test_invalid_hybrid_search_parameters(self, create_params, expected_error_msg): - with self.assertRaises(ValueError) as context: - connection_params = QdrantConnectionParameters(url="http://localhost:6333") - hybrid_search_params = create_params() - search_params = QdrantSearchParameters( - collection_name="test_collection", search_strategy=hybrid_search_params - ) - _ = QdrantSearchEnrichmentHandler( - connection_parameters=connection_params, search_parameters=search_params - ) - self.assertIn(expected_error_msg, str(context.exception)) + @parameterized.expand([ + ( + lambda: QdrantHybridSearchParameters( + dense=None, + sparse=QdrantSparseSearchParameters(vector_name="sparse"), ), + "Both dense and sparse search params are required", + ), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(vector_name="dense"), sparse= + None), + "Both dense and sparse search params are required", + ), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(vector_name="dense"), + sparse=QdrantSparseSearchParameters(vector_name="sparse"), + limit=-1, ), + "Search limit must be positive, got -1", + ), + ]) + def test_invalid_hybrid_search_parameters( + self, create_params, expected_error_msg): + with self.assertRaises(ValueError) as context: + connection_params = QdrantConnectionParameters( + url="http://localhost:6333") + hybrid_search_params = create_params() + search_params = QdrantSearchParameters( + collection_name="test_collection", + search_strategy=hybrid_search_params) + _ = QdrantSearchEnrichmentHandler( + connection_parameters=connection_params, + search_parameters=search_params) + self.assertIn(expected_error_msg, str(context.exception)) if __name__ == "__main__": - unittest.main() + unittest.main() From 0dbf29cfba7f74d9864942e3a112fe89cfc799e8 Mon Sep 17 00:00:00 2001 From: Anush008 Date: Mon, 28 Jul 2025 14:39:02 +0530 Subject: [PATCH 3/4] test: Fix text in error message Signed-off-by: Anush008 --- .../ml/rag/enrichment/qdrant_search_it_test.py | 15 ++++----------- .../ml/rag/enrichment/qdrant_search_test.py | 4 +--- 2 files changed, 5 insertions(+), 14 deletions(-) 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 index d239a9e16ecf..acf5389398e8 100644 --- 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 @@ -98,7 +98,6 @@ def assert_chunks_equivalent( expected_chunks, check_dense=True, check_sparse=True, - check_id_order=True, ): actual_sorted = sorted(actual_chunks, key=lambda c: c.id) expected_sorted = sorted(expected_chunks, key=lambda c: c.id) @@ -116,10 +115,7 @@ def assert_chunks_equivalent( assert "enrichment_data" in actual.metadata actual_data = actual.metadata["enrichment_data"] expected_data = expected.metadata["enrichment_data"] - if check_id_order: - assert actual_data["id"] == expected_data["id"] - else: - assert set(actual_data["id"]) == set(expected_data["id"]) + assert actual_data["id"] == expected_data["id"] for a, e in zip(actual_data["payload"], expected_data["payload"]): assert a["content"] == e["content"] assert a["metadata"] == e["metadata"] @@ -371,8 +367,7 @@ def test_dense_search_values(self): actual, expected_chunks, check_dense=True, - check_sparse=False, - check_id_order=True, ), + check_sparse=False, ), ) def test_sparse_search_values(self): @@ -475,8 +470,7 @@ def test_sparse_search_values(self): actual, expected_chunks, check_dense=False, - check_sparse=True, - check_id_order=True, ), + check_sparse=True, ), ) def test_hybrid_search_values(self): @@ -548,8 +542,7 @@ def test_hybrid_search_values(self): actual, expected_chunks, check_dense=True, - check_sparse=True, - check_id_order=False, ), + check_sparse=True, ), ) 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 index acb90bc69490..52f2365c90e2 100644 --- a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py @@ -108,9 +108,7 @@ def test_missing_dense_embedding(self): connection_parameters=connection_params, search_parameters=search_params) _ = handler._dense_search([chunk]) - self.assertIn( - "missing dense embedding required for dense search", - str(context.exception)) + self.assertIn("missing dense embedding", str(context.exception)) class TestQdrantSparseSearchEnrichment(unittest.TestCase): From 08b08225e7adde408b4d2bb23221577b58f6c2e5 Mon Sep 17 00:00:00 2001 From: Anush008 Date: Mon, 28 Jul 2025 23:21:15 +0530 Subject: [PATCH 4/4] docs: CHANGES.md Signed-off-by: Anush008 --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index ea7090446d00..0b35f8d4b7fd 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -73,6 +73,7 @@ ## New Features / Improvements * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* [Qdrant](https://qdrant.tech/) vector search enrichment handler (Python) ([#35216](https://github.com/apache/beam/pull/35686)). ## Breaking Changes