feat: Add extraction of default and descriptions for fields in Pydantic models

Author: Mahhheshh

py/src/braintrust/parameters.py

@@ -34,6 +34,20 @@ def _pydantic_to_json_schema(model: Any) -> dict[str, Any]:
         raise ValueError(f"Cannot convert {model} to JSON schema - not a pydantic model")
 
 
+def _extract_pydantic_fields(schema: dict[str, dict[str, str]]) -> dict[str, list[dict[str, str]]]:
+    """Extract pydantic fields default and description metadata"""
+    flatten_defaults = []
+    flatten_description = []
+    for field_name, field_metadata in schema.items():
+        flatten_defaults.append({
+            field_name: field_metadata.get("default")
+        })
+        flatten_description.append({field_name: field_metadata.get("description")})
+    return {
+        "default": flatten_defaults,
+        "description": flatten_description
+    }
+
 def validate_parameters(
     parameters: dict[str, Any],
     parameter_schema: EvalParameters,
@@ -143,10 +157,13 @@ def parameters_to_json_schema(parameters: EvalParameters) -> dict[str, Any]:
         else:
             # Pydantic model
             try:
+                pydantic_schema = _pydantic_to_json_schema(schema)
+                parameter_values = _extract_pydantic_fields(pydantic_schema.get("properties", {}))
                 result[name] = {
                     "type": "data",
-                    "schema": _pydantic_to_json_schema(schema),
-                    # TODO: Extract default and description from pydantic model
+                    "schema": pydantic_schema,
+                    "default": parameter_values.get("default"),
+                    "description": parameter_values.get("description")
                 }
             except ValueError:
                 # Not a pydantic model, skip

py/src/braintrust/test_parameters.py

@@ -0,0 +1,181 @@
+from unittest.mock import MagicMock
+
+import pytest
+from braintrust.parameters import _extract_pydantic_fields, parameters_to_json_schema
+
+DEFAULT_VALUE = "You are a helpful assistant."
+DEFAULT_DESCRIPTION = "System prompt for the model"
+SCHEMA_TITLE = "SystemPrompt"
+PARAM_NAME = "system_prompt"
+
+def test_extract_single_field_with_default_and_description():
+    schema = {
+        "value": {"type": "string", "default": DEFAULT_VALUE, "description": DEFAULT_DESCRIPTION},
+    }
+    result = _extract_pydantic_fields(schema)
+    assert result == {
+        "default": [{"value": DEFAULT_VALUE}],
+        "description": [{"value": DEFAULT_DESCRIPTION}],
+    }
+
+
+def test_extract_single_field_missing_default_and_description():
+    schema = {
+        "value": {"type": "string"},
+    }
+    result = _extract_pydantic_fields(schema)
+    assert result == {
+        "default": [{"value": None}],
+        "description": [{"value": None}],
+    }
+
+
+def test_extract_multi_field():
+    schema = {
+        "temperature": {"type": "number", "default": 0.7, "description": "Sampling temperature"},
+        "max_tokens": {"type": "integer", "default": 1024, "description": "Maximum tokens to generate"},
+    }
+    result = _extract_pydantic_fields(schema)
+    assert result == {
+        "default": [
+            {"temperature": 0.7},
+            {"max_tokens": 1024},
+        ],
+        "description": [
+            {"temperature": "Sampling temperature"},
+            {"max_tokens": "Maximum tokens to generate"},
+        ],
+    }
+
+
+def test_extract_multi_field_partial_metadata():
+    schema = {
+        "temperature": {"type": "number", "default": 0.7},
+        "max_tokens": {"type": "integer", "description": "Maximum tokens to generate"},
+    }
+    result = _extract_pydantic_fields(schema)
+    assert result == {
+        "default": [
+            {"temperature": 0.7},
+            {"max_tokens": None},
+        ],
+        "description": [
+            {"temperature": None},
+            {"max_tokens": "Maximum tokens to generate"},
+        ],
+    }
+
+
+def test_extract_empty_schema():
+    result = _extract_pydantic_fields({})
+    assert result == {"default": [], "description": []}
+
+
+@pytest.fixture
+def v2_model():
+    def _make(default=DEFAULT_VALUE, description=DEFAULT_DESCRIPTION):
+        model = MagicMock()
+        model.model_json_schema.return_value = {
+            "title": SCHEMA_TITLE,
+            "type": "object",
+            "properties": {"value": {"type": "string", "default": default, "description": description}},
+        }
+        del model.get
+        return model
+    return _make
+
+
+@pytest.fixture
+def v1_model():
+    def _make(default=DEFAULT_VALUE):
+        model = MagicMock()
+        del model.model_json_schema
+        model.schema.return_value = {
+            "title": SCHEMA_TITLE,
+            "type": "object",
+            "properties": {"value": {"type": "string", "default": default}},
+        }
+        del model.get
+        return model
+    return _make
+
+
+@pytest.fixture
+def v2_multi_field_model():
+    def _make():
+        model = MagicMock()
+        model.model_json_schema.return_value = {
+            "title": "ModelConfig",
+            "type": "object",
+            "properties": {
+                "temperature": {"type": "number", "default": 0.7, "description": "Sampling temperature"},
+                "max_tokens": {"type": "integer", "default": 1024, "description": "Maximum tokens to generate"},
+            },
+        }
+        del model.get
+        return model
+    return _make
+
+
+@pytest.fixture
+def v1_multi_field_model():
+    def _make():
+        model = MagicMock()
+        del model.model_json_schema
+        model.schema.return_value = {
+            "title": "ModelConfig",
+            "type": "object",
+            "properties": {
+                "temperature": {"type": "number", "default": 0.7},
+                "max_tokens": {"type": "integer", "default": 1024},
+            },
+        }
+        del model.get
+        return model
+    return _make
+
+
+def test_pydantic_v2_model(v2_model):
+    schema = parameters_to_json_schema({PARAM_NAME: v2_model()})
+
+    assert schema[PARAM_NAME]["type"] == "data"
+    assert schema[PARAM_NAME]["default"] == [{"value": DEFAULT_VALUE}]
+    assert schema[PARAM_NAME]["description"] == [{"value": DEFAULT_DESCRIPTION}]
+
+
+def test_pydantic_v1_model(v1_model):
+    schema = parameters_to_json_schema({PARAM_NAME: v1_model()})
+
+    assert schema[PARAM_NAME]["type"] == "data"
+    assert schema[PARAM_NAME]["default"] == [{"value": DEFAULT_VALUE}]
+    assert schema[PARAM_NAME]["description"] == [{"value": None}]
+
+
+def test_pydantic_v2_multi_field_model(v2_multi_field_model):
+    schema = parameters_to_json_schema({PARAM_NAME: v2_multi_field_model()})
+
+    assert schema[PARAM_NAME]["type"] == "data"
+    assert schema[PARAM_NAME]["schema"]["title"] == "ModelConfig"
+    assert schema[PARAM_NAME]["default"] == [
+        {"temperature": 0.7},
+        {"max_tokens": 1024},
+    ]
+    assert schema[PARAM_NAME]["description"] == [
+        {"temperature": "Sampling temperature"},
+        {"max_tokens": "Maximum tokens to generate"},
+    ]
+
+
+def test_pydantic_v1_multi_field_model(v1_multi_field_model):
+    schema = parameters_to_json_schema({PARAM_NAME: v1_multi_field_model()})
+
+    assert schema[PARAM_NAME]["type"] == "data"
+    assert schema[PARAM_NAME]["schema"]["title"] == "ModelConfig"
+    assert schema[PARAM_NAME]["default"] == [
+        {"temperature": 0.7},
+        {"max_tokens": 1024},
+    ]
+    assert schema[PARAM_NAME]["description"] == [
+        {"temperature": None},
+        {"max_tokens": None},
+    ]