@arizeai/phoenix-client 1.0.2 → 1.2.0

package/src/experiments/runExperiment.ts

@@ -1,5 +1,8 @@
-import { Dataset, Example } from "../types/datasets";
+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 type {
   Evaluator,
   Experiment,
@@ -9,18 +12,16 @@ import type {
   ExperimentTask,
   RanExperiment,
 } from "../types/experiments";
-import { promisifyResult } from "../utils/promisifyResult";
-import invariant from "tiny-invariant";
-import { pluralize } from "../utils/pluralize";
-import { ClientFn } from "../types/core";
-import { getDatasetBySelector } from "../utils/getDatasetBySelector";
 import { type Logger } from "../types/logger";
+import { getDatasetBySelector } from "../utils/getDatasetBySelector";
+import { pluralize } from "../utils/pluralize";
+import { promisifyResult } from "../utils/promisifyResult";
+import { AnnotatorKind } from "../types/annotations";
 
 /**
  * Parameters for running an experiment.
  *
  * @experimental This feature is not complete, and will change in the future.
- * @deprecated This function will be un-marked as deprecated once the experimental feature flag is removed.
  */
 export type RunExperimentParams = ClientFn & {
   /**
@@ -28,6 +29,14 @@ export type RunExperimentParams = ClientFn & {
    * Defaults to the dataset name + a timestamp
    */
   experimentName?: string;
+  /**
+   * The description of the experiment
+   */
+  experimentDescription?: string;
+  /**
+   * Experiment metadata
+   */
+  experimentMetadata?: Record<string, unknown>;
   /**
    * The dataset to run the experiment on
    */
@@ -40,10 +49,6 @@ export type RunExperimentParams = ClientFn & {
    * The evaluators to use
    */
   evaluators?: Evaluator[];
-  /**
-   * The number of repetitions to run
-   */
-  repetitions?: number;
   /**
    * The project under which the experiment task traces are recorded
    */
@@ -56,43 +61,96 @@ export type RunExperimentParams = ClientFn & {
    * Whether to record the experiment results
    */
   record?: boolean;
+  /**
+   * The number of dataset examples to run in parallel
+   */
+  concurrency?: number;
+  /**
+   * Whether or not to run the experiment as a dry run. If a number is privided, n examples will be run.
+   * @default false
+   */
+  dryRun?: number | boolean;
 };
 
 /**
  * Run an experiment.
  *
+ * @example
+ * ```ts
+ * import { asEvaluator, runExperiment } from "@phoenix/client/experiments";
+ *
+ * const experiment = await runExperiment({
+ *   dataset: "my-dataset",
+ *   task: async (example) => example.input,
+ *   evaluators: [
+ *     asEvaluator("my-evaluator", "CODE", async (params) => params.output),
+ *   ],
+ * });
+ * ```
+ *
  * @experimental This feature is not complete, and will change in the future.
- * @deprecated This function will be un-marked as deprecated once the experimental feature flag is removed.
  */
 export async function runExperiment({
   experimentName: _experimentName,
+  experimentDescription,
+  experimentMetadata,
   client: _client,
   dataset: _dataset,
   task,
   evaluators,
-  repetitions = 1,
   projectName = "default",
   logger = console,
   record = true,
+  concurrency = 5,
+  dryRun = false,
 }: RunExperimentParams): Promise<RanExperiment> {
+  const isDryRun = typeof dryRun === "number" || dryRun === true;
   const client = _client ?? createClient();
   const dataset = await getDatasetBySelector({ dataset: _dataset, 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)
+      : dataset.examples.length;
+
   const experimentName =
     _experimentName ?? `${dataset.name}-${new Date().toISOString()}`;
   const experimentParams: ExperimentParameters = {
-    nRepetitions: repetitions,
-    // TODO: Make configurable?
-    nExamples: dataset.examples.length,
-  };
-  const experiment: Experiment = {
-    id: id(),
-    datasetId: dataset.id,
-    datasetVersionId: dataset.versionId,
-    repetitions,
-    projectName,
+    nExamples,
   };
+  let experiment: Experiment;
+  if (isDryRun) {
+    experiment = {
+      id: id(),
+      datasetId: dataset.id,
+      datasetVersionId: dataset.versionId,
+      projectName,
+    };
+  } else {
+    const experimentResponse = await client
+      .POST("/v1/datasets/{dataset_id}/experiments", {
+        params: {
+          path: {
+            dataset_id: dataset.id,
+          },
+        },
+        body: {
+          name: experimentName,
+          description: experimentDescription,
+          metadata: experimentMetadata,
+          project_name: projectName,
+        },
+      })
+      .then((res) => res.data?.data);
+    invariant(experimentResponse, `Failed to create experiment`);
+    experiment = {
+      id: experimentResponse.id,
+      datasetId: dataset.id,
+      datasetVersionId: dataset.versionId,
+      projectName,
+    };
+  }
 
   if (!record) {
     logger.info(
@@ -104,30 +162,25 @@ export async function runExperiment({
     `🧪 Starting experiment "${experimentName}" on dataset "${dataset.id}" with task "${task.name}" and ${evaluators?.length ?? 0} ${pluralize(
       "evaluator",
       evaluators?.length ?? 0
-    )}`
-  );
-
-  logger.info(
-    `🔁 Running ${repetitions} ${pluralize("repetition", repetitions)} of task "${task.name}"`
+    )} and ${concurrency} concurrent runs`
   );
 
   // Run task against all examples, for each repetition
   type ExperimentRunId = string;
   const runs: Record<ExperimentRunId, ExperimentRun> = {};
-  await Promise.all(
-    Array.from({ length: repetitions }, (_, i) =>
-      runTask({
-        repetition: i + 1,
-        experimentId: experiment.id,
-        task,
-        dataset,
-        logger,
-        onComplete: (run) => {
-          runs[run.id] = run;
-        },
-      })
-    )
-  );
+  await runTask({
+    client,
+    experimentId: experiment.id,
+    task,
+    dataset,
+    logger,
+    onComplete: (run) => {
+      runs[run.id] = run;
+    },
+    concurrency,
+    isDryRun,
+    nExamples,
+  });
   logger.info(`✅ Task runs completed`);
 
   const ranExperiment: RanExperiment = {
@@ -141,6 +194,8 @@ export async function runExperiment({
     evaluators: evaluators ?? [],
     client,
     logger,
+    concurrency,
+    dryRun,
   });
   ranExperiment.evaluationRuns = evaluationRuns;
 
@@ -150,40 +205,49 @@ export async function runExperiment({
 }
 
 /**
- * Run a task against all examples in a dataset.
+ * Run a task against n examples in a dataset.
  */
 function runTask({
+  client,
   experimentId,
   task,
   dataset,
-  repetition,
   onComplete,
   logger,
+  concurrency = 5,
+  isDryRun,
+  nExamples,
 }: {
+  /** The client to use */
+  client: PhoenixClient;
   /** The id of the experiment */
   experimentId: string;
   /** The task to run */
   task: ExperimentTask;
   /** The dataset to run the task on */
   dataset: Dataset;
-  /** The repetition number */
-  repetition: number;
   /** A callback to call when the task is complete */
   onComplete: (run: ExperimentRun) => void;
   /** The logger to use */
   logger: Logger;
+  /** The number of examples to run in parallel */
+  concurrency: number;
+  /** Whether to run the task as a dry run */
+  isDryRun: boolean;
+  /** The number of examples to run */
+  nExamples: number;
 }) {
-  logger.info(
-    `🔧 (${repetition}) Running task "${task.name}" on dataset "${dataset.id}"`
-  );
+  logger.info(`🔧 Running task "${task.name}" on dataset "${dataset.id}"`);
   const run = async (example: Example) => {
+    logger.info(
+      `🔧 Running task "${task.name}" on example "${example.id} of dataset "${dataset.id}"`
+    );
     const thisRun: ExperimentRun = {
       id: id(),
-      traceId: id(),
+      traceId: null, // TODO: fill this in once we trace experiments
       experimentId,
       datasetExampleId: example.id,
       startTime: new Date(),
-      repetitionNumber: repetition,
       endTime: new Date(), // will get replaced with actual end time
       output: null,
       error: null,
@@ -199,22 +263,49 @@ function runTask({
       thisRun.error = error instanceof Error ? error.message : "Unknown error";
     }
     thisRun.endTime = new Date();
+    if (!isDryRun) {
+      // Log the run to the server
+      // We log this without awaiting (e.g. best effort)
+      const res = await client.POST("/v1/experiments/{experiment_id}/runs", {
+        params: {
+          path: {
+            experiment_id: experimentId,
+          },
+        },
+        body: {
+          dataset_example_id: example.id,
+          output: thisRun.output,
+          repetition_number: 0,
+          start_time: thisRun.startTime.toISOString(),
+          end_time: thisRun.endTime.toISOString(),
+          trace_id: thisRun.traceId,
+          error: thisRun.error,
+        },
+      });
+      // replace the local run id with the server-assigned id
+      thisRun.id = res.data?.data.id ?? thisRun.id;
+    }
     onComplete(thisRun);
+    return thisRun;
   };
-  return Promise.all(dataset.examples.map(run));
+  const q = queue(run, concurrency);
+  const examplesToUse = dataset.examples.slice(0, nExamples);
+  examplesToUse.forEach((example) => q.push(example));
+  return q.drain();
 }
 
 /**
  * Evaluate an experiment.
  *
  * @experimental This feature is not complete, and will change in the future.
- * @deprecated This function will be un-marked as deprecated once the experimental feature flag is removed.
  */
 export async function evaluateExperiment({
   experiment,
   evaluators,
   client: _client,
   logger,
+  concurrency = 5,
+  dryRun = false,
 }: {
   /**
    * The experiment to evaluate
@@ -227,7 +318,20 @@ export async function evaluateExperiment({
   client?: PhoenixClient;
   /** The logger to use */
   logger: Logger;
+  /** The number of evaluators to run in parallel */
+  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
+   * @default false
+   * */
+  dryRun?: boolean | number;
 }): Promise<RanExperiment> {
+  const isDryRun = typeof dryRun === "number" || dryRun === true;
+  const nRuns =
+    typeof dryRun === "number"
+      ? Math.max(dryRun, Object.keys(experiment.runs).length)
+      : Object.keys(experiment.runs).length;
   const client = _client ?? createClient();
   const dataset = await getDatasetBySelector({
     dataset: experiment.datasetId,
@@ -240,6 +344,8 @@ export async function evaluateExperiment({
   );
   invariant(experiment.runs, `Experiment "${experiment.id}" has no runs`);
 
+  const runsToEvaluate = Object.values(experiment.runs).slice(0, nRuns);
+
   if (evaluators?.length === 0) {
     return {
       ...experiment,
@@ -266,21 +372,47 @@ export async function evaluateExperiment({
   };
 
   // Run evaluators against all runs
-  await Promise.all(
-    evaluators.map((evaluator) =>
-      Promise.all(
-        Object.values(experiment.runs).map((run) =>
-          runEvaluator({
-            evaluator,
-            run,
-            exampleCache: examplesById,
-            onComplete: onEvaluationComplete,
-          })
-        )
-      )
-    )
+  // Flat list of evaluator + run tuples
+  const evaluatorsAndRuns = evaluators.flatMap((evaluator) =>
+    runsToEvaluate.map((run) => ({
+      evaluator,
+      run,
+    }))
   );
-
+  const evaluatorsQueue = queue(
+    async (evaluatorAndRun: { evaluator: Evaluator; run: ExperimentRun }) => {
+      const evalResult = await runEvaluator({
+        evaluator: evaluatorAndRun.evaluator,
+        run: evaluatorAndRun.run,
+        exampleCache: examplesById,
+        onComplete: onEvaluationComplete,
+      });
+      if (!isDryRun) {
+        logger.info(`📝 Logging evaluation ${evalResult.id}`);
+        // Log the evaluation to the server
+        // We log this without awaiting (e.g. best effort)
+        client.POST("/v1/experiment_evaluations", {
+          body: {
+            experiment_run_id: evaluatorAndRun.run.id,
+            name: evaluatorAndRun.evaluator.name,
+            annotator_kind: evaluatorAndRun.evaluator.kind,
+            start_time: evalResult.startTime.toISOString(),
+            end_time: evalResult.endTime.toISOString(),
+            result: {
+              ...evalResult.result,
+            },
+            error: evalResult.error,
+            trace_id: evalResult.traceId,
+          },
+        });
+      }
+    },
+    concurrency
+  );
+  evaluatorsAndRuns.forEach((evaluatorAndRun) =>
+    evaluatorsQueue.push(evaluatorAndRun)
+  );
+  await evaluatorsQueue.drain();
   logger.info(`✅ Evaluation runs completed`);
 
   return {
@@ -293,7 +425,6 @@ export async function evaluateExperiment({
  * Run an evaluator against a run.
  *
  * @experimental This feature is not complete, and will change in the future.
- * @deprecated This function will be un-marked as deprecated once the experimental feature flag is removed.
  */
 async function runEvaluator({
   evaluator,
@@ -311,7 +442,7 @@ async function runEvaluator({
   const evaluate = async () => {
     const thisEval: ExperimentEvaluationRun = {
       id: id(),
-      traceId: id(),
+      traceId: null, // TODO: fill this in once we trace experiments
       experimentRunId: run.id,
       startTime: new Date(),
       endTime: new Date(), // will get replaced with actual end time
@@ -333,6 +464,7 @@ async function runEvaluator({
     }
     thisEval.endTime = new Date();
     onComplete(thisEval);
+    return thisEval;
   };
 
   return evaluate();
@@ -342,18 +474,23 @@ async function runEvaluator({
  * Wrap an evaluator function in an object with a name property.
  *
  * @experimental This feature is not complete, and will change in the future.
- * @deprecated This function will be un-marked as deprecated once the experimental feature flag is removed.
  *
  * @param name - The name of the evaluator.
  * @param evaluate - The evaluator function.
  * @returns The evaluator object.
  */
-export function asEvaluator(
-  name: string,
-  evaluate: Evaluator["evaluate"]
-): Evaluator {
+export function asEvaluator({
+  name,
+  kind,
+  evaluate,
+}: {
+  name: string;
+  kind: AnnotatorKind;
+  evaluate: Evaluator["evaluate"];
+}): Evaluator {
   return {
     name,
+    kind,
     evaluate,
   };
 }

package/src/spans/addSpanAnnotation.ts

@@ -0,0 +1,59 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { SpanAnnotation, toSpanAnnotationData } from "./types";
+
+/**
+ * Parameters to add a span annotation
+ */
+interface AddSpanAnnotationParams extends ClientFn {
+  spanAnnotation: SpanAnnotation;
+}
+
+/**
+ * Add an annotation to a span.
+ *
+ * The annotation can be of type "LLM", "CODE", or "HUMAN" and can include a label, score, and metadata.
+ * If an identifier is provided and an annotation with that identifier already exists, it will be updated.
+ *
+ * @param params - The parameters to add a span annotation
+ * @returns The ID of the created or updated annotation
+ *
+ * @example
+ * ```ts
+ * const result = await addSpanAnnotation({
+ *   spanAnnotation: {
+ *     spanId: "123abc",
+ *     name: "quality_score",
+ *     label: "good",
+ *     score: 0.95,
+ *     annotatorKind: "LLM",
+ *     identifier: "custom_id_123",
+ *     metadata: {
+ *       model: "gpt-4"
+ *     }
+ *   }
+ * });
+ * ```
+ */
+export async function addSpanAnnotation({
+  client: _client,
+  spanAnnotation,
+}: AddSpanAnnotationParams): Promise<{ id: string }> {
+  const client = _client ?? createClient();
+
+  const { data, error } = await client.POST("/v1/span_annotations", {
+    body: {
+      data: [toSpanAnnotationData(spanAnnotation)],
+    },
+  });
+
+  if (error) {
+    throw new Error(`Failed to add span annotation: ${error}`);
+  }
+
+  if (!data?.data?.[0]?.id) {
+    throw new Error("No annotation ID returned from server");
+  }
+
+  return data.data[0];
+}

package/src/spans/index.ts

@@ -0,0 +1,2 @@
+export * from "./addSpanAnnotation";
+export * from "./logSpanAnnotations";

package/src/spans/logSpanAnnotations.ts

@@ -0,0 +1,71 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { SpanAnnotation, toSpanAnnotationData } from "./types";
+
+/**
+ * Parameters to log multiple span annotations
+ */
+interface LogSpanAnnotationsParams extends ClientFn {
+  /**
+   * The span annotations to log
+   */
+  spanAnnotations: SpanAnnotation[];
+}
+
+/**
+ * Log multiple span annotations in a single request.
+ *
+ * Each annotation can be of type "LLM", "CODE", or "HUMAN" and can include a label, score, and metadata.
+ * If an identifier is provided and an annotation with that identifier already exists, it will be updated.
+ *
+ * @param params - The parameters to log span annotations
+ * @returns The IDs of the created or updated annotations
+ *
+ * @example
+ * ```ts
+ * const results = await logSpanAnnotations({
+ *   spanAnnotations: [
+ *     {
+ *       spanId: "123abc",
+ *       name: "quality_score",
+ *       label: "good",
+ *       score: 0.95,
+ *       annotatorKind: "LLM",
+ *       identifier: "custom_id_123",
+ *       metadata: {
+ *         model: "gpt-4"
+ *       }
+ *     },
+ *     {
+ *       spanId: "456def",
+ *       name: "sentiment",
+ *       label: "positive",
+ *       score: 0.8,
+ *       annotatorKind: "CODE"
+ *     }
+ *   ]
+ * });
+ * ```
+ */
+export async function logSpanAnnotations({
+  client: _client,
+  spanAnnotations,
+}: LogSpanAnnotationsParams): Promise<{ id: string }[]> {
+  const client = _client ?? createClient();
+
+  const { data, error } = await client.POST("/v1/span_annotations", {
+    body: {
+      data: spanAnnotations.map(toSpanAnnotationData),
+    },
+  });
+
+  if (error) {
+    throw new Error(`Failed to log span annotations: ${error}`);
+  }
+
+  if (!data?.data?.length) {
+    throw new Error("No annotation IDs returned from server");
+  }
+
+  return data.data;
+}

package/src/spans/types.ts

@@ -0,0 +1,60 @@
+import { paths } from "../__generated__/api/v1";
+
+type SpanAnnotationData =
+  paths["/v1/span_annotations"]["post"]["requestBody"]["content"]["application/json"]["data"][0];
+
+/**
+ * Parameters for a single span annotation
+ */
+export interface SpanAnnotation {
+  /**
+   * The OpenTelemetry Span ID (hex format without 0x prefix)
+   */
+  spanId: string;
+  /**
+   * The name of the annotation
+   */
+  name: string;
+  /**
+   * The label assigned by the annotation
+   */
+  label?: string;
+  /**
+   * The score assigned by the annotation
+   */
+  score?: number;
+  /**
+   * The identifier of the annotation. If provided, the annotation will be updated if it already exists.
+   */
+  identifier?: string;
+  /**
+   * Metadata for the annotation
+   */
+  metadata?: Record<string, unknown>;
+  /**
+   * The kind of annotator used for the annotation
+   * Can be "HUMAN", "LLM", or "CODE"
+   * @default "HUMAN"
+   */
+  annotatorKind?: SpanAnnotationData["annotator_kind"];
+}
+
+/**
+ * Convert a SpanAnnotation to the API format
+ */
+export function toSpanAnnotationData(
+  annotation: SpanAnnotation
+): SpanAnnotationData {
+  return {
+    span_id: annotation.spanId,
+    name: annotation.name,
+    annotator_kind: annotation.annotatorKind ?? "HUMAN",
+    result: {
+      label: annotation.label ?? null,
+      score: annotation.score ?? null,
+      explanation: null,
+    },
+    metadata: annotation.metadata ?? null,
+    identifier: annotation.identifier ?? "",
+  };
+}

package/src/types/annotations.ts

@@ -1 +1,4 @@
-export type AnnotatorKind = "HUMAN" | "LLM";
+import { components } from "../__generated__/api/v1";
+
+export type AnnotatorKind =
+  components["schemas"]["SpanAnnotationData"]["annotator_kind"];