feat: add _create_in_project_group to logger options

py/src/braintrust/framework2.py

@@ -1,7 +1,7 @@
 import dataclasses
 import json
 from collections.abc import Callable, Mapping, Sequence
-from typing import Any, cast, overload
+from typing import TYPE_CHECKING, Any, overload
 
 import slugify
 from braintrust.logger import _internal_get_global_state, api_conn, login
@@ -21,6 +21,10 @@
 from .util import eprint
 
 
+if TYPE_CHECKING:
+    from braintrust.logger import _CreateProjectRequest
+
+
 class ProjectIdCache:
     def __init__(self):
         self._cache: dict[Project, str] = {}
@@ -29,10 +33,10 @@ def __init__(self):
     def get_by_name(self, project_name: str, project_group_name: str | None = None) -> str:
         if project_name not in self._name_cache:
             state = _internal_get_global_state()
-            body: dict[str, Any] = {"name": project_name, "org_name": state.org_name}
+            body: _CreateProjectRequest = {"name": project_name, "org_name": state.org_name}
             if project_group_name is not None:
                 body["project_group_name"] = project_group_name
-            project = state.api_client().projects.post_project(body=cast(Any, body))
+            project = state.api_client().projects.post_project(body=body)
             self._name_cache[project_name] = project["id"]
         return self._name_cache[project_name]
 

py/src/braintrust/logger.py

@@ -28,6 +28,7 @@
 from multiprocessing import cpu_count
 from types import TracebackType
 from typing import (
+    TYPE_CHECKING,
     Any,
     Generic,
     Literal,
@@ -1876,18 +1877,28 @@ def compute_metadata():
     )
 
 
+if TYPE_CHECKING:
+    from .api._generated.models.projects import CreateProject
+
+    # Keep experimental fields local until they are part of the pinned OpenAPI schema.
+    class _CreateProjectRequest(CreateProject, total=False):
+        project_group_name: str
+
+
 def _compute_logger_metadata(
     project_name: str | None = None,
     project_id: str | None = None,
     state: BraintrustState | None = None,
+    project_group_name: str | None = None,
 ):
     state = state or _state
     state.login()
     org_id = state.org_id
     if project_id is None:
