feat: add tags support to SDK builder classes (#33)

Author: AbhiPrasad

py/src/braintrust/cli/push.py

@@ -251,6 +251,8 @@ def _collect_function_function_defs(
         }
         if f.metadata is not None:
             j["metadata"] = f.metadata
+        if f.tags is not None:
+            j["tags"] = list(f.tags)
         if f.parameters is None:
             raise ValueError(f"Function {f.name} has no supplied parameters")
         j["function_schema"] = {

py/src/braintrust/cli/test_push.py

@@ -0,0 +1,101 @@
+"""Tests for push command serialization."""
+
+import pytest
+
+pydantic = pytest.importorskip("pydantic")
+
+from ..framework2 import (
+    global_,
+    projects,
+)
+from .push import _collect_function_function_defs
+
+
+class ToolInput(pydantic.BaseModel):
+    value: int
+
+
+@pytest.fixture(autouse=True)
+def clear_global_state():
+    global_.functions.clear()
+    global_.prompts.clear()
+    yield
+    global_.functions.clear()
+    global_.prompts.clear()
+
+
+class TestPushMetadata:
+    """Tests for metadata in push command serialization."""
+
+    def test_collect_function_function_defs_includes_metadata(self, mock_project_ids):
+        project = projects.create("test-project")
+        metadata = {"version": "1.0", "author": "test"}
+
+        tool = project.tools.create(
+            handler=lambda x: x,
+            name="test-tool",
+            parameters=ToolInput,
+            metadata=metadata,
+        )
+        global_.functions.append(tool)
+
+        functions = []
+        _collect_function_function_defs(mock_project_ids, functions, "bundle-123", "error")
+
+        assert len(functions) == 1
+        assert functions[0]["metadata"] == metadata
+        assert functions[0]["name"] == "test-tool"
+
+    def test_collect_function_function_defs_excludes_metadata_when_none(self, mock_project_ids):
+        project = projects.create("test-project")
+
+        tool = project.tools.create(
+            handler=lambda x: x,
+            name="test-tool",
+            parameters=ToolInput,
+        )
+        global_.functions.append(tool)
+
+        functions = []
+        _collect_function_function_defs(mock_project_ids, functions, "bundle-123", "error")
+
+        assert len(functions) == 1
+        assert "metadata" not in functions[0]
+
+
+class TestPushTags:
+    """Tests for tags in push command serialization."""
+
+    def test_collect_function_function_defs_includes_tags(self, mock_project_ids):
+        project = projects.create("test-project")
+        tags = ["production", "v1"]
+
+        tool = project.tools.create(
+            handler=lambda x: x,
+            name="test-tool",
+            parameters=ToolInput,
+            tags=tags,
+        )
+        global_.functions.append(tool)
+
+        functions = []
+        _collect_function_function_defs(mock_project_ids, functions, "bundle-123", "error")
+
+        assert len(functions) == 1
+        assert functions[0]["tags"] == ["production", "v1"]
+
+    def test_collect_function_function_defs_excludes_tags_when_none(self, mock_project_ids):
+        project = projects.create("test-project")
+
+        tool = project.tools.create(
+            handler=lambda x: x,
+            name="test-tool",
+            parameters=ToolInput,
+        )
+        global_.functions.append(tool)
+
+        functions = []
+        _collect_function_function_defs(mock_project_ids, functions, "bundle-123", "error")
+
+        assert len(functions) == 1
+        assert "tags" not in functions[0]

py/src/braintrust/conftest.py

@@ -1,6 +1,8 @@
 import os
+from unittest.mock import MagicMock
 
 import pytest
+from braintrust.framework2 import ProjectIdCache
 
 
 def _patch_vcr_aiohttp_stubs():
@@ -107,6 +109,13 @@ async def _patched_record_response(cassette, vcr_request, response):
 _patch_vcr_aiohttp_stubs()
 
 
+@pytest.fixture
+def mock_project_ids():
+    mock = MagicMock(spec=ProjectIdCache)
+    mock.get.return_value = "project-123"
+    return mock
+
+
 @pytest.fixture(autouse=True)
 def override_app_url_for_tests():
     """

py/src/braintrust/framework2.py

@@ -1,6 +1,6 @@
 import dataclasses
 import json
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import Any, overload
 
 import slugify
@@ -53,6 +53,7 @@ class CodeFunction:
     returns: Any
     if_exists: IfExists | None
     metadata: dict[str, Any] | None = None
+    tags: Sequence[str] | None = None
 
 
 @dataclasses.dataclass
@@ -69,6 +70,7 @@ class CodePrompt:
     id: str | None
     if_exists: IfExists | None
     metadata: dict[str, Any] | None = None
+    tags: Sequence[str] | None = None
 
     def to_function_definition(self, if_exists: IfExists | None, project_ids: ProjectIdCache) -> dict[str, Any]:
         prompt_data = self.prompt
@@ -102,6 +104,8 @@ def to_function_definition(self, if_exists: IfExists | None, project_ids: Projec
             j["function_type"] = self.function_type
         if self.metadata is not None:
             j["metadata"] = self.metadata
+        if self.tags is not None:
+            j["tags"] = list(self.tags)
 
         return j
 
@@ -124,6 +128,7 @@ def create(
         returns: Any = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
     ) -> CodeFunction:
         """Creates a tool.
 
@@ -136,6 +141,7 @@ def create(
             returns: The tool's output schema, as a Pydantic model.
             if_exists: What to do if the tool already exists.
             metadata: Custom metadata to attach to the tool.
+            tags: A list of tags for the tool.
 
         Returns:
             A handle to the created tool, that can be used in a prompt.
@@ -160,6 +166,7 @@ def create(
             returns=returns,
             if_exists=if_exists,
             metadata=metadata,
+            tags=tags,
         )
         self.project.add_code_function(f)
         return f
@@ -186,6 +193,7 @@ def create(
         tools: list[CodeFunction | SavedFunctionId | ToolFunctionDefinition] | None = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
     ) -> CodePrompt: ...
 
     @overload  # messages only, no prompt
@@ -202,6 +210,7 @@ def create(
         tools: list[CodeFunction | SavedFunctionId | ToolFunctionDefinition] | None = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
     ) -> CodePrompt: ...
 
     def create(
@@ -218,6 +227,7 @@ def create(
         tools: list[CodeFunction | SavedFunctionId | ToolFunctionDefinition] | None = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
     ):
         """Creates a prompt.
 
@@ -233,6 +243,7 @@ def create(
             tools: The tools to use for the prompt.
             if_exists: What to do if the prompt already exists.
             metadata: Custom metadata to attach to the prompt.
+            tags: A list of tags for the prompt.
         """
         self._task_counter += 1
         if not name:
@@ -282,6 +293,7 @@ def create(
             id=id,
             if_exists=if_exists,
             metadata=metadata,
+            tags=tags,
         )
         self.project.add_prompt(p)
         return p
@@ -304,6 +316,7 @@ def create(
         description: str | None = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
         handler: Callable[..., Any],
         parameters: Any,
         returns: Any = None,
@@ -319,6 +332,7 @@ def create(
         description: str | None = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
         prompt: str,
         model: str,
         params: ModelParams | None = None,
@@ -336,6 +350,7 @@ def create(
         description: str | None = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
         messages: list[ChatCompletionMessageParam],
         model: str,
         params: ModelParams | None = None,
@@ -351,6 +366,7 @@ def create(
         description: str | None = None,
         if_exists: IfExists | None = None,
         metadata: dict[str, Any] | None = None,
+        tags: Sequence[str] | None = None,
         # Code scorer params.
         handler: Callable[..., Any] | None = None,
         parameters: Any = None,
@@ -371,6 +387,7 @@ def create(
             description: The description of the scorer.
             if_exists: What to do if the scorer already exists.
             metadata: Custom metadata to attach to the scorer.
+            tags: A list of tags for the scorer.
 
             The remaining args are mutually exclusive; that is,
             the function will only accept args from one of the following overloads.
@@ -410,6 +427,7 @@ def create(
                 returns=returns,
                 if_exists=if_exists,
                 metadata=metadata,
+                tags=tags,
             )
             self.project.add_code_function(f)
             return f
@@ -449,6 +467,7 @@ def create(
                 id=None,
                 if_exists=if_exists,
                 metadata=metadata,
+                tags=tags,
             )
             self.project.add_prompt(p)
             return p