diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index 9978b53f..d7dceb8b 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -203,6 +203,8 @@ class BraintrustEnv: DEFAULT_BATCH_SIZE = EnvVar("BRAINTRUST_DEFAULT_BATCH_SIZE", EnvParser.INT) NUM_RETRIES = EnvVar("BRAINTRUST_NUM_RETRIES", EnvParser.INT) QUEUE_SIZE = EnvVar("BRAINTRUST_QUEUE_SIZE", EnvParser.INT) + ENVIRONMENT_TYPE = EnvVar("BRAINTRUST_ENVIRONMENT_TYPE", EnvParser.STRING) + ENVIRONMENT_NAME = EnvVar("BRAINTRUST_ENVIRONMENT_NAME", EnvParser.STRING) QUEUE_DROP_LOGGING_PERIOD = EnvVar("BRAINTRUST_QUEUE_DROP_LOGGING_PERIOD", EnvParser.FLOAT) FAILED_PUBLISH_PAYLOADS_DIR = EnvVar("BRAINTRUST_FAILED_PUBLISH_PAYLOADS_DIR", EnvParser.STRING) ALL_PUBLISH_PAYLOADS_DIR = EnvVar("BRAINTRUST_ALL_PUBLISH_PAYLOADS_DIR", EnvParser.STRING) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 2d42fe49..f560874b 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -89,6 +89,7 @@ from .serializable_data_class import SerializableDataClass from .span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3 from .span_identifier_v4 import SpanComponentsV4 +from .span_origin import SpanOriginEnvironment, detect_environment, merge_span_origin_context from .span_types import SpanTypeAttribute from .types import Metadata from .types._eval import ExperimentDatasetEvent @@ -446,6 +447,7 @@ def __init__(self): self.current_span: contextvars.ContextVar[Span] = contextvars.ContextVar( "braintrust_current_span", default=NOOP_SPAN ) + self.span_origin_environment: SpanOriginEnvironment | None = detect_environment() # Context manager is dynamically selected based on current environment self._context_manager = None @@ -1890,6 +1892,7 @@ def init_logger( force_login: bool = False, set_current: bool = True, state: BraintrustState | None = None, + environment: SpanOriginEnvironment | None = None, ) -> "Logger": """ Create a new logger in a specified project. If the project does not exist, it will be created. @@ -1907,6 +1910,7 @@ def init_logger( """ state = state or _state + state.span_origin_environment = detect_environment(environment) compute_metadata_args = dict(project_name=project, project_id=project_id) link_args = { @@ -4586,8 +4590,11 @@ def __init__( span_attributes=dict(**{"type": type, "name": name, **span_attributes}, exec_counter=exec_counter), created=datetime.datetime.now(datetime.timezone.utc).isoformat(), ) - if caller_location: - internal_data["context"] = caller_location + internal_data["context"] = merge_span_origin_context( + caller_location or {}, + "braintrust-python-logger", + self.state.span_origin_environment, + ) # TODO: can be simplified after `event` is typed. id = event.pop("id", None) diff --git a/py/src/braintrust/otel/__init__.py b/py/src/braintrust/otel/__init__.py index 3c0aee73..f8f5c121 100644 --- a/py/src/braintrust/otel/__init__.py +++ b/py/src/braintrust/otel/__init__.py @@ -1,9 +1,11 @@ +import json import logging import os import warnings from urllib.parse import urljoin from braintrust.env import BraintrustEnv +from braintrust.span_origin import SpanOriginEnvironment, detect_environment, merge_span_origin_context INSTALL_ERR_MSG = ( @@ -64,6 +66,32 @@ def _forward_on_ending(processor, span) -> None: on_ending(span) +class _SpanWithAttributes: + def __init__(self, span, attributes): + self._span = span + self.attributes = attributes + + def __getattr__(self, name): + return getattr(self._span, name) + + +def _with_span_origin_attributes(span, environment): + attributes = dict(getattr(span, "attributes", {}) or {}) + existing_context = {} + raw_context = attributes.get("braintrust.context_json") + if isinstance(raw_context, str) and raw_context: + try: + parsed = json.loads(raw_context) + if isinstance(parsed, dict): + existing_context = parsed + except Exception: + existing_context = {} + attributes["braintrust.context_json"] = json.dumps( + merge_span_origin_context(existing_context, "braintrust-python-otel", environment) + ) + return _SpanWithAttributes(span, attributes) + + class AISpanProcessor: """ A span processor that filters spans to only export filtered telemetry. @@ -256,6 +284,7 @@ def add_braintrust_span_processor( filter_ai_spans: bool = False, custom_filter=None, headers: dict[str, str] | None = None, + environment: SpanOriginEnvironment | None = None, ): processor = BraintrustSpanProcessor( api_key=api_key, @@ -264,6 +293,7 @@ def add_braintrust_span_processor( filter_ai_spans=filter_ai_spans, custom_filter=custom_filter, headers=headers, + environment=environment, ) tracer_provider.add_span_processor(processor) @@ -291,6 +321,7 @@ def __init__( filter_ai_spans: bool = False, custom_filter=None, headers: dict[str, str] | None = None, + environment: SpanOriginEnvironment | None = None, SpanProcessor: type | None = None, ): """ @@ -305,6 +336,7 @@ def __init__( headers: Additional headers to include in requests. SpanProcessor: Optional span processor class (BatchSpanProcessor or SimpleSpanProcessor). Defaults to BatchSpanProcessor. """ + self._environment = detect_environment(environment) # Create the exporter # Convert api_url to the full endpoint URL that OtelExporter expects exporter_url = None @@ -386,7 +418,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) + self._processor.on_end(_with_span_origin_attributes(span, self._environment)) def _on_ending(self, span): """Forward pre-end hook when the wrapped processor supports it.""" diff --git a/py/src/braintrust/span_origin.py b/py/src/braintrust/span_origin.py new file mode 100644 index 00000000..842247a8 --- /dev/null +++ b/py/src/braintrust/span_origin.py @@ -0,0 +1,123 @@ +import importlib.metadata +import os +from typing import Any, TypedDict + +from .env import BraintrustEnv + + +class SpanOriginEnvironment(TypedDict, total=False): + type: str + name: str + + +def detect_environment( + explicit: SpanOriginEnvironment | None = None, +) -> SpanOriginEnvironment | None: + if explicit: + return explicit + + env_type = BraintrustEnv.ENVIRONMENT_TYPE.get(None, use_dotenv=True) + env_name = BraintrustEnv.ENVIRONMENT_NAME.get(None, use_dotenv=True) + if env_type or env_name: + return { + **({"type": env_type} if env_type else {}), + **({"name": env_name} if env_name else {}), + } + + ci = _first_present( + { + "GITHUB_ACTIONS": "github_actions", + "GITLAB_CI": "gitlab_ci", + "CIRCLECI": "circleci", + "BUILDKITE": "buildkite", + "JENKINS_URL": "jenkins", + "JENKINS_HOME": "jenkins", + "TF_BUILD": "azure_pipelines", + "TEAMCITY_VERSION": "teamcity", + "TRAVIS": "travis", + "BITBUCKET_BUILD_NUMBER": "bitbucket", + } + ) + if ci: + return {"type": "ci", "name": ci} + if os.environ.get("CI"): + return {"type": "ci", "name": "ci"} + + server = _first_present( + { + "VERCEL": "vercel", + "NETLIFY": "netlify", + } + ) + if server: + return {"type": "server", "name": server} + + if os.environ.get("ECS_CONTAINER_METADATA_URI") or os.environ.get("ECS_CONTAINER_METADATA_URI_V4"): + return {"type": "server", "name": "ecs"} + aws_execution_env = os.environ.get("AWS_EXECUTION_ENV") + if aws_execution_env: + if aws_execution_env.startswith("AWS_ECS_"): + return {"type": "server", "name": "ecs"} + if aws_execution_env.startswith("AWS_Lambda_"): + return {"type": "server", "name": "aws_lambda"} + if os.environ.get("AWS_LAMBDA_FUNCTION_NAME"): + return {"type": "server", "name": "aws_lambda"} + + server = _first_present( + { + "K_SERVICE": "cloud_run", + "FUNCTION_TARGET": "gcp_functions", + "KUBERNETES_SERVICE_HOST": "kubernetes", + "DYNO": "heroku", + "FLY_APP_NAME": "fly", + "RAILWAY_ENVIRONMENT": "railway", + "RENDER_SERVICE_NAME": "render", + } + ) + if server: + return {"type": "server", "name": server} + + return _deployment_mode(os.environ.get("PYTHON_ENV")) + + +def merge_span_origin_context( + context: dict[str, Any] | None, + instrumentation_name: str, + environment: SpanOriginEnvironment | None, +) -> dict[str, Any]: + merged = dict(context or {}) + span_origin = dict(merged.get("span_origin") or {}) + span_origin = { + "name": "braintrust.sdk.python", + "version": _sdk_version(), + "instrumentation": {"name": instrumentation_name}, + **({"environment": environment} if environment else {}), + **span_origin, + } + merged["span_origin"] = span_origin + return merged + + +def _sdk_version() -> str: + try: + return importlib.metadata.version("braintrust") + except importlib.metadata.PackageNotFoundError: + return "unknown" + + +def _first_present(mapping: dict[str, str]) -> str | None: + for key, name in mapping.items(): + if os.environ.get(key): + return name + return None + + +def _deployment_mode(value: str | None) -> SpanOriginEnvironment | None: + if not value: + return None + normalized = value.lower() + if normalized in ("production", "staging"): + return {"type": "server", "name": normalized} + if normalized in ("development", "local"): + return {"type": "local", "name": normalized} + return None diff --git a/py/src/braintrust/test_otel.py b/py/src/braintrust/test_otel.py index 65d116ab..73e5b1a1 100644 --- a/py/src/braintrust/test_otel.py +++ b/py/src/braintrust/test_otel.py @@ -1,4 +1,5 @@ # pylint: disable=not-context-manager +import json import sys import pytest @@ -112,6 +113,87 @@ def test_otel_exporter_uses_env_braintrust_api_key(tmp_path): assert exporter._headers["Authorization"] == "Bearer file-api-key" +def test_braintrust_span_processor_merges_span_origin_with_context_json_set_after_start(): + 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 + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + memory_exporter = InMemorySpanExporter() + provider = TracerProvider() + processor = BraintrustSpanProcessor(api_key="test-api-key", parent="project_name:test") + processor._processor = SimpleSpanProcessor(memory_exporter) + provider.add_span_processor(processor) + tracer = provider.get_tracer("test_tracer") + + try: + with tracer.start_as_current_span("late-context") as span: + span.set_attribute("braintrust.context_json", json.dumps({"metadata": {"source": "late-attribute"}})) + + provider.force_flush() + spans = memory_exporter.get_finished_spans() + assert len(spans) == 1 + context = json.loads(spans[0].attributes["braintrust.context_json"]) + assert context["metadata"]["source"] == "late-attribute" + assert context["span_origin"]["name"] == "braintrust.sdk.python" + assert context["span_origin"]["instrumentation"]["name"] == "braintrust-python-otel" + assert "braintrust.environment.type" not in spans[0].attributes + assert "braintrust.environment.name" not in spans[0].attributes + finally: + provider.shutdown() + + +def test_detect_environment_classifies_aws_ecs_before_lambda(monkeypatch): + from braintrust.span_origin import detect_environment + + monkeypatch.delenv("GITHUB_ACTIONS", raising=False) + monkeypatch.delenv("GITLAB_CI", raising=False) + monkeypatch.delenv("CIRCLECI", raising=False) + monkeypatch.delenv("BUILDKITE", raising=False) + monkeypatch.delenv("CI", raising=False) + monkeypatch.delenv("BRAINTRUST_ENVIRONMENT_TYPE", raising=False) + monkeypatch.delenv("BRAINTRUST_ENVIRONMENT_NAME", raising=False) + monkeypatch.delenv("ECS_CONTAINER_METADATA_URI", raising=False) + monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False) + monkeypatch.setenv("AWS_EXECUTION_ENV", "AWS_ECS_FARGATE") + + assert detect_environment() == {"type": "server", "name": "ecs"} + + +def test_detect_environment_preserves_explicit_name_without_type(monkeypatch): + from braintrust.span_origin import detect_environment + + monkeypatch.delenv("GITHUB_ACTIONS", raising=False) + monkeypatch.delenv("GITLAB_CI", raising=False) + monkeypatch.delenv("CIRCLECI", raising=False) + monkeypatch.delenv("BUILDKITE", raising=False) + monkeypatch.delenv("CI", raising=False) + monkeypatch.delenv("BRAINTRUST_ENVIRONMENT_TYPE", raising=False) + monkeypatch.setenv("BRAINTRUST_ENVIRONMENT_NAME", "staging") + + assert detect_environment() == {"name": "staging"} + + +def test_detect_environment_classifies_lambda_when_lambda_specific(monkeypatch): + from braintrust.span_origin import detect_environment + + monkeypatch.delenv("GITHUB_ACTIONS", raising=False) + monkeypatch.delenv("GITLAB_CI", raising=False) + monkeypatch.delenv("CIRCLECI", raising=False) + monkeypatch.delenv("BUILDKITE", raising=False) + monkeypatch.delenv("CI", raising=False) + monkeypatch.delenv("BRAINTRUST_ENVIRONMENT_TYPE", raising=False) + monkeypatch.delenv("BRAINTRUST_ENVIRONMENT_NAME", raising=False) + monkeypatch.delenv("ECS_CONTAINER_METADATA_URI", raising=False) + monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False) + monkeypatch.setenv("AWS_EXECUTION_ENV", "AWS_Lambda_python3.12") + + assert detect_environment() == {"type": "server", "name": "aws_lambda"} + + 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")