From 7fd6c4ab80513c2336adc261eab58158cd5ff53a Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 26 May 2026 16:50:15 +0200 Subject: [PATCH 1/4] feat: Upwards-recursively read `.env.braintrust` containing `BRAINTRUST_API_KEY` --- py/src/braintrust/btx/span_fetcher.py | 6 +- py/src/braintrust/logger.py | 9 +- py/src/braintrust/otel/__init__.py | 71 ++++++++++--- py/src/braintrust/test_logger.py | 11 ++ py/src/braintrust/test_otel.py | 49 ++++++++- py/src/braintrust/test_util.py | 139 +++++++++++++++++++++++++- py/src/braintrust/util.py | 95 ++++++++++++++++++ 7 files changed, 359 insertions(+), 21 deletions(-) diff --git a/py/src/braintrust/btx/span_fetcher.py b/py/src/braintrust/btx/span_fetcher.py index 519b909c..de3925fd 100644 --- a/py/src/braintrust/btx/span_fetcher.py +++ b/py/src/braintrust/btx/span_fetcher.py @@ -13,6 +13,8 @@ import requests +from braintrust.util import get_braintrust_api_key + _BACKOFF_SECONDS = 30 _MAX_TOTAL_WAIT_SECONDS = 600 @@ -157,9 +159,9 @@ def _fetch_once(root_span_id: str, project_id: str, num_expected: int) -> list[d def _require_api_key() -> str: - key = os.environ.get("BRAINTRUST_API_KEY") + key = get_braintrust_api_key() if not key: - raise ValueError("BRAINTRUST_API_KEY environment variable is not set") + raise ValueError("BRAINTRUST_API_KEY is not set in the environment or nearest .env.braintrust file") return key diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index a908050d..c4d13297 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -88,6 +88,7 @@ coalesce, encode_uri_component, eprint, + get_braintrust_api_key, get_caller_location, get_signature, mask_api_key, @@ -2190,8 +2191,7 @@ def login_to_state( app_public_url = os.environ.get("BRAINTRUST_APP_PUBLIC_URL", app_url) - if api_key is None: - api_key = os.environ.get("BRAINTRUST_API_KEY") + api_key = get_braintrust_api_key(api_key) org_name = _get_org_name(org_name) @@ -2240,7 +2240,10 @@ def login_to_state( conn.set_token(api_key) if not conn: - raise ValueError("Could not login to Braintrust. You may need to set BRAINTRUST_API_KEY in your environment.") + raise ValueError( + "Could not login to Braintrust. You may need to set BRAINTRUST_API_KEY in your environment " + "or nearest .env.braintrust file." + ) # make_long_lived() allows the connection to retry if it breaks, which we're okay with after # this point because we know the connection _can_ successfully ping. diff --git a/py/src/braintrust/otel/__init__.py b/py/src/braintrust/otel/__init__.py index e6fe7f3e..472a3853 100644 --- a/py/src/braintrust/otel/__init__.py +++ b/py/src/braintrust/otel/__init__.py @@ -3,6 +3,8 @@ import warnings from urllib.parse import urljoin +from braintrust.util import get_braintrust_api_key + INSTALL_ERR_MSG = ( "OpenTelemetry packages are not installed. " @@ -29,6 +31,12 @@ class OTLPSpanExporter: def __init__(self, *args, **kwargs): raise ImportError(INSTALL_ERR_MSG) + def export(self, *args, **kwargs): + raise ImportError(INSTALL_ERR_MSG) + + def force_flush(self, *args, **kwargs): + raise ImportError(INSTALL_ERR_MSG) + class BatchSpanProcessor: def __init__(self, *args, **kwargs): raise ImportError(INSTALL_ERR_MSG) @@ -145,7 +153,7 @@ class OtelExporter(OTLPSpanExporter): a more convenient all-in-one interface. Environment Variables: - - BRAINTRUST_API_KEY: Your Braintrust API key. + - BRAINTRUST_API_KEY: Your Braintrust API key. If unset, the nearest .env.braintrust file is used. - BRAINTRUST_PARENT: Parent identifier (e.g., "project_name:test"). - BRAINTRUST_API_URL: Base URL for Braintrust API (defaults to https://api.braintrust.dev). """ @@ -163,7 +171,7 @@ def __init__( Args: url: OTLP endpoint URL. Defaults to {BRAINTRUST_API_URL}/otel/v1/traces. - api_key: Braintrust API key. Defaults to BRAINTRUST_API_KEY env var. + api_key: Braintrust API key. Defaults to BRAINTRUST_API_KEY env var, then .env.braintrust. parent: Parent identifier (e.g., "project_name:test"). Defaults to BRAINTRUST_PARENT env var. headers: Additional headers to include in requests. **kwargs: Additional arguments passed to OTLPSpanExporter. @@ -173,15 +181,13 @@ def __init__( if not base_url.endswith("/"): base_url += "/" endpoint = url or urljoin(base_url, "otel/v1/traces") - api_key = api_key or os.environ.get("BRAINTRUST_API_KEY") + api_key_arg = api_key + env_api_key = os.environ.get("BRAINTRUST_API_KEY") + if api_key is None: + api_key = env_api_key if env_api_key and env_api_key.strip() else None parent = parent or os.environ.get("BRAINTRUST_PARENT") headers = headers or {} - if not api_key: - raise ValueError( - "API key is required. Provide it via api_key parameter or BRAINTRUST_API_KEY environment variable." - ) - # Default parent if not provided if not parent: parent = "project_name:default-otel-project" @@ -190,10 +196,14 @@ def __init__( "Configure with BRAINTRUST_PARENT environment variable or parent parameter." ) - exporter_headers = { - "Authorization": f"Bearer {api_key}", - **headers, - } + self._braintrust_api_key_arg = api_key_arg + self._braintrust_headers_override_authorization = "Authorization" in headers + self._braintrust_has_api_key = bool(api_key and api_key.strip()) + + exporter_headers = {} + if self._braintrust_has_api_key: + exporter_headers["Authorization"] = f"Bearer {api_key}" + exporter_headers.update(headers) if parent: exporter_headers["x-bt-parent"] = parent @@ -202,6 +212,39 @@ def __init__( super().__init__(endpoint=endpoint, headers=exporter_headers, **kwargs) + def _set_api_key_header(self, api_key: str) -> None: + if not self._braintrust_headers_override_authorization: + authorization = {"Authorization": f"Bearer {api_key}"} + exporter_headers = getattr(self, "_headers", None) + if isinstance(exporter_headers, dict): + exporter_headers.update(authorization) + else: + self._headers = {**dict(exporter_headers or {}), **authorization} + + session = getattr(self, "_session", None) + if session is not None: + session.headers.update(authorization) + self._braintrust_has_api_key = True + + def _ensure_api_key(self) -> None: + if self._braintrust_has_api_key: + return + api_key = get_braintrust_api_key(self._braintrust_api_key_arg) + if not api_key or not api_key.strip(): + raise ValueError("API key is required. Provide it via api_key parameter, BRAINTRUST_API_KEY environment variable, or the nearest .env.braintrust file.") + self._set_api_key_header(api_key) + + def initialize(self) -> None: + self._ensure_api_key() + + def export(self, spans): + self._ensure_api_key() + return super().export(spans) + + def force_flush(self, timeout_millis=30000): + self._ensure_api_key() + return super().force_flush(timeout_millis) + def add_braintrust_span_processor( tracer_provider, @@ -252,7 +295,7 @@ def __init__( Initialize the BraintrustSpanProcessor. Args: - api_key: Braintrust API key. Defaults to BRAINTRUST_API_KEY env var. + api_key: Braintrust API key. Defaults to BRAINTRUST_API_KEY env var, then .env.braintrust. parent: Parent identifier (e.g., "project_name:test"). Defaults to BRAINTRUST_PARENT env var. api_url: Base URL for Braintrust API. Defaults to BRAINTRUST_API_URL env var or https://api.braintrust.dev. filter_ai_spans: Whether to enable AI span filtering. Defaults to False. @@ -340,6 +383,7 @@ def _get_parent_otel_braintrust_parent(self, parent_context): def on_end(self, span): """Forward span end events to the inner processor.""" + self._exporter.initialize() self._processor.on_end(span) def _on_ending(self, span): @@ -352,6 +396,7 @@ def shutdown(self): def force_flush(self, timeout_millis=30000): """Force flush the inner processor.""" + self._exporter.initialize() return self._processor.force_flush(timeout_millis) @property diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 3fc5c833..9ba186c1 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -49,6 +49,17 @@ ) +def test_login_to_state_uses_env_braintrust_api_key(tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text(f"BRAINTRUST_API_KEY={logger.TEST_API_KEY}\n") + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + state = logger.login_to_state(org_name="test-org-name") + + assert state.login_token == logger.TEST_API_KEY + assert state.logged_in is True + + class TestInit(TestCase): def test_init_validation(self): with self.assertRaises(ValueError) as cm: diff --git a/py/src/braintrust/test_otel.py b/py/src/braintrust/test_otel.py index 68da6a72..2a8abc8a 100644 --- a/py/src/braintrust/test_otel.py +++ b/py/src/braintrust/test_otel.py @@ -38,7 +38,7 @@ def test_otel_import_behavior(): assert hasattr(OtelExporter, "__init__") -def test_otel_exporter_creation(): +def test_otel_exporter_creation(tmp_path): """Test OtelExporter creation with and without full OpenTelemetry SDK.""" from braintrust.otel import OtelExporter @@ -58,9 +58,12 @@ def test_otel_exporter_creation(): with pytest.MonkeyPatch.context() as m: m.delenv("BRAINTRUST_API_KEY", raising=False) m.delenv("BRAINTRUST_PARENT", raising=False) + m.chdir(tmp_path) + (tmp_path / ".env.braintrust").write_text("") + exporter = OtelExporter() with pytest.raises(ValueError, match="API key is required"): - OtelExporter() + exporter.force_flush() else: # When SDK is not fully installed, instantiation should raise ImportError with pytest.raises(ImportError, match="OpenTelemetry packages are not installed"): @@ -92,6 +95,48 @@ def test_otel_exporter_with_explicit_params(): assert exporter._headers == expected_headers +def test_otel_exporter_uses_env_braintrust_api_key(tmp_path): + if not _check_otel_installed(): + pytest.skip("OpenTelemetry SDK not fully installed, skipping test") + + from braintrust.otel import OtelExporter + + with pytest.MonkeyPatch.context() as m: + m.delenv("BRAINTRUST_API_KEY", raising=False) + m.chdir(tmp_path) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-api-key\n") + + exporter = OtelExporter(parent="project_name:test") + exporter.force_flush() + + assert exporter._headers["Authorization"] == "Bearer file-api-key" + + +def test_braintrust_span_processor_missing_key_raises_on_span_end(tmp_path): + if not _check_otel_installed(): + pytest.skip("OpenTelemetry SDK not fully installed, skipping test") + + from braintrust.otel import BraintrustSpanProcessor + from opentelemetry.sdk.trace import TracerProvider + + with pytest.MonkeyPatch.context() as m: + m.delenv("BRAINTRUST_API_KEY", raising=False) + m.chdir(tmp_path) + (tmp_path / ".env.braintrust").write_text("") + + provider = TracerProvider() + processor = BraintrustSpanProcessor() + provider.add_span_processor(processor) + tracer = provider.get_tracer("test_tracer") + + try: + with pytest.raises(ValueError, match="API key is required"): + with tracer.start_as_current_span("test_span"): + pass + finally: + provider.shutdown() + + def test_otel_exporter_no_parent(caplog): if not _check_otel_installed(): pytest.skip("OpenTelemetry SDK not fully installed, skipping test") diff --git a/py/src/braintrust/test_util.py b/py/src/braintrust/test_util.py index 86bec1c1..d046de10 100644 --- a/py/src/braintrust/test_util.py +++ b/py/src/braintrust/test_util.py @@ -1,8 +1,145 @@ +import os import unittest import pytest -from .util import LazyValue, mask_api_key, merge_dicts_with_paths +from .util import LazyValue, get_braintrust_api_key, mask_api_key, merge_dicts_with_paths, parse_env_var_float + + +class TestParseEnvVarFloat: + """Tests for parse_env_var_float helper.""" + + def test_returns_default_when_env_not_set(self): + assert parse_env_var_float("NONEXISTENT_VAR_12345", 42.0) == 42.0 + + def test_parses_valid_float(self): + os.environ["TEST_FLOAT"] = "123.45" + try: + assert parse_env_var_float("TEST_FLOAT", 0.0) == 123.45 + finally: + del os.environ["TEST_FLOAT"] + + def test_returns_default_for_nan(self): + os.environ["TEST_FLOAT"] = "nan" + try: + assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 + finally: + del os.environ["TEST_FLOAT"] + + def test_returns_default_for_inf(self): + os.environ["TEST_FLOAT"] = "inf" + try: + assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 + finally: + del os.environ["TEST_FLOAT"] + + def test_returns_default_for_negative_inf(self): + os.environ["TEST_FLOAT"] = "-inf" + try: + assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 + finally: + del os.environ["TEST_FLOAT"] + + def test_returns_default_for_empty_string(self): + os.environ["TEST_FLOAT"] = "" + try: + assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 + finally: + del os.environ["TEST_FLOAT"] + + def test_returns_default_for_invalid_string(self): + os.environ["TEST_FLOAT"] = "not_a_number" + try: + assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 + finally: + del os.environ["TEST_FLOAT"] + + def test_allows_negative_values(self): + os.environ["TEST_FLOAT"] = "-5.5" + try: + assert parse_env_var_float("TEST_FLOAT", 0.0) == -5.5 + finally: + del os.environ["TEST_FLOAT"] + + +class TestBraintrustApiKeyLookup: + def test_explicit_api_key_wins(self, tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("BRAINTRUST_API_KEY", "env-key") + + assert get_braintrust_api_key("explicit-key") == "explicit-key" + + def test_nonblank_environment_wins(self, tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("BRAINTRUST_API_KEY", "env-key") + + assert get_braintrust_api_key() == "env-key" + + def test_blank_environment_falls_back_to_file(self, tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("BRAINTRUST_API_KEY", " ") + + assert get_braintrust_api_key() == "file-key" + + def test_uses_nearest_parent_file(self, tmp_path, monkeypatch): + nested = tmp_path / "packages" / "app" + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") + (tmp_path / "packages" / ".env.braintrust").write_text("BRAINTRUST_API_KEY=package-key\n") + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert get_braintrust_api_key() == "package-key" + + @pytest.mark.parametrize("contents", ["OTHER=value\n", 'BRAINTRUST_API_KEY=" "\n']) + def test_nearest_file_is_boundary_without_nonblank_key(self, tmp_path, monkeypatch, contents): + nested = tmp_path / "packages" / "app" + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") + (tmp_path / "packages" / ".env.braintrust").write_text(contents) + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert get_braintrust_api_key() is None + + def test_unreadable_nearest_file_is_boundary(self, tmp_path, monkeypatch): + nested = tmp_path / "packages" / "app" + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") + (tmp_path / "packages" / ".env.braintrust").mkdir() + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert get_braintrust_api_key() is None + + def test_searches_cwd_plus_64_parents(self, tmp_path, monkeypatch): + segments = [f"d{i}" for i in range(65)] + nested = tmp_path.joinpath(*segments) + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=too-high\n") + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert get_braintrust_api_key() is None + + (tmp_path / segments[0] / ".env.braintrust").write_text("BRAINTRUST_API_KEY=boundary-key\n") + + assert get_braintrust_api_key() == "boundary-key" + + def test_supports_dotenv_syntax_and_does_not_mutate_environment(self, tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text( + 'export BRAINTRUST_API_KEY="quoted-key" # comment\nOTHER=value\n' + ) + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + monkeypatch.delenv("OTHER", raising=False) + + assert get_braintrust_api_key() == "quoted-key" + assert os.environ.get("BRAINTRUST_API_KEY") is None + assert os.environ.get("OTHER") is None class TestLazyValue(unittest.TestCase): diff --git a/py/src/braintrust/util.py b/py/src/braintrust/util.py index 0f4d0927..c86238a9 100644 --- a/py/src/braintrust/util.py +++ b/py/src/braintrust/util.py @@ -1,5 +1,9 @@ import inspect +import io import json +import math +import os +import shlex import sys import threading import urllib.parse @@ -10,8 +14,99 @@ from requests import HTTPError, Response +def parse_env_var_float(name: str, default: float) -> float: + """Parse a float from an environment variable, returning default if invalid. + + Returns the default value if the env var is missing, empty, not a valid + float, NaN, or infinity. + """ + value = os.environ.get(name) + if value is None: + return default + try: + result = float(value) + if math.isnan(result) or math.isinf(result): + return default + return result + except (ValueError, TypeError): + return default + + GLOBAL_PROJECT = "Global" BT_IS_ASYNC_ATTRIBUTE = "_BT_IS_ASYNC" +BRAINTRUST_ENV_FILE = ".env.braintrust" +BRAINTRUST_ENV_SEARCH_PARENT_LIMIT = 64 + + +def _parse_braintrust_api_key_dotenv(contents: str) -> str | None: + try: + from dotenv import dotenv_values + + parsed = dotenv_values(stream=io.StringIO(contents), interpolate=False) + value = parsed.get("BRAINTRUST_API_KEY") + return value if value and value.strip() else None + except ImportError: + pass + except Exception: + return None + + for line in contents.splitlines(): + stripped = line.lstrip() + if not stripped or stripped.startswith("#"): + continue + if stripped.startswith("export "): + stripped = stripped[len("export ") :].lstrip() + if "=" not in stripped: + continue + + key, value = stripped.split("=", 1) + if key.strip() != "BRAINTRUST_API_KEY": + continue + + lexer = shlex.shlex(value.lstrip(), posix=True) + lexer.whitespace_split = True + lexer.commenters = "#" + try: + parts = list(lexer) + except ValueError: + return None + if not parts: + return None + api_key = parts[0] + return api_key if api_key.strip() else None + + return None + + +def get_braintrust_api_key(api_key: str | None = None) -> str | None: + if api_key is not None: + return api_key + + env_api_key = os.environ.get("BRAINTRUST_API_KEY") + if env_api_key and env_api_key.strip(): + return env_api_key + + try: + directory = os.getcwd() + except OSError: + return None + + for _ in range(BRAINTRUST_ENV_SEARCH_PARENT_LIMIT + 1): + env_path = os.path.join(directory, BRAINTRUST_ENV_FILE) + try: + with open(env_path, encoding="utf-8") as f: + return _parse_braintrust_api_key_dotenv(f.read()) + except FileNotFoundError: + pass + except OSError: + return None + + parent = os.path.dirname(directory) + if parent == directory: + break + directory = parent + + return None def get_signature(fn: Callable) -> inspect.Signature: From 16d4e1b5a55a6594fb98043e46b45954b61061a7 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 26 May 2026 19:45:37 +0200 Subject: [PATCH 2/4] lint --- py/src/braintrust/btx/span_fetcher.py | 1 - py/src/braintrust/otel/__init__.py | 4 +++- py/src/braintrust/test_util.py | 4 +--- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/py/src/braintrust/btx/span_fetcher.py b/py/src/braintrust/btx/span_fetcher.py index de3925fd..2d05a2e4 100644 --- a/py/src/braintrust/btx/span_fetcher.py +++ b/py/src/braintrust/btx/span_fetcher.py @@ -12,7 +12,6 @@ from typing import Any import requests - from braintrust.util import get_braintrust_api_key diff --git a/py/src/braintrust/otel/__init__.py b/py/src/braintrust/otel/__init__.py index 472a3853..e5ef40b4 100644 --- a/py/src/braintrust/otel/__init__.py +++ b/py/src/braintrust/otel/__init__.py @@ -231,7 +231,9 @@ def _ensure_api_key(self) -> None: return api_key = get_braintrust_api_key(self._braintrust_api_key_arg) if not api_key or not api_key.strip(): - raise ValueError("API key is required. Provide it via api_key parameter, BRAINTRUST_API_KEY environment variable, or the nearest .env.braintrust file.") + raise ValueError( + "API key is required. Provide it via api_key parameter, BRAINTRUST_API_KEY environment variable, or the nearest .env.braintrust file." + ) self._set_api_key_header(api_key) def initialize(self) -> None: diff --git a/py/src/braintrust/test_util.py b/py/src/braintrust/test_util.py index d046de10..badfa794 100644 --- a/py/src/braintrust/test_util.py +++ b/py/src/braintrust/test_util.py @@ -130,9 +130,7 @@ def test_searches_cwd_plus_64_parents(self, tmp_path, monkeypatch): assert get_braintrust_api_key() == "boundary-key" def test_supports_dotenv_syntax_and_does_not_mutate_environment(self, tmp_path, monkeypatch): - (tmp_path / ".env.braintrust").write_text( - 'export BRAINTRUST_API_KEY="quoted-key" # comment\nOTHER=value\n' - ) + (tmp_path / ".env.braintrust").write_text('export BRAINTRUST_API_KEY="quoted-key" # comment\nOTHER=value\n') monkeypatch.chdir(tmp_path) monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) monkeypatch.delenv("OTHER", raising=False) From 688d449d0461ad90fab4cc0f01fbe6b2e9ddac6c Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 28 May 2026 15:45:06 -0400 Subject: [PATCH 3/4] cleanup --- py/src/braintrust/test_util.py | 58 +--------------------------------- py/src/braintrust/util.py | 19 ----------- 2 files changed, 1 insertion(+), 76 deletions(-) diff --git a/py/src/braintrust/test_util.py b/py/src/braintrust/test_util.py index badfa794..9f1b826c 100644 --- a/py/src/braintrust/test_util.py +++ b/py/src/braintrust/test_util.py @@ -3,63 +3,7 @@ import pytest -from .util import LazyValue, get_braintrust_api_key, mask_api_key, merge_dicts_with_paths, parse_env_var_float - - -class TestParseEnvVarFloat: - """Tests for parse_env_var_float helper.""" - - def test_returns_default_when_env_not_set(self): - assert parse_env_var_float("NONEXISTENT_VAR_12345", 42.0) == 42.0 - - def test_parses_valid_float(self): - os.environ["TEST_FLOAT"] = "123.45" - try: - assert parse_env_var_float("TEST_FLOAT", 0.0) == 123.45 - finally: - del os.environ["TEST_FLOAT"] - - def test_returns_default_for_nan(self): - os.environ["TEST_FLOAT"] = "nan" - try: - assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 - finally: - del os.environ["TEST_FLOAT"] - - def test_returns_default_for_inf(self): - os.environ["TEST_FLOAT"] = "inf" - try: - assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 - finally: - del os.environ["TEST_FLOAT"] - - def test_returns_default_for_negative_inf(self): - os.environ["TEST_FLOAT"] = "-inf" - try: - assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 - finally: - del os.environ["TEST_FLOAT"] - - def test_returns_default_for_empty_string(self): - os.environ["TEST_FLOAT"] = "" - try: - assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 - finally: - del os.environ["TEST_FLOAT"] - - def test_returns_default_for_invalid_string(self): - os.environ["TEST_FLOAT"] = "not_a_number" - try: - assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0 - finally: - del os.environ["TEST_FLOAT"] - - def test_allows_negative_values(self): - os.environ["TEST_FLOAT"] = "-5.5" - try: - assert parse_env_var_float("TEST_FLOAT", 0.0) == -5.5 - finally: - del os.environ["TEST_FLOAT"] +from .util import LazyValue, get_braintrust_api_key, mask_api_key, merge_dicts_with_paths class TestBraintrustApiKeyLookup: diff --git a/py/src/braintrust/util.py b/py/src/braintrust/util.py index c86238a9..e3f39377 100644 --- a/py/src/braintrust/util.py +++ b/py/src/braintrust/util.py @@ -1,7 +1,6 @@ import inspect import io import json -import math import os import shlex import sys @@ -14,24 +13,6 @@ from requests import HTTPError, Response -def parse_env_var_float(name: str, default: float) -> float: - """Parse a float from an environment variable, returning default if invalid. - - Returns the default value if the env var is missing, empty, not a valid - float, NaN, or infinity. - """ - value = os.environ.get(name) - if value is None: - return default - try: - result = float(value) - if math.isnan(result) or math.isinf(result): - return default - return result - except (ValueError, TypeError): - return default - - GLOBAL_PROJECT = "Global" BT_IS_ASYNC_ATTRIBUTE = "_BT_IS_ASYNC" BRAINTRUST_ENV_FILE = ".env.braintrust" From caedb2f2a69f96e745d6727a99f0f6b53dfa6f82 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 28 May 2026 15:57:00 -0400 Subject: [PATCH 4/4] use new env helpers --- py/src/braintrust/btx/span_fetcher.py | 4 +- py/src/braintrust/env.py | 90 ++++++++++++++++++++++++--- py/src/braintrust/logger.py | 3 +- py/src/braintrust/otel/__init__.py | 4 +- py/src/braintrust/test_env.py | 72 +++++++++++++++++++++ py/src/braintrust/test_util.py | 81 +----------------------- py/src/braintrust/util.py | 76 ---------------------- 7 files changed, 159 insertions(+), 171 deletions(-) diff --git a/py/src/braintrust/btx/span_fetcher.py b/py/src/braintrust/btx/span_fetcher.py index 2d05a2e4..4af00adc 100644 --- a/py/src/braintrust/btx/span_fetcher.py +++ b/py/src/braintrust/btx/span_fetcher.py @@ -12,7 +12,7 @@ from typing import Any import requests -from braintrust.util import get_braintrust_api_key +from braintrust.env import BraintrustEnv _BACKOFF_SECONDS = 30 @@ -158,7 +158,7 @@ def _fetch_once(root_span_id: str, project_id: str, num_expected: int) -> list[d def _require_api_key() -> str: - key = get_braintrust_api_key() + key = BraintrustEnv.API_KEY.get(None, use_dotenv=True) if not key: raise ValueError("BRAINTRUST_API_KEY is not set in the environment or nearest .env.braintrust file") return key diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index a4933c3e..19aee1dd 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -1,5 +1,7 @@ +import io import math import os +import shlex from collections.abc import Callable from dataclasses import dataclass from enum import Enum @@ -9,6 +11,8 @@ T = TypeVar("T") EnvValue = bool | float | int | str _Parser = Callable[[str], EnvValue | None] +BRAINTRUST_ENV_FILE = ".env.braintrust" +BRAINTRUST_ENV_SEARCH_PARENT_LIMIT = 64 def parse_float(value: str) -> float | None: @@ -48,9 +52,9 @@ def parse_bool(value: str) -> bool | None: def parse_string(value: str) -> str | None: """Parse a string environment variable. - Empty strings are treated as unset so callers fall back to their default. + Empty or whitespace-only strings are treated as unset so callers fall back to their default. """ - return value or None + return value if value.strip() else None class EnvParser(Enum): @@ -68,18 +72,86 @@ class EnvVar: name: str parser: EnvParser - def get(self, default: T) -> T: - value = os.environ.get(self.name) + def get(self, default: T, *, use_dotenv: bool = False) -> T: + parsed = self._parse_value(os.environ.get(self.name)) + if parsed is not None: + return cast(T, parsed) + + if use_dotenv: + parsed = self._get_dotenv_value() + if parsed is not None: + return cast(T, parsed) + + return default + + def _parse_value(self, value: str | None) -> EnvValue | None: if value is None: - return default + return None + return self.parser.parser(value) + + def _get_dotenv_value(self) -> EnvValue | None: + try: + directory = os.getcwd() + except OSError: + return None + + for _ in range(BRAINTRUST_ENV_SEARCH_PARENT_LIMIT + 1): + env_path = os.path.join(directory, BRAINTRUST_ENV_FILE) + try: + with open(env_path, encoding="utf-8") as f: + return self._parse_dotenv_contents(f.read()) + except FileNotFoundError: + pass + except OSError: + return None + + parent = os.path.dirname(directory) + if parent == directory: + break + directory = parent - parsed = self.parser.parser(value) - if parsed is None: - return default - return cast(T, parsed) + return None + + def _parse_dotenv_contents(self, contents: str) -> EnvValue | None: + try: + from dotenv import dotenv_values + + parsed = dotenv_values(stream=io.StringIO(contents), interpolate=False) + return self._parse_value(parsed.get(self.name)) + except ImportError: + pass + except Exception: + return None + + for line in contents.splitlines(): + stripped = line.lstrip() + if not stripped or stripped.startswith("#"): + continue + if stripped.startswith("export "): + stripped = stripped[len("export ") :].lstrip() + if "=" not in stripped: + continue + + key, value = stripped.split("=", 1) + if key.strip() != self.name: + continue + + lexer = shlex.shlex(value.lstrip(), posix=True) + lexer.whitespace_split = True + lexer.commenters = "#" + try: + parts = list(lexer) + except ValueError: + return None + if not parts: + return None + return self._parse_value(parts[0]) + + return None class BraintrustEnv: + API_KEY = EnvVar("BRAINTRUST_API_KEY", EnvParser.STRING) HTTP_TIMEOUT = EnvVar("BRAINTRUST_HTTP_TIMEOUT", EnvParser.FLOAT) SYNC_FLUSH = EnvVar("BRAINTRUST_SYNC_FLUSH", EnvParser.BOOL) MAX_REQUEST_SIZE = EnvVar("BRAINTRUST_MAX_REQUEST_SIZE", EnvParser.INT) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index c4d13297..6fff8304 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -88,7 +88,6 @@ coalesce, encode_uri_component, eprint, - get_braintrust_api_key, get_caller_location, get_signature, mask_api_key, @@ -2191,7 +2190,7 @@ def login_to_state( app_public_url = os.environ.get("BRAINTRUST_APP_PUBLIC_URL", app_url) - api_key = get_braintrust_api_key(api_key) + api_key = api_key or BraintrustEnv.API_KEY.get(None, use_dotenv=True) org_name = _get_org_name(org_name) diff --git a/py/src/braintrust/otel/__init__.py b/py/src/braintrust/otel/__init__.py index e5ef40b4..3c0aee73 100644 --- a/py/src/braintrust/otel/__init__.py +++ b/py/src/braintrust/otel/__init__.py @@ -3,7 +3,7 @@ import warnings from urllib.parse import urljoin -from braintrust.util import get_braintrust_api_key +from braintrust.env import BraintrustEnv INSTALL_ERR_MSG = ( @@ -229,7 +229,7 @@ def _set_api_key_header(self, api_key: str) -> None: def _ensure_api_key(self) -> None: if self._braintrust_has_api_key: return - api_key = get_braintrust_api_key(self._braintrust_api_key_arg) + api_key = self._braintrust_api_key_arg or BraintrustEnv.API_KEY.get(None, use_dotenv=True) if not api_key or not api_key.strip(): raise ValueError( "API key is required. Provide it via api_key parameter, BRAINTRUST_API_KEY environment variable, or the nearest .env.braintrust file." diff --git a/py/src/braintrust/test_env.py b/py/src/braintrust/test_env.py index 7d983902..a4b750be 100644 --- a/py/src/braintrust/test_env.py +++ b/py/src/braintrust/test_env.py @@ -1,3 +1,5 @@ +import pytest + from .env import BraintrustEnv, EnvParser, EnvVar, parse_bool, parse_float, parse_int, parse_string @@ -27,6 +29,7 @@ def test_parse_bool(self): def test_parse_string(self): assert parse_string("value") == "value" assert parse_string("") is None + assert parse_string(" ") is None class TestEnvVar: @@ -53,6 +56,75 @@ def test_default_is_supplied_by_call_site(self, monkeypatch): class TestBraintrustEnv: + def test_api_key_nonblank_environment_wins(self, tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("BRAINTRUST_API_KEY", "env-key") + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) == "env-key" + + def test_api_key_blank_environment_falls_back_to_file(self, tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("BRAINTRUST_API_KEY", " ") + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) == "file-key" + + def test_api_key_uses_nearest_parent_file(self, tmp_path, monkeypatch): + nested = tmp_path / "packages" / "app" + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") + (tmp_path / "packages" / ".env.braintrust").write_text("BRAINTRUST_API_KEY=package-key\n") + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) == "package-key" + + @pytest.mark.parametrize("contents", ["OTHER=value\n", 'BRAINTRUST_API_KEY=" "\n']) + def test_api_key_nearest_file_is_boundary_without_nonblank_key(self, tmp_path, monkeypatch, contents): + nested = tmp_path / "packages" / "app" + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") + (tmp_path / "packages" / ".env.braintrust").write_text(contents) + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) is None + + def test_api_key_unreadable_nearest_file_is_boundary(self, tmp_path, monkeypatch): + nested = tmp_path / "packages" / "app" + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") + (tmp_path / "packages" / ".env.braintrust").mkdir() + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) is None + + def test_api_key_searches_cwd_plus_64_parents(self, tmp_path, monkeypatch): + segments = [f"d{i}" for i in range(65)] + nested = tmp_path.joinpath(*segments) + nested.mkdir(parents=True) + (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=too-high\n") + monkeypatch.chdir(nested) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) is None + + (tmp_path / segments[0] / ".env.braintrust").write_text("BRAINTRUST_API_KEY=boundary-key\n") + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) == "boundary-key" + + def test_api_key_supports_dotenv_syntax_and_does_not_mutate_environment(self, tmp_path, monkeypatch): + (tmp_path / ".env.braintrust").write_text('export BRAINTRUST_API_KEY="quoted-key" # comment\nOTHER=value\n') + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) + monkeypatch.delenv("OTHER", raising=False) + + assert BraintrustEnv.API_KEY.get(None, use_dotenv=True) == "quoted-key" + assert EnvVar("BRAINTRUST_API_KEY", EnvParser.STRING).get(None) is None + assert EnvVar("OTHER", EnvParser.STRING).get(None) is None + def test_centralized_env_definitions_are_lazy(self, monkeypatch): monkeypatch.delenv("BRAINTRUST_HTTP_TIMEOUT", raising=False) assert BraintrustEnv.HTTP_TIMEOUT.get(60.0) == 60.0 diff --git a/py/src/braintrust/test_util.py b/py/src/braintrust/test_util.py index 9f1b826c..86bec1c1 100644 --- a/py/src/braintrust/test_util.py +++ b/py/src/braintrust/test_util.py @@ -1,87 +1,8 @@ -import os import unittest import pytest -from .util import LazyValue, get_braintrust_api_key, mask_api_key, merge_dicts_with_paths - - -class TestBraintrustApiKeyLookup: - def test_explicit_api_key_wins(self, tmp_path, monkeypatch): - (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") - monkeypatch.chdir(tmp_path) - monkeypatch.setenv("BRAINTRUST_API_KEY", "env-key") - - assert get_braintrust_api_key("explicit-key") == "explicit-key" - - def test_nonblank_environment_wins(self, tmp_path, monkeypatch): - (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") - monkeypatch.chdir(tmp_path) - monkeypatch.setenv("BRAINTRUST_API_KEY", "env-key") - - assert get_braintrust_api_key() == "env-key" - - def test_blank_environment_falls_back_to_file(self, tmp_path, monkeypatch): - (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=file-key\n") - monkeypatch.chdir(tmp_path) - monkeypatch.setenv("BRAINTRUST_API_KEY", " ") - - assert get_braintrust_api_key() == "file-key" - - def test_uses_nearest_parent_file(self, tmp_path, monkeypatch): - nested = tmp_path / "packages" / "app" - nested.mkdir(parents=True) - (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") - (tmp_path / "packages" / ".env.braintrust").write_text("BRAINTRUST_API_KEY=package-key\n") - monkeypatch.chdir(nested) - monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) - - assert get_braintrust_api_key() == "package-key" - - @pytest.mark.parametrize("contents", ["OTHER=value\n", 'BRAINTRUST_API_KEY=" "\n']) - def test_nearest_file_is_boundary_without_nonblank_key(self, tmp_path, monkeypatch, contents): - nested = tmp_path / "packages" / "app" - nested.mkdir(parents=True) - (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") - (tmp_path / "packages" / ".env.braintrust").write_text(contents) - monkeypatch.chdir(nested) - monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) - - assert get_braintrust_api_key() is None - - def test_unreadable_nearest_file_is_boundary(self, tmp_path, monkeypatch): - nested = tmp_path / "packages" / "app" - nested.mkdir(parents=True) - (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=root-key\n") - (tmp_path / "packages" / ".env.braintrust").mkdir() - monkeypatch.chdir(nested) - monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) - - assert get_braintrust_api_key() is None - - def test_searches_cwd_plus_64_parents(self, tmp_path, monkeypatch): - segments = [f"d{i}" for i in range(65)] - nested = tmp_path.joinpath(*segments) - nested.mkdir(parents=True) - (tmp_path / ".env.braintrust").write_text("BRAINTRUST_API_KEY=too-high\n") - monkeypatch.chdir(nested) - monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) - - assert get_braintrust_api_key() is None - - (tmp_path / segments[0] / ".env.braintrust").write_text("BRAINTRUST_API_KEY=boundary-key\n") - - assert get_braintrust_api_key() == "boundary-key" - - def test_supports_dotenv_syntax_and_does_not_mutate_environment(self, tmp_path, monkeypatch): - (tmp_path / ".env.braintrust").write_text('export BRAINTRUST_API_KEY="quoted-key" # comment\nOTHER=value\n') - monkeypatch.chdir(tmp_path) - monkeypatch.delenv("BRAINTRUST_API_KEY", raising=False) - monkeypatch.delenv("OTHER", raising=False) - - assert get_braintrust_api_key() == "quoted-key" - assert os.environ.get("BRAINTRUST_API_KEY") is None - assert os.environ.get("OTHER") is None +from .util import LazyValue, mask_api_key, merge_dicts_with_paths class TestLazyValue(unittest.TestCase): diff --git a/py/src/braintrust/util.py b/py/src/braintrust/util.py index e3f39377..0f4d0927 100644 --- a/py/src/braintrust/util.py +++ b/py/src/braintrust/util.py @@ -1,8 +1,5 @@ import inspect -import io import json -import os -import shlex import sys import threading import urllib.parse @@ -15,79 +12,6 @@ GLOBAL_PROJECT = "Global" BT_IS_ASYNC_ATTRIBUTE = "_BT_IS_ASYNC" -BRAINTRUST_ENV_FILE = ".env.braintrust" -BRAINTRUST_ENV_SEARCH_PARENT_LIMIT = 64 - - -def _parse_braintrust_api_key_dotenv(contents: str) -> str | None: - try: - from dotenv import dotenv_values - - parsed = dotenv_values(stream=io.StringIO(contents), interpolate=False) - value = parsed.get("BRAINTRUST_API_KEY") - return value if value and value.strip() else None - except ImportError: - pass - except Exception: - return None - - for line in contents.splitlines(): - stripped = line.lstrip() - if not stripped or stripped.startswith("#"): - continue - if stripped.startswith("export "): - stripped = stripped[len("export ") :].lstrip() - if "=" not in stripped: - continue - - key, value = stripped.split("=", 1) - if key.strip() != "BRAINTRUST_API_KEY": - continue - - lexer = shlex.shlex(value.lstrip(), posix=True) - lexer.whitespace_split = True - lexer.commenters = "#" - try: - parts = list(lexer) - except ValueError: - return None - if not parts: - return None - api_key = parts[0] - return api_key if api_key.strip() else None - - return None - - -def get_braintrust_api_key(api_key: str | None = None) -> str | None: - if api_key is not None: - return api_key - - env_api_key = os.environ.get("BRAINTRUST_API_KEY") - if env_api_key and env_api_key.strip(): - return env_api_key - - try: - directory = os.getcwd() - except OSError: - return None - - for _ in range(BRAINTRUST_ENV_SEARCH_PARENT_LIMIT + 1): - env_path = os.path.join(directory, BRAINTRUST_ENV_FILE) - try: - with open(env_path, encoding="utf-8") as f: - return _parse_braintrust_api_key_dotenv(f.read()) - except FileNotFoundError: - pass - except OSError: - return None - - parent = os.path.dirname(directory) - if parent == directory: - break - directory = parent - - return None def get_signature(fn: Callable) -> inspect.Signature: