diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index f08afb668..3f03b8b6e 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -1,4 +1,5 @@ import io +import json import logging import math import os @@ -36,6 +37,33 @@ def resolve_org_name(org_name: str | None = None) -> str | None: return os.getenv("BRAINTRUST_ORG_NAME") +def resolve_project_mapping(project_name: str | None, project_id: str | None) -> tuple[str | None, str | None]: + """Fill a missing project name or ID from BRAINTRUST_PROJECT_ID_TO_NAME.""" + + raw_mappings = os.getenv("BRAINTRUST_PROJECT_ID_TO_NAME") + if not raw_mappings or (project_name is None and project_id is None): + return project_name, project_id + + try: + mappings = json.loads(raw_mappings) + except (TypeError, ValueError): + _logger.warning("Ignoring invalid BRAINTRUST_PROJECT_ID_TO_NAME: expected a JSON object of project IDs to names") + return project_name, project_id + + if not isinstance(mappings, dict) or not all( + isinstance(key, str) and isinstance(value, str) for key, value in mappings.items() + ): + _logger.warning("Ignoring invalid BRAINTRUST_PROJECT_ID_TO_NAME: expected a JSON object of project IDs to names") + return project_name, project_id + + if project_id is not None and project_name is None: + project_name = mappings.get(project_id) + elif project_name is not None and project_id is None: + project_id = next((mapped_id for mapped_id, name in mappings.items() if name == project_name), None) + + return project_name, project_id + + def parse_float(value: str) -> float | None: """Parse a finite float from a string.""" try: diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 2e7a7adf4..bb3b03e21 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -68,7 +68,7 @@ VALID_SOURCES, ) from .env import DEFAULT_APP_URL as _DEFAULT_APP_URL -from .env import BraintrustEnv, resolve_app_url, resolve_org_name +from .env import BraintrustEnv, resolve_app_url, resolve_org_name, resolve_project_mapping from .generated_types import ( AttachmentReference, AttachmentStatus, @@ -1933,6 +1933,7 @@ def init_logger( state = state or _state state.span_origin_environment = detect_environment(environment) + project, project_id = resolve_project_mapping(project, project_id) compute_metadata_args = dict(project_name=project, project_id=project_id) link_args = { diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 8e0bef213..f87bb2ae6 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -1451,6 +1451,39 @@ def test_span_link_logged_out_org_name_env_vars(with_memory_logger, monkeypatch) ) +def test_init_logger_resolves_project_from_env_mapping(with_memory_logger, monkeypatch): + monkeypatch.setenv( + "BRAINTRUST_PROJECT_ID_TO_NAME", + '{"project-id":"project-name"}', + ) + + logger = init_logger(project_id="project-id") + + assert logger._compute_metadata_args == {"project_name": "project-name", "project_id": "project-id"} + + +def test_init_logger_resolves_project_id_from_env_mapping(with_memory_logger, monkeypatch): + monkeypatch.setenv( + "BRAINTRUST_PROJECT_ID_TO_NAME", + '{"project-id":"project-name"}', + ) + + logger = init_logger(project="project-name") + + assert logger._compute_metadata_args == {"project_name": "project-name", "project_id": "project-id"} + + +def test_init_logger_explicit_project_values_override_env_mapping(with_memory_logger, monkeypatch): + monkeypatch.setenv( + "BRAINTRUST_PROJECT_ID_TO_NAME", + '{"project-id":"project-name"}', + ) + + logger = init_logger(project="explicit-name", project_id="explicit-id") + + assert logger._compute_metadata_args == {"project_name": "explicit-name", "project_id": "explicit-id"} + + def test_span_project_id_logged_in(with_memory_logger, with_simulate_login): logger = init_logger( project="test-project",