chore: import upstream snapshot with attribution

sdks/python/tests/unit/message_processing/batching/test_batch_manager.py

@@ -0,0 +1,116 @@
+from unittest import mock
+import time
+
+from opik.message_processing import messages
+from opik.message_processing.batching import batch_manager
+from opik.message_processing.batching import batchers
+
+from ....testlib import fake_message_factory
+
+NOT_USED = None
+
+
+def test_batch_manager__messages_processing_methods():
+    integers_batcher = mock.Mock()
+    strings_batcher = mock.Mock()
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+
+    assert tested.message_supports_batching("a-string")
+    assert tested.message_supports_batching(42)
+    assert not tested.message_supports_batching(float(42.5))
+
+    tested.process_message(42)
+    integers_batcher.add.assert_called_once_with(42)
+
+    tested.process_message("a-string")
+    strings_batcher.add.assert_called_once_with("a-string")
+
+
+def test_batch_manager__all_batchers_are_empty__batch_manager_is_empty():
+    integers_batcher = mock.Mock()
+    integers_batcher.is_empty.return_value = True
+    strings_batcher = mock.Mock()
+    strings_batcher.is_empty.return_value = True
+
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+
+    assert tested.is_empty()
+    strings_batcher.is_empty.assert_called_once()
+    integers_batcher.is_empty.assert_called_once()
+
+
+def test_batch_manager__at_least_one_batcher_is_not_empty__batch_manager_is_not_empty():
+    integers_batcher = mock.Mock()
+    integers_batcher.is_empty.return_value = True
+    strings_batcher = mock.Mock()
+    strings_batcher.is_empty.return_value = False
+
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+
+    assert not tested.is_empty()
+    strings_batcher.is_empty.assert_called_once()
+    integers_batcher.is_empty.assert_called_once()
+
+
+def test_batch_manager__flush_is_called__all_batchers_are_flushed():
+    integers_batcher = mock.Mock()
+    strings_batcher = mock.Mock()
+
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+    tested.flush()
+    integers_batcher.flush.assert_called_once()
+    strings_batcher.flush.assert_called_once()
+
+
+def test_batch_manager__flush_then_stop__accumulated_data_is_flushed_and_thread_is_stopped():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    span_messages_batch = fake_message_factory.fake_span_create_message_batch(
+        count=10, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    example_span_batcher = batchers.CreateSpanMessageBatcher(
+        flush_callback=flush_callback,
+        max_batch_size=42,
+        flush_interval_seconds=0.1,
+        batch_memory_limit_mb=5,
+    )
+    tested = batch_manager.BatchManager(
+        {messages.CreateSpanMessage: example_span_batcher}
+    )
+
+    tested.start()
+    time.sleep(0.1)
+
+    for span_message in span_messages_batch:
+        tested.process_message(span_message)
+
+    # BatchManager.stop() is a pure signal now — callers must flush first if
+    # they want pending batches pushed downstream.
+    tested.flush()
+    tested.stop()
+
+    assert len(collected_messages) >= 2

sdks/python/tests/unit/message_processing/batching/test_experiment_items_batcher.py

@@ -0,0 +1,194 @@
+from unittest import mock
+import time
+
+
+from opik.message_processing.batching import batchers
+from opik.message_processing import messages
+
+NOT_USED = None
+
+
+def test_create_experiment_items_batch_message_batcher__exactly_max_batch_size_reached__batch_is_flushed():
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    experiment_item_batch_messages = [
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_1",
+                    experiment_id="exp1",
+                    trace_id="trace1",
+                    dataset_item_id="item1",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_2",
+                    experiment_id="exp1",
+                    trace_id="trace2",
+                    dataset_item_id="item2",
+                ),
+            ]
+        ),
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_3",
+                    experiment_id="exp1",
+                    trace_id="trace3",
+                    dataset_item_id="item3",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_4",
+                    experiment_id="exp1",
+                    trace_id="trace4",
+                    dataset_item_id="item4",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_5",
+                    experiment_id="exp1",
+                    trace_id="trace5",
+                    dataset_item_id="item5",
+                ),
+            ]
+        ),
+    ]
+
+    for experiment_items_batch in experiment_item_batch_messages:
+        batcher.add(experiment_items_batch)
+    assert batcher.is_empty()
+
+    # Verify that flush callback was called with all items combined
+    flush_callback.assert_called_once()
+    call_args = flush_callback.call_args[0][0]
+    assert isinstance(call_args, messages.CreateExperimentItemsBatchMessage)
+    assert len(call_args.batch) == 5
+    assert call_args.supports_batching is False
+
+
+def test_create_experiment_items_batch_message_batcher__more_than_max_batch_size__multiple_flushes():
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 3
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    # Add 2 messages
+    batcher.add(
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_1",
+                    experiment_id="exp1",
+                    trace_id="trace1",
+                    dataset_item_id="item1",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_2",
+                    experiment_id="exp1",
+                    trace_id="trace2",
+                    dataset_item_id="item2",
+                ),
+            ]
+        )
+    )
+    assert flush_callback.call_count == 0
+
+    # Add 2 more messages (total 4, exceeds max of 3)
+    batcher.add(
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_3",
+                    experiment_id="exp1",
+                    trace_id="trace3",
+                    dataset_item_id="item3",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_4",
+                    experiment_id="exp1",
+                    trace_id="trace4",
+                    dataset_item_id="item4",
+                ),
+            ]
+        )
+    )
+
+    # Should have flushed the first batch of 3
+    assert flush_callback.call_count == 1
+    first_call_args = flush_callback.call_args_list[0][0][0]
+    assert len(first_call_args.batch) == 3
+
+    # The batcher should have 1 item left
+    assert not batcher.is_empty()
+    assert batcher.size() == 1
+
+
+def test_create_experiment_items_batch_message_batcher__flush_interval_reached__batch_is_flushed():
+    flush_callback = mock.Mock()
+
+    FLUSH_INTERVAL = 0.1  # 100ms
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=1000,  # High enough to not trigger size-based flush
+        flush_callback=flush_callback,
+        flush_interval_seconds=FLUSH_INTERVAL,
+    )
+
+    # Add a message
+    batcher.add(
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_1",
+                    experiment_id="exp1",
+                    trace_id="trace1",
+                    dataset_item_id="item1",
+                ),
+            ]
+        )
+    )
+
+    # Batcher should not be empty yet
+    assert not batcher.is_empty()
+    assert not batcher.is_ready_to_flush()
+
+    # Wait for flush interval
+    time.sleep(FLUSH_INTERVAL + 0.05)
+
+    # Now it should be ready to flush
+    assert batcher.is_ready_to_flush()
+
+    # Manual flush
+    batcher.flush()
+
+    # Should have flushed
+    assert flush_callback.call_count == 1
+    assert batcher.is_empty()
+
+
+def test_create_experiment_items_batch_message_batcher__empty_batch__nothing_flushed():
+    flush_callback = mock.Mock()
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=5,
+        flush_callback=flush_callback,
+        flush_interval_seconds=1.0,
+    )
+
+    # Flush empty batcher
+    batcher.flush()
+
+    # Should not have called the callback
+    flush_callback.assert_not_called()

