Skip to content
4 changes: 3 additions & 1 deletion py/src/braintrust/functions/invoke.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from sseclient import SSEClient

from .._generated_types import FunctionTypeEnum
from ..bt_json import bt_dumps
from ..logger import Exportable, _internal_get_global_state, get_span_parent_object, login, proxy_conn
from ..util import response_raise_for_status
from .constants import INVOKE_API_VERSION
Expand Down Expand Up @@ -201,7 +202,8 @@ def invoke(
if org_name is not None:
headers["x-bt-org-name"] = org_name

resp = proxy_conn().post("function/invoke", json=request, headers=headers, stream=stream)
request_json = bt_dumps(request)
resp = proxy_conn().post("function/invoke", data=request_json, headers=headers, stream=stream)
if resp.status_code == 500:
raise BraintrustInvokeError(resp.text)

Expand Down
57 changes: 56 additions & 1 deletion py/src/braintrust/functions/test_invoke.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
"""Tests for the invoke module, particularly init_function."""

import json
from unittest.mock import MagicMock, patch

from braintrust.functions.invoke import init_function
import pytest
from braintrust.functions.invoke import init_function, invoke
from braintrust.logger import _internal_get_global_state, _internal_reset_global_state


Expand Down Expand Up @@ -59,3 +62,55 @@ def test_init_function_permanently_disables_cache(self):
# Try to start again - should still be disabled because of explicit disable
state.span_cache.start()
assert state.span_cache.disabled is True


def _invoke_with_messages(messages):
"""Call invoke() with mocked proxy_conn; return the parsed request body."""
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = {}
mock_conn = MagicMock()
mock_conn.post.return_value = mock_resp

with (
patch("braintrust.functions.invoke.login"),
patch("braintrust.functions.invoke.get_span_parent_object") as mock_parent,
patch("braintrust.functions.invoke.proxy_conn", return_value=mock_conn),
):
mock_parent.return_value.export.return_value = "span-export"
invoke(project_name="test-project", slug="test-fn", messages=messages)

kwargs = mock_conn.post.call_args.kwargs
assert "data" in kwargs, "invoke must use data= (bt_dumps) not json= (json.dumps) (see issue 38)"
assert "json" not in kwargs
return json.loads(kwargs["data"])


def test_invoke_serializes_openai_messages():
openai_chat = pytest.importorskip("openai.types.chat")
msg = openai_chat.ChatCompletionMessage(role="assistant", content="The answer is X.")
parsed = _invoke_with_messages([msg])
assert isinstance(parsed, dict) and parsed


def test_invoke_serializes_anthropic_messages():
anthropic_types = pytest.importorskip("anthropic.types")
msg = anthropic_types.Message(
id="msg_123",
type="message",
role="assistant",
content=[anthropic_types.TextBlock(type="text", text="The answer is X.")],
model="claude-3-5-sonnet-20241022",
stop_reason="end_turn",
stop_sequence=None,
usage=anthropic_types.Usage(input_tokens=10, output_tokens=20),
)
parsed = _invoke_with_messages([msg])
assert isinstance(parsed, dict) and parsed


def test_invoke_serializes_google_messages():
google_types = pytest.importorskip("google.genai.types")
msg = google_types.Content(role="model", parts=[google_types.Part(text="The answer is X.")])
parsed = _invoke_with_messages([msg])
assert isinstance(parsed, dict) and parsed
Loading