Add saved parameters support to Python SDK

Author: joshuawootonn

py/setup.py

@@ -18,6 +18,7 @@
     "chevron",
     "tqdm",
     "exceptiongroup>=1.2.0",
+    "jsonschema",
     "python-dotenv",
     "sseclient-py",
     "python-slugify",

py/src/braintrust/cli/eval.py

@@ -22,6 +22,7 @@
     set_thread_pool_max_workers,
 )
 from ..logger import Dataset
+from ..parameters import RemoteEvalParameters
 from ..util import eprint
 
 
@@ -131,6 +132,12 @@ async def run_evaluator_task(evaluator, position, opts: EvaluatorOpts):
         if isinstance(evaluator.data, Dataset):
             dataset = evaluator.data
 
+        parameters = None
+        if RemoteEvalParameters.is_parameters(evaluator.parameters) and evaluator.parameters.id is not None:
+            parameters = {"id": evaluator.parameters.id}
+            if evaluator.parameters.version is not None:
+                parameters["version"] = evaluator.parameters.version
+
         # NOTE: This code is duplicated with _EvalCommon in py/src/braintrust/framework.py.
         # Make sure to update those arguments if you change this.
         experiment = init_experiment(
@@ -147,6 +154,7 @@ async def run_evaluator_task(evaluator, position, opts: EvaluatorOpts):
             git_metadata_settings=evaluator.git_metadata_settings,
             repo_info=evaluator.repo_info,
             dataset=dataset,
+            parameters=parameters,
         )
 
     try:

py/src/braintrust/cli/push.py

