diff --git a/py/src/braintrust/devserver/server.py b/py/src/braintrust/devserver/server.py index a141fc79..90ceb2e8 100644 --- a/py/src/braintrust/devserver/server.py +++ b/py/src/braintrust/devserver/server.py @@ -370,6 +370,8 @@ def scorer_fn(input, output, expected, metadata): "strict": True, } headers = {"Accept": "application/json"} + if state.org_name: + headers["x-bt-org-name"] = state.org_name if project_id: headers["x-bt-project-id"] = project_id result = state.proxy_conn().post("function/invoke", json=request, headers=headers) diff --git a/py/src/braintrust/devserver/test_server.py b/py/src/braintrust/devserver/test_server.py new file mode 100644 index 00000000..b983e47d --- /dev/null +++ b/py/src/braintrust/devserver/test_server.py @@ -0,0 +1,30 @@ +from unittest.mock import MagicMock + +import pytest +from braintrust.test_helpers import has_devserver_installed + + +def test_ui_scorer_forwards_org_name_to_function_invoke(): + if not has_devserver_installed(): + pytest.skip("Devserver dependencies not installed (requires .[cli])") + + from braintrust.devserver.server import make_scorer + + response = MagicMock() + response.json.return_value = {"score": 1} + connection = MagicMock() + connection.post.return_value = response + state = MagicMock() + state.org_name = "test-org" + state.proxy_conn.return_value = connection + state.current_span.get.return_value.export.return_value = {} + + scorer = make_scorer(state, "ui-scorer", {"function_id": "function-id"}, "project-id") + + assert scorer("input", "output", "expected", {}) == {"score": 1} + assert connection.post.call_args.args == ("function/invoke",) + assert connection.post.call_args.kwargs["headers"] == { + "Accept": "application/json", + "x-bt-org-name": "test-org", + "x-bt-project-id": "project-id", + }