[python] support uploading examples with attachments and running evals on examples with attachments

python/bench/upload_examples_bench.py

@@ -0,0 +1,124 @@
+import statistics
+import time
+from typing import Dict
+from uuid import uuid4
+from langsmith.schemas import DataType, ExampleCreateWithAttachments
+import sys
+sys.path.append('./../langsmith')
+from client import Client
+
+def create_large_json(length: int) -> Dict:
+    """Create a large JSON object for benchmarking purposes."""
+    large_array = [
+        {
+            "index": i,
+            "data": f"This is element number {i}",
+            "nested": {"id": i, "value": f"Nested value for element {i}"},
+        }
+        for i in range(length)
+    ]
+
+    return {
+        "name": "Huge JSON" + str(uuid4()),
+        "description": "This is a very large JSON object for benchmarking purposes.",
+        "array": large_array,
+        "metadata": {
+            "created_at": "2024-10-22T19:00:00Z",
+            "author": "Python Program",
+            "version": 1.0,
+        },
+    }
+
+
+def create_example_data(dataset_id: str, json_size: int) -> Dict:
+    """Create a single example data object."""
+    return ExampleCreateWithAttachments(**{
+        "dataset_id": dataset_id,
+        "inputs": create_large_json(json_size),
+        "outputs": create_large_json(json_size),
+    })
+
+DATASET_NAME = "TEST DATASET"
+def benchmark_example_uploading(num_examples: int, json_size: int, samples: int = 1) -> Dict:
+    """
+    Benchmark run creation with specified parameters.
+    Returns timing statistics.
+    """
+    multipart_timings, old_timings = [], []
+
+
+    for _ in range(samples):
+        client = Client(api_url="https://dev.api.smith.langchain.com")
+
+        if client.has_dataset(dataset_name=DATASET_NAME):
+            client.delete_dataset(dataset_name=DATASET_NAME)
+
+        dataset = client.create_dataset(
+            DATASET_NAME,
+            description="Test dataset for multipart example upload",
+            data_type=DataType.kv,
+        )
+        examples = [create_example_data(dataset.id, json_size) for i in range(num_examples)]
+
+        # Old method
+        old_start = time.perf_counter()
+        inputs=[e.inputs for e in examples]
+        outputs=[e.outputs for e in examples]
+        # the create_examples endpoint fails above 20mb
+        try:
+            client.create_examples(inputs=inputs,
+                                outputs=outputs,dataset_id=dataset.id)
+            old_elapsed = time.perf_counter() - old_start
+        except:
+            old_elapsed = 1000000
+
+        # New method
+        multipart_start = time.perf_counter()
+        client.upsert_examples_multipart(upserts=examples)
+        multipart_elapsed = time.perf_counter() - multipart_start
+
+        multipart_timings.append(multipart_elapsed)
+        old_timings.append(old_elapsed)
+
+    return {
+        "old": {
+            "mean": statistics.mean(old_timings),
+            "median": statistics.median(old_timings),
+            "stdev": statistics.stdev(old_timings) if len(old_timings) > 1 else 0,
+            "min": min(old_timings),
+            "max": max(old_timings),
+        },
+        "new": {
+            "mean": statistics.mean(multipart_timings),
+            "median": statistics.median(multipart_timings),
+            "stdev": statistics.stdev(multipart_timings) if len(multipart_timings) > 1 else 0,
+            "min": min(multipart_timings),
+            "max": max(multipart_timings),
+        }
+    }
+
+json_size = 1000
+num_examples = 1000
+
+def main(json_size: int, num_examples: int):
+    """
+    Run benchmarks with different combinations of parameters and report results.
+    """
+    results = benchmark_example_uploading(num_examples=num_examples, json_size=json_size)
+
+    print(f"\nBenchmark Results for {num_examples} examples with JSON size {json_size}:")
+    print("-" * 60)
+    print(f"{'Metric':<15} {'Old Method':>20} {'New Method':>20}")
+    print("-" * 60)
+
+    metrics = ['mean', 'median', 'stdev', 'min', 'max']
+    for metric in metrics:
+        print(f"{metric:<15} {results['old'][metric]:>20.4f} {results['new'][metric]:>20.4f}")
+
+    print("-" * 60)
+    print(f"{'Throughput':<15} {num_examples / results['old']['mean']:>20.2f} {num_examples / results['new']['mean']:>20.2f}")
+    print("(examples/second)")
+
+
+if __name__ == "__main__":
+    main(json_size, num_examples)

