Skip to content
Closed
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
28 changes: 28 additions & 0 deletions py/src/braintrust/env.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import io
import json
import logging
import math
import os
Expand Down Expand Up @@ -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:
Expand Down
3 changes: 2 additions & 1 deletion py/src/braintrust/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 = {
Expand Down
33 changes: 33 additions & 0 deletions py/src/braintrust/test_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading