Skip to content

Commit af844ca

Browse files
authored
feat: Add classifier function push builder (#497)
## Summary - add a project-level classifier builder for code-backed saved functions - register classifiers through the existing function push registry with `type_="classifier"` - add focused coverage for normal and lazy-load function registry registration ## Validation - `UV_CACHE_DIR=/private/tmp/uv-cache UV_PYTHON_INSTALL_DIR=/private/tmp/uv-python uv run --group test pytest src/braintrust/test_framework2.py -k classifier` - `UV_CACHE_DIR=/private/tmp/uv-cache UV_PYTHON_INSTALL_DIR=/private/tmp/uv-python uv run --group test python -m py_compile src/braintrust/framework2.py src/braintrust/test_framework2.py`
1 parent 25de685 commit af844ca

2 files changed

Lines changed: 106 additions & 1 deletion

File tree

py/src/braintrust/framework2.py

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,65 @@ def create(
544544
return p
545545

546546

547+
class ClassifierBuilder:
548+
"""Builder to create a classifier in Braintrust."""
549+
550+
def __init__(self, project: "Project"):
551+
self.project = project
552+
self._task_counter = 0
553+
554+
def create(
555+
self,
556+
*,
557+
handler: Callable[..., Any],
558+
name: str | None = None,
559+
slug: str | None = None,
560+
description: str | None = None,
561+
parameters: Any = None,
562+
returns: Any = None,
563+
if_exists: IfExists | None = None,
564+
metadata: Metadata | None = None,
565+
tags: Sequence[str] | None = None,
566+
) -> CodeFunction:
567+
"""Creates a classifier.
568+
569+
Args:
570+
handler: The function that is called when the classifier is used.
571+
name: The name of the classifier.
572+
slug: A unique identifier for the classifier.
573+
description: The description of the classifier.
574+
parameters: The classifier's input schema, as a Pydantic model.
575+
returns: The classifier's output schema, as a Pydantic model.
576+
if_exists: What to do if the classifier already exists.
577+
metadata: Custom metadata to attach to the classifier.
578+
tags: A list of tags for the classifier.
579+
"""
580+
self._task_counter += 1
581+
if name is None or len(name) == 0:
582+
if handler.__name__ and handler.__name__ != "<lambda>":
583+
name = handler.__name__
584+
else:
585+
name = f"Classifier {self._task_counter}"
586+
if slug is None or len(slug) == 0:
587+
slug = slugify.slugify(name)
588+
589+
f = CodeFunction(
590+
project=self.project,
591+
handler=handler,
592+
name=name,
593+
slug=slug,
594+
type_="classifier",
595+
description=description,
596+
parameters=parameters,
597+
returns=returns,
598+
if_exists=if_exists,
599+
metadata=metadata,
600+
tags=tags,
601+
)
602+
self.project.add_code_function(f)
603+
return f
604+
605+
547606
class Project:
548607
"""A handle to a Braintrust project."""
549608

@@ -553,6 +612,7 @@ def __init__(self, name: str):
553612
self.prompts = PromptBuilder(self)
554613
self.parameters = ParametersBuilder(self)
555614
self.scorers = ScorerBuilder(self)
615+
self.classifiers = ClassifierBuilder(self)
556616

557617
self._publishable_code_functions: list[CodeFunction] = []
558618
self._publishable_prompts: list[CodePrompt] = []

py/src/braintrust/test_framework2.py

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@
55

66
import pytest
77

8-
from .framework2 import ProjectIdCache, projects
8+
from .framework import _set_lazy_load
9+
from .framework2 import ProjectIdCache, global_, projects
910

1011

1112
# Check if pydantic is available
@@ -145,6 +146,50 @@ def test_llm_scorer_with_metadata(self):
145146
assert scorer.name == "llm-scorer"
146147

147148

149+
class TestClassifierFunctionPush:
150+
"""Tests for classifier function push registration."""
151+
152+
def test_code_classifier_registers_as_code_function(self):
153+
project = projects.create("test-project")
154+
155+
def classify(output: str):
156+
return {"name": "category", "id": output}
157+
158+
classifier = project.classifiers.create(
159+
handler=classify,
160+
name="test-classifier",
161+
)
162+
163+
assert classifier.type_ == "classifier"
164+
assert classifier.name == "test-classifier"
165+
assert classifier.slug == "test-classifier"
166+
assert project._publishable_code_functions == [classifier]
167+
168+
def test_lazy_code_classifier_registration_uses_functions_registry(self):
169+
global_.functions.clear()
170+
global_.prompts.clear()
171+
global_.parameters.clear()
172+
173+
try:
174+
with _set_lazy_load(True):
175+
project = projects.create("test-project")
176+
177+
def classify(output: str):
178+
return {"name": "category", "id": output}
179+
180+
classifier = project.classifiers.create(
181+
handler=classify,
182+
name="test-classifier",
183+
)
184+
185+
assert global_.functions == [classifier]
186+
assert global_.functions[0].type_ == "classifier"
187+
finally:
188+
global_.functions.clear()
189+
global_.prompts.clear()
190+
global_.parameters.clear()
191+
192+
148193
class TestCodeFunctionTags:
149194
"""Tests for CodeFunction tags support."""
150195

0 commit comments

Comments
 (0)