diff --git a/py/src/braintrust/functions/invoke.py b/py/src/braintrust/functions/invoke.py index 5c566c3f..f0b1c3c0 100644 --- a/py/src/braintrust/functions/invoke.py +++ b/py/src/braintrust/functions/invoke.py @@ -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 @@ -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) diff --git a/py/src/braintrust/functions/test_invoke.py b/py/src/braintrust/functions/test_invoke.py index c38e2e10..264217e2 100644 --- a/py/src/braintrust/functions/test_invoke.py +++ b/py/src/braintrust/functions/test_invoke.py @@ -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 @@ -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