sdks/python: remove the examples from this PR

Author: mohamedawnallah

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/transforms/enrichment_test.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 #
 # Licensed to the Apache Software Foundation (ASF) under one or more
 # contributor license agreements.  See the NOTICE file distributed with
@@ -14,28 +15,74 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+# pytype: skip-file
+# pylint: disable=line-too-long
 
-import logging
 import unittest
+from io import StringIO
 
-import apache_beam as beam
+import mock
 
-# pylint: disable=ungrouped-imports
+# pylint: disable=unused-import
 try:
-  from apache_beam.transforms.enrichment import cross_join
+  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.')
+  raise unittest.SkipTest('RequestResponseIO dependencies are not installed')
 
 
-class TestEnrichmentTransform(unittest.TestCase):
-  def test_cross_join(self):
-    left = {'id': 1, 'key': 'city'}
-    right = {'id': 1, 'value': 'durham'}
-    expected = beam.Row(id=1, key='city', value='durham')
-    output = cross_join(left, right)
-    self.assertEqual(expected, output)
+def validate_enrichment_with_bigtable():
+  expected = '''[START enrichment_with_bigtable]
+Row(sale_id=1, customer_id=1, product_id=1, quantity=1, product={'product_id': '1', 'product_name': 'pixel 5', 'product_stock': '2'})
+Row(sale_id=3, customer_id=3, product_id=2, quantity=3, product={'product_id': '2', 'product_name': 'pixel 6', 'product_stock': '4'})
+Row(sale_id=5, customer_id=5, product_id=4, quantity=2, product={'product_id': '4', 'product_name': 'pixel 8', 'product_stock': '10'})
+  [END enrichment_with_bigtable]'''.splitlines()[1:-1]
+  return expected
+
+
+def validate_enrichment_with_vertex_ai():
+  expected = '''[START enrichment_with_vertex_ai]
+Row(user_id='2963', product_id=14235, sale_price=15.0, age=12.0, state='1', gender='1', country='1')
+Row(user_id='21422', product_id=11203, sale_price=12.0, age=12.0, state='0', gender='0', country='0')
+Row(user_id='20592', product_id=8579, sale_price=9.0, age=12.0, state='2', gender='1', country='2')
+  [END enrichment_with_vertex_ai]'''.splitlines()[1:-1]
+  return expected
+
+
+def validate_enrichment_with_vertex_ai_legacy():
+  expected = '''[START enrichment_with_vertex_ai_legacy]
+Row(entity_id='movie_01', title='The Shawshank Redemption', genres='Drama')
+Row(entity_id='movie_02', title='The Shining', genres='Horror')
+Row(entity_id='movie_04', title='The Dark Knight', genres='Action')
+  [END enrichment_with_vertex_ai_legacy]'''.splitlines()[1:-1]
+  return expected
+
+
+@mock.patch('sys.stdout', new_callable=StringIO)
+class EnrichmentTest(unittest.TestCase):
+  def test_enrichment_with_bigtable(self, mock_stdout):
+    enrichment_with_bigtable()
+    output = mock_stdout.getvalue().splitlines()
+    expected = validate_enrichment_with_bigtable()
+    self.assertEqual(output, expected)
+
+  def test_enrichment_with_vertex_ai(self, mock_stdout):
+    enrichment_with_vertex_ai()
+    output = mock_stdout.getvalue().splitlines()
+    expected = validate_enrichment_with_vertex_ai()
+
+    for i in range(len(expected)):
+      self.assertEqual(set(output[i].split(',')), set(expected[i].split(',')))
+
+  def test_enrichment_with_vertex_ai_legacy(self, mock_stdout):
+    enrichment_with_vertex_ai_legacy()
+    output = mock_stdout.getvalue().splitlines()
+    expected = validate_enrichment_with_vertex_ai_legacy()
+    self.maxDiff = None
+    self.assertEqual(output, expected)
 
 
 if __name__ == '__main__':
-  logging.getLogger().setLevel(logging.INFO)
   unittest.main()