python/langsmith/client.py

@@ -82,6 +82,7 @@
     _SIZE_LIMIT_BYTES,
 )
 from langsmith._internal._multipart import (
+    MultipartPart,
     MultipartPartsAndContext,
     join_multipart_parts_and_context,
 )
@@ -3369,6 +3370,134 @@ def create_example_from_run(
             created_at=created_at,
         )
 
+    def upsert_examples_multipart(
+        self,
+        *,
+        upserts: List[ls_schemas.ExampleCreateWithAttachments] = None,
+    ) -> dict: # Should we create an object for the return type - like UpsertExamplesResponse?
+        """Upsert examples."""
+        if not (self.info.instance_flags or {}).get(
+                "examples_multipart_enabled", False
+            ):
+            raise ValueError("Your LangChain version does not allow using the multipart examples endpoint, please update to the latest version.")
+
+        if upserts is None:
+            upserts = []
+        parts: list[MultipartPart] = []
+
+        for example in upserts:
+            if example.id is not None:
+                example_id = str(example.id)
+            else:
+                example_id = str(uuid.uuid4())
+
+            example_body = {
+                "dataset_id": example.dataset_id,
+                "created_at": example.created_at,
+            }
+            if example.metadata is not None:
+                example_body["metadata"] = example.metadata
+            if example.split is not None:
+                example_body["split"] = example.split
+            valb = _dumps_json(example_body)
+
+            (
+                parts.append(
+                    (
+                        f"{example_id}",
+                        (
+                            None,
+                            valb,
+                            "application/json",
+                            {},
+                        ),
+                    )
+                ),
+            )
+
+            inputsb = _dumps_json(example.inputs)
+            outputsb = _dumps_json(example.outputs)
+
+            (
+                parts.append(
+                    (
+                        f"{example_id}.inputs",
+                        (
+                            None,
+                            inputsb,
+                            "application/json",
+                            {},
+                        ),
+                    )
+                ),
+            )
+
+            (
+                parts.append(
+                    (
+                        f"{example_id}.outputs",
+                        (
+                            None,
+                            outputsb,
+                            "application/json",
+                            {},
+                        ),
+                    )
+                ),
+            )
+
+            if example.attachments:
+                for name, attachment in example.attachments.items():
+                    if isinstance(attachment, tuple):
+                        mime_type, data = attachment
+                        (
+                            parts.append(
+                                (
+                                    f"{example_id}.attachment.{name}",
+                                    (
+                                        None,
+                                        data,
+                                        f"{mime_type}; length={len(data)}",
+                                        {},
+                                    ),
+                                )
+                            ),
+                        )
+                    else:
+                        (
+                            parts.append(
+                                (
+                                    f"{example_id}.attachment.{name}",
+                                    (
+                                        None,
+                                        attachment.data,
+                                        f"{attachment.mime_type}; length={len(attachment.data)}",
+                                        {},
+                                    ),
+                                )
+                            ),
+                        )
+
+        encoder = rqtb_multipart.MultipartEncoder(parts, boundary=BOUNDARY)
+        if encoder.len <= 20_000_000:  # ~20 MB
+            data = encoder.to_string()
+        else:
+            data = encoder
+
+        response = self.request_with_retries(
+            "POST",
+            "/v1/platform/examples/multipart",
+            request_kwargs={
+                "data": data,
+                "headers": {
+                    **self._headers,
+                    "Content-Type": encoder.content_type,
+                },
+            },
+        )
+        ls_utils.raise_for_status_with_text(response)
+        return response.json()
+
     def create_examples(
         self,
         *,

python/langsmith/schemas.py

@@ -89,6 +89,12 @@ class ExampleCreate(ExampleBase):
     split: Optional[Union[str, List[str]]] = None
 
 
+class ExampleCreateWithAttachments(ExampleCreate):
+    """Example create with attachments."""
+
+    attachments: Optional[Attachments] = None
+
+
 class Example(ExampleBase):
     """Example model."""
 
@@ -695,6 +701,8 @@ class LangSmithInfo(BaseModel):
     license_expiration_time: Optional[datetime] = None
     """The time the license will expire."""
     batch_ingest_config: Optional[BatchIngestConfig] = None
+    """The instance flags."""
+    instance_flags: dict[str, Any] = None
 
 
 Example.update_forward_refs()

python/tests/integration_tests/test_client.py

@@ -20,8 +20,9 @@
 from requests_toolbelt import MultipartEncoder, MultipartEncoderMonitor
 
 from langsmith.client import ID_TYPE, Client
-from langsmith.schemas import DataType
+from langsmith.schemas import DataType, ExampleCreateWithAttachments
 from langsmith.utils import (
+    LangSmithNotFoundError,
     LangSmithConnectionError,
     LangSmithError,
     get_env_var,
@@ -368,6 +369,77 @@ def test_error_surfaced_invalid_uri(uri: str) -> None:
     with pytest.raises(LangSmithConnectionError):
         client.create_run("My Run", inputs={"text": "hello world"}, run_type="llm")
 
+# NEED TO FIX ONCE CHANGES PUSH TO PROD
+def test_upsert_examples_multipart() -> None:
+    """Test upserting examples with attachments via multipart endpoint."""
+    dataset_name = "__test_upsert_examples_multipart" + uuid4().hex[:4]
+    langchain_client = Client(api_url="https://dev.api.smith.langchain.com")
+    if langchain_client.has_dataset(dataset_name=dataset_name):
+        langchain_client.delete_dataset(dataset_name=dataset_name)
+
+    dataset = langchain_client.create_dataset(
+        dataset_name,
+        description="Test dataset for multipart example upload",
+        data_type=DataType.kv,
+    )
+
+    # Test example with all fields
+    example_id = uuid4()
+    example_1 = ExampleCreateWithAttachments(
+        id=example_id,
+        dataset_id=dataset.id,
+        inputs={"text": "hello world"},
+        outputs={"response": "greeting"},
+        attachments={
+            "test_file": ("text/plain", b"test content"),
+        },
+    )
+    # Test example without id
+    example_2 = ExampleCreateWithAttachments(
+        dataset_id=dataset.id,
+        inputs={"text": "foo bar"},
+        outputs={"response": "baz"},
+        attachments={
+            "my_file": ("text/plain", b"more test content"),
+        },
+    )
+
+    created_examples = langchain_client.upsert_examples_multipart(upserts=[example_1, example_2])
+    assert created_examples['count'] == 2
+
+    created_example_1 = langchain_client.read_example(created_examples['example_ids'][0])
+    assert created_example_1.inputs["text"] == "hello world"
+    assert created_example_1.outputs["response"] == "greeting"
+
+    created_example_2 = langchain_client.read_example(created_examples['example_ids'][1])
+    assert created_example_2.inputs["text"] == "foo bar"
+    assert created_example_2.outputs["response"] == "baz"
+
+    # make sure examples were sent to the correct dataset
+    all_examples_in_dataset = [example for example in langchain_client.list_examples(dataset_id=dataset.id)]
+    assert len(all_examples_in_dataset) == 2
+
+    # Test that adding invalid example fails - even if valid examples are added alongside
+    example_3 = ExampleCreateWithAttachments(
+        dataset_id=uuid4(), # not a real dataset
+        inputs={"text": "foo bar"},
+        outputs={"response": "baz"},
+        attachments={
+            "my_file": ("text/plain", b"more test content"),
+        },
+    )
+
+    with pytest.raises(LangSmithNotFoundError):
+        langchain_client.upsert_examples_multipart(upserts=[example_3])
+
+    all_examples_in_dataset = [example for example in langchain_client.list_examples(dataset_id=dataset.id)]
+    assert len(all_examples_in_dataset) == 2
+
+    # Throw type errors when not passing ExampleCreateWithAttachments
+    with pytest.raises(AttributeError):
+        langchain_client.upsert_examples_multipart(upserts=[{"foo":"bar"}])
+
+    langchain_client.delete_dataset(dataset_name=dataset_name)
 
 def test_create_dataset(langchain_client: Client) -> None:
     dataset_name = "__test_create_dataset" + uuid4().hex[:4]