feat(functions): add async function invocation

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
+

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.

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):
     """

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"

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"

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