Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions py/src/braintrust/btx/span_fetcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from typing import Any

import requests
from braintrust.env import BraintrustEnv


_BACKOFF_SECONDS = 30
Expand Down Expand Up @@ -157,9 +158,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 = BraintrustEnv.API_KEY.get(None, use_dotenv=True)
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


Expand Down
90 changes: 81 additions & 9 deletions py/src/braintrust/env.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand All @@ -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)
Expand Down
8 changes: 5 additions & 3 deletions py/src/braintrust/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -2190,8 +2190,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 = api_key or BraintrustEnv.API_KEY.get(None, use_dotenv=True)

org_name = _get_org_name(org_name)

Expand Down Expand Up @@ -2240,7 +2239,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.
Expand Down
73 changes: 60 additions & 13 deletions py/src/braintrust/otel/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
import warnings
from urllib.parse import urljoin

from braintrust.env import BraintrustEnv


INSTALL_ERR_MSG = (
"OpenTelemetry packages are not installed. "
Expand All @@ -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)
Expand Down Expand Up @@ -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).
"""
Expand All @@ -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.
Expand All @@ -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"
Expand All @@ -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
Expand All @@ -202,6 +212,41 @@ 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 = 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."
)
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,
Expand Down Expand Up @@ -252,7 +297,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.
Expand Down Expand Up @@ -340,6 +385,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):
Expand All @@ -352,6 +398,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
Expand Down
Loading
Loading