rag: reuse chunk_approximately_equals util as prep for openai integration

Author: moawnallah

sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py

@@ -27,6 +27,7 @@
 from apache_beam.ml.rag.types import Chunk
 from apache_beam.ml.rag.types import Content
 from apache_beam.ml.rag.types import Embedding
+from apache_beam.ml.rag.test_utils import TestHelpers
 from apache_beam.ml.transforms.base import MLTransform
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
@@ -40,19 +41,6 @@
   SENTENCE_TRANSFORMERS_AVAILABLE = False
 
 
-def chunk_approximately_equals(expected, actual):
-  """Compare embeddings allowing for numerical differences."""
-  if not isinstance(expected, Chunk) or not isinstance(actual, Chunk):
-    return False
-
-  return (
-      expected.id == actual.id and expected.metadata == actual.metadata and
-      expected.content == actual.content and
-      len(expected.embedding.dense_embedding) == len(
-          actual.embedding.dense_embedding) and
-      all(isinstance(x, float) for x in actual.embedding.dense_embedding))
-
-
 @pytest.mark.uses_transformers
 @unittest.skipIf(
     not SENTENCE_TRANSFORMERS_AVAILABLE, "sentence-transformers not available")
@@ -105,7 +93,8 @@ def test_embedding_pipeline(self):
           with_transform(embedder))
 
       assert_that(
-          embeddings, equal_to(expected, equals_fn=chunk_approximately_equals))
+          embeddings,
+          equal_to(expected, equals_fn=TestHelpers.chunk_approximately_equals))
 
 
 if __name__ == '__main__':

sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py

@@ -21,6 +21,7 @@
 import unittest
 
 import apache_beam as beam
+from apache_beam.ml.rag.test_utils import TestHelpers
 from apache_beam.ml.rag.types import Chunk
 from apache_beam.ml.rag.types import Content
 from apache_beam.ml.rag.types import Embedding
@@ -39,19 +40,6 @@
   VERTEX_AI_AVAILABLE = False
 
 
-def chunk_approximately_equals(expected, actual):
-  """Compare embeddings allowing for numerical differences."""
-  if not isinstance(expected, Chunk) or not isinstance(actual, Chunk):
-    return False
-
-  return (
-      expected.id == actual.id and expected.metadata == actual.metadata and
-      expected.content == actual.content and
-      len(expected.embedding.dense_embedding) == len(
-          actual.embedding.dense_embedding) and
-      all(isinstance(x, float) for x in actual.embedding.dense_embedding))
-
-
 @unittest.skipIf(
     not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available")
 class VertexAITextEmbeddingsTest(unittest.TestCase):
@@ -104,7 +92,8 @@ def test_embedding_pipeline(self):
           with_transform(embedder))
 
       assert_that(
-          embeddings, equal_to(expected, equals_fn=chunk_approximately_equals))
+          embeddings,
+          equal_to(expected, equals_fn=TestHelpers.chunk_approximately_equals))
 
 
 if __name__ == '__main__':

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

@@ -80,6 +80,19 @@ def find_free_port():
       # Return the port number assigned by OS.
       return s.getsockname()[1]
 
+  @staticmethod
+  def chunk_approximately_equals(expected, actual):
+    """Compare embeddings allowing for numerical differences."""
+    if not isinstance(expected, Chunk) or not isinstance(actual, Chunk):
+      return False
+
+    return (
+        expected.id == actual.id and expected.metadata == actual.metadata and
+        expected.content == actual.content and
+        len(expected.embedding.dense_embedding) == len(
+            actual.embedding.dense_embedding) and
+        all(isinstance(x, float) for x in actual.embedding.dense_embedding))
+
 
 class CustomMilvusContainer(MilvusContainer):
   """Custom Milvus container with configurable ports and environment setup.
@@ -407,7 +420,3 @@ def assert_chunks_equivalent(
           # Validate field metadata.
           err_msg = f"Field Metadata doesn't match for chunk {actual.id}"
           assert a_f['metadata'] == e_f['metadata'], err_msg
-
-
-if __name__ == '__main__':
-  unittest.main()