Port JS Parameters implementation to Python

Author: joshuawootonn

py/src/braintrust/cli/push.py

@@ -268,6 +268,13 @@ def _collect_prompt_function_defs(
         functions.append(p.to_function_definition(if_exists, project_ids))
 
 
+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(
@@ -306,6 +313,8 @@ def run(args):
         _collect_function_function_defs(project_ids, functions, bundle_id, args.if_exists)
     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

@@ -26,7 +26,7 @@
 
 from ..framework import EvalAsync, EvalScorer, Evaluator, ExperimentSummary, SSEProgressEvent
 from ..generated_types import FunctionId
-from ..logger import BraintrustState, bt_iscoroutinefunction
+from ..logger import BraintrustState, RemoteEvalParameters, bt_iscoroutinefunction
 from ..parameters import parameters_to_json_schema, validate_parameters
 from ..span_identifier_v4 import parse_parent
 from .auth import AuthorizationMiddleware
@@ -79,6 +79,40 @@ async def index(request: Request) -> PlainTextResponse:
     return PlainTextResponse("Hello, world!")
 
 
+def _ensure_input_field(row: Any) -> Any:
+    if isinstance(row, dict) and "input" not in row:
+        return {**row, "input": None}
+    return row
+
+
+def _serialize_parameters_container(parameters: Any) -> dict[str, Any]:
+    if parameters is None:
+        return {}
+
+    if RemoteEvalParameters.is_parameters(parameters):
+        return {
+            "type": "braintrust.parameters",
+            "schema": dict(parameters.schema),
+            "source": {
+                "parametersId": parameters.id,
+                "slug": parameters.slug,
+                "name": parameters.name,
+                "projectId": parameters.project_id,
+                "version": parameters.version,
+            },
+        }
+
+    schema = parameters_to_json_schema(parameters)
+    if schema:
+        return {
+            "type": "braintrust.staticParameters",
+            "schema": schema,
+            "source": None,
+        }
+
+    return {}
+
+
 async def list_evaluators(request: Request) -> JSONResponse:
     # Get the authenticated context
     ctx = getattr(request.state, "ctx", None)
@@ -93,7 +127,7 @@ 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_parameters_container(evaluator.parameters),
             "scores": [{"name": getattr(score, "name", f"score_{i}")} for i, score in enumerate(evaluator.scores)],
         }
 
@@ -130,12 +164,20 @@ async def run_eval(request: Request) -> JSONResponse | StreamingResponse:
     if not evaluator:
         return JSONResponse({"error": f"Evaluator '{eval_data['name']}' not found"}, status_code=404)
 
-    # Get the dataset if data is provided
-    try:
-        dataset = await get_dataset(state, eval_data["data"])
-    except Exception as e:
-        print(f"Error loading dataset: {e}", file=sys.stderr)
-        return JSONResponse({"error": f"Failed to load dataset: {str(e)}"}, status_code=400)
+    # Get the dataset if data is provided, otherwise fall back to the evaluator's own data
+    dataset = None
+    raw_data = eval_data.get("data")
+    if raw_data is not None:
+        try:
+            dataset = await get_dataset(state, raw_data)
+        except Exception as e:
+            print(f"Error loading dataset from request, falling back to evaluator data: {e}", file=sys.stderr)
+
+    if dataset is None:
+        dataset = evaluator.data
+
+    if isinstance(dataset, list):
+        dataset = [_ensure_input_field(row) for row in dataset]
 
     # Validate parameters if provided
     validated_parameters = None

py/src/braintrust/framework.py

@@ -34,6 +34,7 @@
     Experiment,
     ExperimentSummary,
     Metadata,
+    RemoteEvalParameters,
     ScoreSummary,
     Span,
     _ExperimentDatasetEvent,
@@ -42,7 +43,7 @@
     stringify_exception,
 )
 from .logger import init as _init_experiment
-from .parameters import EvalParameters
+from .parameters import EvalParameters, validate_parameters
 from .resource_manager import ResourceManager
 from .score import Score, is_score, is_scorer
 from .serializable_data_class import SerializableDataClass
@@ -438,7 +439,7 @@ 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.
@@ -674,7 +675,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,
@@ -803,7 +804,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,
@@ -930,7 +931,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,
@@ -1390,6 +1391,17 @@ def get_other_fields(s):
     scorer_names = [_scorer_name(scorer, i) for i, scorer in enumerate(scorers)]
     unhandled_scores = scorer_names
 
+    resolved_parameters: dict[str, Any] | None = None
+    if evaluator.parameters is not None:
+        if RemoteEvalParameters.is_parameters(evaluator.parameters):
+            resolved_parameters = validate_parameters({}, evaluator.parameters)
+        elif isinstance(evaluator.parameters, dict):
+            resolved_parameters = validate_parameters({}, evaluator.parameters)
+        else:
+            raise ValueError(
+                "Invalid evaluator.parameters. Expected an EvalParameters schema or RemoteEvalParameters."
+            )
+
     async def run_evaluator_task(datum, trial_index=0):
         if isinstance(datum, dict):
             datum = EvalCase.from_dict(datum)
