diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 1e66fced..6952ebe2 100644 --- a/py/src/braintrust/__init__.py +++ b/py/src/braintrust/__init__.py @@ -87,3 +87,4 @@ def is_equal(expected, output): from .sandbox import register_sandbox as register_sandbox from .util import BT_IS_ASYNC_ATTRIBUTE as BT_IS_ASYNC_ATTRIBUTE from .util import MarkAsyncWrapper as MarkAsyncWrapper +from .workflow_eval import * diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 70cf41b6..bd0c300a 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -173,6 +173,18 @@ class SpanInternalOptions(TypedDict, total=False): `context.span_origin.instrumentation.name`. Set by SDK integrations (`openai-auto`, `anthropic-auto`, etc.).""" + initial_span_write_as_merge: bool + """Emit the initial row as a merge for an SDK-owned deterministic span.""" + + skip_initial_span_write: bool + """Rehydrate an existing span without emitting an initial row.""" + + span_id: str + """SDK-controlled span ID.""" + + root_span_id: str + """SDK-controlled root span ID.""" + T = TypeVar("T") TMapping = TypeVar("TMapping", bound=Mapping[str, Any]) @@ -3091,6 +3103,33 @@ def flush(): _state.global_bg_logger().flush() +def _internal_start_span_with_initial_merge( + name: str, + *, + parent: str, + span_id: str, + root_span_id: str, + state: BraintrustState | None = None, + type: SpanTypeAttribute | None = None, + span_attributes: SpanAttributes | Mapping[str, Any] | None = None, + **event: Any, +) -> Span: + """Start a deterministic SDK-owned span whose first row is merge-safe.""" + return start_span( + name=name, + parent=parent, + state=state, + type=type, + span_attributes=span_attributes, + internal={ + "initial_span_write_as_merge": True, + "span_id": span_id, + "root_span_id": root_span_id, + }, + **event, + ) + + def _check_org_info(state, org_info, org_name): if len(org_info) == 0: raise ValueError("This user is not part of any organizations.") @@ -3986,6 +4025,26 @@ def update_span(exported: str, **event: Any) -> None: ) +def _internal_resume_span(exported: str, state: BraintrustState | None = None) -> Span: + """Rehydrate an exported root span so SDK work can continue later.""" + components = SpanComponentsV4.from_str(exported) + if not components.row_id or not components.span_id or not components.root_span_id: + raise ValueError("Only exported root spans can be resumed") + return SpanImpl( + parent_object_type=components.object_type, + parent_object_id=LazyValue(_span_components_to_object_id_lambda(components), use_mutex=False), + parent_compute_object_metadata_args=components.compute_object_metadata_args, + parent_span_ids=None, + event={"id": components.row_id}, + propagated_event=components.propagated_event, + span_id=components.span_id, + root_span_id=components.root_span_id, + state=state, + lookup_span_parent=False, + internal={"skip_initial_span_write": True}, + ) + + @dataclasses.dataclass class ParentSpanIds: span_id: str @@ -4744,7 +4803,8 @@ def __init__( span_attributes=dict(**{"type": type, "name": name, **span_attributes}, exec_counter=exec_counter), created=datetime.datetime.now(datetime.timezone.utc).isoformat(), ) - self._instrumentation = (internal or {}).get("instrumentation") or "braintrust-python-logger" + internal = internal or {} + self._instrumentation = internal.get("instrumentation") or "braintrust-python-logger" internal_data["context"] = merge_span_origin_context( caller_location or {}, self._instrumentation, @@ -4759,8 +4819,8 @@ def __init__( # Resolve all span IDs (span_id, root_span_id, span_parents) span_ids = _resolve_span_ids( - span_id=span_id, - root_span_id=root_span_id, + span_id=internal.get("span_id", span_id), + root_span_id=internal.get("root_span_id", root_span_id), parent_span_ids=parent_span_ids, lookup_span_parent=lookup_span_parent, id_generator=self.state.id_generator, @@ -4772,8 +4832,9 @@ def __init__( # The first log is a replacement, but subsequent logs to the same span # object will be merges. - self._is_merge = False - self.log_internal(event=event, internal_data=internal_data) + self._is_merge = internal.get("initial_span_write_as_merge", False) + if not internal.get("skip_initial_span_write", False): + self.log_internal(event=event, internal_data=internal_data) self._is_merge = True @property diff --git a/py/src/braintrust/test_workflow_eval.py b/py/src/braintrust/test_workflow_eval.py new file mode 100644 index 00000000..c7f6fd8b --- /dev/null +++ b/py/src/braintrust/test_workflow_eval.py @@ -0,0 +1,991 @@ +"""Tests for the experimental workflow eval API.""" + +import asyncio +from unittest.mock import patch + +import pytest + +from . import workflow_eval as workflow_eval_module +from .logger import BraintrustState, Dataset, ObjectMetadata, ProjectDatasetMetadata +from .test_helpers import init_test_exp, with_memory_logger, with_simulate_login # noqa: F401 +from .util import LazyValue +from .workflow_eval import ( + WorkflowEvalCompletedResult, + WorkflowEvalMemoryStore, + WorkflowEvalRedisStore, + WorkflowEvalWaitingResult, + WorkflowScorer, + WorkflowScorerResult, + WorkflowSubmissionCompletionPoll, + WorkflowSubmissionCompletionWebhook, + WorkflowSubmissionPoll, + WorkflowTask, + WorkflowTaskResult, + define_workflow_eval, +) + + +@pytest.mark.asyncio +async def test_memory_store_is_atomic_and_copies_values(): + store = WorkflowEvalMemoryStore() + value = bytearray(b"first") + + first = await store.get_or_set("key", value) + value[:] = b"other" + second = await store.get_or_set("key", b"second") + + assert first.created is True + assert first.value == b"first" + assert second.created is False + assert second.value == b"first" + assert await store.read("key") == b"first" + assert await store.get_set_size("key") == 0 + await asyncio.gather(*(store.add_to_set("key", str(i % 3)) for i in range(30))) + assert await store.get_set_size("key") == 3 + assert await store.read("key") == b"first" + + +class _SyncRedis: + def __init__(self): + self.values = {} + self.calls = [] + self.sets = {} + self.expirations = {} + + def eval(self, script, numkeys, key, member, ttl_ms): + assert "SADD" in script and "PEXPIRE" in script + assert numkeys == 1 + self.sets.setdefault(key, set()).add(member) + self.expirations[key] = ttl_ms + return 1 + + def scard(self, key): + return len(self.sets.get(key, set())) + + def get(self, key): + return self.values.get(key) + + def set(self, key, value, **kwargs): + self.calls.append((key, kwargs)) + old = self.values.get(key) + if kwargs.get("nx") and old is not None: + return old if kwargs.get("get") else None + self.values[key] = value + return old if kwargs.get("get") else True + + +class _AsyncRedis(_SyncRedis): + async def eval(self, *args): + return super().eval(*args) + + async def scard(self, key): + return super().scard(key) + + async def get(self, key): + return super().get(key) + + async def set(self, key, value, **kwargs): + return super().set(key, value, **kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_type", [_SyncRedis, _AsyncRedis]) +async def test_redis_store_supports_sync_and_async_redis_py(client_type): + client = client_type() + store = WorkflowEvalRedisStore(client, key_prefix="test:", ttl_ms=1234) + + await store.write("one", b"value") + assert await store.read("one") == b"value" + assert (await store.get_or_set("two", b"first")).created is True + existing = await store.get_or_set("two", b"second") + + assert existing.created is False + assert existing.value == b"first" + assert all(key.startswith("test:") for key, _ in client.calls) + assert all(options["px"] == 1234 for _, options in client.calls) + assert await store.get_set_size("progress") == 0 + await store.add_to_set("progress", "a") + await store.add_to_set("progress", "a") + await store.add_to_set("progress", "b") + assert await store.get_set_size("progress") == 2 + assert client.sets == {"test:progress": {"a", "b"}} + assert client.expirations == {"test:progress": 1234} + + +@pytest.mark.asyncio +async def test_ordinary_workflow_eval_completes_locally_and_is_idempotent(): + task_calls = [] + score_calls = [] + + def task(input, hooks): + task_calls.append((input, hooks.trial_index)) + hooks.metadata["task"] = True + hooks.tags = ["updated"] + return input * 2 + + def scorer(output, expected, metadata, tags): + score_calls.append((output, metadata, tags)) + return output == expected + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 2, "expected": 4, "metadata": {"source": "test"}}], + task=task, + scores=[scorer], + trial_count=2, + ) + + result = await workflow_eval.start(no_send_logs=True) + repeated = await workflow_eval.status(result.run_id) + + assert isinstance(result, WorkflowEvalCompletedResult) + assert isinstance(repeated, WorkflowEvalCompletedResult) + assert result.summary.scores["scorer"].score == 1 + assert sorted(task_calls) == [(2, 0), (2, 1)] + assert score_calls == [ + (4, {"source": "test", "task": True}, ["updated"]), + (4, {"source": "test", "task": True}, ["updated"]), + ] + + +@pytest.mark.asyncio +async def test_ordinary_only_runs_do_not_require_case_ids_and_get_new_run_ids(): + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"input": 1}], + task=lambda value: value, + scores=[lambda: 1], + ) + + first = await workflow_eval.start(no_send_logs=True) + second = await workflow_eval.start(no_send_logs=True) + + assert first.status == "completed" + assert second.status == "completed" + assert first.run_id != second.run_id + + +@pytest.mark.asyncio +async def test_poll_advances_task_then_mixed_scorers_without_polling_new_submissions(): + submitted = {} + ready = set() + calls = [] + + async def task_submit(item, context): + calls.append(("submit-task", context.submission_id)) + submitted[context.submission_id] = item + return {"provider_id": context.submission_id} + + async def task_collect(_submission, context): + calls.append(("collect-task", context.submission_id)) + return WorkflowTaskResult(output=submitted[context.submission_id].input * 2) + + async def score_submit(item, context): + calls.append(("submit-score", context.submission_id)) + submitted[context.submission_id] = item + return {"provider_id": context.submission_id} + + async def score_collect(_submission, context): + calls.append(("collect-score", context.submission_id)) + item = submitted[context.submission_id] + return WorkflowScorerResult(score=item.output == item.expected) + + async def poll(_submission, context): + calls.append(("poll", context.submission_id)) + return WorkflowSubmissionPoll("complete" if context.submission_id in ready else "pending") + + task = WorkflowTask( + submit=task_submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=task_collect, + ) + submission_score = WorkflowScorer( + name="submission", + submit=score_submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=score_collect, + ) + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i, "expected": i * 2} for i in range(3)], + task=task, + scores=[lambda output, expected: output == expected, submission_score], + ) + + started = await workflow_eval.start(no_send_logs=True) + assert isinstance(started, WorkflowEvalWaitingResult) + assert started.pending.poll == 3 + task_submission_ids = set(submitted) + ready.update(task_submission_ids) + + after_tasks = await workflow_eval.poll(started.run_id) + assert isinstance(after_tasks, WorkflowEvalWaitingResult) + assert after_tasks.pending.poll == 3 + score_submission_ids = set(submitted) - task_submission_ids + assert score_submission_ids + assert not any(call == ("poll", submission_id) for submission_id in score_submission_ids for call in calls) + + ready.update(score_submission_ids) + completed = await workflow_eval.poll(started.run_id) + assert isinstance(completed, WorkflowEvalCompletedResult) + assert completed.summary.scores["scorer_0"].score == 1 + assert completed.summary.scores["submission"].score == 1 + calls_before_status = list(calls) + assert (await workflow_eval.status(started.run_id)).status == "completed" + assert calls == calls_before_status + + +@pytest.mark.asyncio +async def test_failed_poll_raises_and_can_be_retried(): + failure = RuntimeError("provider failed") + ready = False + + async def poll(_submission, _context): + return WorkflowSubmissionPoll("complete") if ready else WorkflowSubmissionPoll("failed", error=failure) + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1}], + task=WorkflowTask( + submit=lambda item, context: {"id": context.submission_id}, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda _submission, _context: WorkflowTaskResult(output=1), + ), + ) + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="provider failed") as raised: + await workflow.poll(started.run_id) + assert raised.value is failure + assert (await workflow.status(started.run_id)).pending.poll == 1 + ready = True + assert (await workflow.poll(started.run_id)).status == "completed" + + +@pytest.mark.asyncio +async def test_webhook_result_can_be_matched_by_external_id(): + submitted = {} + + async def submit(item, context): + submitted[context.submission_id] = item + return {"id": f"external-{context.submission_id}"} + + async def collect(_submission, context): + return WorkflowTaskResult(output=submitted[context.submission_id].input) + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": "ok"}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=collect, + ), + ) + started = await workflow_eval.start(no_send_logs=True) + submission_id = next(iter(submitted)) + + completed = await workflow_eval.process_submission_result(started.run_id, external_id=f"external-{submission_id}") + + assert isinstance(completed, WorkflowEvalCompletedResult) + assert ( + await workflow_eval.process_submission_result(started.run_id, submission_id=submission_id) + ).status == "completed" + + +@pytest.mark.asyncio +async def test_collect_rejects_array_results(): + submission_ids = [] + + async def submit(_item, context): + submission_ids.append(context.submission_id) + return {"id": context.submission_id} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "a", "input": 1}, {"id": "b", "input": 2}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: [WorkflowTaskResult(output=1)], + ), + ) + started = await workflow_eval.start(no_send_logs=True) + + with pytest.raises(TypeError, match="single WorkflowTaskResult"): + await workflow_eval.process_submission_result(started.run_id, submission_id=submission_ids[0]) + + +@pytest.mark.asyncio +async def test_case_ids_must_be_stable_and_unique(): + submission_task = WorkflowTask( + submit=lambda _item, _context: {"id": "unused"}, + completion=WorkflowSubmissionCompletionPoll(lambda _submission, _context: WorkflowSubmissionPoll("pending")), + collect=lambda _submission, _context: [], + ) + missing = define_workflow_eval( + "project", store=WorkflowEvalMemoryStore(), data=[{"input": 1}], task=submission_task + ) + with pytest.raises(ValueError, match="non-empty id"): + await missing.start(no_send_logs=True) + + duplicate = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "same", "input": 1}, {"id": "same", "input": 2}], + task=submission_task, + ) + with pytest.raises(ValueError, match="duplicate"): + await duplicate.start(no_send_logs=True) + + +@pytest.mark.asyncio +async def test_case_persistence_does_not_deepcopy_inputs(): + class SerializableWithoutDeepcopy: + def __deepcopy__(self, _memo): + raise AssertionError("input was deep-copied") + + def model_dump(self, **_kwargs): + return {"value": "serialized"} + + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "pending"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": SerializableWithoutDeepcopy()}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("pending") + ), + collect=lambda _submission, _context: [], + ), + ) + + result = await workflow_eval.start(no_send_logs=True) + + assert result.status == "waiting" + assert submitted[0].input == {"value": "serialized"} + + +@pytest.mark.asyncio +async def test_concurrent_webhooks_claim_downstream_submission_once(): + task_items = [] + collects_started = 0 + both_collecting = asyncio.Event() + score_submissions = 0 + + async def submit_task(item, _context): + task_items.append(item) + return {"id": "task-provider"} + + async def collect_task(_submission, _context): + nonlocal collects_started + collects_started += 1 + if collects_started == 2: + both_collecting.set() + await both_collecting.wait() + return WorkflowTaskResult(output=task_items[0].input * 2) + + async def submit_score(_item, _context): + nonlocal score_submissions + score_submissions += 1 + return {"id": "score-provider"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 2, "expected": 4}], + task=WorkflowTask( + submit=submit_task, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=collect_task, + ), + scores=[ + WorkflowScorer( + name="exact", + submit=submit_score, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: [], + ) + ], + ) + started = await workflow_eval.start(no_send_logs=True) + + await asyncio.gather( + workflow_eval.process_submission_result(started.run_id, external_id="task-provider"), + workflow_eval.process_submission_result(started.run_id, external_id="task-provider"), + ) + + assert score_submissions == 1 + current = await workflow_eval.status(started.run_id) + assert current.status == "waiting" + assert current.pending.webhook == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy_ids", [False, True]) +async def test_workflow_logging_uses_stable_spans_and_resume_metadata( + monkeypatch, with_memory_logger, with_simulate_login, legacy_ids +): + if legacy_ids: + monkeypatch.setenv("BRAINTRUST_LEGACY_IDS", "true") + else: + monkeypatch.delenv("BRAINTRUST_LEGACY_IDS", raising=False) + local_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1, "expected": 2}], + task=lambda value: value + 1, + scores=[lambda output, expected: output == expected], + ) + local_summary = (await local_eval.start(no_send_logs=True)).summary + experiment = init_test_exp("workflow", "project") + monkeypatch.setattr(workflow_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + submitted = [] + trace_configurations = [] + classifier_trace_configurations = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "task"} + + def scorer(output, expected, trace): + trace_configurations.append(trace.get_configuration()) + return output == expected + + def classifier(output, trace): + classifier_trace_configurations.append(trace.get_configuration()) + return {"id": "positive"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1, "expected": 2}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda _submission, _context: WorkflowTaskResult(output=submitted[0].input + 1), + ), + scores=[scorer], + classifiers=[classifier], + experiment_name="workflow", + ) + + waiting = await workflow_eval.start() + result = await workflow_eval.poll(waiting.run_id) + logs = with_memory_logger.pop() + + assert result.status == "completed" + assert len(logs) == 4 + roots = [row for row in logs if not row["span_parents"]] + assert len(roots) == 1 + assert roots[0]["metadata"]["workflow_eval"] == { + "run_id": result.run_id, + "case_id": "case", + "trial_index": 0, + } + assert trace_configurations == [ + { + "object_type": "experiment", + "object_id": experiment.id, + "root_span_id": roots[0]["root_span_id"], + } + ] + assert classifier_trace_configurations == trace_configurations + assert len({row["span_id"] for row in logs}) == 4 + await workflow_eval.status(result.run_id) + assert with_memory_logger.pop() == [] + + +@pytest.mark.asyncio +async def test_workflow_logging_flushes_before_persisting_log_markers( + monkeypatch, with_memory_logger, with_simulate_login +): + local_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": "case", "input": 1}], + task=lambda value: value, + ) + local_summary = (await local_eval.start(no_send_logs=True)).summary + experiment = init_test_exp("workflow", "project") + monkeypatch.setattr(workflow_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + + flush_count = 0 + marker_flush_counts = [] + original_flush = with_memory_logger.flush + + def flush(*args, **kwargs): + nonlocal flush_count + flush_count += 1 + return original_flush(*args, **kwargs) + + monkeypatch.setattr(with_memory_logger, "flush", flush) + + class RecordingStore(WorkflowEvalMemoryStore): + async def write(self, key, value): + if "/task-log/" in key or "/score-log-" in key or "/classification-log-" in key: + marker_flush_counts.append(flush_count) + await super().write(key, value) + + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "task"} + + workflow_eval = define_workflow_eval( + "project", + store=RecordingStore(), + data=[{"id": "case", "input": 1, "expected": 1}], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda _submission, _context: WorkflowTaskResult(output=submitted[0].input), + ), + scores=[lambda output, expected: output == expected], + classifiers=[lambda output: {"id": "positive"}], + experiment_name="workflow", + ) + + waiting = await workflow_eval.start() + completed = await workflow_eval.poll(waiting.run_id) + + assert completed.status == "completed" + assert len(marker_flush_counts) == 3 + assert all(count > 0 for count in marker_flush_counts) + + +@pytest.mark.asyncio +async def test_workflow_dataset_rows_preserve_dataset_origin(monkeypatch, with_memory_logger, with_simulate_login): + project_metadata = ObjectMetadata(id="test-project", name="test-project", full_info={}) + dataset_metadata = ObjectMetadata(id="active-dataset", name="test-dataset", full_info={}) + dataset = Dataset( + lazy_metadata=LazyValue( + lambda: ProjectDatasetMetadata(project=project_metadata, dataset=dataset_metadata), + use_mutex=False, + ), + state=BraintrustState(), + ) + row = { + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + "input": 1, + } + local_summary = ( + await define_workflow_eval( + "project", store=WorkflowEvalMemoryStore(), data=[{"input": 1}], task=lambda value: value + ).start(no_send_logs=True) + ).summary + experiment = init_test_exp("workflow", "project") + monkeypatch.setattr(workflow_eval_module, "init_experiment", lambda **_kwargs: experiment) + monkeypatch.setattr(experiment, "summarize", lambda **_kwargs: local_summary) + + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": "task"} + + workflow_eval = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=dataset, + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda _submission, _context: WorkflowTaskResult(output=submitted[0].input), + ), + experiment_name="workflow", + ) + + with patch.object(dataset, "_refetch", return_value=[row]): + waiting = await workflow_eval.start() + await workflow_eval.poll(waiting.run_id) + + root = next(log for log in with_memory_logger.pop() if not log["span_parents"]) + assert root["origin"] == { + "object_type": "dataset", + "object_id": "active-dataset", + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + } + + +@pytest.mark.asyncio +async def test_only_completed_cases_start_scorers_and_classifiers(): + ready = set() + submitted = {} + local_scores = [] + classifications = [] + + async def submit(item, context): + submitted[context.submission_id] = item + return {"id": context.submission_id} + + async def poll(submission, _context): + return WorkflowSubmissionPoll("complete" if submission["id"] in ready else "pending") + + async def score(output): + local_scores.append(output) + return 1 + + async def classify(output): + classifications.append(output) + return {"id": "positive"} + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(3)], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda submission, _context: WorkflowTaskResult(output=submitted[submission["id"]].input), + ), + scores=[ + score, + WorkflowScorer( + name="remote", + submit=submit, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda _submission, _context: WorkflowScorerResult(score=1), + ), + ], + classifiers=[classify], + ) + started = await workflow.start(no_send_logs=True) + task_ids = set(submitted) + ready.add(next(key for key, item in submitted.items() if item.input == 1)) + partial = await workflow.poll(started.run_id) + assert partial.pending.poll == 3 # two tasks and the ready case's scorer + assert local_scores == [1] + assert classifications == [1] + score_ids = set(submitted) - task_ids + assert len(score_ids) == 1 + assert submitted[next(iter(score_ids))].output == 1 + ready.update(score_ids) + partial = await workflow.poll(started.run_id) + assert partial.status == "waiting" + assert partial.pending.poll == 2 + assert local_scores == [1] + ready.update(task_ids) + await workflow.poll(started.run_id) + ready.update(submitted) + completed = await workflow.poll(started.run_id) + assert completed.status == "completed" + assert sorted(local_scores) == [0, 1, 2] + assert sorted(classifications) == [0, 1, 2] + assert completed.summary.scores["remote"].score == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("max_concurrency", [None, 1, 3]) +async def test_provider_callback_concurrency_is_bounded(max_concurrency): + active = {"submit": 0, "poll": 0, "collect": 0} + peak = active.copy() + + async def callback(stage): + active[stage] += 1 + peak[stage] = max(peak[stage], active[stage]) + await asyncio.sleep(0.001) + active[stage] -= 1 + + async def submit(item, context): + await callback("submit") + return {"id": context.submission_id, "input": item.input} + + async def poll(_submission, _context): + await callback("poll") + return WorkflowSubmissionPoll("complete") + + async def collect(submission, _context): + await callback("collect") + return WorkflowTaskResult(output=submission["input"]) + + options = {} if max_concurrency is None else {"max_concurrency": max_concurrency} + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(25)], + task=WorkflowTask(submit=submit, completion=WorkflowSubmissionCompletionPoll(poll), collect=collect), + **options, + ) + started = await workflow.start(no_send_logs=True) + assert started.pending.poll == 25 + assert (await workflow.poll(started.run_id)).status == "completed" + limit = max_concurrency or 10 + assert peak["submit"] == limit + assert peak["poll"] == limit + assert 0 < peak["collect"] <= limit + assert all(value == 0 for value in active.values()) + + +@pytest.mark.parametrize("max_concurrency", [0, -1, 1.5, True, "2", None]) +def test_invalid_concurrency_is_rejected(max_concurrency): + with pytest.raises(ValueError, match="max_concurrency must be a positive integer"): + define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[], + task=lambda value: value, + max_concurrency=max_concurrency, + ) + + +@pytest.mark.asyncio +async def test_poll_error_does_not_prevent_independent_case_progress(): + scores = [] + + async def poll(submission, _context): + if submission["input"] == 0: + raise RuntimeError("provider unavailable") + return WorkflowSubmissionPoll("complete") + + async def score(output): + scores.append(output) + return 1 + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(3)], + task=WorkflowTask( + submit=lambda item, _context: {"input": item.input}, + completion=WorkflowSubmissionCompletionPoll(poll), + collect=lambda submission, _context: WorkflowTaskResult(output=submission["input"]), + ), + scores=[score], + max_concurrency=1, + ) + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="provider unavailable"): + await workflow.poll(started.run_id) + assert scores == [1, 2] + assert (await workflow.status(started.run_id)).pending.poll == 1 + + +@pytest.mark.asyncio +async def test_submission_error_waits_for_independent_submissions(): + calls = [] + run_ids = [] + + async def submit(item, context): + run_ids.append(context.run_id) + if item.input == 0: + raise RuntimeError("submission failed") + await asyncio.sleep(0) + calls.append(item.input) + return {"input": item.input} + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[{"id": str(i), "input": i} for i in range(3)], + task=WorkflowTask( + submit=submit, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("pending") + ), + collect=lambda submission, _context: WorkflowTaskResult(output=submission["input"]), + ), + max_concurrency=1, + ) + with pytest.raises(RuntimeError, match="submission failed"): + await workflow.start(no_send_logs=True) + assert calls == [1, 2] + assert (await workflow.status(run_ids[0])).pending.poll == 2 + + +@pytest.mark.asyncio +async def test_webhooks_resume_fresh_definitions_and_validate_both_locators(): + store = WorkflowEvalMemoryStore() + submitted = {} + collected = [] + + async def submit(item, context): + submitted[item.input] = context.submission_id + return {"id": item.input} + + async def collect(submission, _context): + collected.append(submission["id"]) + await asyncio.sleep(0) + return WorkflowTaskResult(output=submission["id"]) + + def definition(): + return define_workflow_eval( + "project", + store=store, + data=[{"id": value, "input": value} for value in ("a", "b")], + task=WorkflowTask( + submit=submit, + collect=collect, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + ), + ) + + started = await definition().start(no_send_logs=True) + workflow = definition() + for locators in [ + {"submission_id": submitted["a"], "external_id": "b"}, + {"submission_id": submitted["a"], "external_id": "unknown"}, + {"submission_id": "unknown", "external_id": "a"}, + ]: + with pytest.raises(ValueError, match="different submissions"): + await workflow.process_submission_result(started.run_id, **locators) + with pytest.raises(ValueError, match="requires submission_id or external_id"): + await workflow.process_submission_result(started.run_id) + with pytest.raises(ValueError, match="No submission matches"): + await workflow.process_submission_result(started.run_id, external_id="unknown") + await asyncio.gather( + workflow.process_submission_result(started.run_id, submission_id=submitted["a"]), + workflow.process_submission_result(started.run_id, external_id="b"), + ) + assert (await workflow.status(started.run_id)).status == "completed" + assert sorted(collected) == ["a", "b"] + await workflow.process_submission_result(started.run_id, external_id="a") + assert sorted(collected) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_workflow_trials_preserve_parameters_metadata_and_tags(): + task_items = [] + score_items = [] + + async def submit_task(item, context): + task_items.append(item) + return {"id": context.submission_id, "trial": item.trial_index} + + async def submit_score(item, _context): + score_items.append(item) + return {"id": item.id} + + workflow = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + case_id=lambda datum: "case", + data=[{"input": 1, "expected": 2, "metadata": {"original": True}, "tags": ["old"], "trial_count": 2}], + task=WorkflowTask( + submit=submit_task, + completion=WorkflowSubmissionCompletionPoll( + lambda _submission, _context: WorkflowSubmissionPoll("complete") + ), + collect=lambda submission, _context: WorkflowTaskResult( + output=submission["trial"], metadata={"new": True}, tags=["updated"] + ), + ), + scores=[ + WorkflowScorer( + name="remote", + submit=submit_score, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]), + collect=lambda _submission, _context: WorkflowScorerResult(score=1), + ) + ], + ) + started = await workflow.start({"temperature": 0.5}, no_send_logs=True) + assert {item.id for item in task_items} == {"case:trial:0", "case:trial:1"} + assert all(item.parameters == {"temperature": 0.5} for item in task_items) + assert sorted(item.trial_index for item in task_items) == [0, 1] + assert (await workflow.poll(started.run_id)).pending.webhook == 2 + assert all(item.metadata == {"original": True, "new": True} for item in score_items) + assert all(item.tags == ["updated"] and item.expected == 2 for item in score_items) + for item in score_items: + result = await workflow.process_submission_result(started.run_id, external_id=item.id) + assert result.status == "completed" + + +@pytest.mark.asyncio +async def test_ordinary_tasks_with_workflow_scorers_and_empty_data(): + submitted = [] + + async def submit(item, _context): + submitted.append(item) + return {"id": item.id} + + scorer = WorkflowScorer( + name="remote", + submit=submit, + completion=WorkflowSubmissionCompletionPoll(lambda _submission, _context: WorkflowSubmissionPoll("complete")), + collect=lambda _submission, _context: WorkflowScorerResult(score=1), + ) + for data in [[], [{"id": "case", "input": 1}]]: + workflow = define_workflow_eval( + "project", store=WorkflowEvalMemoryStore(), data=data, task=lambda value: value + 1, scores=[scorer] + ) + started = await workflow.start(no_send_logs=True) + if data: + assert started.pending.poll == 1 + assert submitted[0].output == 2 + else: + assert started.status == "completed" + assert (await workflow.poll(started.run_id)).status == "completed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["poll", "webhook"]) +async def test_replay_repairs_interrupted_completion_progress(mode): + class InterruptedStore(WorkflowEvalMemoryStore): + interrupted = False + + async def add_to_set(self, key, member): + if key.endswith("/complete") and not self.interrupted: + self.interrupted = True + raise RuntimeError("interrupted progress") + await super().add_to_set(key, member) + + collects = [] + + async def collect(submission, _context): + collects.append(submission["id"]) + return WorkflowTaskResult(output=1) + + workflow = define_workflow_eval( + "project", + store=InterruptedStore(), + data=[{"id": value, "input": value} for value in ("a", "b")], + task=WorkflowTask( + submit=lambda item, _context: {"id": item.input}, + collect=collect, + completion=WorkflowSubmissionCompletionWebhook(lambda submission, _context: submission["id"]) + if mode == "webhook" + else WorkflowSubmissionCompletionPoll( + lambda submission, _context: WorkflowSubmissionPoll( + "complete" if submission["id"] == "a" else "pending" + ) + ), + ), + ) + started = await workflow.start(no_send_logs=True) + with pytest.raises(RuntimeError, match="interrupted progress"): + if mode == "webhook": + await workflow.process_submission_result(started.run_id, external_id="a") + else: + await workflow.poll(started.run_id) + if mode == "webhook": + result = await workflow.process_submission_result(started.run_id, external_id="a") + else: + result = await workflow.poll(started.run_id) + assert result.pending.poll + result.pending.webhook == 1 + assert collects == ["a"] diff --git a/py/src/braintrust/type_tests/test_workflow_eval.py b/py/src/braintrust/type_tests/test_workflow_eval.py new file mode 100644 index 00000000..a43b7683 --- /dev/null +++ b/py/src/braintrust/type_tests/test_workflow_eval.py @@ -0,0 +1,75 @@ +"""Static and runtime type coverage for the experimental workflow eval API.""" + +from typing import TypedDict + +import pytest +from braintrust import ( + EvalCase, + WorkflowEval, + WorkflowEvalMemoryStore, + WorkflowScorer, + WorkflowScorerItem, + WorkflowScorerResult, + WorkflowSubmissionCompletionPoll, + WorkflowSubmissionContext, + WorkflowSubmissionPoll, + WorkflowTask, + WorkflowTaskItem, + WorkflowTaskResult, + define_workflow_eval, +) + + +class Submission(TypedDict): + id: str + + +async def submit_task(item: WorkflowTaskItem[str, str], context: WorkflowSubmissionContext) -> Submission: + assert item.id + return {"id": context.submission_id} + + +async def collect_task(submission: Submission, context: WorkflowSubmissionContext) -> WorkflowTaskResult[int]: + return WorkflowTaskResult(output=len(submission["id"] + context.run_id)) + + +async def submit_score(item: WorkflowScorerItem[str, int, str], context: WorkflowSubmissionContext) -> Submission: + assert item.id + return {"id": context.submission_id} + + +async def collect_score(submission: Submission, context: WorkflowSubmissionContext) -> WorkflowScorerResult: + assert submission["id"] == context.submission_id + return WorkflowScorerResult(score=1) + + +async def poll_submission(submission: Submission, context: WorkflowSubmissionContext) -> WorkflowSubmissionPoll: + assert submission["id"] == context.submission_id + return WorkflowSubmissionPoll(status="pending") + + +task: WorkflowTask[str, int, str, Submission] = WorkflowTask( + submit=submit_task, + completion=WorkflowSubmissionCompletionPoll(poll_submission), + collect=collect_task, +) +score: WorkflowScorer[str, int, str, Submission] = WorkflowScorer( + name="score", + submit=submit_score, + completion=WorkflowSubmissionCompletionPoll(poll_submission), + collect=collect_score, +) +workflow_eval: WorkflowEval[str, int, str] = define_workflow_eval( + "project", + store=WorkflowEvalMemoryStore(), + data=[EvalCase(id="case", input="input", expected="expected")], + task=task, + scores=[score], +) + + +@pytest.mark.asyncio +async def test_workflow_eval_types_at_runtime(): + result = await workflow_eval.start(no_send_logs=True) + assert result.status == "waiting" + assert result.pending.poll == 1 diff --git a/py/src/braintrust/workflow_eval.py b/py/src/braintrust/workflow_eval.py new file mode 100644 index 00000000..01d5ae55 --- /dev/null +++ b/py/src/braintrust/workflow_eval.py @@ -0,0 +1,1442 @@ +"""Experimental workflow evaluations with one asynchronous provider submission per case/trial.""" + +import asyncio +import base64 +import dataclasses +import functools +import hashlib +import inspect +import threading +import uuid +from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence +from typing import Any, Generic, Literal, Protocol, TypeVar, cast + +from .bt_json import bt_dumps, bt_loads +from .env import BraintrustEnv +from .framework import ( + BaseExperiment, + DictEvalHooks, + EvalCase, + EvalClassifier, + EvalData, + EvalResult, + EvalScorer, + EvalTask, + Evaluator, + OneOrMoreScores, + _build_classification_span_output, + _build_span_metadata, + _classifier_name, + _get_persisted_base_experiment_id, + _normalize_score, + _scorer_name, + _validate_classification_result, + _validated_object_reference, + await_or_run, + build_local_summary, + call_user_fn, + init_experiment, + run_evaluator, +) +from .git_fields import GitMetadataSettings, RepoInfo +from .logger import ( + NOOP_SPAN, + BraintrustState, + Dataset, + Experiment, + ExperimentSummary, + Metadata, + Span, + _internal_get_global_state, + _internal_resume_span, + _internal_start_span_with_initial_merge, + span_components_to_object_id, +) +from .logger import init as _init_experiment +from .parameters import EvalParameters, RemoteEvalParameters, ValidatedParameters, validate_parameters +from .score import Classification, Score, ScoreLike, is_score, is_scorer +from .span_identifier_v3 import span_object_type_v3_to_typed_string +from .span_identifier_v4 import SpanComponentsV4 +from .span_types import SpanTypeAttribute +from .trace import LocalTrace +from .util import get_signature, merge_dicts + + +Input = TypeVar("Input") +Output = TypeVar("Output") +Expected = TypeVar("Expected") +SubmissionData = TypeVar("SubmissionData") + +DEFAULT_MAX_CONCURRENCY = 10 +DEFAULT_REDIS_TTL_MS = 1_000 * 60 * 60 * 24 * 7 +_SCHEMA_PREFIX = "workflow-eval/python/v1" + + +@dataclasses.dataclass(frozen=True) +class WorkflowSubmissionContext: + """Identifiers supplied to a workflow submission processor callback.""" + + run_id: str + submission_id: str + + +@dataclasses.dataclass(frozen=True) +class WorkflowSubmissionPoll: + """The current state returned by a polling completion callback.""" + + status: Literal["pending", "complete", "failed"] + error: Any = None + + +@dataclasses.dataclass(frozen=True) +class WorkflowSubmissionCompletionPoll(Generic[SubmissionData]): + """Configures a submission processor whose provider is checked by polling.""" + + poll: Callable[ + [SubmissionData, WorkflowSubmissionContext], WorkflowSubmissionPoll | Awaitable[WorkflowSubmissionPoll] + ] + mode: Literal["poll"] = dataclasses.field(default="poll", init=False) + + +@dataclasses.dataclass(frozen=True) +class WorkflowSubmissionCompletionWebhook(Generic[SubmissionData]): + """Configures a submission processor completed by an incoming webhook.""" + + get_external_id: Callable[[SubmissionData, WorkflowSubmissionContext], str | Awaitable[str]] + mode: Literal["webhook"] = dataclasses.field(default="webhook", init=False) + + +WorkflowSubmissionCompletion = ( + WorkflowSubmissionCompletionPoll[SubmissionData] | WorkflowSubmissionCompletionWebhook[SubmissionData] +) + + +@dataclasses.dataclass(frozen=True) +class WorkflowTaskItem(Generic[Input, Expected]): + """One case/trial passed to a workflow task submission callback.""" + + id: str + input: Input + expected: Expected | None + metadata: Metadata + tags: list[str] | None + parameters: ValidatedParameters | None + trial_index: int + + +@dataclasses.dataclass(frozen=True) +class WorkflowTaskResult(Generic[Output]): + """A collected result for one task item.""" + + output: Output + metadata: Metadata | None = None + tags: list[str] | None = None + + +@dataclasses.dataclass(frozen=True) +class WorkflowScorerItem(Generic[Input, Output, Expected]): + """One completed case/trial passed to a workflow scorer submission callback.""" + + id: str + input: Input + output: Output + expected: Expected | None + metadata: Metadata + tags: list[str] | None + trial_index: int + + +@dataclasses.dataclass(frozen=True) +class WorkflowScorerResult: + """A collected score for one scorer item.""" + + score: OneOrMoreScores + + +@dataclasses.dataclass(frozen=True) +class WorkflowTask(Generic[Input, Output, Expected, SubmissionData]): + """Submits one asynchronous provider operation per case/trial and collects its task result.""" + + submit: Callable[ + [WorkflowTaskItem[Input, Expected], WorkflowSubmissionContext], SubmissionData | Awaitable[SubmissionData] + ] + completion: WorkflowSubmissionCompletion[SubmissionData] + collect: Callable[ + [SubmissionData, WorkflowSubmissionContext], WorkflowTaskResult[Output] | Awaitable[WorkflowTaskResult[Output]] + ] + + +@dataclasses.dataclass(frozen=True) +class WorkflowScorer(Generic[Input, Output, Expected, SubmissionData]): + """Submits one asynchronous provider operation per case/trial and collects its score.""" + + name: str + submit: Callable[ + [WorkflowScorerItem[Input, Output, Expected], WorkflowSubmissionContext], + SubmissionData | Awaitable[SubmissionData], + ] + completion: WorkflowSubmissionCompletion[SubmissionData] + collect: Callable[ + [SubmissionData, WorkflowSubmissionContext], WorkflowScorerResult | Awaitable[WorkflowScorerResult] + ] + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("WorkflowScorer name must be a non-empty string") + + +@dataclasses.dataclass(frozen=True) +class WorkflowEvalStoreEntry: + """Result of an atomic workflow-store get-or-set operation.""" + + value: bytes + created: bool + + +class WorkflowEvalStore(Protocol): + """Minimal persistence interface used by workflow evaluations.""" + + async def read(self, key: str) -> bytes | None: ... + + async def write(self, key: str, value: bytes) -> None: ... + + async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: ... + + async def add_to_set(self, key: str, member: str) -> None: + """Atomically add a unique member; retain sets as long as run records.""" + ... + + async def get_set_size(self, key: str) -> int: ... + + +class WorkflowEvalMemoryStore: + """Process-local workflow evaluation state, intended for tests and local runs.""" + + def __init__(self) -> None: + self._values: dict[str, bytes] = {} + self._sets: dict[str, set[str]] = {} + self._lock = threading.Lock() + + async def read(self, key: str) -> bytes | None: + with self._lock: + value = self._values.get(key) + return bytes(value) if value is not None else None + + async def write(self, key: str, value: bytes) -> None: + with self._lock: + self._values[key] = bytes(value) + + async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: + with self._lock: + existing = self._values.get(key) + if existing is not None: + return WorkflowEvalStoreEntry(value=bytes(existing), created=False) + self._values[key] = bytes(value) + return WorkflowEvalStoreEntry(value=bytes(value), created=True) + + async def add_to_set(self, key: str, member: str) -> None: + with self._lock: + self._sets.setdefault(key, set()).add(member) + + async def get_set_size(self, key: str) -> int: + with self._lock: + return len(self._sets.get(key, set())) + + +class WorkflowEvalRedisStore: + """Workflow state backed by an existing sync or async redis-py client.""" + + def __init__( + self, + client: Any, + *, + key_prefix: str = "braintrust-eval:", + ttl_ms: int = DEFAULT_REDIS_TTL_MS, + ) -> None: + if not isinstance(ttl_ms, int) or isinstance(ttl_ms, bool) or ttl_ms < 1: + raise ValueError("WorkflowEvalRedisStore ttl_ms must be a positive integer") + self.client = client + self.key_prefix = key_prefix + self.ttl_ms = ttl_ms + + async def _call(self, method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + if inspect.iscoroutinefunction(method): + return await method(*args, **kwargs) + value = await asyncio.get_running_loop().run_in_executor(None, functools.partial(method, *args, **kwargs)) + if inspect.isawaitable(value): + return await value + return value + + async def read(self, key: str) -> bytes | None: + value = await self._call(self.client.get, f"{self.key_prefix}{key}") + if value is None: + return None + if isinstance(value, bytes): + value = value.decode("ascii") + if not isinstance(value, str): + raise TypeError("WorkflowEvalRedisStore expected GET to return str, bytes, or None") + return base64.b64decode(value) + + async def write(self, key: str, value: bytes) -> None: + encoded = base64.b64encode(value).decode("ascii") + await self._call(self.client.set, f"{self.key_prefix}{key}", encoded, px=self.ttl_ms) + + async def get_or_set(self, key: str, value: bytes) -> WorkflowEvalStoreEntry: + redis_key = f"{self.key_prefix}{key}" + encoded = base64.b64encode(value).decode("ascii") + existing = await self._call(self.client.set, redis_key, encoded, px=self.ttl_ms, nx=True, get=True) + if existing is None: + return WorkflowEvalStoreEntry(value=bytes(value), created=True) + if isinstance(existing, bytes): + existing = existing.decode("ascii") + if not isinstance(existing, str): + raise TypeError("WorkflowEvalRedisStore expected atomic SET to return str, bytes, or None") + return WorkflowEvalStoreEntry(value=base64.b64decode(existing), created=False) + + async def add_to_set(self, key: str, member: str) -> None: + await self._call( + self.client.eval, + "redis.call('SADD', KEYS[1], ARGV[1]); redis.call('PEXPIRE', KEYS[1], ARGV[2]); return 1", + 1, + f"{self.key_prefix}{key}", + member, + self.ttl_ms, + ) + + async def get_set_size(self, key: str) -> int: + return int(await self._call(self.client.scard, f"{self.key_prefix}{key}")) + + +@dataclasses.dataclass(frozen=True) +class WorkflowEvalPending: + """Counts of submitted provider submissions awaiting completion.""" + + poll: int + webhook: int + + +@dataclasses.dataclass(frozen=True) +class WorkflowEvalWaitingResult: + run_id: str + pending: WorkflowEvalPending + status: Literal["waiting"] = dataclasses.field(default="waiting", init=False) + + +@dataclasses.dataclass(frozen=True) +class WorkflowEvalCompletedResult: + run_id: str + pending: WorkflowEvalPending + summary: ExperimentSummary + status: Literal["completed"] = dataclasses.field(default="completed", init=False) + + +WorkflowEvalResult = WorkflowEvalWaitingResult | WorkflowEvalCompletedResult + + +@dataclasses.dataclass(frozen=True) +class _WorkflowEvalConfig(Generic[Input, Output, Expected]): + project_name: str + store: WorkflowEvalStore + data: EvalData[Input, Expected] + task: EvalTask[Input, Output, Expected] | WorkflowTask[Input, Output, Expected, Any] + scores: Sequence[EvalScorer[Input, Output, Expected] | WorkflowScorer[Input, Output, Expected, Any]] + classifiers: Sequence[EvalClassifier[Input, Output, Expected]] + case_id: Callable[[EvalCase[Input, Expected]], str | Awaitable[str]] | None + experiment_name: str | None + trial_count: int + max_concurrency: int + metadata: Metadata | None + tags: Sequence[str] | None + is_public: bool + project_id: str | None + base_experiment_name: str | None + base_experiment_id: str | None + git_metadata_settings: GitMetadataSettings | None + repo_info: RepoInfo | None + description: str | None + summarize_scores: bool + parameters: EvalParameters | RemoteEvalParameters | None + state: BraintrustState | None + + +class WorkflowEval(Generic[Input, Output, Expected]): + """A workflow evaluation definition that can be started and resumed.""" + + def __init__(self, config: _WorkflowEvalConfig[Input, Output, Expected]) -> None: + self._config = config + + async def start( + self, parameters: Mapping[str, Any] | None = None, *, no_send_logs: bool = False + ) -> WorkflowEvalResult: + return await _WorkflowEvalRunner(self._config).start(parameters, no_send_logs=no_send_logs) + + async def status(self, run_id: str) -> WorkflowEvalResult: + return await _WorkflowEvalRunner(self._config).status(run_id) + + async def poll(self, run_id: str) -> WorkflowEvalResult: + return await _WorkflowEvalRunner(self._config).poll(run_id) + + async def process_submission_result( + self, run_id: str, *, submission_id: str | None = None, external_id: str | None = None + ) -> WorkflowEvalResult: + return await _WorkflowEvalRunner(self._config).process_submission_result( + run_id, submission_id=submission_id, external_id=external_id + ) + + +def define_workflow_eval( + project_name: str, + *, + store: WorkflowEvalStore, + data: EvalData[Input, Expected], + task: EvalTask[Input, Output, Expected] | WorkflowTask[Input, Output, Expected, Any], + scores: Sequence[EvalScorer[Input, Output, Expected] | WorkflowScorer[Input, Output, Expected, Any]] | None = None, + classifiers: Sequence[EvalClassifier[Input, Output, Expected]] | None = None, + case_id: Callable[[EvalCase[Input, Expected]], str | Awaitable[str]] | None = None, + experiment_name: str | None = None, + trial_count: int = 1, + max_concurrency: int = DEFAULT_MAX_CONCURRENCY, + metadata: Metadata | None = None, + tags: Sequence[str] | None = None, + is_public: bool = False, + project_id: str | None = None, + base_experiment_name: str | None = None, + base_experiment_id: str | None = None, + git_metadata_settings: GitMetadataSettings | None = None, + repo_info: RepoInfo | None = None, + description: str | None = None, + summarize_scores: bool = True, + parameters: EvalParameters | RemoteEvalParameters | None = None, + state: BraintrustState | None = None, +) -> WorkflowEval[Input, Output, Expected]: + """Define an experimental evaluation that can pause across provider submissions. + + WorkflowTask and WorkflowScorer submit one case/trial and collect one result. + Submission data and collected values must be JSON serializable. + Scorers and classifiers start once their case's task is persisted and logged. + + Each poll checks existing submissions once; newly submitted work waits for a + later invocation. Provider callbacks use max_concurrency (default 10). + Callback failures are raised after independent work advances. + Webhooks resume via process_submission_result with submission_id or external_id. + Collection callbacks must tolerate replay, including concurrent delivery. + """ + if not isinstance(trial_count, int) or isinstance(trial_count, bool) or trial_count < 1: + raise ValueError("trial_count must be a positive integer") + if not isinstance(max_concurrency, int) or isinstance(max_concurrency, bool) or max_concurrency < 1: + raise ValueError("max_concurrency must be a positive integer") + return WorkflowEval( + _WorkflowEvalConfig( + project_name=project_name, + store=store, + data=data, + task=task, + scores=list(scores or []), + classifiers=list(classifiers or []), + case_id=case_id, + experiment_name=experiment_name, + trial_count=trial_count, + max_concurrency=max_concurrency, + metadata=metadata, + tags=tags, + is_public=is_public, + project_id=project_id, + base_experiment_name=base_experiment_name, + base_experiment_id=base_experiment_id, + git_metadata_settings=git_metadata_settings, + repo_info=repo_info, + description=description, + summarize_scores=summarize_scores, + parameters=parameters, + state=state, + ) + ) + + +def _json_bytes(value: Any) -> bytes: + return bt_dumps(value).encode("utf-8") + + +def _decode(value: bytes) -> Any: + return bt_loads(value.decode("utf-8")) + + +def _json_value(value: Any) -> Any: + """Validate and normalize a value at the workflow JSON boundary.""" + return _decode(_json_bytes(value)) + + +def _stable_hex(*parts: str, length: int) -> str: + return hashlib.sha256("\0".join(parts).encode("utf-8")).hexdigest()[:length] + + +def _stable_uuid(*parts: str) -> str: + digest = hashlib.sha256("\0".join(parts).encode("utf-8")).digest() + return str(uuid.UUID(bytes=digest[:16], version=5)) + + +def _stage_kind(prefix: str, name: str) -> str: + return f"{prefix}-{_stable_hex(name, length=32)}" + + +def _scorer_args(case: Mapping[str, Any], task_result: Mapping[str, Any]) -> dict[str, Any]: + datum = case["datum"] + return { + "input": datum["input"], + "output": task_result["output"], + "expected": datum.get("expected"), + "metadata": task_result["metadata"], + "id": case["case_id"], + "tags": task_result.get("tags"), + } + + +def _score_fields(result: ScoreLike) -> dict[str, Any]: + return {key: value for key, value in result.as_dict().items() if key not in ("metadata", "name")} + + +def _raise_callback_errors(errors: Sequence[BaseException]) -> None: + if len(errors) == 1: + raise errors[0] + raise RuntimeError( + "Workflow submission callbacks failed: " + "; ".join(str(error) for error in errors) + ) from errors[0] + + +class _WorkflowEvalRunner(Generic[Input, Output, Expected]): + def __init__(self, config: _WorkflowEvalConfig[Input, Output, Expected]) -> None: + self.config = config + self.store = config.store + self._callback_semaphore = asyncio.Semaphore(config.max_concurrency) + self.eval_name = config.experiment_name or config.project_name + self.definition_key = _stable_hex(config.project_name, self.eval_name, length=32) + + def _key(self, run_id: str, kind: str, identifier: str | None = None) -> str: + suffix = f"/{identifier}" if identifier is not None else "" + return f"{_SCHEMA_PREFIX}/{self.definition_key}/{run_id}/{kind}{suffix}" + + async def _read(self, run_id: str, kind: str, identifier: str | None = None) -> Any | None: + value = await self.store.read(self._key(run_id, kind, identifier)) + return _decode(value) if value is not None else None + + async def _required(self, run_id: str, kind: str, identifier: str | None = None) -> Any: + value = await self._read(run_id, kind, identifier) + if value is None: + target = f" {identifier!r}" if identifier is not None else "" + raise ValueError(f"Unknown workflow evaluation {kind}{target} for run {run_id!r}") + return value + + async def _write(self, run_id: str, kind: str, value: Any, identifier: str | None = None) -> None: + await self.store.write(self._key(run_id, kind, identifier), _json_bytes(value)) + + async def _claim(self, run_id: str, action: str) -> bool: + result = await self.store.get_or_set(self._key(run_id, "claim", action), b"1") + return result.created + + async def _call(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + async with self._callback_semaphore: + return await await_or_run(asyncio.get_running_loop(), fn, *args, **kwargs) + + async def _flush_logs(self) -> None: + state = self.config.state or _internal_get_global_state() + background_logger = state.global_bg_logger() + await asyncio.get_running_loop().run_in_executor(None, background_logger.flush) + + def _parameters(self, run: Mapping[str, Any]) -> ValidatedParameters | None: + raw = run.get("parameters") + if self.config.parameters is None: + return cast(ValidatedParameters | None, raw) + return validate_parameters(raw or {}, self.config.parameters) + + def _experiment(self, run: Mapping[str, Any]) -> Experiment | None: + if run["no_send_logs"]: + return None + experiment_parameters = None + if isinstance(self.config.parameters, RemoteEvalParameters) and self.config.parameters.id is not None: + experiment_parameters = {"id": self.config.parameters.id} + if self.config.parameters.version is not None: + experiment_parameters["version"] = self.config.parameters.version + dataset = self.config.data if isinstance(self.config.data, Dataset) else None + return init_experiment( + project_name=self.config.project_name if self.config.project_id is None else None, + project_id=self.config.project_id, + experiment_name=run["experiment_name"], + description=self.config.description, + metadata=self.config.metadata, + tags=self.config.tags, + is_public=self.config.is_public, + update=True, + base_experiment=self.config.base_experiment_name, + base_experiment_id=self.config.base_experiment_id, + git_metadata_settings=self.config.git_metadata_settings, + repo_info=self.config.repo_info, + dataset=dataset, + parameters=experiment_parameters, + state=self.config.state, + ) + + async def _resolve_data(self, experiment: Experiment | None) -> list[EvalCase[Input, Expected]]: + data: Any = self.config.data + if inspect.isclass(data): + data = data() + if isinstance(data, BaseExperiment): + if experiment is None: + raise ValueError("Cannot use BaseExperiment without sending logs") + base_name = data.name + if base_name is None: + base = experiment.fetch_base_experiment() + if base is None: + raise ValueError("BaseExperiment failed to resolve a base experiment") + base_name = base.name + data = _init_experiment( + project=self.config.project_name if self.config.project_id is None else None, + project_id=self.config.project_id, + experiment=base_name, + open=True, + set_current=False, + state=self.config.state, + ).as_dataset() + elif callable(data) and not isinstance(data, Dataset): + data = await self._call(data) + if inspect.isawaitable(data): + data = await data + + values: list[Any] = [] + if isinstance(data, AsyncIterable): + async for value in data: + values.append(value) + else: + values.extend(data) + return [value if isinstance(value, EvalCase) else EvalCase.from_dict(value) for value in values] + + def _uses_submission_processor(self) -> bool: + return isinstance(self.config.task, WorkflowTask) or any( + isinstance(score, WorkflowScorer) for score in self.config.scores + ) + + async def _save_completed(self, run: Mapping[str, Any], summary: ExperimentSummary) -> WorkflowEvalCompletedResult: + completed = {**run, "status": "completed", "summary": _json_value(summary.as_dict())} + await self._write(run["run_id"], "run", completed) + return WorkflowEvalCompletedResult( + run_id=run["run_id"], pending=WorkflowEvalPending(poll=0, webhook=0), summary=summary + ) + + async def _run_ordinary_evaluator( + self, + run: Mapping[str, Any], + data: list[EvalCase[Input, Expected]], + experiment: Experiment | None, + parameters: ValidatedParameters, + ) -> WorkflowEvalCompletedResult: + evaluator = Evaluator( + project_name=self.config.project_name, + eval_name=self.eval_name, + data=data, + task=cast(EvalTask[Input, Output, Expected], self.config.task), + scores=cast(Sequence[EvalScorer[Input, Output, Expected]], self.config.scores), + classifiers=list(self.config.classifiers), + experiment_name=run["experiment_name"], + metadata=self.config.metadata, + tags=self.config.tags, + trial_count=self.config.trial_count, + max_concurrency=self.config.max_concurrency, + is_public=self.config.is_public, + update=True, + project_id=self.config.project_id, + base_experiment_name=self.config.base_experiment_name, + base_experiment_id=self.config.base_experiment_id, + git_metadata_settings=self.config.git_metadata_settings, + repo_info=self.config.repo_info, + description=self.config.description, + summarize_scores=self.config.summarize_scores, + parameters=self.config.parameters, + parameter_values=cast(dict[str, Any], parameters), + ) + result = await run_evaluator( + experiment, + evaluator, + position=None, + filters=[], + state=self.config.state, + ) + return await self._save_completed(run, result.summary) + + async def _persist_cases(self, run: dict[str, Any], data: Sequence[EvalCase[Input, Expected]]) -> None: + seen_case_ids: set[str] = set() + item_ids: list[str] = [] + for datum in data: + case_id = datum.id + if not case_id and self.config.case_id is not None: + case_id = await self._call(self.config.case_id, datum) + if not isinstance(case_id, str) or not case_id: + raise ValueError("Every workflow evaluation case must have a non-empty id or be assigned by case_id") + if case_id in seen_case_ids: + raise ValueError(f"Workflow evaluation case IDs must be unique; found duplicate {case_id!r}") + seen_case_ids.add(case_id) + + trial_count = datum.trial_count if datum.trial_count is not None else self.config.trial_count + if not isinstance(trial_count, int) or isinstance(trial_count, bool) or trial_count < 1: + raise ValueError(f"trial_count for case {case_id!r} must be a positive integer") + for trial_index in range(trial_count): + item_id = f"{case_id}:trial:{trial_index}" + item_ids.append(item_id) + await self._write( + run["run_id"], + "case", + { + "id": item_id, + "case_id": case_id, + "trial_index": trial_index, + "datum": _json_value( + {field.name: getattr(datum, field.name) for field in dataclasses.fields(datum)} + ), + "metadata": _json_value(dict(datum.metadata or {})), + "tags": list(datum.tags) if datum.tags is not None else None, + }, + item_id, + ) + run["case_ids"] = item_ids + await self._write(run["run_id"], "run", run) + + async def start(self, parameters: Mapping[str, Any] | None, *, no_send_logs: bool) -> WorkflowEvalResult: + validated_parameters = validate_parameters(parameters or {}, self.config.parameters) + run_id = str(uuid.uuid4()) + run = { + "run_id": run_id, + "experiment_name": self.config.experiment_name or f"{self.eval_name}-{run_id}", + "no_send_logs": no_send_logs, + "parameters": _json_value(validated_parameters), + "status": "running", + "case_ids": [], + } + created = await self.store.get_or_set(self._key(run_id, "run"), _json_bytes(run)) + if not created.created: + raise RuntimeError(f"Workflow evaluation run ID collision: {run_id}") + + experiment = self._experiment(run) + data = await self._resolve_data(experiment) + if not self._uses_submission_processor(): + return await self._run_ordinary_evaluator(run, data, experiment, validated_parameters) + + await self._persist_cases(run, data) + return await self._advance(run, experiment=experiment) + + async def status(self, run_id: str) -> WorkflowEvalResult: + run = await self._required(run_id, "run") + if run["status"] == "completed": + return WorkflowEvalCompletedResult( + run_id=run_id, + pending=WorkflowEvalPending(poll=0, webhook=0), + summary=ExperimentSummary.from_dict_deep(run["summary"]), + ) + return await self._waiting_status(run) + + async def poll(self, run_id: str) -> WorkflowEvalResult: + run = await self._required(run_id, "run") + if run["status"] == "completed": + return await self.status(run_id) + # Snapshot first: newly submitted downstream work waits for the next poll call. + submissions = await self._submission_records(run) + + async def poll_submission(submission: dict[str, Any], processor: Any) -> None: + completion = processor.completion + assert isinstance(completion, WorkflowSubmissionCompletionPoll) + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + result = await self._call(completion.poll, submission["submission_data"], context) + if isinstance(result, Mapping): + result = WorkflowSubmissionPoll(**result) + if not isinstance(result, WorkflowSubmissionPoll): + raise TypeError("Submission poll callback must return WorkflowSubmissionPoll") + if result.status not in ("pending", "complete", "failed"): + raise ValueError(f"Submission poll callback returned unsupported status {result.status!r}") + if result.status == "failed": + if isinstance(result.error, BaseException): + raise result.error + raise RuntimeError(str(result.error)) + if result.status == "complete": + await self._collect_submission(run, submission, processor) + + results = await asyncio.gather( + *( + poll_submission(submission, processor) + for submission, processor in submissions + if submission["status"] == "submitted" and submission["mode"] == "poll" + ), + return_exceptions=True, + ) + # Persist independent progress before reporting provider callback failures. + errors = [result for result in results if isinstance(result, BaseException)] + try: + status = await self._advance(run) + except Exception as error: + errors.append(error) + if errors: + _raise_callback_errors(errors) + return status + + async def process_submission_result( + self, run_id: str, *, submission_id: str | None, external_id: str | None + ) -> WorkflowEvalResult: + if not submission_id and not external_id: + raise ValueError("process_submission_result requires submission_id or external_id") + run = await self._required(run_id, "run") + external_submission_id = await self._read(run_id, "external", external_id) if external_id is not None else None + if submission_id and external_submission_id and submission_id != external_submission_id: + raise ValueError("submission_id and external_id identify different submissions") + submission_id = submission_id or external_submission_id + submission = await self._read(run_id, "submission", submission_id) if submission_id else None + if submission is None: + raise ValueError("No submission matches this result") + if external_id is not None and submission.get("external_id") != external_id: + raise ValueError("submission_id and external_id identify different submissions") + if run["status"] == "completed": + return await self.status(run_id) + processor = ( + self.config.task + if submission["kind"] == "task" + else next( + score + for score in self.config.scores + if isinstance(score, WorkflowScorer) and score.name == submission["scorer_name"] + ) + ) + await self._collect_submission(run, submission, processor) + # Only the completed case needs advancing; other cases can still be pending. + return await self._advance(run, item_ids=[submission["item_id"]]) + + def _submission_specs(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: + specs: list[tuple[dict[str, Any], Any]] = [] + for item_id in run["case_ids"]: + if isinstance(self.config.task, WorkflowTask): + specs.append( + ( + { + "id": f"submission-{_stable_hex(run['run_id'], 'task', item_id, length=32)}", + "kind": "task", + "item_id": item_id, + }, + self.config.task, + ) + ) + for scorer in self.config.scores: + if isinstance(scorer, WorkflowScorer): + specs.append( + ( + { + "id": f"submission-{_stable_hex(run['run_id'], 'score', scorer.name, item_id, length=32)}", + "kind": "score", + "scorer_name": scorer.name, + "item_id": item_id, + }, + scorer, + ) + ) + return specs + + async def _submission_records(self, run: Mapping[str, Any]) -> list[tuple[dict[str, Any], Any]]: + records: list[tuple[dict[str, Any], Any]] = [] + for spec, processor in self._submission_specs(run): + record = await self._read(run["run_id"], "submission", spec["id"]) + if record is not None: + await self._persist_submission_progress(run["run_id"], record) + records.append((record, processor)) + return records + + async def _case(self, run_id: str, item_id: str) -> dict[str, Any]: + return cast(dict[str, Any], await self._required(run_id, "case", item_id)) + + async def _task_item( + self, run: Mapping[str, Any], item_id: str, parameters: ValidatedParameters | None + ) -> WorkflowTaskItem[Any, Any]: + case = await self._case(run["run_id"], item_id) + datum = case["datum"] + return WorkflowTaskItem( + id=item_id, + input=datum["input"], + expected=datum.get("expected"), + metadata=dict(case["metadata"]), + tags=case.get("tags"), + parameters=parameters, + trial_index=case["trial_index"], + ) + + async def _scorer_item(self, run: Mapping[str, Any], item_id: str) -> WorkflowScorerItem[Any, Any, Any]: + case = await self._case(run["run_id"], item_id) + task = await self._required(run["run_id"], "task-result", item_id) + datum = case["datum"] + return WorkflowScorerItem( + id=item_id, + input=datum["input"], + output=task["output"], + expected=datum.get("expected"), + metadata=dict(task["metadata"]), + tags=task.get("tags"), + trial_index=case["trial_index"], + ) + + async def _submit_submission(self, run: Mapping[str, Any], spec: dict[str, Any], processor: Any) -> None: + run_id = run["run_id"] + existing = await self._read(run_id, "submission", spec["id"]) + if existing is not None: + await self._persist_submission_progress(run_id, existing) + return + if not await self._claim(run_id, f"submit:{spec['id']}"): + return + parameters = self._parameters(run) + if spec["kind"] == "task": + item = await self._task_item(run, spec["item_id"], parameters) + else: + item = await self._scorer_item(run, spec["item_id"]) + context = WorkflowSubmissionContext(run_id=run_id, submission_id=spec["id"]) + submission_data = await self._call(processor.submit, item, context) + submission_data = _json_value(submission_data) + completion = processor.completion + external_id = None + if isinstance(completion, WorkflowSubmissionCompletionWebhook): + external_id = await self._call(completion.get_external_id, submission_data, context) + if not isinstance(external_id, str) or not external_id.strip(): + raise ValueError("Submission webhook get_external_id must return a non-empty string") + submission = { + **spec, + "submission_data": submission_data, + "external_id": external_id, + "mode": completion.mode, + "status": "submitted", + } + await self._write(run_id, "submission", submission, spec["id"]) + await self._persist_submission_progress(run_id, submission) + + async def _persist_submission_progress(self, run_id: str, submission: Mapping[str, Any]) -> None: + if submission.get("external_id") is not None: + await self._write(run_id, "external", submission["id"], submission["external_id"]) + await self.store.add_to_set(self._key(run_id, "progress", f"{submission['mode']}/submitted"), submission["id"]) + if submission["status"] == "complete": + await self.store.add_to_set( + self._key(run_id, "progress", f"{submission['mode']}/complete"), submission["id"] + ) + + async def _collect_submission(self, run: Mapping[str, Any], submission: dict[str, Any], processor: Any) -> None: + run_id = run["run_id"] + if submission["status"] != "complete": + context = WorkflowSubmissionContext(run_id=run_id, submission_id=submission["id"]) + result = await self._call(processor.collect, submission["submission_data"], context) + result_type = WorkflowTaskResult if submission["kind"] == "task" else WorkflowScorerResult + if isinstance(result, Mapping): + result = result_type(**result) + if not isinstance(result, result_type): + raise TypeError(f"Submission collect callback must return a single {result_type.__name__}") + item_id = submission["item_id"] + if submission["kind"] == "task": + case = await self._case(run_id, item_id) + metadata = {**case["metadata"], **(result.metadata or {})} + tags = result.tags if result.tags is not None else case.get("tags") + await self._write( + run_id, + "task-result", + { + "output": _json_value(result.output), + "metadata": metadata, + "tags": tags, + }, + item_id, + ) + else: + await self._write( + run_id, _stage_kind("score-result", submission["scorer_name"]), _json_value(result.score), item_id + ) + submission["status"] = "complete" + await self._write(run_id, "submission", submission, submission["id"]) + # Replay repairs an interrupted progress update without collecting again. + await self._persist_submission_progress(run_id, submission) + + async def _waiting_status(self, run: Mapping[str, Any]) -> WorkflowEvalWaitingResult: + pending = {} + for mode in ("poll", "webhook"): + completed = await self.store.get_set_size(self._key(run["run_id"], "progress", f"{mode}/complete")) + submitted = await self.store.get_set_size(self._key(run["run_id"], "progress", f"{mode}/submitted")) + pending[mode] = submitted - completed + return WorkflowEvalWaitingResult( + run_id=run["run_id"], pending=WorkflowEvalPending(poll=pending["poll"], webhook=pending["webhook"]) + ) + + def _span_ids(self, run_id: str, item_id: str, stage: str) -> tuple[str, str, str]: + if BraintrustEnv.LEGACY_IDS: + root_span_id = _stable_uuid(run_id, item_id, "root") + span_id = root_span_id if stage == "root" else _stable_uuid(run_id, item_id, stage) + return _stable_uuid(run_id, item_id, f"row:{stage}"), span_id, root_span_id + root_span_id = _stable_hex(run_id, item_id, "root", length=32) + span_id = root_span_id[:16] if stage == "root" else _stable_hex(run_id, item_id, stage, length=16) + row_id = _stable_uuid(run_id, item_id, stage) + return row_id, span_id, root_span_id + + def _start_root(self, experiment: Experiment | None, run: Mapping[str, Any], case: Mapping[str, Any]) -> Span: + if experiment is None: + return NOOP_SPAN + datum = case["datum"] + event_dataset = experiment.dataset or (self.config.data if isinstance(self.config.data, Dataset) else None) + if ( + event_dataset is not None + and isinstance(datum.get("id"), str) + and datum["id"] + and isinstance(datum.get("_xact_id"), str) + and datum["_xact_id"] + ): + origin = { + "object_type": "dataset", + "object_id": event_dataset.id, + "id": datum["id"], + "_xact_id": datum["_xact_id"], + **({"created": datum["created"]} if isinstance(datum.get("created"), str) else {}), + } + else: + origin = _validated_object_reference(datum.get("origin")) + row_id, span_id, root_span_id = self._span_ids(run["run_id"], case["id"], "root") + return _internal_start_span_with_initial_merge( + "eval", + parent=experiment.export(), + span_id=span_id, + root_span_id=root_span_id, + state=self.config.state, + type=SpanTypeAttribute.EVAL, + id=row_id, + input=datum["input"], + expected=datum.get("expected"), + metadata={ + **case["metadata"], + "workflow_eval": { + "run_id": run["run_id"], + "case_id": case["case_id"], + "trial_index": case["trial_index"], + }, + }, + tags=case.get("tags"), + **({"origin": origin} if origin is not None else {}), + ) + + def _start_child( + self, + root: Span, + run_id: str, + item_id: str, + stage: str, + name: str, + span_type: SpanTypeAttribute, + **event: Any, + ) -> Span: + row_id, span_id, _ = self._span_ids(run_id, item_id, stage) + span_attributes: dict[str, Any] = {"type": span_type} + if span_type != SpanTypeAttribute.TASK: + span_attributes["purpose"] = "scorer" + return root.start_span( + name, + span_attributes=span_attributes, + internal={"initial_span_write_as_merge": True, "span_id": span_id}, + id=row_id, + **event, + ) + + async def _run_ordinary_task( + self, + run: Mapping[str, Any], + case: dict[str, Any], + experiment: Experiment | None, + parameters: ValidatedParameters | None, + ) -> None: + run_id = run["run_id"] + item_id = case["id"] + if await self._read(run_id, "task-result", item_id) is not None: + await self._log_task(run, case, experiment) + return + if not await self._claim(run_id, f"task:{item_id}"): + return + datum = case["datum"] + metadata = dict(case["metadata"]) + hooks = DictEvalHooks( + metadata, + expected=datum.get("expected"), + trial_index=case["trial_index"], + tags=case.get("tags"), + parameters=parameters, + ) + root = self._start_root(experiment, run, case) + task = cast(Callable[..., Any], self.config.task) + task_args: list[Any] = [datum["input"]] + try: + if len(get_signature(task).parameters) == 2: + task_args.append(hooks) + except Exception: + pass + with root: + with self._start_child( + root, + run_id, + item_id, + "task", + "task", + SpanTypeAttribute.TASK, + input=datum["input"], + ) as span: + hooks.set_span(span) + output = await self._call(task, *task_args) + span.log(output=output) + tags = list(hooks.tags) if hooks.tags else None + task_result = { + "output": _json_value(output), + "metadata": _json_value(metadata), + "tags": tags, + } + await self._write(run_id, "task-result", task_result, item_id) + root.log(output=output, metadata=metadata, tags=tags) + if root is not NOOP_SPAN: + await self._flush_logs() + await self._write( + run_id, + "task-log", + {"root_span": root.export() if experiment is not None else None}, + item_id, + ) + + async def _log_task(self, run: Mapping[str, Any], case: dict[str, Any], experiment: Experiment | None) -> None: + run_id = run["run_id"] + item_id = case["id"] + if await self._read(run_id, "task-log", item_id) is not None: + return + task_result = await self._required(run_id, "task-result", item_id) + root = self._start_root(experiment, run, case) + with root: + with self._start_child( + root, + run_id, + item_id, + "task", + "task", + SpanTypeAttribute.TASK, + input=case["datum"]["input"], + ) as span: + span.log(output=task_result["output"]) + root.log(output=task_result["output"], metadata=task_result["metadata"], tags=task_result.get("tags")) + if root is not NOOP_SPAN: + await self._flush_logs() + await self._write( + run_id, + "task-log", + {"root_span": root.export() if experiment is not None else None}, + item_id, + ) + + async def _root_for_case(self, run: Mapping[str, Any], item_id: str) -> Span: + task_log = await self._required(run["run_id"], "task-log", item_id) + exported = task_log.get("root_span") + return _internal_resume_span(exported, self.config.state) if exported else NOOP_SPAN + + async def _trace_for_case(self, run: Mapping[str, Any], item_id: str) -> LocalTrace | None: + task_log = await self._required(run["run_id"], "task-log", item_id) + exported = task_log.get("root_span") + if not exported: + return None + components = SpanComponentsV4.from_str(exported) + if not components.root_span_id: + raise ValueError("Persisted workflow evaluation root span is missing its root span ID") + trace_state = self.config.state or _internal_get_global_state() + + async def ensure_spans_flushed() -> None: + await asyncio.get_running_loop().run_in_executor(None, trace_state.flush) + await trace_state.flush_otel() + + return LocalTrace( + object_type=span_object_type_v3_to_typed_string(components.object_type), + object_id=span_components_to_object_id(components), + root_span_id=components.root_span_id, + ensure_spans_flushed=ensure_spans_flushed, + state=trace_state, + ) + + def _prepare_scores(self, raw: Any, name: str) -> list[ScoreLike]: + if isinstance(raw, dict): + raw = _normalize_score(raw, "When returning a dict, it must be a valid Score object.") + if isinstance(raw, Iterable) and not isinstance(raw, (str, bytes, Mapping)): + return [ + _normalize_score(value, "When returning an array of scores, each score must be a valid Score object.") + for value in raw + ] + if is_score(raw): + return [raw] + return [Score(name=name, score=raw)] + + async def _run_ordinary_score(self, run: Mapping[str, Any], case: dict[str, Any], scorer: Any, name: str) -> None: + run_id = run["run_id"] + item_id = case["id"] + result_kind = _stage_kind("score-result", name) + if await self._read(run_id, result_kind, item_id) is None: + if not await self._claim(run_id, f"score:{name}:{item_id}"): + return + task_result = await self._required(run_id, "task-result", item_id) + fn = scorer.eval_async if hasattr(scorer, "eval_async") else scorer + trace = await self._trace_for_case(run, item_id) + raw = await call_user_fn( + asyncio.get_running_loop(), + fn, + **_scorer_args(case, task_result), + trace=trace, + ) + await self._write(run_id, result_kind, _json_value(raw), item_id) + await self._log_score(run, case, name) + + async def _log_score(self, run: Mapping[str, Any], case: dict[str, Any], name: str) -> None: + run_id = run["run_id"] + item_id = case["id"] + log_kind = _stage_kind("score-log", name) + if await self._read(run_id, log_kind, item_id) is not None: + return + raw = await self._required(run_id, _stage_kind("score-result", name), item_id) + results = self._prepare_scores(raw, name) + task_result = await self._required(run_id, "task-result", item_id) + root = await self._root_for_case(run, item_id) + propagated = merge_dicts({**(root.propagated_event or {})}, {"span_attributes": {"purpose": "scorer"}}) + with root: + with self._start_child( + root, + run_id, + item_id, + f"score:{name}", + name, + SpanTypeAttribute.SCORE, + input=_scorer_args(case, task_result), + propagated_event=propagated, + ) as span: + output = ( + {result.name: _score_fields(result) for result in results} + if len(results) != 1 + else _score_fields(results[0]) + ) + scores = {result.name: result.score for result in results} + span.log(output=output, metadata=_build_span_metadata(results), scores=scores) + root.log(scores=scores) + if root is not NOOP_SPAN: + await self._flush_logs() + await self._write(run_id, log_kind, True, item_id) + + async def _run_classifier(self, run: Mapping[str, Any], case: dict[str, Any], classifier: Any, name: str) -> None: + run_id = run["run_id"] + item_id = case["id"] + result_kind = _stage_kind("classification-result", name) + if await self._read(run_id, result_kind, item_id) is None: + if not await self._claim(run_id, f"classification:{name}:{item_id}"): + return + task_result = await self._required(run_id, "task-result", item_id) + trace = await self._trace_for_case(run, item_id) + raw = await call_user_fn( + asyncio.get_running_loop(), + classifier, + **_scorer_args(case, task_result), + trace=trace, + ) + if raw is None: + values: list[Any] = [] + elif isinstance(raw, Iterable) and not isinstance(raw, (str, bytes, Mapping)): + values = list(raw) + else: + values = [raw] + classifications = [_validate_classification_result(value, name) for value in values] + await self._write(run_id, result_kind, [value.as_dict() for value in classifications], item_id) + await self._log_classifier(run, case, name) + + async def _log_classifier(self, run: Mapping[str, Any], case: dict[str, Any], name: str) -> None: + run_id = run["run_id"] + item_id = case["id"] + log_kind = _stage_kind("classification-log", name) + if await self._read(run_id, log_kind, item_id) is not None: + return + raw = await self._required(run_id, _stage_kind("classification-result", name), item_id) + classifications = [Classification.from_dict(value) for value in raw] + task_result = await self._required(run_id, "task-result", item_id) + root = await self._root_for_case(run, item_id) + with root: + with self._start_child( + root, + run_id, + item_id, + f"classification:{name}", + name, + SpanTypeAttribute.CLASSIFIER, + input=_scorer_args(case, task_result), + ) as span: + if classifications: + span.log( + output=_build_classification_span_output(classifications), + metadata=_build_span_metadata(classifications), + ) + grouped: dict[str, list[Any]] = {} + for result in classifications: + grouped.setdefault(cast(str, result.name), []).append(result.as_item()) + root.log(classifications=grouped) + else: + span.log(output={}, metadata=None) + if root is not NOOP_SPAN: + await self._flush_logs() + await self._write(run_id, log_kind, True, item_id) + + async def _advance_case( + self, + run: Mapping[str, Any], + item_id: str, + experiment: Experiment | None, + scorers: Sequence[Any], + scorer_names: Sequence[str], + classifier_names: Sequence[str], + ) -> bool: + run_id = run["run_id"] + case = await self._case(run_id, item_id) + case_run = {**run, "case_ids": [item_id]} + if isinstance(self.config.task, WorkflowTask): + for spec, processor in self._submission_specs(case_run): + if spec["kind"] == "task": + await self._submit_submission(run, spec, processor) + if await self._read(run_id, "task-result", item_id) is None: + return False + await self._log_task(run, case, experiment) + else: + await self._run_ordinary_task(run, case, experiment, self._parameters(run)) + if await self._read(run_id, "task-log", item_id) is None: + return False + + for spec, processor in self._submission_specs(case_run): + if spec["kind"] == "score": + await self._submit_submission(run, spec, processor) + complete = True + for scorer, name in zip(scorers, scorer_names): + if isinstance(scorer, WorkflowScorer): + if await self._read(run_id, _stage_kind("score-result", name), item_id) is None: + complete = False + continue + await self._log_score(run, case, name) + else: + await self._run_ordinary_score(run, case, scorer, name) + if await self._read(run_id, _stage_kind("score-log", name), item_id) is None: + complete = False + for classifier, name in zip(self.config.classifiers, classifier_names): + await self._run_classifier(run, case, classifier, name) + if await self._read(run_id, _stage_kind("classification-log", name), item_id) is None: + complete = False + if complete: + await self.store.add_to_set(self._key(run_id, "progress", "cases"), item_id) + return complete + + async def _advance( + self, run: Mapping[str, Any], *, experiment: Experiment | None = None, item_ids: Sequence[str] | None = None + ) -> WorkflowEvalResult: + run = await self._required(run["run_id"], "run") + run_id = run["run_id"] + if run["status"] == "completed": + return await self.status(run_id) + if experiment is None and not run["no_send_logs"]: + experiment = self._experiment(run) + resolved_scores = [ + score() if inspect.isclass(score) and is_scorer(score) else score for score in self.config.scores + ] + scorer_names = [ + score.name if isinstance(score, WorkflowScorer) else _scorer_name(score, index) + for index, score in enumerate(resolved_scores) + ] + if len(scorer_names) != len(set(scorer_names)): + raise ValueError("Workflow evaluation scorer names must be unique") + classifier_names = [ + _classifier_name(classifier, index) for index, classifier in enumerate(self.config.classifiers) + ] + if len(classifier_names) != len(set(classifier_names)): + raise ValueError("Workflow evaluation classifier names must be unique") + + results = await asyncio.gather( + *( + self._advance_case(run, item_id, experiment, resolved_scores, scorer_names, classifier_names) + for item_id in (item_ids if item_ids is not None else run["case_ids"]) + ), + return_exceptions=True, + ) + errors = [result for result in results if isinstance(result, BaseException)] + if errors: + _raise_callback_errors(errors) + if await self.store.get_set_size(self._key(run_id, "progress", "cases")) != len(run["case_ids"]): + return await self._waiting_status(run) + if not await self._claim(run_id, "finalize"): + return await self.status(run_id) + summary = await self._summary(run, scorer_names, classifier_names, experiment) + return await self._save_completed(run, summary) + + async def _summary( + self, + run: Mapping[str, Any], + scorer_names: Sequence[str], + classifier_names: Sequence[str], + experiment: Experiment | None, + ) -> ExperimentSummary: + if experiment is not None: + comparison_experiment_id = self.config.base_experiment_id + if comparison_experiment_id is None: + comparison_experiment_id = _get_persisted_base_experiment_id(experiment) + return experiment.summarize( + summarize_scores=self.config.summarize_scores, + comparison_experiment_id=comparison_experiment_id, + ) + + results: list[EvalResult[Any, Any, Any]] = [] + for item_id in run["case_ids"]: + case = await self._case(run["run_id"], item_id) + task_result = await self._required(run["run_id"], "task-result", item_id) + scores: dict[str, float | None] = {} + for name in scorer_names: + raw = await self._required(run["run_id"], _stage_kind("score-result", name), item_id) + for score in self._prepare_scores(raw, name): + scores[score.name] = score.score + classifications: dict[str, list[Any]] = {} + for name in classifier_names: + raw = await self._required(run["run_id"], _stage_kind("classification-result", name), item_id) + for value in raw: + classification = Classification.from_dict(value) + classifications.setdefault(cast(str, classification.name), []).append(classification.as_item()) + datum = case["datum"] + results.append( + EvalResult( + input=datum["input"], + output=task_result["output"], + scores=scores, + classifications=classifications or None, + expected=datum.get("expected"), + metadata=task_result["metadata"], + tags=task_result.get("tags"), + ) + ) + evaluator = Evaluator( + project_name=self.config.project_name, + eval_name=self.eval_name, + data=[], + task=cast(Any, lambda value: value), + scores=[], + experiment_name=run["experiment_name"], + metadata=self.config.metadata, + ) + return build_local_summary(evaluator, cast(Any, results)) + + +__all__ = [ + "WorkflowSubmissionCompletionPoll", + "WorkflowSubmissionCompletionWebhook", + "WorkflowSubmissionContext", + "WorkflowSubmissionPoll", + "WorkflowScorer", + "WorkflowScorerItem", + "WorkflowScorerResult", + "WorkflowTask", + "WorkflowTaskItem", + "WorkflowTaskResult", + "WorkflowEval", + "WorkflowEvalCompletedResult", + "WorkflowEvalMemoryStore", + "WorkflowEvalPending", + "WorkflowEvalRedisStore", + "WorkflowEvalResult", + "WorkflowEvalStore", + "WorkflowEvalStoreEntry", + "WorkflowEvalWaitingResult", + "define_workflow_eval", +]