From 152ec59849ec81622738c443765f6b5da91ce44e Mon Sep 17 00:00:00 2001
From: isaac hershenson <ihershenson@hmc.edu>
Date: Wed, 13 Nov 2024 11:23:09 -0800
Subject: [PATCH] integration test skeleton

---
 python/langsmith/client.py                    | 11 ++-
 python/tests/integration_tests/test_client.py | 68 ++++++++++++++++++-
 python/tests/unit_tests/test_client.py        |  6 +-
 3 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/python/langsmith/client.py b/python/langsmith/client.py
index dca29d9fc..7e823573d 100644
--- a/python/langsmith/client.py
+++ b/python/langsmith/client.py
@@ -3370,12 +3370,19 @@ def create_example_from_run(
             created_at=created_at,
         )
 
-    def upsert_example_multipart(
+    def upsert_examples_multipart(
         self,
         *,
-        upserts: List[ls_schemas.ExampleCreateWithAttachments] = None,
+        upserts: List[ls_schemas.ExampleCreateWithAttachments] = [],
     ) -> None:
         """Upsert examples."""
+        # not sure if the below checks are necessary
+        if not isinstance(upserts, list):
+            raise TypeError(f"upserts must be a list, got {type(upserts)}")
+        for item in upserts:
+            if not isinstance(item, ls_schemas.ExampleCreateWithAttachments):
+                raise TypeError(f"Each item must be ExampleCreateWithAttachments, got {type(item)}")
+            
         parts: list[MultipartPart] = []
 
         for example in upserts:
diff --git a/python/tests/integration_tests/test_client.py b/python/tests/integration_tests/test_client.py
index 57a6e2171..bfe0d818a 100644
--- a/python/tests/integration_tests/test_client.py
+++ b/python/tests/integration_tests/test_client.py
@@ -20,7 +20,7 @@
 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 (
     LangSmithConnectionError,
     LangSmithError,
@@ -369,6 +369,72 @@ def test_error_surfaced_invalid_uri(uri: str) -> None:
         client.create_run("My Run", inputs={"text": "hello world"}, run_type="llm")
 
 
+@pytest.mark.parametrize("uri", ["http://dev.api.smith.langchain.com"])
+def test_upsert_examples_multipart(uri: str) -> None:
+    """Test upserting examples with attachments via multipart endpoint."""
+    dataset_name = "__test_upsert_examples_multipart" + uuid4().hex[:4]
+    langchain_client = Client(api_url=uri, api_key="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")
+    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"),
+        },
+    )
+
+    langchain_client.upsert_examples_multipart([example_1, example_2])
+    
+    created_example = langchain_client.read_example(example_id)
+    assert created_example.inputs["text"] == "hello world"
+    assert created_example.outputs["response"] == "greeting"
+
+    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"),
+        },
+    )
+
+    # will this throw an error? idk need to test
+    langchain_client.upsert_examples_multipart([example_2, example_3]) # don't add example_1 because of explicit id
+
+    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(TypeError):
+        langchain_client.upsert_examples_multipart([{"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]
     if langchain_client.has_dataset(dataset_name=dataset_name):
diff --git a/python/tests/unit_tests/test_client.py b/python/tests/unit_tests/test_client.py
index edda5dd09..4b68ce368 100644
--- a/python/tests/unit_tests/test_client.py
+++ b/python/tests/unit_tests/test_client.py
@@ -417,8 +417,8 @@ def test_create_run_mutate(
 
 
 @mock.patch("langsmith.client.requests.Session")
-def test_upsert_example_multipart(mock_session_cls: mock.Mock) -> None:
-    """Test that upsert_example_multipart sends correct multipart data."""
+def test_upsert_examples_multipart(mock_session_cls: mock.Mock) -> None:
+    """Test that upsert_examples_multipart sends correct multipart data."""
     mock_session = MagicMock()
     mock_response = MagicMock()
     mock_response.status_code = 200
@@ -447,7 +447,7 @@ def test_upsert_example_multipart(mock_session_cls: mock.Mock) -> None:
             ),
         },
     )
-    client.upsert_example_multipart(upserts=[example])
+    client.upsert_examples_multipart(upserts=[example])
 
     # Verify the request
     assert mock_session.request.call_count == 2  # we always make a call to /info