From 31e8a8c1de809fd97337c4e17b66f9b029480e42 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 5 Aug 2026 11:13:50 -0400 Subject: [PATCH] feat(functions): add async function invocation Add invoke_async() with the same configuration and typed stream overloads as invoke(). Non-streaming calls return the decoded result without blocking the event loop; streaming calls return BraintrustStream with async iteration and final_value_async() support. Bridge the existing synchronous HTTP and SSE transports through worker threads, preserving pending stream reads across caller cancellation. Cover the API with VCR-backed invocation, stream, cancellation, and type tests. --- ...voke_async_encodes_body_as_utf8_bytes.yaml | 29 +++++ py/src/braintrust/functions/invoke.py | 117 ++++++++++++++++++ py/src/braintrust/functions/stream.py | 55 +++++++- py/src/braintrust/functions/test_invoke.py | 25 ++++ py/src/braintrust/functions/test_stream.py | 69 +++++++++++ py/src/braintrust/type_tests/test_invoke.py | 12 ++ 6 files changed, 302 insertions(+), 5 deletions(-) create mode 100644 py/src/braintrust/functions/cassettes/test_invoke_async_encodes_body_as_utf8_bytes.yaml create mode 100644 py/src/braintrust/functions/test_stream.py create mode 100644 py/src/braintrust/type_tests/test_invoke.py diff --git a/py/src/braintrust/functions/cassettes/test_invoke_async_encodes_body_as_utf8_bytes.yaml b/py/src/braintrust/functions/cassettes/test_invoke_async_encodes_body_as_utf8_bytes.yaml new file mode 100644 index 000000000..64123a0a9 --- /dev/null +++ b/py/src/braintrust/functions/cassettes/test_invoke_async_encodes_body_as_utf8_bytes.yaml @@ -0,0 +1,29 @@ +interactions: +- request: + body: '{"api_version": 1, "input": {"text": "result \u2014 excellent"}, "metadata": + null, "parent": "", "project_name": "test-project", "slug": "test-fn", "stream": + false, "tags": null}' + headers: + Accept: + - application/json + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Type: + - application/json + User-Agent: + - python-requests/2.32.5 + method: POST + uri: https://proxy.braintrust.ai/function/invoke + response: + body: + string: '{"output": "result \u2014 excellent"}' + headers: + Content-Type: + - application/json; charset=utf-8 + status: + code: 200 + message: OK +version: 1 + diff --git a/py/src/braintrust/functions/invoke.py b/py/src/braintrust/functions/invoke.py index ded7ed46a..08bf65594 100644 --- a/py/src/braintrust/functions/invoke.py +++ b/py/src/braintrust/functions/invoke.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import Sequence from typing import Any, Literal, TypedDict, TypeVar, overload @@ -230,6 +231,122 @@ def invoke( return resp.json() +@overload +async def invoke_async( + # the permutations of arguments for a function id + function_id: str | None = None, + version: str | None = None, + prompt_session_id: str | None = None, + prompt_session_function_id: str | None = None, + project_name: str | None = None, + project_id: str | None = None, + slug: str | None = None, + global_function: str | None = None, + function_type: FunctionTypeEnum | None = None, + # arguments to the function + input: Any = None, + messages: Sequence[Any] | None = None, + metadata: Metadata | None = None, + tags: Sequence[str] | None = None, + parent: Exportable | str | None = None, + stream: Literal[False] | None = None, + mode: ModeType | None = None, + strict: bool | None = None, + overrides: dict[str, Any] | None = None, + org_name: str | None = None, + api_key: str | None = None, + app_url: str | None = None, + force_login: bool = False, +) -> T: ... + + +@overload +async def invoke_async( + # the permutations of arguments for a function id + function_id: str | None = None, + version: str | None = None, + prompt_session_id: str | None = None, + prompt_session_function_id: str | None = None, + project_name: str | None = None, + project_id: str | None = None, + slug: str | None = None, + global_function: str | None = None, + function_type: FunctionTypeEnum | None = None, + # arguments to the function + input: Any = None, + messages: Sequence[Any] | None = None, + metadata: Metadata | None = None, + tags: Sequence[str] | None = None, + parent: Exportable | str | None = None, + stream: Literal[True] = True, + mode: ModeType | None = None, + strict: bool | None = None, + overrides: dict[str, Any] | None = None, + org_name: str | None = None, + api_key: str | None = None, + app_url: str | None = None, + force_login: bool = False, +) -> BraintrustStream: ... + + +async def invoke_async( + # the permutations of arguments for a function id + function_id: str | None = None, + version: str | None = None, + prompt_session_id: str | None = None, + prompt_session_function_id: str | None = None, + project_name: str | None = None, + project_id: str | None = None, + slug: str | None = None, + global_function: str | None = None, + function_type: FunctionTypeEnum | None = None, + # arguments to the function + input: Any = None, + messages: Sequence[Any] | None = None, + metadata: Metadata | None = None, + tags: Sequence[str] | None = None, + parent: Exportable | str | None = None, + stream: bool = False, + mode: ModeType | None = None, + strict: bool | None = None, + overrides: dict[str, Any] | None = None, + org_name: str | None = None, + api_key: str | None = None, + app_url: str | None = None, + force_login: bool = False, +) -> BraintrustStream | T: + """Asynchronously invoke a Braintrust function without blocking the event loop. + + When ``stream=True``, the returned ``BraintrustStream`` supports ``async for`` + and ``final_value_async()``. + """ + return await asyncio.to_thread( + invoke, + function_id=function_id, + version=version, + prompt_session_id=prompt_session_id, + prompt_session_function_id=prompt_session_function_id, + project_name=project_name, + project_id=project_id, + slug=slug, + global_function=global_function, + function_type=function_type, + input=input, + messages=messages, + metadata=metadata, + tags=tags, + parent=parent, + stream=stream, + mode=mode, + strict=strict, + overrides=overrides, + org_name=org_name, + api_key=api_key, + app_url=app_url, + force_login=force_login, + ) + + def init_function(project_name: str, slug: str, version: str | None = None): """ Creates a function that can be used as either a task or scorer in the Eval framework. diff --git a/py/src/braintrust/functions/stream.py b/py/src/braintrust/functions/stream.py index d3a6b84d2..36ca86540 100644 --- a/py/src/braintrust/functions/stream.py +++ b/py/src/braintrust/functions/stream.py @@ -5,11 +5,12 @@ with utility methods to make them easy to log and convert into various formats. """ +import asyncio import dataclasses import json -from collections.abc import Generator, Iterable +from collections.abc import AsyncIterator, Generator, Iterable, Iterator from itertools import tee -from typing import Literal +from typing import Literal, cast from sseclient import SSEClient @@ -82,12 +83,21 @@ class BraintrustInvokeError(ValueError): BraintrustStreamChunk = ( BraintrustTextChunk | BraintrustJsonChunk | BraintrustErrorChunk | BraintrustConsoleChunk | BraintrustProgressChunk ) +_STREAM_END = object() +_UNSET = object() + + +def _next_or_end(stream: Iterator[BraintrustStreamChunk]) -> BraintrustStreamChunk | object: + try: + return next(stream) + except StopIteration: + return _STREAM_END class BraintrustStream: """ A Braintrust stream. This is a wrapper around a generator of `BraintrustStreamChunk`, - with utility methods to make them easy to log and convert into various formats. + with synchronous and asynchronous iteration utilities for logging and conversion. """ def __init__(self, base_stream: SSEClient | list[BraintrustStreamChunk]): @@ -101,7 +111,9 @@ def __init__(self, base_stream: SSEClient | list[BraintrustStreamChunk]): self.stream: Iterable[BraintrustStreamChunk] = self._parse_sse_stream(base_stream) else: self.stream = base_stream - self._memoized_final_value = None + self._async_iterator: Iterator[BraintrustStreamChunk] | None = None + self._pending_async_read: asyncio.Task[BraintrustStreamChunk | object] | None = None + self._memoized_final_value = _UNSET def _parse_sse_stream(self, sse_client: SSEClient) -> Generator[BraintrustStreamChunk, None, None]: """ @@ -149,6 +161,8 @@ def copy(self): """ current_stream = self.stream self.stream, new_stream = tee(current_stream) + if self._async_iterator is not None: + self._async_iterator = iter(self.stream) return BraintrustStream(new_stream) def final_value(self): @@ -163,10 +177,14 @@ def final_value(self): Returns: The final value of the stream. """ - if self._memoized_final_value is None: + if self._memoized_final_value is _UNSET: self._memoized_final_value = parse_stream(self) return self._memoized_final_value + async def final_value_async(self): + """Consume the stream asynchronously and return its combined final value.""" + return await asyncio.to_thread(self.final_value) + def __iter__(self): """ Iterate over the stream chunks. @@ -176,6 +194,33 @@ def __iter__(self): """ yield from self.stream + def __aiter__(self) -> AsyncIterator[BraintrustStreamChunk]: + return self + + async def __anext__(self) -> BraintrustStreamChunk: + if self._async_iterator is None: + self._async_iterator = iter(self.stream) + self.stream = self._async_iterator + if self._pending_async_read is None: + self._pending_async_read = asyncio.create_task(asyncio.to_thread(_next_or_end, self._async_iterator)) + + pending_read = self._pending_async_read + try: + chunk = await asyncio.shield(pending_read) + except asyncio.CancelledError: + if pending_read.cancelled(): + self._pending_async_read = None + raise + except BaseException: + self._pending_async_read = None + raise + else: + self._pending_async_read = None + + if chunk is _STREAM_END: + raise StopAsyncIteration + return cast(BraintrustStreamChunk, chunk) + def parse_stream(stream: BraintrustStream): """ diff --git a/py/src/braintrust/functions/test_invoke.py b/py/src/braintrust/functions/test_invoke.py index 650f5c819..09ba6304f 100644 --- a/py/src/braintrust/functions/test_invoke.py +++ b/py/src/braintrust/functions/test_invoke.py @@ -1,9 +1,11 @@ """Tests for the invoke module, particularly init_function.""" +import inspect import json from unittest.mock import MagicMock, patch import pytest +from braintrust import invoke_async from braintrust.functions.invoke import init_function, invoke from braintrust.logger import TEST_API_KEY, _internal_get_global_state, _internal_reset_global_state @@ -88,6 +90,10 @@ def _invoke_with_messages(messages): return json.loads(data.decode("utf-8")) +def test_invoke_async_signature_matches_invoke(): + assert inspect.signature(invoke_async).parameters == inspect.signature(invoke).parameters + + def test_invoke_serializes_openai_messages(): openai_chat = pytest.importorskip("openai.types.chat") msg = openai_chat.ChatCompletionMessage(role="assistant", content="The answer is X.") @@ -146,3 +152,22 @@ def test_invoke_encodes_body_as_utf8_bytes(monkeypatch): api_key=TEST_API_KEY, ) assert result["output"] == f"result {em_dash} excellent" + + +@pytest.mark.asyncio +@pytest.mark.vcr +async def test_invoke_async_encodes_body_as_utf8_bytes(monkeypatch): + """The async API returns the same decoded response as invoke().""" + monkeypatch.delenv("BRAINTRUST_PROXY_URL", raising=False) + monkeypatch.delenv("BRAINTRUST_API_URL", raising=False) + _internal_reset_global_state() + + em_dash = "\u2014" + result = await invoke_async( + project_name="test-project", + slug="test-fn", + input={"text": f"result {em_dash} excellent"}, + parent="", + api_key=TEST_API_KEY, + ) + assert result["output"] == f"result {em_dash} excellent" diff --git a/py/src/braintrust/functions/test_stream.py b/py/src/braintrust/functions/test_stream.py new file mode 100644 index 000000000..a23cfbc03 --- /dev/null +++ b/py/src/braintrust/functions/test_stream.py @@ -0,0 +1,69 @@ +"""Tests for synchronous and asynchronous Braintrust streams.""" + +import asyncio +import queue +import threading + +import pytest +from braintrust.functions.stream import BraintrustStream, BraintrustTextChunk + + +def test_braintrust_stream_supports_sync_for(): + stream = BraintrustStream( + [ + BraintrustTextChunk(data="Hello"), + BraintrustTextChunk(data=" world"), + ] + ) + + assert [chunk.data for chunk in stream] == ["Hello", " world"] + + +@pytest.mark.asyncio +async def test_braintrust_stream_supports_async_for(): + stream = BraintrustStream( + [ + BraintrustTextChunk(data="Hello"), + BraintrustTextChunk(data=" world"), + ] + ) + + chunks = [chunk async for chunk in stream] + + assert [chunk.data for chunk in chunks] == ["Hello", " world"] + + +@pytest.mark.asyncio +async def test_braintrust_stream_cancellation_preserves_pending_chunk(): + chunk_queue: queue.Queue[BraintrustTextChunk] = queue.Queue() + read_started = threading.Event() + first_chunk = BraintrustTextChunk(data="first") + + def chunks(): + read_started.set() + yield chunk_queue.get(timeout=2) + yield BraintrustTextChunk(data="second") + + stream = BraintrustStream(chunks()) # type: ignore[arg-type] + pending_read = asyncio.create_task(anext(stream)) + assert await asyncio.wait_for(asyncio.to_thread(read_started.wait), timeout=1) + + pending_read.cancel() + with pytest.raises(asyncio.CancelledError): + await pending_read + + chunk_queue.put(first_chunk) + assert await anext(stream) is first_chunk + + +@pytest.mark.asyncio +async def test_braintrust_stream_final_value_async(): + stream = BraintrustStream( + [ + BraintrustTextChunk(data="Hello"), + BraintrustTextChunk(data=" world"), + ] + ) + + assert await stream.final_value_async() == "Hello world" + assert await stream.final_value_async() == "Hello world" diff --git a/py/src/braintrust/type_tests/test_invoke.py b/py/src/braintrust/type_tests/test_invoke.py new file mode 100644 index 000000000..9af7f6df8 --- /dev/null +++ b/py/src/braintrust/type_tests/test_invoke.py @@ -0,0 +1,12 @@ +"""Type-check tests for asynchronous function invocation.""" + +from braintrust import BraintrustStream, invoke_async + + +async def _check_invoke_async_return_types() -> None: + output: dict[str, str] = await invoke_async(project_name="project", slug="function", input={}) + stream: BraintrustStream = await invoke_async(project_name="project", slug="function", input={}, stream=True) + + assert output is not None + async for chunk in stream: + assert chunk is not None