diff --git a/py/src/braintrust/framework2.py b/py/src/braintrust/framework2.py index 781dc2ca..61652c87 100644 --- a/py/src/braintrust/framework2.py +++ b/py/src/braintrust/framework2.py @@ -128,6 +128,7 @@ class CodeParameters: schema: EvalParameters if_exists: IfExists | None metadata: Mapping[str, Any] | None = None + tags: Sequence[str] | None = None def to_function_definition(self, if_exists: IfExists | None, project_ids: ProjectIdCache) -> dict[str, Any]: schema = parameters_to_json_schema(self.schema) @@ -146,6 +147,8 @@ def to_function_definition(self, if_exists: IfExists | None, project_ids: Projec } if self.metadata is not None: j["metadata"] = self.metadata + if self.tags is not None: + j["tags"] = list(self.tags) return j @@ -353,6 +356,7 @@ def create( description: str | None = None, if_exists: IfExists | None = None, metadata: Mapping[str, Any] | None = None, + tags: Sequence[str] | None = None, ) -> EvalParameters: if slug is None or len(slug) == 0: slug = slugify.slugify(name) @@ -365,6 +369,7 @@ def create( schema=schema, if_exists=if_exists, metadata=metadata, + tags=tags, ) self.project.add_parameters(parameters) return schema diff --git a/py/src/braintrust/parameters.py b/py/src/braintrust/parameters.py index c8b71ddc..13c84a2c 100644 --- a/py/src/braintrust/parameters.py +++ b/py/src/braintrust/parameters.py @@ -68,10 +68,12 @@ class RemoteEvalParameters(SerializableDataClass): version: str | None schema: ParametersSchema data: dict[str, Any] + tags: list[str] | None = None @classmethod def from_function_row(cls, row: dict[str, Any]) -> "RemoteEvalParameters": function_data = row.get("function_data") or {} + tags = row.get("tags") return cls( id=row.get("id"), project_id=row.get("project_id"), @@ -80,6 +82,7 @@ def from_function_row(cls, row: dict[str, Any]) -> "RemoteEvalParameters": version=row.get("_xact_id"), schema=function_data.get("__schema") or {}, data=function_data.get("data") or {}, + tags=list(tags) if tags is not None else None, ) def validate(self, data: Any) -> bool: diff --git a/py/src/braintrust/test_framework2.py b/py/src/braintrust/test_framework2.py index d8745a4a..da874f02 100644 --- a/py/src/braintrust/test_framework2.py +++ b/py/src/braintrust/test_framework2.py @@ -331,3 +331,34 @@ class RequiredParam(BaseModel): } assert func_def["function_data"]["__schema"]["properties"]["model"]["x-bt-type"] == "model" assert "required_text" not in func_def["function_data"]["data"] + + +def test_project_parameters_create_includes_tags(mock_project_ids): + project = projects.create("test-project") + project.parameters.create( + name="test-parameters", + schema={"model": {"type": "model", "default": "gpt-5-mini"}}, + tags=["production", "v1"], + ) + + parameters = project._publishable_parameters + assert len(parameters) == 1 + + func_def = parameters[0].to_function_definition(None, mock_project_ids) + + assert func_def["tags"] == ["production", "v1"] + + +def test_project_parameters_create_excludes_tags_when_none(mock_project_ids): + project = projects.create("test-project") + project.parameters.create( + name="test-parameters", + schema={"model": {"type": "model", "default": "gpt-5-mini"}}, + ) + + parameters = project._publishable_parameters + assert len(parameters) == 1 + + func_def = parameters[0].to_function_definition(None, mock_project_ids) + + assert "tags" not in func_def diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 3fc5c833..58dace90 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -331,6 +331,7 @@ def test_load_parameters_returns_remote_object(self): "name": "Saved parameters", "slug": "saved-parameters", "_xact_id": "v1", + "tags": ["production", "v1"], "function_data": { "type": "parameters", "data": {"prefix": "hello"}, @@ -353,6 +354,7 @@ def test_load_parameters_returns_remote_object(self): assert isinstance(parameters, RemoteEvalParameters) assert parameters.id == "params-123" assert parameters.version == "v1" + assert parameters.tags == ["production", "v1"] assert parameters.data == {"prefix": "hello"} assert ( logger._state._parameters_cache.get( diff --git a/py/src/braintrust/type_tests/test_parameters_schema.py b/py/src/braintrust/type_tests/test_parameters_schema.py index 920f2816..452c4e23 100644 --- a/py/src/braintrust/type_tests/test_parameters_schema.py +++ b/py/src/braintrust/type_tests/test_parameters_schema.py @@ -26,6 +26,6 @@ def test_parameters_create_accepts_discriminated_schema_entries() -> None: }, } - created = project.parameters.create(name="test-parameters", schema=schema) + created = project.parameters.create(name="test-parameters", schema=schema, tags=["production", "v1"]) assert created is schema