sdks/python: address Danny's feedback (2)

Author: moawnallah

sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py

@@ -170,10 +170,17 @@ def enrichment_with_milvus():
       output_fields=["id", "content", "domain", "cost", "metadata"],
       round_decimal=2)
 
+  # MilvusCollectionLoadParameters is optional and provides fine-grained control
+  # over how collections are loaded into memory. For simple use cases or when
+  # getting started, this parameter can be omitted to use default loading
+  # behavior. Consider using it in resource-constrained environments to optimize
+  # memory usage and query performance.
   collection_load_parameters = MilvusCollectionLoadParameters()
 
   milvus_search_handler = MilvusSearchEnrichmentHandler(
-      connection_parameters, search_parameters, collection_load_parameters)
+      connection_parameters=connection_parameters,
+      search_parameters=search_parameters,
+      collection_load_parameters=collection_load_parameters)
   with beam.Pipeline() as p:
     _ = (
         p

sdks/python/apache_beam/ml/rag/enrichment/milvus_search.py

@@ -229,28 +229,41 @@ def __post_init__(self):
 
 
 @dataclass
-class HybridSearchNamespace:
-  """Namespace containing all parameters for hybrid search operations.
+class HybridSearchParameters:
+  """Parameters for hybrid (vector + keyword) search operations.
 
   Args:
     vector: Parameters for the vector search component.
     keyword: Parameters for the keyword search component.
-    hybrid: Parameters for combining the vector and keyword results.
+    ranker: Ranker for combining vector and keyword search results.
+      Example: RRFRanker(k=100).
+    limit: Maximum number of results to return per query. Defaults to 3 search
+      results.
+    kwargs: Optional keyword arguments for additional hybrid search parameters.
+      Enables forward compatibility.
   """
   vector: VectorSearchParameters
   keyword: KeywordSearchParameters
-  hybrid: HybridSearchParameters
+  ranker: MilvusBaseRanker
+  limit: int = 3
+  kwargs: Dict[str, Any] = field(default_factory=dict)
 
   def __post_init__(self):
-    if not self.vector or not self.keyword or not self.hybrid:
+    if not self.vector or not self.keyword:
       raise ValueError(
-          "Vector, keyword, and hybrid search parameters must be provided for "
+          "Vector and keyword search parameters must be provided for "
           "hybrid search")
 
+    if not self.ranker:
+      raise ValueError("Ranker must be provided for hybrid search")
+
+    if self.limit <= 0:
+      raise ValueError(f"Search limit must be positive, got {self.limit}")
+
 
 SearchStrategyType = Union[VectorSearchParameters,
                            KeywordSearchParameters,
-                           HybridSearchNamespace]
+                           HybridSearchParameters]
 
 
 @dataclass
@@ -315,6 +328,22 @@ class MilvusCollectionLoadParameters:
   kwargs: Dict[str, Any] = field(default_factory=dict)
 
 
+@dataclass
+class MilvusSearchResult:
+  """Search result from Milvus per chunk.
+
+  Args:
+    id: List of entity IDs returned from the search. Can be either string or
+      integer IDs.
+    distance: List of distances/similarity scores for each returned entity.
+    fields: List of dictionaries containing additional field values for each
+      entity. Each dictionary corresponds to one returned entity.
+  """
+  id: List[Union[str, int]] = field(default_factory=list)
+  distance: List[float] = field(default_factory=list)
+  fields: List[Dict[str, Any]] = field(default_factory=list)
+
+
 InputT, OutputT = Union[Chunk, List[Chunk]], List[Tuple[Chunk, Dict[str, Any]]]
 
 
@@ -343,8 +372,8 @@ def __init__(
       self,
       connection_parameters: MilvusConnectionParameters,
       search_parameters: MilvusSearchParameters,
-      collection_load_parameters: MilvusCollectionLoadParameters,
       *,
+      collection_load_parameters: Optional[MilvusCollectionLoadParameters],
       min_batch_size: int = 1,
       max_batch_size: int = 1000,
       **kwargs):
@@ -360,7 +389,7 @@ def __init__(
       milvus_handler = MilvusSearchEnrichmentHandler(
         connection_paramters,
         search_parameters,
-        collection_load_parameters,
+        collection_load_parameters=collection_load_parameters,
         min_batch_size=10,
         max_batch_size=100)
 
@@ -371,8 +400,8 @@ def __init__(
       search_parameters (MilvusSearchParameters): Configuration for search
         operations, including collection name, search strategy, and output
         fields.
-      collection_load_parameters (MilvusCollectionLoadParameters): Parameters
-        controlling how collections are loaded into memory, which can
+      collection_load_parameters (Optional[MilvusCollectionLoadParameters]):
+        Parameters controlling how collections are loaded into memory, which can
         significantly impact resource usage and performance.
       min_batch_size (int): Minimum number of elements to batch together when
         querying Milvus. Default is 1 (no batching when max_batch_size is 1).
@@ -390,22 +419,25 @@ def __init__(
     self._connection_parameters = connection_parameters
     self._search_parameters = search_parameters
     self._collection_load_parameters = collection_load_parameters
-    self.kwargs = kwargs
+    if not self._collection_load_parameters:
+      self._collection_load_parameters = MilvusCollectionLoadParameters()
     self._batching_kwargs = {
         'min_batch_size': min_batch_size, 'max_batch_size': max_batch_size
     }
+    self.kwargs = kwargs
     self.join_fn = join_fn
     self.use_custom_types = True
 
   def __enter__(self):
-    connectionParams = unpack_dataclass_with_kwargs(self._connection_parameters)
-    loadCollectionParams = unpack_dataclass_with_kwargs(
+    connection_params = unpack_dataclass_with_kwargs(
+        self._connection_parameters)
+    collection_load_params = unpack_dataclass_with_kwargs(
         self._collection_load_parameters)
-    self._client = MilvusClient(**connectionParams)
+    self._client = MilvusClient(**connection_params)
     self._client.load_collection(
         collection_name=self.collection_name,
         partition_names=self.partition_names,
-        **loadCollectionParams)
+        **collection_load_params)
 
   def __call__(self, request: Union[Chunk, List[Chunk]], *args,
                **kwargs) -> List[Tuple[Chunk, Dict[str, Any]]]:
@@ -414,18 +446,18 @@ def __call__(self, request: Union[Chunk, List[Chunk]], *args,
     return self._get_call_response(reqs, search_result)
 
   def _search_documents(self, chunks: List[Chunk]):
-    if isinstance(self.search_strategy, HybridSearchNamespace):
+    if isinstance(self.search_strategy, HybridSearchParameters):
       data = self._get_hybrid_search_data(chunks)
-      hybrid_search_params = unpack_dataclass_with_kwargs(
-          self.search_strategy.hybrid)
       return self._client.hybrid_search(
           collection_name=self.collection_name,
           partition_names=self.partition_names,
           output_fields=self.output_fields,
           timeout=self.timeout,
           round_decimal=self.round_decimal,
           reqs=data,
-          **hybrid_search_params)
+          ranker=self.search_strategy.ranker,
+          limit=self.search_strategy.limit,
+          **self.search_strategy.kwargs)
     elif isinstance(self.search_strategy, VectorSearchParameters):
       data = list(map(self._get_vector_search_data, chunks))
       vector_search_params = unpack_dataclass_with_kwargs(self.search_strategy)
@@ -497,14 +529,14 @@ def _get_call_response(
     for i in range(len(chunks)):
       chunk = chunks[i]
       hits: Hits = search_result[i]
-      result = defaultdict(list)
+      result = MilvusSearchResult()
       for i in range(len(hits)):
         hit: Hit = hits[i]
         normalized_fields = self._normalize_milvus_fields(hit.fields)
-        result["id"].append(hit.id)
-        result["distance"].append(hit.distance)
-        result["fields"].append(normalized_fields)
-      response.append((chunk, result))
+        result.id.append(hit.id)
+        result.distance.append(hit.distance)
+        result.fields.append(normalized_fields)
+      response.append((chunk, result.__dict__))
     return response
 
   def _normalize_milvus_fields(self, fields: Dict[str, Any]):

sdks/python/apache_beam/ml/rag/enrichment/milvus_search_it_test.py

@@ -63,7 +63,6 @@
       MilvusCollectionLoadParameters,
       VectorSearchParameters,
       KeywordSearchParameters,
-      HybridSearchNamespace,
       HybridSearchParameters,
       VectorSearchMetrics,
       KeywordSearchMetrics)
@@ -523,7 +522,9 @@ def test_invalid_query_on_non_existent_collection(self):
     collection_load_parameters = MilvusCollectionLoadParameters()
 
     handler = MilvusSearchEnrichmentHandler(
-        self._connection_params, search_parameters, collection_load_parameters)
+        self._connection_params,
+        search_parameters,
+        collection_load_parameters=collection_load_parameters)
 
     with self.assertRaises(Exception) as context:
       with TestPipeline() as p:
@@ -549,7 +550,9 @@ def test_invalid_query_on_non_existent_field(self):
     collection_load_parameters = MilvusCollectionLoadParameters()
 
     handler = MilvusSearchEnrichmentHandler(
-        self._connection_params, search_parameters, collection_load_parameters)
+        self._connection_params,
+        search_parameters,
+        collection_load_parameters=collection_load_parameters)
 
     with self.assertRaises(Exception) as context:
       with TestPipeline() as p:
@@ -569,7 +572,9 @@ def test_empty_input_chunks(self):
     collection_load_parameters = MilvusCollectionLoadParameters()
 
     handler = MilvusSearchEnrichmentHandler(
-        self._connection_params, search_parameters, collection_load_parameters)
+        self._connection_params,
+        search_parameters,
+        collection_load_parameters=collection_load_parameters)
 
     expected_chunks = []
 
@@ -1187,16 +1192,14 @@ def test_hybrid_search(self):
         search_params=addition_keyword_search_params)
 
     hybrid_search_parameters = HybridSearchParameters(
-        ranker=RRFRanker(1), limit=1)
-
-    hybrid_search_ns = HybridSearchNamespace(
         vector=vector_search_parameters,
         keyword=keyword_search_parameters,
-        hybrid=hybrid_search_parameters)
+        ranker=RRFRanker(1),
+        limit=1)
 
     search_parameters = MilvusSearchParameters(
         collection_name=MILVUS_IT_CONFIG["collection_name"],
-        search_strategy=hybrid_search_ns,
+        search_strategy=hybrid_search_parameters,
         output_fields=["id", "content", "metadata"],
         round_decimal=1)
 
@@ -1306,10 +1309,6 @@ def assert_chunks_equivalent(
     assert 'enrichment_data' in actual.metadata, err_msg
 
     # For enrichment_data, ensure consistent ordering of results.
-    # If "expected" has values for enrichment_data but actual doesn't, that's
-    # acceptable since vector search results can vary based on many factors
-    # including implementation details, vector database state, and slight
-    # variations in similarity calculations.
     actual_data = actual.metadata['enrichment_data']
     expected_data = expected.metadata['enrichment_data']
 

sdks/python/apache_beam/ml/rag/enrichment/milvus_search_test.py

@@ -30,7 +30,6 @@
       VectorSearchParameters,
       KeywordSearchParameters,
       HybridSearchParameters,
-      HybridSearchNamespace,
       MilvusBaseRanker,
       unpack_dataclass_with_kwargs)
 except ImportError as e:
@@ -285,47 +284,39 @@ class TestMilvusHybridSearchEnrichment(unittest.TestCase):
   """Tests specific to hybrid search functionality"""
 
   @parameterized.expand([
-      # Missing vector in hybrid search namespace.
+      # Missing vector in hybrid search parameters.
       (
-          lambda: HybridSearchNamespace(
+          lambda: HybridSearchParameters(
               vector=None,  # type: ignore[arg-type]
               keyword=KeywordSearchParameters(anns_field="sparse_embedding"),
-              hybrid=HybridSearchParameters(ranker=MockRanker())),
-          "Vector, keyword, and hybrid search parameters must be provided for "
-          "hybrid search"
+              ranker=MockRanker()),
+          "Vector and keyword search parameters must be provided for hybrid "
+          "search"
       ),
-      # Missing keyword in hybrid search namespace.
+      # Missing keyword in hybrid search parameters.
       (
-          lambda: HybridSearchNamespace(
+          lambda: HybridSearchParameters(
               vector=VectorSearchParameters(anns_field="embedding"),
               keyword=None,  # type: ignore[arg-type]
-              hybrid=HybridSearchParameters(ranker=MockRanker())),
-          "Vector, keyword, and hybrid search parameters must be provided for "
-          "hybrid search"
-      ),
-      # Missing hybrid in hybrid search namespace.
-      (
-          lambda: HybridSearchNamespace(
-              vector=VectorSearchParameters(anns_field="embedding"),
-              keyword=KeywordSearchParameters(anns_field="sparse_embedding"),
-              hybrid=None),  # type: ignore[arg-type]
-          "Vector, keyword, and hybrid search parameters must be provided for "
-          "hybrid search"
+              ranker=MockRanker()),
+          "Vector and keyword search parameters must be provided for hybrid "
+          "search"
       ),
       # Missing ranker in hybrid search parameters.
       (
-          lambda: HybridSearchNamespace(
+          lambda: HybridSearchParameters(
               vector=VectorSearchParameters(anns_field="embedding"),
               keyword=KeywordSearchParameters(anns_field="sparse_embedding"),
-              hybrid=HybridSearchParameters(ranker=None)),  # type: ignore[arg-type]
+              ranker=None),  # type: ignore[arg-type]
           "Ranker must be provided for hybrid search"
       ),
       # Negative limit in hybrid search parameters.
       (
-          lambda: HybridSearchNamespace(
+          lambda: HybridSearchParameters(
               vector=VectorSearchParameters(anns_field="embedding"),
               keyword=KeywordSearchParameters(anns_field="sparse_embedding"),
-              hybrid=HybridSearchParameters(ranker=MockRanker(), limit=-1)),
+              ranker=MockRanker(),
+              limit=-1),
           "Search limit must be positive, got -1"
       ),
   ])
@@ -334,10 +325,10 @@ def test_invalid_search_parameters(self, create_params, expected_error_msg):
     with self.assertRaises(ValueError) as context:
       connection_params = MilvusConnectionParameters(
           uri="http://localhost:19530")
-      hybrid_search_namespace = create_params()
+      hybrid_search_params = create_params()
       search_params = MilvusSearchParameters(
           collection_name="test_collection",
-          search_strategy=hybrid_search_namespace)
+          search_strategy=hybrid_search_params)
       collection_load_params = MilvusCollectionLoadParameters()
 
       _ = MilvusSearchEnrichmentHandler(