@@ -271,7 +271,6 @@ def _collect_prompt_function_defs(
     for p in global_.prompts:
         functions.append(p.to_function_definition(if_exists, project_ids))
 
-
 def _collect_evaluator_defs(
     project_ids: ProjectIdCache,
     functions: list[dict[str, Any]],
@@ -322,6 +321,13 @@ def _collect_evaluator_defs(
         )
 
 
+def _collect_parameters_function_defs(
+    project_ids: ProjectIdCache, functions: list[dict[str, Any]], if_exists: IfExists
+) -> None:
+    for p in global_.parameters:
+        functions.append(p.to_function_definition(if_exists, project_ids))
+
+
 def run(args):
     """Runs the braintrust push subcommand."""
     login(
@@ -379,6 +385,8 @@ def run(args):
 
     if len(global_.prompts) > 0:
         _collect_prompt_function_defs(project_ids, functions, args.if_exists)
+    if len(global_.parameters) > 0:
+        _collect_parameters_function_defs(project_ids, functions, args.if_exists)
 
     if len(functions) > 0:
         api_conn().post_json("insert-functions", {"functions": functions})

py/src/braintrust/devserver/server.py

@@ -28,7 +28,7 @@
 from ..framework import EvalAsync, EvalScorer, Evaluator, ExperimentSummary, SSEProgressEvent
 from ..generated_types import FunctionId
 from ..logger import BraintrustState, bt_iscoroutinefunction
-from ..parameters import parameters_to_json_schema, validate_parameters
+from ..parameters import serialize_remote_eval_parameters_container, validate_parameters
 from ..span_identifier_v4 import parse_parent
 from .auth import AuthorizationMiddleware
 from .cache import cached_login
@@ -41,6 +41,42 @@
 _all_evaluators: dict[str, Evaluator[Any, Any]] = {}
 
 
+class _ParameterOverrideHooks:
+    def __init__(self, hooks: Any, parameters: dict[str, Any]):
+        self._hooks = hooks
+        self._parameters = parameters
+
+    @property
+    def metadata(self):
+        return self._hooks.metadata
+
+    @property
+    def expected(self):
+        return self._hooks.expected
+
+    @property
+    def span(self):
+        return self._hooks.span
+
+    @property
+    def trial_index(self):
+        return self._hooks.trial_index
+
+    @property
+    def tags(self):
+        return self._hooks.tags
+
+    @property
+    def parameters(self):
+        return self._parameters
+
+    def report_progress(self, progress):
+        return self._hooks.report_progress(progress)
+
+    def meta(self, **info: Any):
+        return self._hooks.meta(**info)
+
+
 class CheckAuthorizedMiddleware(BaseHTTPMiddleware):
     def __init__(self, app, allowed_org_name: str | None = None):
         super().__init__(app)
@@ -95,7 +131,9 @@ async def list_evaluators(request: Request) -> JSONResponse:
     evaluator_list = {}
     for name, evaluator in _all_evaluators.items():
         evaluator_list[name] = {
-            "parameters": parameters_to_json_schema(evaluator.parameters) if evaluator.parameters else {},
+            "parameters": (
+                serialize_remote_eval_parameters_container(evaluator.parameters) if evaluator.parameters else None
+            ),
             "scores": [{"name": getattr(score, "name", f"score_{i}")} for i, score in enumerate(evaluator.scores)],
         }
 
@@ -155,11 +193,12 @@ async def run_eval(request: Request) -> JSONResponse | StreamingResponse:
     sse_queue = SSEQueue()
 
     async def task(input, hooks):
+        task_hooks = hooks if validated_parameters is None else _ParameterOverrideHooks(hooks, validated_parameters)
         if bt_iscoroutinefunction(evaluator.task):
-            result = await evaluator.task(input, hooks)
+            result = await evaluator.task(input, task_hooks)
         else:
-            result = evaluator.task(input, hooks)
-        hooks.report_progress(
+            result = evaluator.task(input, task_hooks)
+        task_hooks.report_progress(
             {
                 "format": "code",
                 "output_type": "completion",
@@ -186,10 +225,9 @@ def stream_fn(event: SSEProgressEvent):
     if parent:
         parent = parse_parent(parent)
 
-    # Override evaluator parameters with validated ones if provided
-    eval_kwargs = {k: v for (k, v) in evaluator.__dict__.items() if k not in ["eval_name", "project_name"]}
-    if validated_parameters is not None:
-        eval_kwargs["parameters"] = validated_parameters
+    eval_kwargs = {
+        k: v for (k, v) in evaluator.__dict__.items() if k not in ["eval_name", "project_name", "parameter_values"]
+    }
 
     try:
         eval_task = asyncio.create_task(

py/src/braintrust/framework.py

@@ -42,7 +42,7 @@
     stringify_exception,
 )
 from .logger import init as _init_experiment
-from .parameters import EvalParameters
+from .parameters import EvalParameters, RemoteEvalParameters, is_eval_parameter_schema, validate_parameters
 from .resource_manager import ResourceManager
 from .score import Score, is_score, is_scorer
 from .serializable_data_class import SerializableDataClass
@@ -439,12 +439,14 @@ class Evaluator(Generic[Input, Output]):
     Whether to summarize the scores of the experiment after it has run.
     """
 
-    parameters: EvalParameters | None = None
+    parameters: EvalParameters | RemoteEvalParameters | None = None
     """
     A set of parameters that will be passed to the evaluator.
     Can be used to define prompts or other configurable values.
     """
 
+    parameter_values: dict[str, Any] | None = None
+
 
 @dataclasses.dataclass
 class EvalResultWithSummary(SerializableDataClass, Generic[Input, Output]):
@@ -675,7 +677,7 @@ def _EvalCommon(
     summarize_scores: bool,
     no_send_logs: bool,
     error_score_handler: ErrorScoreHandler | None = None,
-    parameters: EvalParameters | None = None,
+    parameters: EvalParameters | RemoteEvalParameters | None = None,
     on_start: Callable[[ExperimentSummary], None] | None = None,
     stream: Callable[[SSEProgressEvent], None] | None = None,
     parent: str | None = None,
@@ -741,6 +743,12 @@ async def make_empty_summary():
         if isinstance(evaluator.data, Dataset):
             dataset = evaluator.data
 
+        experiment_parameters = None
+        if RemoteEvalParameters.is_parameters(evaluator.parameters) and evaluator.parameters.id is not None:
+            experiment_parameters = {"id": evaluator.parameters.id}
+            if evaluator.parameters.version is not None:
+                experiment_parameters["version"] = evaluator.parameters.version
+
         # NOTE: This code is duplicated with run_evaluator_task in py/src/braintrust/cli/eval.py.
         # Make sure to update those arguments if you change this.
         experiment = None
@@ -759,6 +767,7 @@ async def make_empty_summary():
                 git_metadata_settings=evaluator.git_metadata_settings,
                 repo_info=evaluator.repo_info,
                 dataset=dataset,
+                parameters=experiment_parameters,
                 state=state,
             )
 
@@ -804,7 +813,7 @@ async def EvalAsync(
     description: str | None = None,
     summarize_scores: bool = True,
     no_send_logs: bool = False,
-    parameters: EvalParameters | None = None,
+    parameters: EvalParameters | RemoteEvalParameters | None = None,
     on_start: Callable[[ExperimentSummary], None] | None = None,
     stream: Callable[[SSEProgressEvent], None] | None = None,
     parent: str | None = None,
@@ -931,7 +940,7 @@ def Eval(
     description: str | None = None,
     summarize_scores: bool = True,
     no_send_logs: bool = False,
-    parameters: EvalParameters | None = None,
+    parameters: EvalParameters | RemoteEvalParameters | None = None,
     on_start: Callable[[ExperimentSummary], None] | None = None,
     stream: Callable[[SSEProgressEvent], None] | None = None,
     parent: str | None = None,
@@ -1392,6 +1401,15 @@ def get_other_fields(s):
     scorer_names = [_scorer_name(scorer, i) for i, scorer in enumerate(scorers)]
     unhandled_scores = scorer_names
 
+    if evaluator.parameter_values is not None:
+        resolved_evaluator_parameters = evaluator.parameter_values
+    elif RemoteEvalParameters.is_parameters(evaluator.parameters):
+        resolved_evaluator_parameters = validate_parameters({}, evaluator.parameters)
+    elif is_eval_parameter_schema(evaluator.parameters):
+        resolved_evaluator_parameters = validate_parameters({}, evaluator.parameters)
+    else:
+        resolved_evaluator_parameters = evaluator.parameters
+
     async def run_evaluator_task(datum, trial_index=0):
         if isinstance(datum, dict):
             datum = EvalCase.from_dict(datum)
@@ -1451,7 +1469,7 @@ def report_progress(event: TaskProgressEvent):
                     trial_index=trial_index,
                     tags=tags,
                     report_progress=report_progress,
-                    parameters=evaluator.parameters,
+                    parameters=resolved_evaluator_parameters,
                 )
 
                 # Check if the task takes a hooks argument

py/src/braintrust/framework2.py

@@ -16,6 +16,7 @@
     SavedFunctionId,
     ToolFunctionDefinition,
 )
+from .parameters import EvalParameters, get_default_data_from_parameters_schema, parameters_to_json_schema
 from .util import eprint
 
 
@@ -40,6 +41,7 @@ class _GlobalState:
     def __init__(self):
         self.functions: list[CodeFunction] = []
         self.prompts: list[CodePrompt] = []
+        self.parameters: list["CodeParameters"] = []
 
 
 global_ = _GlobalState()
@@ -116,6 +118,36 @@ def to_function_definition(self, if_exists: IfExists | None, project_ids: Projec
         return j
 
 
+@dataclasses.dataclass
+class CodeParameters:
+    project: "Project"
+    name: str
+    slug: str
+    description: str | None
+    schema: EvalParameters
+    if_exists: IfExists | None
+    metadata: dict[str, Any] | None = None
+
+    def to_function_definition(self, if_exists: IfExists | None, project_ids: ProjectIdCache) -> dict[str, Any]:
+        schema = parameters_to_json_schema(self.schema)
+        j: dict[str, Any] = {
+            "project_id": project_ids.get(self.project),
+            "name": self.name,
+            "slug": self.slug,
+            "description": self.description or "",
+            "function_type": "parameters",
+            "function_data": {
+                "type": "parameters",
+                "data": get_default_data_from_parameters_schema(schema),
+                "__schema": schema,
+            },
+            "if_exists": self.if_exists if self.if_exists is not None else if_exists,
+        }
+        if self.metadata is not None:
+            j["metadata"] = self.metadata
+        return j
+
+
 class ToolBuilder:
     """Builder to create a tool in Braintrust."""
 
@@ -305,6 +337,38 @@ def create(
         return p
 
 
+class ParametersBuilder:
+    """Builder to create saved parameters in Braintrust."""
+
+    def __init__(self, project: "Project"):
+        self.project = project
+
+    def create(
+        self,
+        *,
+        name: str,
+        schema: EvalParameters,
+        slug: str | None = None,
+        description: str | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> EvalParameters:
+        if slug is None or len(slug) == 0:
+            slug = slugify.slugify(name)
+
+        parameters = CodeParameters(
+            project=self.project,
+            name=name,
+            slug=slug,
+            description=description,
+            schema=schema,
+            if_exists=if_exists,
+            metadata=metadata,
+        )
+        self.project.add_parameters(parameters)
+        return schema
+
+
 class ScorerBuilder:
     """Builder to create a scorer in Braintrust."""
 
@@ -486,10 +550,12 @@ def __init__(self, name: str):
         self.name = name
         self.tools = ToolBuilder(self)
         self.prompts = PromptBuilder(self)
+        self.parameters = ParametersBuilder(self)
         self.scorers = ScorerBuilder(self)
 
         self._publishable_code_functions: list[CodeFunction] = []
         self._publishable_prompts: list[CodePrompt] = []
+        self._publishable_parameters: list[CodeParameters] = []
 
     def add_code_function(self, fn: CodeFunction):
         self._publishable_code_functions.append(fn)
@@ -501,6 +567,11 @@ def add_prompt(self, prompt: CodePrompt):
         if _is_lazy_load():
             global_.prompts.append(prompt)
 
+    def add_parameters(self, parameters: CodeParameters):
+        self._publishable_parameters.append(parameters)
+        if _is_lazy_load():
+            global_.parameters.append(parameters)
+
     def publish(self):
         if _is_lazy_load():
             eprint(f"{bcolors.WARNING}publish() is a no-op when running `braintrust push`.{bcolors.ENDC}")
@@ -518,6 +589,8 @@ def publish(self):
         for prompt in self._publishable_prompts:
             prompt_definition = prompt.to_function_definition(None, project_id_cache)
             definitions.append(prompt_definition)
+        for parameters in self._publishable_parameters:
+            definitions.append(parameters.to_function_definition(None, project_id_cache))
         return api_conn().post_json("insert-functions", {"functions": definitions})