sdks/python/apache_beam/ml/rag/chunking/base_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for apache_beam.ml.rag.chunking.base."""

import unittest
from typing import Any
from typing import Dict
from typing import Optional

import pytest

import apache_beam as beam
from apache_beam.ml.rag.chunking.base import ChunkIdFn
from apache_beam.ml.rag.chunking.base import ChunkingTransformProvider
from apache_beam.ml.rag.types import Chunk
from apache_beam.ml.rag.types import Content
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to


class WordSplitter(beam.DoFn):
  def process(self, element):
    words = element['text'].split()
    for i, word in enumerate(words):
      yield Chunk(
          content=Content(text=word),
          index=i,
          metadata={'source': element['source']})


class InvalidChunkingProvider(ChunkingTransformProvider):
  def __init__(self, chunk_id_fn: Optional[ChunkIdFn] = None):
    super().__init__(chunk_id_fn=chunk_id_fn)


class MockChunkingProvider(ChunkingTransformProvider):
  def __init__(self, chunk_id_fn: Optional[ChunkIdFn] = None):
    super().__init__(chunk_id_fn=chunk_id_fn)

  def get_splitter_transform(
      self
  ) -> beam.PTransform[beam.PCollection[Dict[str, Any]],
                       beam.PCollection[Chunk]]:
    return beam.ParDo(WordSplitter())


def chunk_equals(expected, actual):
  """Custom equality function for Chunk objects."""
  if not isinstance(expected, Chunk) or not isinstance(actual, Chunk):
    return False
  # Don't compare IDs since they're randomly generated
  return (
      expected.index == actual.index and expected.content == actual.content and
      expected.metadata == actual.metadata)


def id_equals(expected, actual):
  """Custom equality function for Chunk object id's."""
  if not isinstance(expected, Chunk) or not isinstance(actual, Chunk):
    return False
  return (expected.id == actual.id)


@pytest.mark.uses_transformers
class ChunkingTransformProviderTest(unittest.TestCase):
  def setUp(self):
    self.test_doc = {'text': 'hello world test', 'source': 'test.txt'}

  def test_doesnt_override_get_text_splitter_transform(self):
    provider = InvalidChunkingProvider()
    with self.assertRaises(NotImplementedError):
      provider.get_splitter_transform()

  def test_chunking_transform(self):
    """Test the complete chunking transform."""
    provider = MockChunkingProvider()

    with TestPipeline() as p:
      chunks = (
          p
          | beam.Create([self.test_doc])
          | provider.get_ptransform_for_processing())

      expected = [
          Chunk(
              content=Content(text="hello"),
              index=0,
              metadata={'source': 'test.txt'}),
          Chunk(
              content=Content(text="world"),
              index=1,
              metadata={'source': 'test.txt'}),
          Chunk(
              content=Content(text="test"),
              index=2,
              metadata={'source': 'test.txt'})
      ]

      assert_that(chunks, equal_to(expected, equals_fn=chunk_equals))

  def test_custom_chunk_id_fn(self):
    """Test the a custom chink id function."""
    def source_index_id_fn(chunk: Chunk):
      return f"{chunk.metadata['source']}_{chunk.index}"

    provider = MockChunkingProvider(chunk_id_fn=source_index_id_fn)

    with TestPipeline() as p:
      chunks = (
          p
          | beam.Create([self.test_doc])
          | provider.get_ptransform_for_processing())

      expected = [
          Chunk(content=Content(text="hello"), id="test.txt_0"),
          Chunk(content=Content(text="world"), id="test.txt_1"),
          Chunk(content=Content(text="test"), id="test.txt_2")
      ]

      assert_that(chunks, equal_to(expected, equals_fn=id_equals))


if __name__ == '__main__':
  unittest.main()