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.
 
@@ -2546,18 +2548,28 @@ def create_dataset(
         Dataset
             The created dataset.
         """
-        dataset = ls_schemas.DatasetCreate(
-            name=dataset_name,
-            description=description,
-            data_type=data_type,
-        )
+        dataset: Dict[str, Any] = {
+            "name": dataset_name,
+            "data_type": data_type.value,
+            "created_at": datetime.datetime.now().isoformat(),
+        }
+        if description is not None:
+            dataset["description"] = description
+
+        if inputs_schema is not None:
+            dataset["inputs_schema_definition"] = inputs_schema
+
+        if outputs_schema is not None:
+            dataset["outputs_schema_definition"] = outputs_schema
+
         response = self.request_with_retries(
             "POST",
             "/datasets",
             headers={**self._headers, "Content-Type": "application/json"},
-            data=dataset.json(),
+            data=orjson.dumps(dataset),
         )
         ls_utils.raise_for_status_with_text(response)
+
         return ls_schemas.Dataset(
             **response.json(),
             _host_url=self._host_url,

python/langsmith/schemas.py

@@ -135,13 +135,6 @@ class Config:
         frozen = True
 
 
-class DatasetCreate(DatasetBase):
-    """Dataset create model."""
-
-    id: Optional[UUID] = None
-    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
-
-
 class Dataset(DatasetBase):
     """Dataset ORM model."""
 
@@ -151,6 +144,8 @@ class Dataset(DatasetBase):
     example_count: Optional[int] = None
     session_count: Optional[int] = None
     last_session_start_time: Optional[datetime] = None
+    inputs_schema: Optional[Dict[str, Any]] = None
+    outputs_schema: Optional[Dict[str, Any]] = None
     _host_url: Optional[str] = PrivateAttr(default=None)
     _tenant_id: Optional[UUID] = PrivateAttr(default=None)
     _public_path: Optional[str] = PrivateAttr(default=None)
@@ -163,6 +158,12 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         """Initialize a Dataset object."""
+        if "inputs_schema_definition" in kwargs:
+            kwargs["inputs_schema"] = kwargs.pop("inputs_schema_definition")
+
+        if "outputs_schema_definition" in kwargs:
+            kwargs["outputs_schema"] = kwargs.pop("outputs_schema_definition")
+
         super().__init__(**kwargs)
         self._host_url = _host_url
         self._tenant_id = _tenant_id

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]