sdks/python: remove examples from this PR

Author: moawnallah

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

@@ -116,75 +116,3 @@ def enrichment_with_vertex_ai_legacy():
         | "Enrich W/ Vertex AI" >> Enrichment(vertex_ai_handler)
         | "Print" >> beam.Map(print))
   # [END enrichment_with_vertex_ai_legacy]
-
-
-def enrichment_with_milvus():
-  # [START enrichment_with_milvus]
-  import os
-  import apache_beam as beam
-  from apache_beam.ml.rag.types import Content
-  from apache_beam.ml.rag.types import Chunk
-  from apache_beam.ml.rag.types import Embedding
-  from apache_beam.transforms.enrichment import Enrichment
-  from apache_beam.ml.rag.enrichment.milvus_search import (
-      MilvusSearchEnrichmentHandler,
-      MilvusConnectionParameters,
-      MilvusSearchParameters,
-      MilvusCollectionLoadParameters,
-      VectorSearchParameters,
-      VectorSearchMetrics)
-
-  uri = os.environ.get("MILVUS_VECTOR_DB_URI")
-  user = os.environ.get("MILVUS_VECTOR_DB_USER")
-  password = os.environ.get("MILVUS_VECTOR_DB_PASSWORD")
-  db_id = os.environ.get("MILVUS_VECTOR_DB_ID")
-  token = os.environ.get("MILVUS_VECTOR_DB_TOKEN")
-  collection_name = os.environ.get("MILVUS_VECTOR_DB_COLLECTION_NAME")
-
-  data = [
-      Chunk(
-          id="query1",
-          embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]),
-          content=Content())
-  ]
-
-  connection_parameters = MilvusConnectionParameters(
-      uri, user, password, db_id, token)
-
-  # The first condition (language == "en") excludes documents in other
-  # languages. Initially, this gives us two documents. After applying the second
-  # condition (cost < 50), only the first document returns in search results.
-  filter_expr = 'metadata["language"] == "en" AND cost < 50'
-
-  search_params = {"metric_type": VectorSearchMetrics.COSINE.value, "nprobe": 1}
-
-  vector_search_params = VectorSearchParameters(
-      anns_field="dense_embedding_cosine",
-      limit=3,
-      filter=filter_expr,
-      search_params=search_params)
-
-  search_parameters = MilvusSearchParameters(
-      collection_name=collection_name,
-      search_strategy=vector_search_params,
-      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=connection_parameters,
-      search_parameters=search_parameters,
-      collection_load_parameters=collection_load_parameters)
-  with beam.Pipeline() as p:
-    _ = (
-        p
-        | "Create" >> beam.Create(data)
-        | "Enrich W/ Milvus" >> Enrichment(milvus_search_handler)
-        | "Print" >> beam.Map(print))
-  # [END enrichment_with_milvus]

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

@@ -18,27 +18,16 @@
 # pytype: skip-file
 # pylint: disable=line-too-long
 
-import os
 import unittest
 from io import StringIO
 
 import mock
-import pytest
 
 # pylint: disable=unused-import
 try:
-  from apache_beam.examples.snippets.transforms.elementwise.enrichment import (
-      enrichment_with_bigtable,
-      enrichment_with_vertex_ai,
-      enrichment_with_vertex_ai_legacy,
-      enrichment_with_milvus)
-  from apache_beam.ml.rag.enrichment.milvus_search import (
-      MilvusConnectionParameters)
-  from apache_beam.ml.rag.enrichment.milvus_search_it_test import (
-      MilvusEnrichmentTestHelper,
-      MilvusDBContainerInfo,
-      parse_chunk_strings,
-      assert_chunks_equivalent)
+  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_bigtable, \
+  enrichment_with_vertex_ai_legacy
+  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_vertex_ai
   from apache_beam.io.requestresponse import RequestResponseIO
 except ImportError:
   raise unittest.SkipTest('RequestResponseIO dependencies are not installed')
@@ -71,15 +60,7 @@ def validate_enrichment_with_vertex_ai_legacy():
   return expected
 
 
-def validate_enrichment_with_milvus():
-  expected = '''[START enrichment_with_milvus]
-Chunk(content=Content(text=None), id='query1', index=0, metadata={'enrichment_data': defaultdict(<class 'list'>, {'id': [1], 'distance': [1.0], 'fields': [{'content': 'This is a test document', 'cost': 49, 'domain': 'medical', 'id': 1, 'metadata': {'language': 'en'}}]})}, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3], sparse_embedding=None))
-  [END enrichment_with_milvus]'''.splitlines()[1:-1]
-  return expected
-
-
 @mock.patch('sys.stdout', new_callable=StringIO)
-@pytest.mark.uses_testcontainer
 class EnrichmentTest(unittest.TestCase):
   def test_enrichment_with_bigtable(self, mock_stdout):
     enrichment_with_bigtable()
@@ -102,63 +83,6 @@ def test_enrichment_with_vertex_ai_legacy(self, mock_stdout):
     self.maxDiff = None
     self.assertEqual(output, expected)
 
-  def test_enrichment_with_milvus(self, mock_stdout):
-    milvus_db = None
-    try:
-      milvus_db = EnrichmentTestHelpers.pre_milvus_enrichment()
-      enrichment_with_milvus()
-      output = mock_stdout.getvalue().splitlines()
-      expected = validate_enrichment_with_milvus()
-      self.maxDiff = None
-      output = parse_chunk_strings(output)
-      expected = parse_chunk_strings(expected)
-      assert_chunks_equivalent(output, expected)
-    except Exception as e:
-      self.fail(f"Test failed with unexpected error: {e}")
-    finally:
-      if milvus_db:
-        EnrichmentTestHelpers.post_milvus_enrichment(milvus_db)
-
-
-class EnrichmentTestHelpers:
-  @staticmethod
-  def pre_milvus_enrichment() -> MilvusDBContainerInfo:
-    # Create Milvus db container and make sure it is up and running.
-    db = MilvusEnrichmentTestHelper.start_db_container()
-
-    # Construct connection parameters.
-    connection_params = MilvusConnectionParameters(
-        uri=db.uri,
-        user=db.user,
-        password=db.password,
-        db_id=db.id,
-        token=db.token)
-
-    # Initialize db with data required for testing.
-    collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
-        connection_params)
-
-    # Setup environment variables for db and collection configuration. This will
-    # be used downstream by the milvus enrichment handler.
-    os.environ['MILVUS_VECTOR_DB_URI'] = db.uri
-    os.environ['MILVUS_VECTOR_DB_USER'] = db.user
-    os.environ['MILVUS_VECTOR_DB_PASSWORD'] = db.password
-    os.environ['MILVUS_VECTOR_DB_ID'] = db.id
-    os.environ['MILVUS_VECTOR_DB_TOKEN'] = db.token
-    os.environ['MILVUS_VECTOR_DB_COLLECTION_NAME'] = collection_name
-
-    return db
-
-  @staticmethod
-  def post_milvus_enrichment(db: MilvusDBContainerInfo):
-    MilvusEnrichmentTestHelper.stop_db_container(db)
-    os.environ.pop('MILVUS_VECTOR_DB_URI', None)
-    os.environ.pop('MILVUS_VECTOR_DB_USER', None)
-    os.environ.pop('MILVUS_VECTOR_DB_PASSWORD', None)
-    os.environ.pop('MILVUS_VECTOR_DB_ID', None)
-    os.environ.pop('MILVUS_VECTOR_DB_TOKEN', None)
-    os.environ.pop('MILVUS_VECTOR_DB_COLLECTION_NAME', None)
-
 
 if __name__ == '__main__':
   unittest.main()