sdks/python: fix linting and formatting issues

Author: moawnallah

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

@@ -28,21 +28,20 @@
 
 from google.protobuf.json_format import MessageToDict
 
-from apache_beam.ml.rag.types import Chunk
-from apache_beam.ml.rag.types import Embedding
-from apache_beam.transforms.enrichment import EnrichmentSourceHandler
 from pymilvus import AnnSearchRequest
 from pymilvus import Hit
 from pymilvus import Hits
 from pymilvus import MilvusClient
 from pymilvus import SearchResult
 
+from apache_beam.ml.rag.types import Chunk
+from apache_beam.ml.rag.types import Embedding
+from apache_beam.transforms.enrichment import EnrichmentSourceHandler
+
 from apache_beam.ml.rag.types import Chunk
 from apache_beam.ml.rag.types import Embedding
 from apache_beam.ml.rag.utils import (
-  MilvusHelpers,
-  MilvusConnectionConfig,
-  unpack_dataclass_with_kwargs)
+    MilvusHelpers, MilvusConnectionConfig, unpack_dataclass_with_kwargs)
 from apache_beam.transforms.enrichment import EnrichmentSourceHandler
 
 

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

@@ -38,11 +38,11 @@
 from apache_beam.ml.rag.types import Content
 from apache_beam.ml.rag.types import Embedding
 from apache_beam.ml.rag.utils import (
-    MilvusConnectionConfig,
-    unpack_dataclass_with_kwargs)
+    MilvusConnectionConfig, unpack_dataclass_with_kwargs)
 from apache_beam.ml.rag.test_utils import (
     VectorDBContainerInfo,
-    MilvusTestHelpers,)
+    MilvusTestHelpers,
+)
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 
@@ -124,6 +124,7 @@ def _construct_index_params():
 
   return index_params
 
+
 @dataclass
 class MilvusITDataConstruct:
   id: int
@@ -139,6 +140,7 @@ class MilvusITDataConstruct:
   def __getitem__(self, key):
     return getattr(self, key)
 
+
 MILVUS_ENRICHMENT_IT_CONFIG = {
     "collection_name": "docs_catalog",
     "fields": [
@@ -228,70 +230,71 @@ def __getitem__(self, key):
     }
 }
 
+
 def initialize_db_with_data(connc_params: MilvusConnectionConfig):
