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

Author: Mahhheshh

py/src/braintrust/parameters.py

@@ -34,6 +34,26 @@ 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, Any]]) -> tuple[Any, Any] | tuple[dict[str, Any], dict[str, Any]]:
+    """Extract pydantic fields default and description metadata"""
+    flatten_defaults = {}
+    flatten_description = {}
+    schema_items = schema.items()
+    if len(schema_items) == 1:
+        for _, field_metadata in schema_items:
+            return (
+                field_metadata.get("default"),
+                field_metadata.get("description")
+            )
+
+    for field_name, field_metadata in schema_items:
+        flatten_defaults[field_name] = field_metadata.get("default")
+        flatten_description[field_name] = field_metadata.get("description")
+    return (
+        flatten_defaults,
+        flatten_description
+    )
+
 def validate_parameters(
     parameters: dict[str, Any],
     parameter_schema: EvalParameters,
@@ -143,10 +163,13 @@ def parameters_to_json_schema(parameters: EvalParameters) -> dict[str, Any]:
         else:
             # Pydantic model
             try:
+                pydantic_schema = _pydantic_to_json_schema(schema)
+                model_defaults, model_descriptions = _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": model_defaults,
+                    "description": model_descriptions
                 }
             except ValueError:
                 # Not a pydantic model, skip

py/src/braintrust/test_pydantic_parameters.py

@@ -0,0 +1,150 @@
+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},
+    }
+    defaults, descriptions = _extract_pydantic_fields(schema)
+    assert defaults == DEFAULT_VALUE
+    assert descriptions == DEFAULT_DESCRIPTION
+
+
+def test_extract_single_field_missing_default_and_description():
+    schema = {
+        "value": {"type": "string"},
+    }
+    defaults, descriptions = _extract_pydantic_fields(schema)
+    assert defaults is None
+    assert descriptions is 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"},
+    }
+    defaults, descriptions = _extract_pydantic_fields(schema)
+    assert defaults == {"temperature": 0.7, "max_tokens": 1024}
+    assert descriptions == {"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"},
+    }
+    defaults, descriptions = _extract_pydantic_fields(schema)
+    assert defaults == {"temperature": 0.7, "max_tokens": None}
+    assert descriptions == {"temperature": None, "max_tokens": "Maximum tokens to generate"}
+
+
+def test_extract_empty_schema():
+    defaults, descriptions = _extract_pydantic_fields({})
+    assert defaults == {}
+    assert descriptions == {}
+
+
+@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"] == DEFAULT_VALUE
+    assert schema[PARAM_NAME]["description"] == 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"] == DEFAULT_VALUE
+    assert schema[PARAM_NAME]["description"] is 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}