sdks/python/tests/unit/message_processing/batching/test_feedback_scores_batcher.py

@@ -0,0 +1,165 @@
+from unittest import mock
+import time
+
+import pytest
+
+from opik.message_processing.batching import batchers
+from opik.message_processing import messages
+
+NOT_USED = None
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class, batch_message_class",
+    [
+        (
+            batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+            messages.AddSpanFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+            messages.AddTraceFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+            messages.AddThreadsFeedbackScoresBatchMessage,
+        ),
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__exactly_max_batch_size_reached__batch_is_flushed(
+    message_batcher_class,
+    batch_message_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    add_feedback_score_batch_messages = [
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[1, 2]),
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[3, 4, 5]),
+    ]  # batcher doesn't care about the content
+
+    for feedback_scores_batch in add_feedback_score_batch_messages:
+        batcher.add(feedback_scores_batch)
+    assert batcher.is_empty()
+
+    flush_callback.assert_called_once_with(
+        batch_message_class(batch=[1, 2, 3, 4, 5], supports_batching=False)
+    )
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class,batch_message_class",
+    [
+        (
+            batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+            messages.AddSpanFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+            messages.AddTraceFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+            messages.AddThreadsFeedbackScoresBatchMessage,
+        ),
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__more_than_max_batch_size_items_added__one_batch_flushed__some_data_remains_in_batcher(
+    message_batcher_class,
+    batch_message_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    add_feedback_score_batch_messages = [
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[1, 2]),
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[3, 4, 5, 6]),
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[7, 8]),
+    ]  # batcher doesn't care about the content
+
+    for feedback_scores_batch in add_feedback_score_batch_messages:
+        batcher.add(feedback_scores_batch)
+
+    assert not batcher.is_empty()
+    flush_callback.assert_called_once_with(
+        batch_message_class(batch=[1, 2, 3, 4, 5], supports_batching=False)
+    )
+    flush_callback.reset_mock()
+
+    batcher.flush()
+    flush_callback.assert_called_once_with(
+        batch_message_class(batch=[6, 7, 8], supports_batching=False)
+    )
+    assert batcher.is_empty()
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+        batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+        batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__batcher_doesnt_have_items__flush_is_called__flush_callback_NOT_called(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    batcher.flush()
+    flush_callback.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+        batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+        batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__ready_to_flush_returns_True__is_flush_interval_passed(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+    FLUSH_INTERVAL = 0.1
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=FLUSH_INTERVAL,
+    )
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()
+    batcher.flush()
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()

sdks/python/tests/unit/message_processing/batching/test_flushing_thread.py

@@ -0,0 +1,174 @@
+import threading
+import time
+from unittest import mock
+from opik.message_processing import messages
+from opik.message_processing.batching import (
+    batch_manager,
+    batchers,
+    flushing_thread,
+)
+from ....testlib import fake_message_factory
+
+
+def test_flushing_thread__batcher_is_flushed__every_time_flush_interval_time_passes():
+    flush_callback = mock.Mock()
+    FLUSH_INTERVAL = 0.2
+    very_big_batch_size = float("inf")
+    batcher = batchers.CreateSpanMessageBatcher(
+        flush_callback=flush_callback,
+        max_batch_size=very_big_batch_size,
+        flush_interval_seconds=FLUSH_INTERVAL,
+    )
+    manager = batch_manager.BatchManager(
+        message_to_batcher_mapping={messages.CreateSpanMessage: batcher},
+    )
+
+    spans_messages = fake_message_factory.fake_span_create_message_batch(
+        count=2, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    manager.start()
+    try:
+        manager.process_message(spans_messages[0])
+        flush_callback.assert_not_called()
+
+        time.sleep(FLUSH_INTERVAL + 0.1)
+        # flush interval has passed after batcher was created, batcher is ready to be flushed
+        # (0.1 is added because a thread probation interval is 0.1, and it's already made its first check)
+        flush_callback.assert_called_once()
+
+        flush_callback.reset_mock()
+
+        manager.process_message(spans_messages[1])
+        time.sleep(FLUSH_INTERVAL)
+        # flush interval has passed after a previous flush, batcher is ready to be flushed again
+        flush_callback.assert_called_once()
+    finally:
+        manager.stop(flush=False)
+
+
+def test_flushing_thread__concurrent_adds_during_flush__no_messages_dropped():
+    """Regression test for OPIK-6444.
+
+    The FlushingThread used to call ``batcher.flush()`` without holding the
+    BatchManager lock, while ``BatchManager.process_message`` holds it for
+    ``add()``. Messages appended during the post-iteration tail of
+    ``_create_batches_from_accumulated_messages()`` (notably during
+    ``sequence_splitter.split_into_batches()``) were not in the returned
+    batches, but were wiped by ``self._accumulated_messages = []``.
+    """
+    delivered_trace_ids = set()
+    delivered_lock = threading.Lock()
+
+    def flush_callback(batch_message):
+        with delivered_lock:
+            for trace in batch_message.batch:
+                delivered_trace_ids.add(trace.id)
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        flush_callback=flush_callback,
+        max_batch_size=10_000,
+        flush_interval_seconds=0.005,
+    )
+    manager = batch_manager.BatchManager(
+        message_to_batcher_mapping={messages.CreateTraceMessage: batcher},
+    )
+    # Probe aggressively to maximise the race window in the buggy code path.
+    manager._flushing_thread._probe_interval_seconds = 0.001
+    manager.start()
+
+    WRITERS = 4
+    ADDS_PER_WRITER = 100
+    PAYLOAD_SIZE = 50 * fake_message_factory.ONE_KILOBYTE
+    added_trace_ids = []
+    added_lock = threading.Lock()
+
+    def writer():
+        trace_messages = fake_message_factory.fake_create_trace_message_batch(
+            count=ADDS_PER_WRITER,
+            approximate_trace_size=PAYLOAD_SIZE,
+            has_ended=True,
+        )
+        for message in trace_messages:
+            manager.process_message(message)
+            with added_lock:
+                added_trace_ids.append(message.trace_id)
+
+    writer_threads = [threading.Thread(target=writer) for _ in range(WRITERS)]
+    for writer_thread in writer_threads:
+        writer_thread.start()
+    for writer_thread in writer_threads:
+        writer_thread.join()
+
+    manager.stop(flush=True)
+
+    assert set(added_trace_ids) == delivered_trace_ids, (
+        f"dropped {len(set(added_trace_ids) - delivered_trace_ids)} of "
+        f"{len(added_trace_ids)} trace messages"
+    )
+
+
+def test_flushing_thread__exception_in_flush_callable__thread_keeps_running():
+    """Regression test for OPIK-6444.
+
+    ``FlushingThread.run`` had no try/except, so any exception in the
+    periodic callable silently killed the daemon thread, after which
+    periodic flushing stopped.
+    """
+    call_counter = {"n": 0}
+
+    def flaky_callable():
+        call_counter["n"] += 1
+        if call_counter["n"] == 1:
+            raise RuntimeError("simulated transient failure")
+
+    tested = flushing_thread.FlushingThread(
+        flush_callable=flaky_callable, probe_interval_seconds=0.01
+    )
+    tested.start()
+    try:
+        time.sleep(0.2)
+    finally:
+        tested.close()
+        tested.join(timeout=2)
+
+    assert call_counter["n"] >= 2, "FlushingThread stopped after the first failure"
+
+
+def test_batch_manager_flush_ready__failing_batcher__remaining_batchers_still_flushed():
+    """A failing flush on one batcher must not skip the others in the same tick."""
+    good_callback = mock.Mock()
+    bad_callback = mock.Mock(side_effect=RuntimeError("boom"))
+
+    bad_batcher = batchers.CreateSpanMessageBatcher(
+        flush_callback=bad_callback,
+        max_batch_size=float("inf"),
+        flush_interval_seconds=0.0,
+    )
+    good_batcher = batchers.CreateTraceMessageBatcher(
+        flush_callback=good_callback,
+        max_batch_size=float("inf"),
+        flush_interval_seconds=0.0,
+    )
+    manager = batch_manager.BatchManager(
+        message_to_batcher_mapping={
+            messages.CreateSpanMessage: bad_batcher,
+            messages.CreateTraceMessage: good_batcher,
+        },
+    )
+
+    span = fake_message_factory.fake_span_create_message_batch(
+        count=1, approximate_span_size=fake_message_factory.ONE_KILOBYTE
+    )[0]
+    trace = fake_message_factory.fake_create_trace_message_batch(
+        count=1,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=True,
+    )[0]
+    manager.process_message(span)
+    manager.process_message(trace)
+
+    manager.flush_ready()
+
+    bad_callback.assert_called_once()
+    good_callback.assert_called_once()

sdks/python/tests/unit/message_processing/batching/test_sequence_splitter.py

@@ -0,0 +1,134 @@
+import json
+
+import pytest
+
+from opik.message_processing.batching import sequence_splitter
+
+from ....testlib import fake_message_factory
+
+
+ONE_MEGABYTE_OBJECT_A = fake_message_factory.LongStr(
+    "a" * fake_message_factory.ONE_MEGABYTE
+)
+ONE_MEGABYTE_OBJECT_B = fake_message_factory.LongStr(
+    "b" * fake_message_factory.ONE_MEGABYTE
+)
+ONE_MEGABYTE_OBJECT_C = fake_message_factory.LongStr(
+    "c" * fake_message_factory.ONE_MEGABYTE
+)
+
+
+def test_split_list_into_batches__by_size_only():
+    items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    batches = sequence_splitter.split_into_batches(items, max_length=4)
+
+    assert batches == [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10]]
+
+
+def test_split_list_into_batches__by_memory_only():
+    items = [ONE_MEGABYTE_OBJECT_A] * 2 + [ONE_MEGABYTE_OBJECT_B] * 2
+    batches = sequence_splitter.split_into_batches(items, max_payload_size_MB=3.5)
+
+    assert batches == [
+        [ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_B],
+        [ONE_MEGABYTE_OBJECT_B],
+    ]
+
+
+def test_split_list_into_batches__by_memory_and_by_size():
+    FOUR_MEGABYTE_OBJECT_C = [ONE_MEGABYTE_OBJECT_C] * 4
+    items = (
+        [ONE_MEGABYTE_OBJECT_A] * 2
+        + [FOUR_MEGABYTE_OBJECT_C]
+        + [ONE_MEGABYTE_OBJECT_B] * 2
+    )
+    batches = sequence_splitter.split_into_batches(
+        items, max_length=2, max_payload_size_MB=3.5
+    )
+
+    # Object C comes before object A because if item is bigger than the max payload size
+    # it is immediately added to the result batches list before batch which is currently accumulating
+    assert batches == [
+        [FOUR_MEGABYTE_OBJECT_C],
+        [ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_A],
+        [ONE_MEGABYTE_OBJECT_B, ONE_MEGABYTE_OBJECT_B],
+    ]
+
+
+def test_split_list_into_batches__empty_list():
+    items = []
+    batches = sequence_splitter.split_into_batches(
+        items, max_length=3, max_payload_size_MB=3.5
+    )
+
+    assert batches == []
+
+
+def test_split_list_into_batches__multiple_large_objects():
+    items = [ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_B, ONE_MEGABYTE_OBJECT_C]
+    batches = sequence_splitter.split_into_batches(
+        items, max_length=2, max_payload_size_MB=0.5
+    )
+
+    assert batches == [
+        [ONE_MEGABYTE_OBJECT_A],
+        [ONE_MEGABYTE_OBJECT_B],
+        [ONE_MEGABYTE_OBJECT_C],
+    ]
+
+
+@pytest.mark.parametrize(
+    "test_input",
+    [
+        ("string"),  # Basic string
+        (""),  # Empty string
+        ("Unicode👋🏼"),  # String with non-UTF8 characters
+        (123.76),  # Float
+        (123456789),  # Integer
+        (True),  # Boolean
+        (False),  # Boolean
+        ([1, 2, 3]),  # List
+        ({"key": "value"}),  # Dictionary
+        (None),  # null
+    ],
+)
+def test_get_json_size_basic(test_input):
+    expected = len(
+        json.dumps(test_input, separators=(",", ":"), ensure_ascii=False).encode(
+            "utf-8"
+        )
+    )
+    assert sequence_splitter._get_json_size(test_input) == expected
+
+
+def test_get_json_size_complex_nested():
+    test_input = {
+        "string": "simple value",
+        "number": 42,
+        "boolean": True,
+        "null": None,
+        "nested_dict": {
+            "level1": {
+                "level2": {
+                    "level3": "deep value",
+                    "numbers": [1, 2, 3],
+                    "mixed_list": [True, "string", 42.0, None],
+                },
+                "sibling2": {
+                    "data": [{"id": 1, "name": "item1"}, {"id": 2, "name": "item2"}]
+                },
+            }
+        },
+        "array_of_arrays": [[1, 2, 3], ["a", "b", "c"], [{"x": 1}, {"y": 2}]],
+        "mixed_nesting": {
+            "lists": [[1, 2], [3, 4]],
+            "dicts": {"a": {"b": {"c": "value"}}, "x": {"y": {"z": 100}}},
+        },
+    }
+
+    expected = len(
+        json.dumps(test_input, separators=(",", ":"), ensure_ascii=False).encode(
+            "utf-8"
+        )
+    )
+    assert sequence_splitter._get_json_size(test_input) == expected