@@ -1449,7 +1461,7 @@ def report_progress(event: TaskProgressEvent):
                     trial_index=trial_index,
                     tags=tags,
                     report_progress=report_progress,
-                    parameters=evaluator.parameters,
+                    parameters=resolved_parameters,
                 )
 
                 # Check if the task takes a hooks argument

py/src/braintrust/framework2.py

@@ -16,6 +16,7 @@
     SavedFunctionId,
     ToolFunctionDefinition,
 )
+from .parameters import EvalParameters, _pydantic_to_json_schema
 from .util import eprint
 
 
@@ -34,6 +35,7 @@ class _GlobalState:
     def __init__(self):
         self.functions: list[CodeFunction] = []
         self.prompts: list[CodePrompt] = []
+        self.parameters: list[CodeParameters] = []
 
 
 global_ = _GlobalState()
@@ -287,6 +289,161 @@ def create(
         return p
 
 
+def _maybe_serialize_prompt_default(default: Any) -> Any:
+    as_dict = getattr(default, "as_dict", None)
+    if callable(as_dict):
+        return as_dict()
+    return default
+
+
+def _pydantic_instance_to_plain(value: Any) -> Any:
+    if hasattr(value, "model_dump"):
+        return value.model_dump()
+    if hasattr(value, "dict"):
+        return value.dict()
+    return value
+
+
+def _is_single_field_value_model(model: Any) -> bool:
+    fields = getattr(model, "__fields__", None) or getattr(model, "model_fields", {})
+    return isinstance(fields, dict) and len(fields) == 1 and "value" in fields
+
+
+def _maybe_set_default_from_pydantic_model(model: Any, schema_obj: dict[str, Any]) -> dict[str, Any]:
+    if "default" in schema_obj:
+        return schema_obj
+    try:
+        instance = model()
+    except Exception:
+        return schema_obj
+
+    if _is_single_field_value_model(model) and hasattr(instance, "value"):
+        return {**schema_obj, "default": _pydantic_instance_to_plain(getattr(instance, "value"))}
+
+    return {**schema_obj, "default": _pydantic_instance_to_plain(instance)}
+
+
+def serialize_eval_parameters_to_parameters_schema(parameters: EvalParameters) -> dict[str, Any]:
+    properties: dict[str, Any] = {}
+    required: list[str] = []
+
+    for name, schema in parameters.items():
+        if isinstance(schema, dict) and schema.get("type") == "prompt":
+            prompt_schema: dict[str, Any] = {"type": "object", "x-bt-type": "prompt"}
+
+            description = schema.get("description")
+            if description is not None:
+                prompt_schema["description"] = description
+
+            default_value = schema.get("default")
+            if default_value is not None:
+                prompt_schema["default"] = _maybe_serialize_prompt_default(default_value)
+            else:
+                required.append(name)
+
+            properties[name] = prompt_schema
+            continue
+
+        if schema is None:
+            raise ValueError(f"Parameter '{name}' has no schema")
+
+        if not (hasattr(schema, "model_json_schema") or hasattr(schema, "schema")):
+            raise ValueError(
+                f"Invalid schema for parameter '{name}'. Expected a pydantic model (v1 or v2) or a prompt parameter."
+            )
+
+        schema_obj = _pydantic_to_json_schema(schema)
+        if _is_single_field_value_model(schema):
+            value_schema = schema_obj.get("properties", {}).get("value")
+            if not isinstance(value_schema, dict):
+                raise ValueError(f"Invalid pydantic schema for parameter '{name}': missing properties.value")
+            parameter_schema = _maybe_set_default_from_pydantic_model(schema, value_schema)
+        else:
+            parameter_schema = _maybe_set_default_from_pydantic_model(schema, schema_obj)
+
+        properties[name] = parameter_schema
+        if "default" not in parameter_schema:
+            required.append(name)
+
+    out: dict[str, Any] = {"type": "object", "properties": properties, "additionalProperties": True}
+    if required:
+        out["required"] = required
+    return out
+
+
+def get_default_data_from_parameters_schema(schema: dict[str, Any]) -> dict[str, Any]:
+    properties = schema.get("properties")
+    if not isinstance(properties, dict):
+        return {}
+
+    return {k: v["default"] for k, v in properties.items() if isinstance(v, dict) and "default" in v}
+
+
+@dataclasses.dataclass
+class CodeParameters:
+    """Parameters defined in code, with metadata."""
+
+    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 = serialize_eval_parameters_to_parameters_schema(self.schema)
+        j: dict[str, Any] = {
+            "project_id": project_ids.get(self.project),
+            "name": self.name,
+            "slug": self.slug,
+            "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.description is not None:
+            j["description"] = self.description
+        if self.metadata is not None:
+            j["metadata"] = self.metadata
+        return j
+
+
+class ParametersBuilder:
+    """Builder to create parameters in Braintrust."""
+
+    def __init__(self, project: "Project"):
+        self.project = project
+
+    def create(
+        self,
+        *,
+        name: str,
+        slug: str | None = None,
+        description: str | None = None,
+        schema: EvalParameters,
+        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."""
 
@@ -461,10 +618,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)
@@ -476,6 +635,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}")