feat: support few shot in js (#933)

Please review closely because I am not a JS dev

Adds two few shot related methods for indexing and finding similar
examples to the js langsmith client. Mirrors implementations in python
here:
https://github.com/langchain-ai/langsmith-sdk/pull/925/files#diff-619399e55cdfbfd5de40fb5e942606b34a9d5e16c289ae44e02d7a4a46d82e4fR14

js/package.json

@@ -276,4 +276,4 @@
     },
     "./package.json": "./package.json"
   }
-}
+}

js/src/client.ts

@@ -27,6 +27,7 @@ import {
   RunCreate,
   RunUpdate,
   ScoreType,
+  ExampleSearch,
   TimeDelta,
   TracerSession,
   TracerSessionResult,
@@ -1892,7 +1893,14 @@ export class Client {
     {
       description,
       dataType,
-    }: { description?: string; dataType?: DataType } = {}
+      inputsSchema,
+      outputsSchema,
+    }: {
+      description?: string;
+      dataType?: DataType;
+      inputsSchema?: KVMap;
+      outputsSchema?: KVMap;
+    } = {}
   ): Promise<Dataset> {
     const body: KVMap = {
       name,
@@ -1901,6 +1909,12 @@ export class Client {
     if (dataType) {
       body.data_type = dataType;
     }
+    if (inputsSchema) {
+      body.inputs_schema_definition = inputsSchema;
+    }
+    if (outputsSchema) {
+      body.outputs_schema_definition = outputsSchema;
+    }
     const response = await this.caller.call(fetch, `${this.apiUrl}/datasets`, {
       method: "POST",
       headers: { ...this.headers, "Content-Type": "application/json" },
@@ -2148,6 +2162,104 @@ export class Client {
     await response.json();
   }
 
+  public async indexDataset({
+    datasetId,
+    datasetName,
+    tag,
+  }: {
+    datasetId?: string;
+    datasetName?: string;
+    tag?: string;
+  }): Promise<void> {
+    let datasetId_ = datasetId;
+    if (!datasetId_ && !datasetName) {
+      throw new Error("Must provide either datasetName or datasetId");
+    } else if (datasetId_ && datasetName) {
+      throw new Error("Must provide either datasetName or datasetId, not both");
+    } else if (!datasetId_) {
+      const dataset = await this.readDataset({ datasetName });
+      datasetId_ = dataset.id;
+    }
+    assertUuid(datasetId_);
+
+    const data = {
+      tag: tag,
+    };
+    const response = await this.caller.call(
+      fetch,
+      `${this.apiUrl}/datasets/${datasetId_}/index`,
+      {
+        method: "POST",
+        headers: { ...this.headers, "Content-Type": "application/json" },
+        body: JSON.stringify(data),
+        signal: AbortSignal.timeout(this.timeout_ms),
+        ...this.fetchOptions,
+      }
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to index dataset ${datasetId_}: ${response.status} ${response.statusText}`
+      );
+    }
+    await response.json();
+  }
+
+  /**
+   * Lets you run a similarity search query on a dataset.
+   *
+   * Requires the dataset to be indexed. Please see the `indexDataset` method to set up indexing.
+   *
+   * @param inputs      The input on which to run the similarity search. Must have the
+   *                    same schema as the dataset.
+   *
+   * @param datasetId   The dataset to search for similar examples.
+   *
+   * @param limit       The maximum number of examples to return. Will return the top `limit` most
+   *                    similar examples in order of most similar to least similar. If no similar
+   *                    examples are found, random examples will be returned.
+   *
+   * @returns           A list of similar examples.
+   *
+   *
+   * @example
+   * dataset_id = "123e4567-e89b-12d3-a456-426614174000"
+   * inputs = {"text": "How many people live in Berlin?"}
+   * limit = 5
+   * examples = await client.similarExamples(inputs, dataset_id, limit)
+   */
+  public async similarExamples(
+    inputs: KVMap,
+    datasetId: string,
+    limit: number
+  ): Promise<ExampleSearch[]> {
+    const data = {
+      limit: limit,
+      inputs: inputs,
+    };
+
+    assertUuid(datasetId);
+    const response = await this.caller.call(
+      fetch,
+      `${this.apiUrl}/datasets/${datasetId}/search`,
+      {
+        method: "POST",
+        headers: { ...this.headers, "Content-Type": "application/json" },
+        body: JSON.stringify(data),
+        signal: AbortSignal.timeout(this.timeout_ms),
+        ...this.fetchOptions,
+      }
+    );
+
+    if (!response.ok) {
+      throw new Error(
+        `Failed to fetch similar examples: ${response.status} ${response.statusText}`
+      );
+    }
+
+    const result = await response.json();
+    return result["examples"] as ExampleSearch[];
+  }
+
   public async createExample(
     inputs: KVMap,
     outputs: KVMap,

js/src/schemas.ts

@@ -251,6 +251,11 @@ export interface ExampleUpdate {
 export interface ExampleUpdateWithId extends ExampleUpdate {
   id: string;
 }
+
+export interface ExampleSearch extends BaseExample {
+  id: string;
+}
+
 export interface BaseDataset {
   name: string;
   description: string;

js/src/tests/few_shot.int.test.ts

@@ -0,0 +1,65 @@
+import { KVMap, ExampleSearch } from "../schemas.js";
+import { Client } from "../index.js";
+import { v4 as uuidv4 } from "uuid";
+
+const TESTING_DATASET_NAME = `test_dataset_few_shot_js_${uuidv4()}`;
+
+test("evaluate can evaluate", async () => {
+  const client = new Client();
+
+  const schema: KVMap = {
+    type: "object",
+    properties: {
+      name: { type: "string" },
+    },
+    required: ["name"],
+    additionalProperties: false,
+  };
+
+  const has_dataset = await client.hasDataset({
+    datasetName: TESTING_DATASET_NAME,
+  });
+  if (has_dataset === true) {
+    await client.deleteDataset({ datasetName: TESTING_DATASET_NAME });
+  }
+
+  const dataset = await client.createDataset(TESTING_DATASET_NAME, {
+    description:
+      "For testing purposed. Is created & deleted for each test run.",
+    inputsSchema: schema,
+  });
+
+  // create examples
+  const res = await client.createExamples({
+    inputs: [{ name: "foo" }, { name: "bar" }],
+    outputs: [{ output: 2 }, { output: 3 }],
+    datasetName: TESTING_DATASET_NAME,
+  });
+  if (res.length !== 2) {
+    throw new Error("Failed to create examples");
+  }
+
+  await client.indexDataset({ datasetId: dataset.id });
+
+  let i = 0;
+  let examples: ExampleSearch[] = [];
+  while (i < 10) {
+    examples = await client.similarExamples(
+      { name: "foo" },
+      dataset.id,
+      // specify limit of 5 so you return all examples
+      5
+    );
+    if (examples.length === 2) {
+      break;
+    }
+
+    // sleep for one second
+    await new Promise((r) => setTimeout(r, 1000));
+    i++;
+  }
+
+  expect(examples.length).toBe(2);
+  expect(examples[0].inputs).toEqual({ name: "foo" });
+  expect(examples[1].inputs).toEqual({ name: "bar" });
+});