-    # Open the connection to the milvus db.
-    config = unpack_dataclass_with_kwargs(connc_params)
-    config["alias"] = f"milvus_conn_{uuid.uuid4().hex[:8]}"
-    client = MilvusClient(**config)
-
-    # Configure schema.
-    field_schemas: List[FieldSchema] = cast(
-        List[FieldSchema], MILVUS_ENRICHMENT_IT_CONFIG["fields"])
-    schema = CollectionSchema(
-        fields=field_schemas,
-        functions=MILVUS_ENRICHMENT_IT_CONFIG["functions"])
-
-    # Create collection with the schema.
-    collection_name = MILVUS_ENRICHMENT_IT_CONFIG["collection_name"]
-    index_function: Callable[[], IndexParams] = cast(
-        Callable[[], IndexParams], MILVUS_ENRICHMENT_IT_CONFIG["index"])
-    client.create_collection(
-        collection_name=collection_name,
-        schema=schema,
-        index_params=index_function())
-
-    # Assert that collection was created.
-    collection_error = f"Expected collection '{collection_name}' to be created."
-    assert client.has_collection(collection_name), collection_error
-
-    # Gather all fields we have excluding 'sparse_embedding_bm25' special field.
-    fields = [field.name for field in field_schemas]
-
-    # Prep data for indexing. Currently we can't insert sparse vectors for BM25
-    # sparse embedding field as it would be automatically generated by Milvus
-    # through the registered BM25 function.
-    data_ready_to_index = []
-    for doc in MILVUS_ENRICHMENT_IT_CONFIG["corpus"]:
-        item = {}
-        for field in fields:
-            if field.startswith("dense_embedding"):
-                item[field] = doc["dense_embedding"]
-            elif field == "sparse_embedding_inner_product":
-                item[field] = doc["sparse_embedding"]
-            elif field == "sparse_embedding_bm25":
-                # It is automatically generated by Milvus from content field.
-                continue
-            else:
-                item[field] = doc[field]
-        data_ready_to_index.append(item)
-
-    # Index data.
-    result = client.insert(
-        collection_name=collection_name, data=data_ready_to_index)
-
-    # Assert that the intended data has been properly indexed.
-    insertion_err = f'failed to insert the {result["insert_count"]} data points'
-    assert result["insert_count"] == len(data_ready_to_index), insertion_err
-
-    # Release the collection from memory. It will be loaded lazily when the
-    # enrichment handler is invoked.
-    client.release_collection(collection_name)
-
-    # Close the connection to the Milvus database, as no further preparation
-    # operations are needed  before executing the enrichment handler.
-    client.close()
-
-    return collection_name
+  # Open the connection to the milvus db.
+  config = unpack_dataclass_with_kwargs(connc_params)
+  config["alias"] = f"milvus_conn_{uuid.uuid4().hex[:8]}"
+  client = MilvusClient(**config)
+
+  # Configure schema.
+  field_schemas: List[FieldSchema] = cast(
+      List[FieldSchema], MILVUS_ENRICHMENT_IT_CONFIG["fields"])
+  schema = CollectionSchema(
+      fields=field_schemas, functions=MILVUS_ENRICHMENT_IT_CONFIG["functions"])
+
+  # Create collection with the schema.
+  collection_name = MILVUS_ENRICHMENT_IT_CONFIG["collection_name"]
+  index_function: Callable[[], IndexParams] = cast(
+      Callable[[], IndexParams], MILVUS_ENRICHMENT_IT_CONFIG["index"])
+  client.create_collection(
+      collection_name=collection_name,
+      schema=schema,
+      index_params=index_function())
+
+  # Assert that collection was created.
+  collection_error = f"Expected collection '{collection_name}' to be created."
+  assert client.has_collection(collection_name), collection_error
+
+  # Gather all fields we have excluding 'sparse_embedding_bm25' special field.
+  fields = [field.name for field in field_schemas]
+
+  # Prep data for indexing. Currently we can't insert sparse vectors for BM25
+  # sparse embedding field as it would be automatically generated by Milvus
+  # through the registered BM25 function.
+  data_ready_to_index = []
+  for doc in MILVUS_ENRICHMENT_IT_CONFIG["corpus"]:
+    item = {}
+    for field in fields:
+      if field.startswith("dense_embedding"):
+        item[field] = doc["dense_embedding"]
+      elif field == "sparse_embedding_inner_product":
+        item[field] = doc["sparse_embedding"]
+      elif field == "sparse_embedding_bm25":
+        # It is automatically generated by Milvus from content field.
+        continue
+      else:
+        item[field] = doc[field]
+    data_ready_to_index.append(item)
+
+  # Index data.
+  result = client.insert(
+      collection_name=collection_name, data=data_ready_to_index)
+
+  # Assert that the intended data has been properly indexed.
+  insertion_err = f'failed to insert the {result["insert_count"]} data points'
+  assert result["insert_count"] == len(data_ready_to_index), insertion_err
+
+  # Release the collection from memory. It will be loaded lazily when the
+  # enrichment handler is invoked.
+  client.release_collection(collection_name)
+
+  # Close the connection to the Milvus database, as no further preparation
+  # operations are needed  before executing the enrichment handler.
+  client.close()
+
+  return collection_name
+
 
 @pytest.mark.uses_testcontainer
 @unittest.skipUnless(
@@ -401,9 +404,8 @@ def test_empty_input_chunks(self):
     with TestPipeline(is_integration_test=True) as p:
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
-          result,
-          lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
-            actual, expected_chunks))
+          result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
+              actual, expected_chunks))
 
   def test_filtered_search_with_cosine_similarity_and_batching(self):
     test_chunks = [
@@ -530,9 +532,8 @@ def test_filtered_search_with_cosine_similarity_and_batching(self):
     with TestPipeline(is_integration_test=True) as p:
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
-          result,
-          lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
-            actual, expected_chunks))
+          result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
+              actual, expected_chunks))
 
   def test_filtered_search_with_bm25_full_text_and_batching(self):
     test_chunks = [
@@ -636,9 +637,8 @@ def test_filtered_search_with_bm25_full_text_and_batching(self):
     with TestPipeline(is_integration_test=True) as p:
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
-          result,
-          lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
-            actual, expected_chunks))
+          result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
+              actual, expected_chunks))
 
   def test_vector_search_with_euclidean_distance(self):
     test_chunks = [
@@ -778,9 +778,8 @@ def test_vector_search_with_euclidean_distance(self):
     with TestPipeline(is_integration_test=True) as p:
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
-          result,
-          lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
-            actual, expected_chunks))
+          result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
+              actual, expected_chunks))
 
   def test_vector_search_with_inner_product_similarity(self):
     test_chunks = [
@@ -919,9 +918,8 @@ def test_vector_search_with_inner_product_similarity(self):
     with TestPipeline(is_integration_test=True) as p:
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
-          result,
-          lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
-            actual, expected_chunks))
+          result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
+              actual, expected_chunks))
 
   def test_keyword_search_with_inner_product_sparse_embedding(self):
     test_chunks = [
@@ -985,9 +983,8 @@ def test_keyword_search_with_inner_product_sparse_embedding(self):
     with TestPipeline(is_integration_test=True) as p:
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
-          result,
-          lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
-            actual, expected_chunks))
+          result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
+              actual, expected_chunks))
 
   def test_hybrid_search(self):
     test_chunks = [
@@ -1059,10 +1056,9 @@ def test_hybrid_search(self):
     with TestPipeline(is_integration_test=True) as p:
       result = (p | beam.Create(test_chunks) | Enrichment(handler))
       assert_that(
-          result,
-          lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
-            actual, expected_chunks))
+          result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent(
+              actual, expected_chunks))
 
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

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

