sdks/python: fix linting issues

Author: mohamedawnallah

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

@@ -52,10 +52,9 @@
       ConnectionConfig,
       CloudSQLConnectionConfig,
       ExternalSQLDBConnectionConfig)
-  from apache_beam.ml.rag.enrichment.milvus_search import (
-      MilvusConnectionParameters)
-  from apache_beam.ml.rag.enrichment.milvus_search_it_test import (
-      MilvusEnrichmentTestHelper, MilvusDBContainerInfo)
+  from apache_beam.ml.rag.enrichment.milvus_search import MilvusConnectionParameters
+  from apache_beam.ml.rag.test_utils import MilvusTestHelpers
+  from apache_beam.ml.rag.test_utils import VectorDBContainerInfo
   from apache_beam.ml.rag.test_utils import MilvusTestHelpers
   from apache_beam.ml.rag.utils import parse_chunk_strings
   from apache_beam.io.requestresponse import RequestResponseIO
@@ -261,7 +260,7 @@ def sql_test_context(is_cloudsql: bool, db_adapter: DatabaseTypeAdapter):
   @staticmethod
   @contextmanager
   def milvus_test_context():
-    db: Optional[MilvusDBContainerInfo] = None
+    db: Optional[VectorDBContainerInfo] = None
     try:
       db = EnrichmentTestHelpers.pre_milvus_enrichment()
       yield
@@ -374,16 +373,16 @@ def post_sql_enrichment_test(res: CloudSQLEnrichmentTestDataConstruct):
       os.environ.pop('GOOGLE_CLOUD_SQL_DB_TABLE_ID', None)
 
   @staticmethod
-  def pre_milvus_enrichment() -> MilvusDBContainerInfo:
+  def pre_milvus_enrichment() -> VectorDBContainerInfo:
     try:
-      db = MilvusEnrichmentTestHelper.start_db_container()
+      db = MilvusTestHelpers.start_db_container()
       connection_params = MilvusConnectionParameters(
           uri=db.uri,
           user=db.user,
           password=db.password,
           db_id=db.id,
           token=db.token)
-      collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
+      collection_name = MilvusTestHelpers.initialize_db_with_data(
           connection_params)
     except Exception as e:
       raise TestContainerStartupError(
@@ -401,9 +400,9 @@ def pre_milvus_enrichment() -> MilvusDBContainerInfo:
     return db
 
   @staticmethod
-  def post_milvus_enrichment(db: MilvusDBContainerInfo):
+  def post_milvus_enrichment(db: VectorDBContainerInfo):
     try:
-      MilvusEnrichmentTestHelper.stop_db_container(db)
+      MilvusTestHelpers.stop_db_container(db)
     except Exception as e:
       raise TestContainerTeardownError(
           f"Milvus container failed to tear down: {str(e)}")

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

@@ -47,9 +47,6 @@
   from pymilvus.exceptions import MilvusException
   from pymilvus.milvus_client import IndexParams
 
-  from apache_beam.transforms.enrichment import Enrichment
-  from apache_beam.ml.rag.test_utils import MilvusTestHelpers
-  from apache_beam.ml.rag.test_utils import VectorDBContainerInfo
   from apache_beam.ml.rag.enrichment.milvus_search import HybridSearchParameters
   from apache_beam.ml.rag.enrichment.milvus_search import KeywordSearchMetrics
   from apache_beam.ml.rag.enrichment.milvus_search import KeywordSearchParameters
@@ -59,6 +56,9 @@
   from apache_beam.ml.rag.enrichment.milvus_search import MilvusSearchParameters
   from apache_beam.ml.rag.enrichment.milvus_search import VectorSearchMetrics
   from apache_beam.ml.rag.enrichment.milvus_search import VectorSearchParameters
+  from apache_beam.ml.rag.test_utils import MilvusTestHelpers
+  from apache_beam.ml.rag.test_utils import VectorDBContainerInfo
+  from apache_beam.transforms.enrichment import Enrichment
 except ImportError as e:
   raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}')
 
@@ -235,77 +235,6 @@ def __getitem__(self, key):
 }
 
 
-def initialize_db_with_data(connc_params: MilvusConnectionParameters):
-  # Open the connection to the milvus db with retry.
-  def create_client():
-    return MilvusClient(**connc_params.__dict__)
-
-  client = retry_with_backoff(
-      create_client,
-      max_retries=3,
-      retry_delay=1.0,
-      operation_name="Test Milvus client connection",
-      exception_types=(MilvusException, ))
-
-  # Configure schema.
-  field_schemas: List[FieldSchema] = cast(
-      List[FieldSchema], MILVUS_IT_CONFIG["fields"])
-  schema = CollectionSchema(
-      fields=field_schemas, functions=MILVUS_IT_CONFIG["functions"])
-
-  # Create collection with the schema.
-  collection_name = MILVUS_IT_CONFIG["collection_name"]
-  index_function: Callable[[], IndexParams] = cast(
-      Callable[[], IndexParams], MILVUS_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 = list(map(lambda field: field.name, 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_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 the 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.require_docker_in_docker
 @unittest.skipUnless(
     platform.system() == "Linux",
@@ -329,7 +258,8 @@ def setUpClass(cls):
         db_name=cls._db.id,
         token=cls._db.token)
     cls._collection_load_params = MilvusCollectionLoadParameters()
-    cls._collection_name = initialize_db_with_data(cls._connection_params)
+    cls._collection_name = initialize_db_with_data(
+        cls._connection_params, MILVUS_IT_CONFIG)
 
   @classmethod
   def tearDownClass(cls):

sdks/python/apache_beam/ml/rag/test_utils.py

@@ -21,15 +21,24 @@
 import socket
 import tempfile
 from dataclasses import dataclass
+from typing import Callable
 from typing import List
 from typing import Optional
+from typing import cast
 
 import yaml
 from testcontainers.core.config import testcontainers_config
 from testcontainers.core.generic import DbContainer
 from testcontainers.milvus import MilvusContainer
+from pymilvus import CollectionSchema
+from pymilvus import FieldSchema
+from pymilvus import MilvusClient
+from pymilvus.exceptions import MilvusException
+from pymilvus.milvus_client import IndexParams
 
 from apache_beam.ml.rag.types import Chunk
+from apache_beam.ml.rag.enrichment.milvus_search import MilvusConnectionParameters
+from apache_beam.ml.rag.utils import retry_with_backoff
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -180,6 +189,76 @@ def stop_db_container(db_info: VectorDBContainerInfo):
     db_info.container.stop()
     _LOGGER.info("milvus db container stopped successfully.")
 
+  def initialize_db_with_data(
+      connc_params: MilvusConnectionParameters, config: dict):
+    # Open the connection to the milvus db with retry.
+    def create_client():
+      return MilvusClient(**connc_params.__dict__)
+
+    client = retry_with_backoff(
+        create_client,
+        max_retries=3,
+        retry_delay=1.0,
+        operation_name="Test Milvus client connection",
+        exception_types=(MilvusException, ))
+
+    # Configure schema.
+    field_schemas: List[FieldSchema] = cast(List[FieldSchema], config["fields"])
+    schema = CollectionSchema(
+        fields=field_schemas, functions=config["functions"])
+
+    # Create collection with the schema.
+    collection_name = config["collection_name"]
+    index_function: Callable[[], IndexParams] = cast(
+        Callable[[], IndexParams], 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 = list(map(lambda field: field.name, 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 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 the 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
+
   @staticmethod
   @contextlib.contextmanager
   def create_user_yaml(service_port: int, max_vector_field_num=5):