sdks/python/tests/unit/message_processing/batching/test_span_batcher.py

@@ -0,0 +1,236 @@
+from unittest import mock
+import time
+
+import pytest
+
+from opik.message_processing.batching import batchers
+from opik.message_processing import messages
+from ....testlib import fake_message_factory
+
+NOT_USED = mock.sentinel.NOT_USED
+
+
+def test_create_span_message_batcher__exactly_max_batch_size_reached__batch_is_flushed():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=10,
+    )
+
+    assert batcher.is_empty()
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=MAX_BATCH_SIZE, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    for span_message in span_messages:
+        batcher.add(span_message)
+    assert batcher.is_empty()
+
+    assert len(collected_messages) == 1
+
+
+def test_create_span_message_batcher__more_than_max_batch_size_items_added__one_batch_flushed__some_data_remains_in_batcher():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=10,
+    )
+
+    assert batcher.is_empty()
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=MAX_BATCH_SIZE + 2,
+        approximate_span_size=fake_message_factory.ONE_MEGABYTE,
+    )
+
+    for span_message in span_messages:
+        batcher.add(span_message)
+
+    assert not batcher.is_empty()
+    assert len(collected_messages) == 1
+
+    batcher.flush()
+    assert len(collected_messages) == 2
+
+
+def test_create_span_message_batcher__split_message_into_batches__size_limit_reached():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=3,
+    )
+
+    assert batcher.is_empty()
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=2 * 2, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    for span_message in span_messages:
+        batcher.add(span_message)
+
+    assert not batcher.is_empty()
+    batcher.flush()
+
+    # we created 4 messages about 1 MB each, and batcher is limited to 3 MB - thus we expect 2 batches
+    assert len(collected_messages) == 2
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.CreateSpanMessageBatcher,
+        batchers.CreateTraceMessageBatcher,
+    ],
+)
+def test_create_message_batcher__batcher_doesnt_have_items__flush_is_called__flush_callback_NOT_called(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=10,
+    )
+
+    assert batcher.is_empty()
+    batcher.flush()
+    flush_callback.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.CreateSpanMessageBatcher,
+        batchers.CreateTraceMessageBatcher,
+    ],
+)
+def test_create_message_batcher__ready_to_flush_returns_True__is_flush_interval_passed(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+    FLUSH_INTERVAL = 0.1
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=FLUSH_INTERVAL,
+        batch_memory_limit_mb=10,
+    )
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()
+    batcher.flush()
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()
+
+
+def test_create_span_message_batcher__add_duplicated_span__previous_span_is_removed_from_batcher():
+    """
+    Tests the handling of duplicate spans in a span message batcher, specifically focusing on ensuring
+    that newer ones replace older duplicated spans.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+    assert batcher.is_empty()
+
+    span_messages = []
+    span_messages += fake_message_factory.fake_span_create_message_batch(
+        count=1,
+        approximate_span_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+    span_messages += fake_message_factory.fake_span_create_message_batch(
+        count=1, approximate_span_size=fake_message_factory.ONE_KILOBYTE, has_ended=True
+    )
+
+    batcher.add(span_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second span to be considered the same as the first one
+    span_messages[1].span_id = span_messages[0].span_id
+    batcher.add(span_messages[1])
+
+    # assert that the first span is removed from the batcher
+    assert batcher.size() == 1
+
+    # check that the second message is in the batch
+    batcher.flush()
+    assert len(batches) == 1
+    assert len(batches[0].batch) == 1
+    assert batches[0].batch[0].id == span_messages[1].span_id
+
+
+def test_create_span_message_batcher__add_duplicated_span__previous_span_is_not_removed_from_batcher__not_ended_span():
+    """
+    Tests the behavior of adding a duplicate span to a span message batcher and ensures
+    that the previous span is not removed when adding a duplicate span if both are not ended.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+    assert batcher.is_empty()
+
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=2,
+        approximate_span_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+
+    batcher.add(span_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second span to be considered the same as the first one
+    span_messages[1].span_id = span_messages[0].span_id
+    batcher.add(span_messages[1])
+
+    # assert that both spans are in the batcher
+    assert batcher.size() == 2

sdks/python/tests/unit/message_processing/batching/test_trace_batcher.py

@@ -0,0 +1,123 @@
+from unittest import mock
+
+from opik.message_processing import messages
+from opik.message_processing.batching import batchers
+
+from ....testlib import fake_message_factory
+
+NOT_USED = mock.sentinel.NOT_USED
+
+
+def test_create_trace_message_batcher__split_message_into_batches__size_limit_reached():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=3,
+    )
+
+    assert batcher.is_empty()
+    trace_messages = fake_message_factory.fake_create_trace_message_batch(
+        count=2 * 2, approximate_trace_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    for trace_message in trace_messages:
+        batcher.add(trace_message)
+
+    assert not batcher.is_empty()
+    batcher.flush()
+
+    # we created 4 messages about 1 MB each, and batcher is limited to 3 MB - thus we expect 2 batches
+    assert len(collected_messages) == 2
+
+
+def test_create_trace_message_batcher__add_duplicated_trace__previous_trace_is_removed_from_batcher():
+    """
+    Tests the handling of duplicate traces in a trace message batcher, specifically focusing on ensuring
+    that newer ones replace older duplicated traces.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+
+    trace_messages = []
+    trace_messages += fake_message_factory.fake_create_trace_message_batch(
+        count=1,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+    trace_messages += fake_message_factory.fake_create_trace_message_batch(
+        count=1,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=True,
+    )
+
+    batcher.add(trace_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second trace to be considered the same as the first one
+    trace_messages[1].trace_id = trace_messages[0].trace_id
+
+    # add the second trace to the batcher and check that the first one is removed from the batcher
+    batcher.add(trace_messages[1])
+    assert batcher.size() == 1
+
+    batcher.flush()
+    assert len(batches) == 1
+    assert len(batches[0].batch) == 1
+    assert batches[0].batch[0].id == trace_messages[1].trace_id
+
+
+def test_create_trace_message_batcher__add_duplicated_trace__previous_trace_is_not_removed_from_batcher__not_ended_trace():
+    """
+    Tests the behavior of adding a duplicate trace to a trace message batcher and ensures
+    that the previous trace is not removed when adding a duplicate trace if both are not ended.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+
+    trace_messages = fake_message_factory.fake_create_trace_message_batch(
+        count=2,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+
+    batcher.add(trace_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second trace to be considered the same as the first one
+    trace_messages[1].trace_id = trace_messages[0].trace_id
+
+    # add the second trace to the batcher and check that the first one is removed from the batcher
+    batcher.add(trace_messages[1])
+    assert batcher.size() == 2