@@ -23,8 +23,7 @@
   from apache_beam.ml.rag.types import Embedding
   from apache_beam.ml.rag.types import Content
   from apache_beam.ml.rag.utils import (
-    MilvusConnectionConfig,
-    unpack_dataclass_with_kwargs)
+      MilvusConnectionConfig, unpack_dataclass_with_kwargs)
   from apache_beam.ml.rag.enrichment.milvus_search import (
       MilvusSearchEnrichmentHandler,
       MilvusSearchParameters,
@@ -81,8 +80,7 @@ def test_invalid_connection_parameters(self):
   def test_invalid_search_parameters(self, create_params, expected_error_msg):
     """Test validation errors for invalid general search parameters."""
     with self.assertRaises(ValueError) as context:
-      connection_params = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       search_params = create_params()
       collection_load_params = MilvusCollectionLoadParameters()
 
@@ -138,8 +136,7 @@ class TestMilvusVectorSearchEnrichment(unittest.TestCase):
   def test_invalid_search_parameters(self, create_params, expected_error_msg):
     """Test validation errors for invalid vector search parameters."""
     with self.assertRaises(ValueError) as context:
-      connection_params = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       vector_search_params = create_params()
       search_params = MilvusSearchParameters(
           collection_name="test_collection",
@@ -157,8 +154,7 @@ 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 = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       vector_search_params = VectorSearchParameters(anns_field="embedding")
       search_params = MilvusSearchParameters(
           collection_name="test_collection",
@@ -194,8 +190,7 @@ class TestMilvusKeywordSearchEnrichment(unittest.TestCase):
   def test_invalid_search_parameters(self, create_params, expected_error_msg):
     """Test validation errors for invalid keyword search parameters."""
     with self.assertRaises(ValueError) as context:
-      connection_params = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       keyword_search_params = create_params()
       search_params = MilvusSearchParameters(
           collection_name="test_collection",
@@ -215,8 +210,7 @@ def test_missing_text_content_and_sparse_embedding(self):
           id=1,
           content=Content(text=None),
           embedding=Embedding(sparse_embedding=None))
-      connection_params = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       vector_search_params = VectorSearchParameters(anns_field="embedding")
       search_params = MilvusSearchParameters(
           collection_name="test_collection",
@@ -241,8 +235,7 @@ def test_missing_text_content_only(self):
           content=Content(text=None),
           embedding=Embedding(
               sparse_embedding=([1, 2, 3, 4], [0.05, 0.41, 0.05, 0.41])))
-      connection_params = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       vector_search_params = VectorSearchParameters(anns_field="embedding")
       search_params = MilvusSearchParameters(
           collection_name="test_collection",
@@ -263,8 +256,7 @@ def test_missing_sparse_embedding_only(self):
           id=1,
           content=Content(text="what is apache beam?"),
           embedding=Embedding(sparse_embedding=None))
-      connection_params = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       vector_search_params = VectorSearchParameters(anns_field="embedding")
       search_params = MilvusSearchParameters(
           collection_name="test_collection",
@@ -323,8 +315,7 @@ class TestMilvusHybridSearchEnrichment(unittest.TestCase):
   def test_invalid_search_parameters(self, create_params, expected_error_msg):
     """Test validation errors for invalid hybrid search parameters."""
     with self.assertRaises(ValueError) as context:
-      connection_params = MilvusConnectionConfig(
-          uri="http://localhost:19530")
+      connection_params = MilvusConnectionConfig(uri="http://localhost:19530")
       hybrid_search_params = create_params()
       search_params = MilvusSearchParameters(
           collection_name="test_collection",