From a5560ce5fe5f5d0d97ee326f0b4976d5071d2547 Mon Sep 17 00:00:00 2001
From: Ankush Gola <ankush.gola@gmail.com>
Date: Wed, 27 Mar 2024 17:38:56 -0700
Subject: [PATCH] fix metadata support

---
 js/package.json                               |  2 +-
 js/src/client.ts                              |  5 ++++-
 js/src/schemas.ts                             |  2 ++
 js/src/tests/client.int.test.ts               | 10 ++++++++++
 python/langsmith/client.py                    |  2 ++
 python/langsmith/schemas.py                   |  1 +
 python/tests/integration_tests/test_client.py |  2 ++
 7 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/js/package.json b/js/package.json
index 968c170a2..fd6c5e0cf 100644
--- a/js/package.json
+++ b/js/package.json
@@ -138,4 +138,4 @@
     },
     "./package.json": "./package.json"
   }
-}
\ No newline at end of file
+}
diff --git a/js/src/client.ts b/js/src/client.ts
index 8bc650296..7797749d9 100644
--- a/js/src/client.ts
+++ b/js/src/client.ts
@@ -218,6 +218,8 @@ export type CreateExampleOptions = {
   datasetName?: string;
   createdAt?: Date;
   exampleId?: string;
+
+  metadata?: KVMap;
 };
 
 type AutoBatchQueueItem = {
@@ -1867,7 +1869,7 @@ export class Client {
   public async createExample(
     inputs: KVMap,
     outputs: KVMap,
-    { datasetId, datasetName, createdAt, exampleId }: CreateExampleOptions
+    { datasetId, datasetName, createdAt, exampleId, metadata }: CreateExampleOptions
   ): Promise<Example> {
     let datasetId_ = datasetId;
     if (datasetId_ === undefined && datasetName === undefined) {
@@ -1886,6 +1888,7 @@ export class Client {
       outputs,
       created_at: createdAt_?.toISOString(),
       id: exampleId,
+      metadata,
     };
 
     const response = await this.caller.call(fetch, `${this.apiUrl}/examples`, {
diff --git a/js/src/schemas.ts b/js/src/schemas.ts
index bcf94ee62..cf207c674 100644
--- a/js/src/schemas.ts
+++ b/js/src/schemas.ts
@@ -57,6 +57,7 @@ export interface BaseExample {
   dataset_id: string;
   inputs: KVMap;
   outputs?: KVMap;
+  metadata?: KVMap;
 }
 
 /**
@@ -222,6 +223,7 @@ export interface ExampleUpdate {
   dataset_id?: string;
   inputs?: KVMap;
   outputs?: KVMap;
+  metadata?: KVMap;
 }
 export interface BaseDataset {
   name: string;
diff --git a/js/src/tests/client.int.test.ts b/js/src/tests/client.int.test.ts
index 144528f33..14e8afc67 100644
--- a/js/src/tests/client.int.test.ts
+++ b/js/src/tests/client.int.test.ts
@@ -446,12 +446,22 @@ test.concurrent(
       { output: "hi there" },
       {
         datasetId: dataset.id,
+        metadata: { key: "value" },
       }
     );
     const exampleValue = await client.readExample(example.id);
     const initialVersion = exampleValue.modified_at;
     expect(exampleValue.inputs.input).toEqual("hello world");
     expect(exampleValue?.outputs?.output).toEqual("hi there");
+    expect(exampleValue?.metadata?.key).toEqual("value");
+
+    // Update the example by modifying the metadata
+    await client.updateExample(example.id, {
+      metadata: { key: "new value" },
+    });
+    const updatedExampleValue = await client.readExample(example.id);
+    expect(updatedExampleValue?.metadata?.key).toEqual("new value");
+
     // Create multiple
     await client.createExamples({
       inputs: [
diff --git a/python/langsmith/client.py b/python/langsmith/client.py
index 0280d390a..51405552c 100644
--- a/python/langsmith/client.py
+++ b/python/langsmith/client.py
@@ -3043,6 +3043,8 @@ def update_example(
             The input values to update.
         outputs : Mapping[str, Any] or None, default=None
             The output values to update.
+        metadata : Dict or None, default=None
+            The metadata to update.
         dataset_id : UUID or None, default=None
             The ID of the dataset to update.
 
diff --git a/python/langsmith/schemas.py b/python/langsmith/schemas.py
index 97e5b7cfa..50947e6d4 100644
--- a/python/langsmith/schemas.py
+++ b/python/langsmith/schemas.py
@@ -104,6 +104,7 @@ class ExampleUpdate(BaseModel):
     dataset_id: Optional[UUID] = None
     inputs: Optional[Dict[str, Any]] = None
     outputs: Optional[Dict[str, Any]] = None
+    metadata: Optional[Dict[str, Any]] = None
 
     class Config:
         """Configuration class for the schema."""
diff --git a/python/tests/integration_tests/test_client.py b/python/tests/integration_tests/test_client.py
index 3ee7c81a0..f6e21f755 100644
--- a/python/tests/integration_tests/test_client.py
+++ b/python/tests/integration_tests/test_client.py
@@ -91,6 +91,7 @@ def test_datasets(langchain_client: Client) -> None:
         example_id=example.id,
         inputs={"col1": "updatedExampleCol1"},
         outputs={"col2": "updatedExampleCol2"},
+        metadata={"foo": "bar"},
     )
     updated_example = langchain_client.read_example(example.id)
     assert updated_example.id == example.id
@@ -98,6 +99,7 @@ def test_datasets(langchain_client: Client) -> None:
     assert updated_example_value.inputs["col1"] == "updatedExampleCol1"
     assert updated_example_value.outputs is not None
     assert updated_example_value.outputs["col2"] == "updatedExampleCol2"
+    assert updated_example_value.metadata["foo"] == "bar"
 
     langchain_client.delete_example(example.id)
     examples2 = list(