@arizeai/phoenix-client 1.3.0 → 2.0.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@arizeai/phoenix-client",
-  "version": "1.3.0",
+  "version": "2.0.0",
   "description": "A client for the Phoenix API",
   "main": "dist/src/index.js",
   "module": "dist/esm/index.js",
@@ -22,6 +22,10 @@
       "import": "./dist/esm/experiments/index.js",
       "require": "./dist/src/experiments/index.js"
     },
+    "./datasets": {
+      "import": "./dist/esm/datasets/index.js",
+      "require": "./dist/src/datasets/index.js"
+    },
     "./utils/*": {
       "import": "./dist/esm/utils/*.js",
       "require": "./dist/src/utils/*.js"

package/src/datasets/appendDatasetExamples.ts

@@ -0,0 +1,55 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { Example, DatasetSelector } from "../types/datasets";
+import invariant from "tiny-invariant";
+
+export type AppendDatasetExamplesParams = ClientFn & {
+  /**
+   * The dataset to append examples to
+   */
+  dataset: DatasetSelector;
+  /**
+   * The examples to append to the dataset
+   */
+  examples: Example[];
+};
+
+export type AppendDatasetExamplesResponse = {
+  datasetId: string;
+  // TODO: respond with the versionId
+  // versionId: string;
+};
+
+/**
+ * Append examples to an existing dataset
+ * @experimental this interface may change in the future
+ */
+export async function appendDatasetExamples({
+  client: _client,
+  dataset,
+  examples,
+}: AppendDatasetExamplesParams): Promise<AppendDatasetExamplesResponse> {
+  const client = _client || createClient();
+  const inputs = examples.map((example) => example.input);
+  const outputs = examples.map((example) => example.output ?? {}); // Treat null as an empty object
+  const metadata = examples.map((example) => example.metadata);
+  const appendResponse = await client.POST("/v1/datasets/upload", {
+    params: {
+      query: {
+        sync: true,
+      },
+    },
+    body: {
+      name: dataset.datasetId,
+      action: "append",
+      inputs,
+      outputs,
+      metadata,
+    },
+  });
+  invariant(appendResponse.data?.data, "Failed to append dataset examples");
+  const datasetId = appendResponse.data.data.dataset_id;
+  return {
+    datasetId,
+  };
+}

package/src/datasets/createDataset.ts

@@ -0,0 +1,60 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { Example } from "../types/datasets";
+import invariant from "tiny-invariant";
+
+export type CreateDatasetParams = ClientFn & {
+  /**
+   * The name of the dataset
+   */
+  name: string;
+  /**
+   * The description of the dataset
+   */
+  description: string;
+  /**
+   * The examples to create in the dataset
+   */
+  examples: Example[];
+};
+
+export type CreateDatasetResponse = {
+  datasetId: string;
+};
+
+/**
+ * Create a new dataset
+ * @experimental this interface may change in the future
+ */
+export async function createDataset({
+  client: _client,
+  name,
+  description,
+  examples,
+}: CreateDatasetParams): Promise<CreateDatasetResponse> {
+  const client = _client || createClient();
+  const inputs = examples.map((example) => example.input);
+  const outputs = examples.map((example) => example.output ?? {}); // Treat null as an empty object
+  const metadata = examples.map((example) => example.metadata);
+  const createDatasetResponse = await client.POST("/v1/datasets/upload", {
+    params: {
+      query: {
+        // TODO: parameterize this
+        sync: true,
+      },
+    },
+    body: {
+      name,
+      description,
+      action: "create",
+      inputs,
+      outputs,
+      metadata,
+    },
+  });
+  invariant(createDatasetResponse.data?.data, "Failed to create dataset");
+  const datasetId = createDatasetResponse.data.data.dataset_id;
+  return {
+    datasetId,
+  };
+}

package/src/datasets/getDataset.ts

@@ -0,0 +1,27 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { DatasetSelector, Dataset } from "../types/datasets";
+import { getDatasetExamples } from "./getDatasetExamples";
+import { getDatasetInfo } from "./getDatasetInfo";
+
+export type GetDatasetParams = ClientFn & {
+  dataset: DatasetSelector;
+};
+
+/**
+ * Get dataset info and the examples from the latest version of the dataset
+ */
+export async function getDataset({
+  client: _client,
+  dataset,
+}: GetDatasetParams): Promise<Dataset> {
+  const client = _client || createClient();
+  const [datasetInfo, datasetExamples] = await Promise.all([
+    getDatasetInfo({ client, dataset }),
+    getDatasetExamples({ client, dataset }),
+  ]);
+  return {
+    ...datasetInfo,
+    ...datasetExamples,
+  };
+}

package/src/datasets/getDatasetExamples.ts

@@ -0,0 +1,34 @@
+import invariant from "tiny-invariant";
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { DatasetSelector, DatasetExamples } from "../types/datasets";
+
+export type GetDatasetExamplesParams = ClientFn & {
+  dataset: DatasetSelector;
+};
+
+/**
+ * Get the latest examples from a dataset
+ */
+export async function getDatasetExamples({
+  client: _client,
+  dataset,
+}: GetDatasetExamplesParams): Promise<DatasetExamples> {
+  const client = _client || createClient();
+  const response = await client.GET("/v1/datasets/{id}/examples", {
+    params: {
+      path: {
+        id: dataset.datasetId,
+      },
+    },
+  });
+  invariant(response.data?.data, "Failed to get dataset examples");
+  const examplesData = response.data.data;
+  return {
+    versionId: examplesData.version_id,
+    examples: examplesData.examples.map((example) => ({
+      ...example,
+      updatedAt: new Date(example.updated_at),
+    })),
+  };
+}

package/src/datasets/getDatasetInfo.ts

@@ -0,0 +1,34 @@
+import invariant from "tiny-invariant";
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { DatasetSelector, DatasetInfo } from "../types/datasets";
+
+export type GetDatasetInfoParams = ClientFn & {
+  dataset: DatasetSelector;
+};
+
+/**
+ * Get an overview of the information in a dataset
+ * Note: this does not include the examples contained in the dataset
+ */
+export async function getDatasetInfo({
+  client: _client,
+  dataset,
+}: GetDatasetInfoParams): Promise<DatasetInfo> {
+  const client = _client || createClient();
+  const datasetResponse = await client.GET("/v1/datasets/{id}", {
+    params: {
+      path: {
+        id: dataset.datasetId,
+      },
+    },
+  });
+  invariant(datasetResponse.data?.data, "Failed to get dataset info");
+  const datasetInfo = datasetResponse.data.data;
+  return {
+    id: datasetInfo.id,
+    name: datasetInfo.name,
+    description: datasetInfo.description || undefined,
+    metadata: datasetInfo.metadata,
+  };
+}

package/src/datasets/index.ts

@@ -0,0 +1,6 @@
+export * from "./createDataset";
+export * from "./getDataset";
+export * from "./getDatasetExamples";
+export * from "./appendDatasetExamples";
+export * from "./getDatasetExamples";
+export * from "./getDatasetInfo";

package/src/datasets/listDatasets.ts

@@ -0,0 +1,37 @@
+import { createClient } from "../client";
+import { DatasetInfo } from "../types/datasets";
+import { ClientFn } from "../types/core";
+import invariant from "tiny-invariant";
+
+export type ListDatasetsParams = ClientFn;
+
+type FullDatasetInfo = DatasetInfo & {
+  startDate: Date;
+  endDate: Date;
+};
+
+/**
+ * List the information about all datasets available to the client.
+ *
+ * @example
+ * ```ts
+ * import { listDatasets } from "@arizeai/phoenix-client/datasets";
+ *
+ * const datasets = await listDatasets({});
+ * console.log(datasets);
+ * ```
+ *
+ * @throws {Error} If the datasets cannot be listed or the response is invalid.
+ */
+export async function listDatasets({
+  client: _client,
+}: ListDatasetsParams): Promise<FullDatasetInfo[]> {
+  const client = _client || createClient();
+  const response = await client.GET("/v1/datasets");
+  invariant(response.data?.data, "Failed to list datasets");
+  return response.data.data.map((dataset) => ({
+    ...dataset,
+    startDate: new Date(dataset.created_at),
+    endDate: new Date(dataset.updated_at),
+  }));
+}

package/src/experiments/getExperiment.ts

@@ -0,0 +1,40 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { ExperimentRunsMap, RanExperiment } from "../types/experiments";
+import { getExperimentInfo } from "./getExperimentInfo";
+import { getExperimentRuns } from "./getExperimentRuns";
+
+export type GetExperimentResultParams = ClientFn & {
+  /**
+   * The experiment ID.
+   */
+  experimentId: string;
+};
+
+/**
+ * A function that gets the result of a experiment.
+ * Fetches the experiment data as well as the runs.
+ */
+export async function getExperiment({
+  client: _client,
+  experimentId,
+}: GetExperimentResultParams): Promise<RanExperiment> {
+  const client = _client || createClient();
+  const [experiment, experimentRuns] = await Promise.all([
+    getExperimentInfo({ client, experimentId }),
+    getExperimentRuns({ client, experimentId }),
+  ]);
+  const experimentRunsMap: ExperimentRunsMap = {
+    runs: experimentRuns.runs.reduce(
+      (acc, run) => {
+        acc[run.id] = run;
+        return acc;
+      },
+      {} as ExperimentRunsMap["runs"]
+    ),
+  };
+  return {
+    ...experiment,
+    ...experimentRunsMap,
+  };
+}

package/src/experiments/getExperimentInfo.ts

@@ -0,0 +1,39 @@
+import invariant from "tiny-invariant";
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { type ExperimentInfo } from "../types/experiments";
+
+export type GetExperimentParams = ClientFn & {
+  /**
+   * The experiment ID
+   */
+  experimentId: string;
+};
+
+/**
+ * Returns an object containing the high-level info about an experiment
+ */
+export async function getExperimentInfo({
+  client: _client,
+  experimentId: experiment_id,
+}: GetExperimentParams): Promise<ExperimentInfo> {
+  const client = _client || createClient();
+  const { data: { data: experimentData } = {} } = await client.GET(
+    "/v1/experiments/{experiment_id}",
+    {
+      params: {
+        path: {
+          experiment_id,
+        },
+      },
+    }
+  );
+  invariant(experimentData, "Failed to get experiment");
+  return {
+    id: experimentData.id,
+    datasetId: experimentData.dataset_id,
+    datasetVersionId: experimentData.dataset_version_id,
+    projectName: experimentData.project_name || "", // This will never happen
+    metadata: experimentData.metadata,
+  };
+}

package/src/experiments/getExperimentRuns.ts

@@ -0,0 +1,45 @@
+import { createClient } from "../client";
+import invariant from "tiny-invariant";
+import { ClientFn } from "../types/core";
+import { ExperimentRun } from "../types/experiments";
+
+export type GetExperimentRunsParams = ClientFn & {
+  /**
+   * The experiment ID.
+   */
+  experimentId: string;
+};
+
+/**
+ * A function that gets the runs (e.g. the results) of a experiment
+ */
+export async function getExperimentRuns({
+  client: _client,
+  experimentId,
+}: GetExperimentRunsParams): Promise<{ runs: ExperimentRun[] }> {
+  const client = _client || createClient();
+  const getRunsPromise = client.GET("/v1/experiments/{experiment_id}/runs", {
+    params: {
+      path: {
+        experiment_id: experimentId,
+      },
+    },
+  });
+  const [experimentRunResponse] = await Promise.all([getRunsPromise]);
+  const { data: { data: experimentRunsData } = {} } = experimentRunResponse;
+  invariant(experimentRunsData, "Failed to retrieve experiment runs");
+  return {
+    runs: experimentRunsData.map((run) => {
+      return {
+        id: run.id,
+        traceId: run.trace_id || null,
+        experimentId: run.experiment_id,
+        datasetExampleId: run.dataset_example_id,
+        startTime: new Date(run.start_time),
+        endTime: new Date(run.end_time),
+        output: run.output as ExperimentRun["output"],
+        error: run.error || null,
+      };
+    }),
+  };
+}

package/src/experiments/index.ts

@@ -1 +1,4 @@
+export * from "./getExperimentInfo";
+export * from "./getExperiment";
+export * from "./getExperimentRuns";
 export * from "./runExperiment";

package/src/experiments/runExperiment.ts

@@ -2,17 +2,23 @@ import { queue } from "async";
 import invariant from "tiny-invariant";
 import { createClient, type PhoenixClient } from "../client";
 import { ClientFn } from "../types/core";
-import { Dataset, Example } from "../types/datasets";
+import {
+  Dataset,
+  DatasetSelector,
+  Example,
+  ExampleWithId,
+} from "../types/datasets";
 import type {
   Evaluator,
-  Experiment,
+  ExperimentInfo,
   ExperimentEvaluationRun,
   ExperimentRun,
+  ExperimentRunID,
   ExperimentTask,
   RanExperiment,
 } from "../types/experiments";
 import { type Logger } from "../types/logger";
-import { getDatasetBySelector } from "../utils/getDatasetBySelector";
+import { getDataset } from "../datasets/getDataset";
 import { pluralize } from "../utils/pluralize";
 import { promisifyResult } from "../utils/promisifyResult";
 import { AnnotatorKind } from "../types/annotations";
@@ -44,12 +50,13 @@ export type RunExperimentParams = ClientFn & {
   experimentDescription?: string;
   /**
    * Experiment metadata
+   * E.x. modelName
    */
   experimentMetadata?: Record<string, unknown>;
   /**
    * The dataset to run the experiment on
    */
-  dataset: Dataset | string | Example[];
+  dataset: DatasetSelector;
   /**
    * The task to run
    */
@@ -112,9 +119,9 @@ export type RunExperimentParams = ClientFn & {
 export async function runExperiment({
   experimentName,
   experimentDescription,
-  experimentMetadata,
+  experimentMetadata = {},
   client: _client,
-  dataset: _dataset,
+  dataset: DatasetSelector,
   task,
   evaluators,
   logger = console,
@@ -125,24 +132,25 @@ export async function runExperiment({
   let provider: NodeTracerProvider | undefined;
   const isDryRun = typeof dryRun === "number" || dryRun === true;
   const client = _client ?? createClient();
-  const dataset = await getDatasetBySelector({ dataset: _dataset, client });
+  const dataset = await getDataset({ dataset: DatasetSelector, client });
   invariant(dataset, `Dataset not found`);
   invariant(dataset.examples.length > 0, `Dataset has no examples`);
   const nExamples =
     typeof dryRun === "number"
-      ? Math.max(dryRun, dataset.examples.length)
+      ? Math.min(dryRun, dataset.examples.length)
       : dataset.examples.length;
 
   let projectName = `${dataset.name}-exp-${new Date().toISOString()}`;
   // initialize the tracer into scope
   let taskTracer: Tracer;
-  let experiment: Experiment;
+  let experiment: ExperimentInfo;
   if (isDryRun) {
     experiment = {
       id: localId(),
       datasetId: dataset.id,
       datasetVersionId: dataset.versionId,
       projectName,
+      metadata: experimentMetadata,
     };
     taskTracer = createNoOpProvider().getTracer("no-op");
   } else {
@@ -165,9 +173,10 @@ export async function runExperiment({
     projectName = experimentResponse.project_name ?? projectName;
     experiment = {
       id: experimentResponse.id,
-      datasetId: dataset.id,
-      datasetVersionId: dataset.versionId,
+      datasetId: experimentResponse.dataset_id,
+      datasetVersionId: experimentResponse.dataset_version_id,
       projectName,
+      metadata: experimentResponse.metadata,
     };
     // Initialize the tracer, now that we have a project name
     const baseUrl = client.config.baseUrl;
@@ -189,16 +198,14 @@ export async function runExperiment({
   }
 
   logger.info(
-    `🧪 Starting experiment "${experimentName}" on dataset "${dataset.id}" with task "${task.name}" and ${evaluators?.length ?? 0} ${pluralize(
+    `🧪 Starting experiment "${experimentName || `<unnamed>`}" on dataset "${dataset.id}" with task "${task.name}" and ${evaluators?.length ?? 0} ${pluralize(
       "evaluator",
       evaluators?.length ?? 0
     )} and ${concurrency} concurrent runs`
   );
 
-  // Run task against all examples, for each repetition
-  type ExperimentRunId = string;
-  const runs: Record<ExperimentRunId, ExperimentRun> = {};
-  await runTask({
+  const runs: Record<ExperimentRunID, ExperimentRun> = {};
+  await runTaskWithExamples({
     client,
     experimentId: experiment.id,
     task,
@@ -242,7 +249,7 @@ export async function runExperiment({
 /**
  * Run a task against n examples in a dataset.
  */
-function runTask({
+function runTaskWithExamples({
   client,
   experimentId,
   task,
@@ -274,9 +281,9 @@ function runTask({
   nExamples: number;
   /** TraceProvider instance that will be used to create spans from task calls */
   tracer: Tracer;
-}) {
+}): Promise<void> {
   logger.info(`🔧 Running task "${task.name}" on dataset "${dataset.id}"`);
-  const run = async (example: Example) => {
+  const run = async (example: ExampleWithId) => {
     return tracer.startActiveSpan(`Task: ${task.name}`, async (span) => {
       logger.info(
         `🔧 Running task "${task.name}" on example "${example.id} of dataset "${dataset.id}"`
@@ -366,13 +373,12 @@ export async function evaluateExperiment({
   experiment,
   evaluators,
   client: _client,
-  logger,
+  logger = console,
   concurrency = 5,
   dryRun = false,
 }: {
   /**
    * The experiment to evaluate
-   * @todo also accept Experiment, and attempt to fetch the runs from the server
    **/
   experiment: RanExperiment;
   /** The evaluators to use */
@@ -380,9 +386,9 @@ export async function evaluateExperiment({
   /** The client to use */
   client?: PhoenixClient;
   /** The logger to use */
-  logger: Logger;
+  logger?: Logger;
   /** The number of evaluators to run in parallel */
-  concurrency: number;
+  concurrency?: number;
   /**
    * Whether to run the evaluation as a dry run
    * If a number is provided, the evaluation will be run for the first n runs
@@ -414,8 +420,8 @@ export async function evaluateExperiment({
     typeof dryRun === "number"
       ? Math.max(dryRun, Object.keys(experiment.runs).length)
       : Object.keys(experiment.runs).length;
-  const dataset = await getDatasetBySelector({
-    dataset: experiment.datasetId,
+  const dataset = await getDataset({
+    dataset: { datasetId: experiment.datasetId },
     client,
   });
   invariant(dataset, `Dataset "${experiment.datasetId}" not found`);

package/src/types/datasets.ts

@@ -1,22 +1,48 @@
 import { Node } from "./core";
 
+/**
+ * A dataset can be identified by its datasetId
+ * TODO: add support for datasetName and datasetVersionId via discriminated union
+ */
+export type DatasetSelector = { datasetId: string };
+
+/**
+ * Overview information about a dataset
+ */
+export interface DatasetInfo extends Node {
+  name: string;
+  description?: string | null;
+  metadata?: Record<string, unknown>;
+}
+
+/**
+ * A dataset's examples
+ */
+export interface DatasetExamples {
+  examples: ExampleWithId[];
+  /**
+   * The version ID of the dataset examples
+   */
+  versionId: string;
+}
+
 /**
  * An example is a record to feed into an AI task
  */
-export interface Example extends Node {
-  id: string;
-  updatedAt: Date;
+export interface Example {
   input: Record<string, unknown>;
   output: Record<string, unknown> | null;
   metadata: Record<string, unknown>;
 }
 
 /**
- * A dataset is a collection of examples for an AI task
+ * An example that has been synced to the server
  */
-export interface Dataset extends Node {
-  id: string;
-  name: string;
-  versionId: string;
-  examples: Example[];
+export interface ExampleWithId extends Example, Node {
+  updatedAt: Date;
 }
+
+/**
+ * A dataset is a collection of examples for an AI task
+ */
+export interface Dataset extends DatasetInfo, DatasetExamples, Node {}

package/src/types/experiments.ts

@@ -5,17 +5,33 @@ import { Example } from "./datasets";
 /**
  * An experiment is a set of task runs on a dataset version
  */
-export interface Experiment extends Node {
+export interface ExperimentInfo extends Node {
   datasetId: string;
   datasetVersionId: string;
   /**
    * The project under which the experiment task traces are recorded
    */
   projectName: string;
+  /**
+   * Metadata about the experiment as an object of key values
+   * e.x. model name
+   */
+  metadata: Record<string, unknown>;
+}
+
+export type ExperimentRunID = string;
+
+/**
+ * A map of an experiment runId to the run
+ */
+export interface ExperimentRunsMap {
+  runs: Record<ExperimentRunID, ExperimentRun>;
 }
 
-export interface RanExperiment extends Experiment {
-  runs: Record<string, ExperimentRun>;
+/**
+ * An experiment that has been run and been recorded on the server
+ */
+export interface RanExperiment extends ExperimentInfo, ExperimentRunsMap {
   evaluationRuns?: ExperimentEvaluationRun[];
 }