-        response = state.api_client().projects.post_project(
-            body={"name": project_name or GLOBAL_PROJECT, "org_name": state.org_name}
-        )
+        body: _CreateProjectRequest = {"name": project_name or GLOBAL_PROJECT, "org_name": state.org_name}
+        if project_group_name is not None:
+            body["project_group_name"] = project_group_name
+        response = state.api_client().projects.post_project(body=body)
         return OrgProjectMetadata(
             org_id=org_id,
             project=ObjectMetadata(id=response["id"], name=response["name"], full_info=dict(response)),
@@ -1915,6 +1926,7 @@ def init_logger(
     set_current: bool = True,
     state: BraintrustState | None = None,
     environment: SpanOriginEnvironment | None = None,
+    _create_in_project_group: str | None = None,
 ) -> "Logger":
     """
     Create a new logger in a specified project. If the project does not exist, it will be created.
@@ -1928,12 +1940,15 @@ def init_logger(
     :param org_name: (Optional) The name of a specific organization to connect to. This is useful if you belong to multiple.
     :param force_login: Login again, even if you have already logged in (by default, the logger will not login if you are already logged in)
     :param set_current: If true (the default), set the global current-experiment to the newly-created one.
+    :param _create_in_project_group: Experimental: the name of an existing project group to create the project in. Existing projects must already belong to the group, otherwise registration fails with a 409. Ignored when project_id is provided.
     :returns: The newly created Logger.
     """
 
     state = state or _state
     state.span_origin_environment = detect_environment(environment)
     compute_metadata_args = dict(project_name=project, project_id=project_id)
+    if project_id is None and _create_in_project_group is not None:
+        compute_metadata_args["project_group_name"] = _create_in_project_group
 
     link_args = {
         "app_url": app_url,
@@ -2662,10 +2677,14 @@ def _current_braintrust_parent(state: BraintrustState | None = None) -> str | No
             if components.object_id:
                 return f"project_id:{components.object_id}"
             meta = components.compute_object_metadata_args or {}
+            # A `project_name:` parent would register the project outside its group.
+            if meta.get("project_group_name"):
+                return f"project_id:{logger.id}"
             name = meta.get("project_name")
             if name:
                 return f"project_name:{name}"
-        except Exception:
+        except Exception as e:
+            logging.warning(f"Failed to resolve braintrust.parent from the current logger: {e}")
             return None
 
     return None
@@ -5207,6 +5226,12 @@ def _get_otel_parent(self):
         if parent_type == SpanObjectTypeV3.PROJECT_LOGS:
             _id = info.get("id")
             _name = info.get("name")
+            # A `project_name:` parent would register the project outside its group.
+            if not _id and (self.parent_compute_object_metadata_args or {}).get("project_group_name"):
+                try:
+                    _id = self.parent_object_id.get()
+                except Exception:
+                    return None
             if _id:
                 return f"project_id:{_id}"
             elif _name:

py/src/braintrust/test_helpers.py

@@ -104,7 +104,7 @@ def init_test_logger(project_name: str):
     l._lazy_metadata = lazy_metadata  # Skip actual login by setting fake metadata directly
 
     # Replace the global _compute_logger_metadata function with a resolved LazyValue
-    def fake_compute_logger_metadata(project_name=None, project_id=None, state=None):
+    def fake_compute_logger_metadata(project_name=None, project_id=None, state=None, project_group_name=None):
         if project_id:
             project_metadata = ObjectMetadata(id=project_id, name=project_name, full_info=dict())
         else:

py/src/braintrust/test_logger.py

@@ -37,6 +37,7 @@
     BraintrustState,
     RemoteEvalParameters,
     _check_org_info,
+    _compute_logger_metadata,
     _extract_attachments,
     parent_context,
     render_message,
@@ -2781,6 +2782,123 @@ def test_experiment_fetch_uses_generated_paginated_fetch(self):
         mock_state.api_conn.assert_not_called()
 
 
+class TestLoggerProjectGroup(TestCase):
+    def setUp(self):
+        # init_test_logger permanently replaces the module-level function with a fake, so restore the real one.
+        metadata_patch = patch.object(logger, "_compute_logger_metadata", _compute_logger_metadata)
+        metadata_patch.start()
+        self.addCleanup(metadata_patch.stop)
+        self.state = MagicMock()
+        self.state.org_name = "test-org"
+        self.state.context_manager = BraintrustState().context_manager
+        self.state.id_generator = logger.id_gen.get_id_generator()
+        self.projects = self.state.api_client.return_value.projects
+        self.projects.post_project.return_value = {"id": "project-id", "name": "project"}
+        self.projects.get_project_id.return_value = {"id": "project-id", "name": "project"}
+
+    def test_registration_body_and_lazy_resolution(self):
+        for project in ["project", None]:
+            for group in ["my-group", None]:
+                with self.subTest(project=project, group=group):
+                    self.projects.reset_mock()
+                    instance = init_logger(
+                        project=project,
+                        _create_in_project_group=group,
+                        state=self.state,
+                        set_current=False,
+                    )
+                    self.projects.post_project.assert_not_called()
+                    self.assertEqual(instance.id, "project-id")
+                    # Second access must not re-register.
+                    self.assertEqual(instance.id, "project-id")
+                    body = {"name": project or "Global", "org_name": "test-org"}
+                    if group is not None:
+                        body["project_group_name"] = group
+                    self.projects.post_project.assert_called_once_with(body=body)
+
+    def test_project_id_takes_precedence(self):
+        for project in ["project", None]:
+            with self.subTest(project=project):
+                self.projects.reset_mock()
+                instance = init_logger(
+                    project=project,
+                    project_id="project-id",
+                    _create_in_project_group="my-group",
+                    state=self.state,
+                    set_current=False,
+                )
+                self.assertNotIn("project_group_name", instance._compute_metadata_args)
+                self.assertEqual(instance.id, "project-id")
+                self.projects.post_project.assert_not_called()
+                if project is None:
+                    self.projects.get_project_id.assert_called_once_with("project-id")
+                else:
+                    self.projects.get_project_id.assert_not_called()
+
+    def test_group_survives_unresolved_exports(self):
+        instance = init_logger(
+            project="project",
+            _create_in_project_group="my-group",
+            state=self.state,
+            set_current=False,
+        )
+        exports = [instance.export(), instance.start_span(name="parent").export()]
+        # Generated by JS's SpanComponentsV4 with the same logger options.
+        js_export = "BAIAeyJjb21wdXRlX29iamVjdF9tZXRhZGF0YV9hcmdzIjp7InByb2plY3RfbmFtZSI6InByb2plY3QiLCJwcm9qZWN0X2dyb3VwX25hbWUiOiJteS1ncm91cCJ9fQ=="
+        exports.append(js_export)
+        self.projects.post_project.assert_not_called()
+        for exported in exports:
+            with self.subTest(exported=exported):
+                self.projects.reset_mock()
+                components = logger.SpanComponentsV4.from_str(exported)
+                self.assertEqual(
+                    components.compute_object_metadata_args,
+                    {
+                        "project_name": "project",
+                        "project_group_name": "my-group",
+                        **({"project_id": None} if exported != js_export else {}),
+                    },
+                )
+                with patch.object(logger, "_state", self.state):
+                    self.assertEqual(logger.span_components_to_object_id(components), "project-id")
+                self.projects.post_project.assert_called_once_with(
+                    body={"name": "project", "org_name": "test-org", "project_group_name": "my-group"}
+                )
+
+    def test_otel_parent_resolves_project_id_when_grouped(self):
+        for group, expected_parent in [("my-group", "project_id:project-id"), (None, "project_name:project")]:
+            with self.subTest(group=group):
+                self.projects.reset_mock()
+                # Fresh loggers so each check starts from an unresolved project id.
+                span_logger, current = (
+                    init_logger(project="project", _create_in_project_group=group, state=self.state, set_current=False)
+                    for _ in range(2)
+                )
+                self.assertEqual(span_logger.start_span(name="parent")._get_otel_parent(), expected_parent)
+                with (
+                    patch.object(logger, "current_experiment", return_value=None),
+                    patch.object(logger, "current_logger", return_value=current),
+                ):
+                    self.assertEqual(logger._current_braintrust_parent(), expected_parent)
+                expected_calls = 0 if group is None else 2
+                self.assertEqual(self.projects.post_project.call_count, expected_calls)
+
+    def test_otel_parent_tolerates_failed_group_registration(self):
+        self.projects.post_project.side_effect = Exception("409 Conflict")
+        instance = init_logger(
+            project="project", _create_in_project_group="my-group", state=self.state, set_current=False
+        )
+        # Must not fall back to a `project_name:` parent, which would register the project outside its group.
+        self.assertIsNone(instance.start_span(name="parent")._get_otel_parent())
+        with (
+            patch.object(logger, "current_experiment", return_value=None),
+            patch.object(logger, "current_logger", return_value=instance),
+            self.assertLogs(level="WARNING") as logs,
+        ):
+            self.assertIsNone(logger._current_braintrust_parent())
+        self.assertIn("409 Conflict", "\n".join(logs.output))
+
+
 class TestProjectGeneratedAPI(TestCase):
     def test_lazy_project_uses_generated_registration_and_lookup(self):
         mock_state = MagicMock()