feat: schema validation in langsmith sdk

python/langsmith/client.py

@@ -2529,6 +2529,8 @@ def create_dataset(
         *,
         description: Optional[str] = None,
         data_type: ls_schemas.DataType = ls_schemas.DataType.kv,
+        inputs_schema: Optional[Dict[str, Any]] = None,
+        outputs_schema: Optional[Dict[str, Any]] = None,
     ) -> ls_schemas.Dataset:
         """Create a dataset in the LangSmith API.
 
@@ -2550,12 +2552,15 @@ def create_dataset(
             name=dataset_name,
             description=description,
             data_type=data_type,
+            inputs_schema=inputs_schema,
+            outputs_schema=outputs_schema,
         )
+
         response = self.request_with_retries(
             "POST",
             "/datasets",
             headers={**self._headers, "Content-Type": "application/json"},
-            data=dataset.json(),
+            data=dataset.json(by_alias=True),
         )
         ls_utils.raise_for_status_with_text(response)
         return ls_schemas.Dataset(

python/langsmith/schemas.py

@@ -128,10 +128,17 @@ class DatasetBase(BaseModel):
     name: str
     description: Optional[str] = None
     data_type: Optional[DataType] = None
+    inputs_schema: Optional[Dict[str, Any]] = Field(
+        None, alias="inputs_schema_definition"
+    )
+    outputs_schema: Optional[Dict[str, Any]] = Field(
+        None, alias="outputs_schema_definition"
+    )
 
     class Config:
         """Configuration class for the schema."""
 
+        allow_population_by_field_name = True
         frozen = True
 
 

python/tests/integration_tests/test_client.py

@@ -13,6 +13,7 @@
 
 import pytest
 from freezegun import freeze_time
+from pydantic import BaseModel
 
 from langsmith.client import ID_TYPE, Client
 from langsmith.schemas import DataType
@@ -312,11 +313,7 @@ def test_error_surfaced_invalid_uri(monkeypatch: pytest.MonkeyPatch, uri: str) -
         client.create_run("My Run", inputs={"text": "hello world"}, run_type="llm")
 
 
-def test_create_dataset(
-    monkeypatch: pytest.MonkeyPatch, langchain_client: Client
-) -> None:
-    """Test persisting runs and adding feedback."""
-    monkeypatch.setenv("LANGCHAIN_ENDPOINT", "https://dev.api.smith.langchain.com")
+def test_create_dataset(langchain_client: Client) -> None:
     dataset_name = "__test_create_dataset" + uuid4().hex[:4]
     if langchain_client.has_dataset(dataset_name=dataset_name):
         langchain_client.delete_dataset(dataset_name=dataset_name)
@@ -360,6 +357,59 @@ def test_create_dataset(
     langchain_client.delete_dataset(dataset_id=dataset.id)
 
 
+def test_dataset_schema_validation(langchain_client: Client) -> None:
+    dataset_name = "__test_create_dataset" + uuid4().hex[:4]
+    if langchain_client.has_dataset(dataset_name=dataset_name):
+        langchain_client.delete_dataset(dataset_name=dataset_name)
+
+    class InputSchema(BaseModel):
+        input: str
+
+    class OutputSchema(BaseModel):
+        output: str
+
+    dataset = langchain_client.create_dataset(
+        dataset_name,
+        data_type=DataType.kv,
+        inputs_schema=InputSchema.model_json_schema(),
+        outputs_schema=OutputSchema.model_json_schema(),
+    )
+
+    # confirm we store the schema from the create request
+    assert dataset.inputs_schema == InputSchema.model_json_schema()
+    assert dataset.outputs_schema == OutputSchema.model_json_schema()
+
+    # create an example that matches the schema, which should succeed
+    langchain_client.create_example(
+        inputs={"input": "hello world"},
+        outputs={"output": "hello"},
+        dataset_id=dataset.id,
+    )
+
+    # create an example that does not match the input schema
+    with pytest.raises(LangSmithError):
+        langchain_client.create_example(
+            inputs={"john": 1},
+            outputs={"output": "hello"},
+            dataset_id=dataset.id,
+        )
+
+    # create an example that does not match the output schema
+    with pytest.raises(LangSmithError):
+        langchain_client.create_example(
+            inputs={"input": "hello world"},
+            outputs={"john": 1},
+            dataset_id=dataset.id,
+        )
+
+    # assert read API includes the schema definition
+    read_dataset = langchain_client.read_dataset(dataset_id=dataset.id)
+    assert read_dataset.inputs_schema == InputSchema.model_json_schema()
+    assert read_dataset.outputs_schema == OutputSchema.model_json_schema()
+
+    langchain_client.delete_dataset(dataset_id=dataset.id)
+
+
 @freeze_time("2023-01-01")
 def test_list_datasets(langchain_client: Client) -> None:
     ds1n = "__test_list_datasets1" + uuid4().hex[:4]