diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 10ccd3cc3..05c5fe12a 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -1429,6 +1429,22 @@ def _internal_get_global_state() -> BraintrustState: return _state +def _resolve_state( + state: "BraintrustState | None", api_key: str | None, app_url: str | None +) -> "BraintrustState": + """Resolve the state to use for a logger/experiment/dataset. + + If an explicit state is provided, use it. Otherwise, if api_key or app_url + is specified, create a new isolated BraintrustState to avoid conflicts with + the global state. If neither is specified, use the global state. + """ + if state is not None: + return state + if api_key is not None or app_url is not None: + return BraintrustState() + return _state + + _internal_reset_global_state() _logger = logging.getLogger("braintrust") @@ -1576,7 +1592,7 @@ def init( :returns: The experiment object. """ - state: BraintrustState = state or _state + state = _resolve_state(state, api_key, app_url) if project is None and project_id is None: raise ValueError("Must specify at least one of project or project_id") @@ -1739,7 +1755,7 @@ def init_dataset( :returns: The dataset object. """ - state = state or _state + state = _resolve_state(state, api_key, app_url) def compute_metadata(): state.login(org_name=org_name, api_key=api_key, app_url=app_url) @@ -1766,11 +1782,12 @@ def compute_metadata(): ) -def _compute_logger_metadata(project_name: str | None = None, project_id: str | None = None): - login() - org_id = _state.org_id +def _compute_logger_metadata( + state: "BraintrustState", project_name: str | None = None, project_id: str | None = None +): + org_id = state.org_id if project_id is None: - response = _state.app_conn().post_json( + response = state.app_conn().post_json( "api/project/register", { "project_name": project_name or GLOBAL_PROJECT, @@ -1783,7 +1800,7 @@ def _compute_logger_metadata(project_name: str | None = None, project_id: str | project=ObjectMetadata(id=resp_project["id"], name=resp_project["name"], full_info=resp_project), ) elif project_name is None: - response = _state.app_conn().get_json("api/project", {"id": project_id}) + response = state.app_conn().get_json("api/project", {"id": project_id}) return OrgProjectMetadata( org_id=org_id, project=ObjectMetadata(id=project_id, name=response["name"], full_info=response) ) @@ -1819,7 +1836,7 @@ def init_logger( :returns: The newly created Logger. """ - state = state or _state + state = _resolve_state(state, api_key, app_url) compute_metadata_args = dict(project_name=project, project_id=project_id) link_args = { @@ -1831,7 +1848,7 @@ def init_logger( def compute_metadata(): state.login(org_name=org_name, api_key=api_key, app_url=app_url, force_login=force_login) - return _compute_logger_metadata(**compute_metadata_args) + return _compute_logger_metadata(state, **compute_metadata_args) # For loggers, enable queue size limit enforcement (bounded queue) state.enforce_queue_size_limit(True) @@ -3400,7 +3417,12 @@ def _span_components_to_object_id_lambda(components: SpanComponentsV4) -> Callab raise Exception("Impossible: compute_object_metadata_args not supported for experiments") elif components.object_type == SpanObjectTypeV3.PROJECT_LOGS: captured_compute_object_metadata_args = components.compute_object_metadata_args - return lambda: _compute_logger_metadata(**captured_compute_object_metadata_args).project.id + + def compute_project_id(): + login() + return _compute_logger_metadata(_state, **captured_compute_object_metadata_args).project.id + + return compute_project_id else: raise Exception(f"Unknown object type: {components.object_type}") diff --git a/py/src/braintrust/test_helpers.py b/py/src/braintrust/test_helpers.py index 7e24bb238..75fad2c1b 100644 --- a/py/src/braintrust/test_helpers.py +++ b/py/src/braintrust/test_helpers.py @@ -103,7 +103,7 @@ def init_test_logger(project_name: str): l._lazy_metadata = lazy_metadata # Skip actual login by setting fake metadata directly # Replace the global _compute_logger_metadata function with a resolved LazyValue - def fake_compute_logger_metadata(project_name=None, project_id=None): + def fake_compute_logger_metadata(state, project_name=None, project_id=None): if project_id: project_metadata = ObjectMetadata(id=project_id, name=project_name, full_info=dict()) else: diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 21382b8c8..8d381e9b0 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -3174,3 +3174,104 @@ def test_multiple_attachment_types_tracked(with_memory_logger, with_simulate_log assert attachment in with_memory_logger.upload_attempts assert json_attachment in with_memory_logger.upload_attempts assert ext_attachment in with_memory_logger.upload_attempts + + +class TestImplicitStateIsolation(TestCase): + """Test that passing api_key or app_url creates an isolated BraintrustState.""" + + def test_init_logger_with_api_key_creates_isolated_state(self): + """Test that init_logger with api_key creates a new isolated state.""" + from braintrust.logger import BraintrustState, _state + + # Create a logger with an explicit api_key + logger_a = init_logger( + project="test-project", + project_id="test-project-id", + api_key="test-api-key-a", + set_current=False, + ) + + # The logger should have its own state, not the global _state + assert logger_a.state is not _state + assert isinstance(logger_a.state, BraintrustState) + + def test_init_logger_with_app_url_creates_isolated_state(self): + """Test that init_logger with app_url creates a new isolated state.""" + from braintrust.logger import BraintrustState, _state + + # Create a logger with an explicit app_url + logger_a = init_logger( + project="test-project", + project_id="test-project-id", + app_url="https://custom.braintrust.dev", + set_current=False, + ) + + # The logger should have its own state, not the global _state + assert logger_a.state is not _state + assert isinstance(logger_a.state, BraintrustState) + + def test_init_logger_without_api_key_uses_global_state(self): + """Test that init_logger without api_key uses the global state.""" + from braintrust.logger import _state + + # Create a logger without api_key or app_url + logger_a = init_logger( + project="test-project", + project_id="test-project-id", + set_current=False, + ) + + # The logger should use the global _state + assert logger_a.state is _state + + def test_multiple_loggers_with_different_api_keys_have_separate_states(self): + """Test that multiple loggers with different api_keys have separate states.""" + # Create two loggers with different api_keys + logger_a = init_logger( + project="test-project-a", + project_id="test-project-id-a", + api_key="test-api-key-a", + set_current=False, + ) + + logger_b = init_logger( + project="test-project-b", + project_id="test-project-id-b", + api_key="test-api-key-b", + set_current=False, + ) + + # Each logger should have its own separate state + assert logger_a.state is not logger_b.state + + def test_init_with_api_key_creates_isolated_state(self): + """Test that init (experiment) with api_key creates a new isolated state.""" + from braintrust.logger import BraintrustState, _state + + # Create an experiment with an explicit api_key + experiment = braintrust.init( + project="test-project", + experiment="test-experiment", + api_key="test-api-key", + set_current=False, + ) + + # The experiment should have its own state, not the global _state + assert experiment.state is not _state + assert isinstance(experiment.state, BraintrustState) + + def test_init_dataset_with_api_key_creates_isolated_state(self): + """Test that init_dataset with api_key creates a new isolated state.""" + from braintrust.logger import BraintrustState, _state + + # Create a dataset with an explicit api_key + dataset = braintrust.init_dataset( + project="test-project", + name="test-dataset", + api_key="test-api-key", + ) + + # The dataset should have its own state, not the global _state + assert dataset.state is not _state + assert isinstance(dataset.state, BraintrustState)