Skip to content

Commit 4033ca8

Browse files
elijah0528claudelforst
authored
feat: add project_group_name to projects.create() (#835)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: lforst <8118419+lforst@users.noreply.github.com>
1 parent b4a204c commit 4033ca8

2 files changed

Lines changed: 48 additions & 7 deletions

File tree

‎py/src/braintrust/framework2.py‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import dataclasses
22
import json
33
from collections.abc import Callable, Mapping, Sequence
4-
from typing import Any, overload
4+
from typing import Any, cast, overload
55

66
import slugify
77
from braintrust.logger import _internal_get_global_state, api_conn, login
@@ -26,16 +26,19 @@ def __init__(self):
2626
self._cache: dict[Project, str] = {}
2727
self._name_cache: dict[str, str] = {}
2828

29-
def get_by_name(self, project_name: str) -> str:
29+
def get_by_name(self, project_name: str, project_group_name: str | None = None) -> str:
3030
if project_name not in self._name_cache:
3131
state = _internal_get_global_state()
32-
project = state.api_client().projects.post_project(body={"name": project_name, "org_name": state.org_name})
32+
body: dict[str, Any] = {"name": project_name, "org_name": state.org_name}
33+
if project_group_name is not None:
34+
body["project_group_name"] = project_group_name
35+
project = state.api_client().projects.post_project(body=cast(Any, body))
3336
self._name_cache[project_name] = project["id"]
3437
return self._name_cache[project_name]
3538

3639
def get(self, project: "Project") -> str:
3740
if project not in self._cache:
38-
self._cache[project] = self.get_by_name(project.name)
41+
self._cache[project] = self.get_by_name(project.name, project.project_group_name)
3942
return self._cache[project]
4043

4144

@@ -607,8 +610,9 @@ def create(
607610
class Project:
608611
"""A handle to a Braintrust project."""
609612

610-
def __init__(self, name: str):
613+
def __init__(self, name: str, project_group_name: str | None = None):
611614
self.name = name
615+
self.project_group_name = project_group_name
612616
self.tools = ToolBuilder(self)
613617
self.prompts = PromptBuilder(self)
614618
self.parameters = ParametersBuilder(self)
@@ -659,8 +663,13 @@ def publish(self):
659663
class ProjectBuilder:
660664
"""Creates handles to Braintrust projects."""
661665

662-
def create(self, name: str) -> Project:
663-
return Project(name)
666+
def create(self, name: str, project_group_name: str | None = None) -> Project:
667+
"""Create a handle to a Braintrust project.
668+
669+
:param name: The name of the project.
670+
:param project_group_name: (Optional) If specified, creates the project inside the project group with this name when the project does not already exist. Requires permission to create projects in that group.
671+
"""
672+
return Project(name, project_group_name=project_group_name)
664673

665674

666675
projects = ProjectBuilder()

‎py/src/braintrust/test_framework2.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,38 @@ def test_project_id_cache_uses_generated_project_registration():
3030
mock_state.app_conn.assert_not_called()
3131

3232

33+
def test_project_id_cache_creates_the_project_in_its_project_group():
34+
mock_state = MagicMock()
35+
mock_state.org_name = "test-org"
36+
mock_state.api_client.return_value.projects.post_project.return_value = {
37+
"id": "generated-project-id",
38+
"name": "test-project",
39+
}
40+
project = projects.create("test-project", project_group_name="my-group")
41+
with patch("braintrust.logger._state", mock_state):
42+
project_id = ProjectIdCache().get(project)
43+
44+
assert project_id == "generated-project-id"
45+
mock_state.api_client.return_value.projects.post_project.assert_called_once_with(
46+
body={"name": "test-project", "org_name": "test-org", "project_group_name": "my-group"}
47+
)
48+
49+
50+
def test_project_id_cache_omits_project_group_name_when_unspecified():
51+
mock_state = MagicMock()
52+
mock_state.org_name = "test-org"
53+
mock_state.api_client.return_value.projects.post_project.return_value = {
54+
"id": "generated-project-id",
55+
"name": "test-project",
56+
}
57+
with patch("braintrust.logger._state", mock_state):
58+
ProjectIdCache().get(projects.create("test-project"))
59+
60+
mock_state.api_client.return_value.projects.post_project.assert_called_once_with(
61+
body={"name": "test-project", "org_name": "test-org"}
62+
)
63+
64+
3365
class TestCodeFunctionMetadata:
3466
"""Tests for CodeFunction metadata support."""
3567

0 commit comments

Comments
 (0)