diff --git a/py/src/braintrust/framework2.py b/py/src/braintrust/framework2.py index 51fb23630..7540659c0 100644 --- a/py/src/braintrust/framework2.py +++ b/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] diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 2e7a7adf4..21155a2dd 100644 --- a/py/src/braintrust/logger.py +++ b/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: diff --git a/py/src/braintrust/test_helpers.py b/py/src/braintrust/test_helpers.py index 585454bc3..8b0086b77 100644 --- a/py/src/braintrust/test_helpers.py +++ b/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: diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 8e0bef213..597ab4f77 100644 --- a/py/src/braintrust/test_logger.py +++ b/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()