diff --git a/langfuse/batch_evaluation.py b/langfuse/batch_evaluation.py index 4d7bf3371..325cd4637 100644 --- a/langfuse/batch_evaluation.py +++ b/langfuse/batch_evaluation.py @@ -838,7 +838,8 @@ async def run_async( This is the main implementation method that orchestrates the entire batch evaluation process: fetching items from `GET /api/public/v2/observations` with cursor pagination, mapping, evaluating, creating scores, and tracking - statistics. + statistics. On completion, queued scores are flushed before returning + without blocking the caller's event loop. Args: mapper: Function to transform `ObservationV2` items to evaluator inputs. @@ -1060,7 +1061,7 @@ async def process_item( if verbose: logger.info("Flushing scores to Langfuse...") - self.client.flush() + await asyncio.to_thread(self.client.flush) if verbose: logger.info( diff --git a/tests/unit/test_batch_evaluation.py b/tests/unit/test_batch_evaluation.py index f6c29f4f6..10470aaf9 100644 --- a/tests/unit/test_batch_evaluation.py +++ b/tests/unit/test_batch_evaluation.py @@ -1,13 +1,18 @@ """Unit tests for batch evaluation on the v2 observations API.""" +import asyncio import inspect import json +import threading from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any, Dict, List, Optional from unittest.mock import MagicMock +import httpx import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from langfuse import Langfuse from langfuse.api import ObservationsV2Response, ObservationV2 @@ -141,6 +146,60 @@ async def test_paginates_with_cursor_and_scores_observations(): client.flush.assert_called_once() +@pytest.mark.asyncio +async def test_flush_keeps_event_loop_responsive_until_scores_are_delivered(): + loop = asyncio.get_running_loop() + loop_progressed = threading.Event() + progress_during_upload = [] + delivered_scores = [] + observations = [make_observation(i) for i in range(2)] + + def respond(request: httpx.Request) -> httpx.Response: + if request.url.path == "/api/public/v2/observations": + return httpx.Response( + 200, + json={ + "data": [ + item.model_dump(mode="json", by_alias=True) + for item in observations + ], + "meta": {}, + }, + ) + assert request.url.path == "/api/public/scores" + # The upload starts when flush() wakes the score consumer. Its caller's + # event loop must remain available while flush waits for delivery. + loop.call_soon_threadsafe(loop_progressed.set) + progress_during_upload.append(loop_progressed.wait(timeout=2)) + delivered_scores.extend(json.loads(request.content)) + return httpx.Response(202, json={"message": "Accepted"}) + + with httpx.Client(transport=httpx.MockTransport(respond)) as httpx_client: + client = Langfuse( + public_key="pk-batch-async-flush", + secret_key="sk-batch-async-flush", + base_url="http://langfuse.test", + httpx_client=httpx_client, + tracer_provider=TracerProvider(), + span_exporter=InMemorySpanExporter(), + flush_at=1000, + flush_interval=60, + ) + try: + result = await run(BatchEvaluationRunner(client)) + + assert result.total_items_processed == 2 + assert result.total_scores_created == 2 + assert progress_during_upload == [True] + assert {score["observationId"] for score in delivered_scores} == { + "obs-0", + "obs-1", + } + assert client._resources._score_queue.unfinished_tasks == 0 + finally: + client.shutdown() + + @pytest.mark.asyncio async def test_mapper_receives_raw_string_io_and_requested_fields(): runner, api, _ = make_runner([make_observation(1)])