@arizeai/phoenix-client 4.0.3 → 4.2.0

package/src/datasets/createOrGetDataset.ts

@@ -0,0 +1,39 @@
+import { createClient } from "../client";
+import { CreateDatasetParams, createDataset } from "./createDataset";
+import { getDatasetInfoByName } from "./getDatasetInfoByName";
+
+export type CreateOrGetDatasetParams = CreateDatasetParams;
+
+export type CreateOrGetDatasetResponse = {
+  datasetId: string;
+};
+
+/**
+ * Given the parameters to create a dataset, this function will either
+ * retrieve an existing dataset by name or create a new one with the provided parameters.
+ *
+ * This is useful in cases where you would like to re-run a pipeline like:
+ * - ensure dataset exists
+ * - create a task
+ * - run experiment
+ * - evaluate experiment
+ * without having to create a new dataset each time.
+ */
+export async function createOrGetDataset({
+  name,
+  description,
+  examples,
+  client: _client,
+}: CreateOrGetDatasetParams): Promise<CreateOrGetDatasetResponse> {
+  const client = _client || createClient();
+  // start by fetching an existing dataset by name, catching any errors that occur
+  try {
+    const dataset = await getDatasetInfoByName({ datasetName: name, client });
+    return {
+      datasetId: dataset.id,
+    };
+  } catch {
+    // If the dataset doesn't exist, create it, falling back to the error handling inside createDataset
+    return await createDataset({ name, description, examples, client });
+  }
+}

package/src/datasets/index.ts

@@ -3,3 +3,4 @@ export * from "./getDataset";
 export * from "./getDatasetExamples";
 export * from "./appendDatasetExamples";
 export * from "./getDatasetInfo";
+export * from "./createOrGetDataset";

package/src/experiments/getExperimentRuns.ts

@@ -2,35 +2,60 @@ import { createClient } from "../client";
 import invariant from "tiny-invariant";
 import { ClientFn } from "../types/core";
 import { ExperimentRun } from "../types/experiments";
+import { components } from "../__generated__/api/v1";
 
 export type GetExperimentRunsParams = ClientFn & {
   /**
    * The experiment ID.
    */
   experimentId: string;
+  /**
+   * The pagination size by which to pull runs
+   * Exposed for controlling the rate at which runs are pulled
+   * @default 100
+   */
+  pageSize?: number;
 };
 
+const DEFAULT_PAGE_SIZE = 100;
+
 /**
- * A function that gets the runs (e.g. the results) of a experiment
+ * A function that gets all the runs (e.g. the results) of a experiment
  */
 export async function getExperimentRuns({
   client: _client,
   experimentId,
+  pageSize = DEFAULT_PAGE_SIZE,
 }: GetExperimentRunsParams): Promise<{ runs: ExperimentRun[] }> {
   const client = _client || createClient();
-  const getRunsPromise = client.GET("/v1/experiments/{experiment_id}/runs", {
-    params: {
-      path: {
-        experiment_id: experimentId,
+
+  // Validate that the parameter is an integer and exit early
+  invariant(
+    Number.isInteger(pageSize) && pageSize > 0,
+    "pageSize must be a positive integer greater than 0"
+  );
+  const runs: ExperimentRun[] = [];
+  let cursor: string | null = null;
+  do {
+    const res: {
+      data?: components["schemas"]["ListExperimentRunsResponseBody"];
+    } = await client.GET("/v1/experiments/{experiment_id}/runs", {
+      params: {
+        path: {
+          experiment_id: experimentId,
+        },
+        query: {
+          cursor,
+          limit: pageSize,
+        },
       },
-    },
-  });
-  const [experimentRunResponse] = await Promise.all([getRunsPromise]);
-  const { data: { data: experimentRunsData } = {} } = experimentRunResponse;
-  invariant(experimentRunsData, "Failed to retrieve experiment runs");
-  return {
-    runs: experimentRunsData.map((run) => {
-      return {
+    });
+    // NB: older versions of phoenix simply don't respond with a cursor and fetch all
+    cursor = res.data?.next_cursor || null;
+    const data = res.data?.data;
+    invariant(data, "Failed to fetch runs");
+    runs.push(
+      ...data.map((run) => ({
         id: run.id,
         traceId: run.trace_id || null,
         experimentId: run.experiment_id,
@@ -39,7 +64,11 @@ export async function getExperimentRuns({
         endTime: new Date(run.end_time),
         output: run.output as ExperimentRun["output"],
         error: run.error || null,
-      };
-    }),
+      }))
+    );
+  } while (cursor != null);
+
+  return {
+    runs,
   };
 }

package/src/experiments/instrumentation.ts

@@ -20,6 +20,7 @@ export function createProvider({
   baseUrl,
   headers,
   useBatchSpanProcessor = true,
+  diagLogLevel,
 }: {
   projectName: string;
   headers: HeadersOptions;
@@ -32,8 +33,15 @@ export function createProvider({
    * The base URL of the Phoenix. Doesn't include the /v1/traces path.
    */
   baseUrl: string;
+  /**
+   * The diag log level to set for the built in DiagConsoleLogger instance.
+   * Omit to disable built in logging.
+   */
+  diagLogLevel?: DiagLogLevel;
 }) {
-  diag.setLogger(new DiagConsoleLogger(), DiagLogLevel.ERROR);
+  if (diagLogLevel) {
+    diag.setLogger(new DiagConsoleLogger(), diagLogLevel);
+  }
 
   const exporter = new OTLPTraceExporter({
     url: `${baseUrl}/v1/traces`,

package/src/experiments/runExperiment.ts

@@ -23,7 +23,12 @@ import { pluralize } from "../utils/pluralize";
 import { promisifyResult } from "../utils/promisifyResult";
 import { AnnotatorKind } from "../types/annotations";
 import { createProvider, createNoOpProvider } from "./instrumentation";
-import { SpanStatusCode, Tracer, trace } from "@opentelemetry/api";
+import {
+  type DiagLogLevel,
+  SpanStatusCode,
+  Tracer,
+  trace,
+} from "@opentelemetry/api";
 import {
   MimeType,
   OpenInferenceSpanKind,
@@ -111,6 +116,11 @@ export type RunExperimentParams = ClientFn & {
    * @default true
    */
   useBatchSpanProcessor?: boolean;
+  /**
+   * Log level to set for the default DiagConsoleLogger when tracing.
+   * Omit to disable default diag logging, or to bring your own.
+   */
+  diagLogLevel?: DiagLogLevel;
 };
 
 /**
@@ -160,6 +170,7 @@ export async function runExperiment({
   setGlobalTracerProvider = true,
   repetitions = 1,
   useBatchSpanProcessor = true,
+  diagLogLevel,
 }: RunExperimentParams): Promise<RanExperiment> {
   // Validation
   assert(
@@ -227,6 +238,7 @@ export async function runExperiment({
       baseUrl,
       headers: client.config.headers ?? {},
       useBatchSpanProcessor,
+      diagLogLevel,
     });
     // Register the provider
     if (setGlobalTracerProvider) {
@@ -298,6 +310,8 @@ export async function runExperiment({
     concurrency,
     dryRun,
     tracerProvider: provider,
+    diagLogLevel,
+    useBatchSpanProcessor,
   });
   ranExperiment.evaluationRuns = evaluationRuns;
 
@@ -468,6 +482,7 @@ export async function evaluateExperiment({
   setGlobalTracerProvider = true,
   useBatchSpanProcessor = true,
   tracerProvider: paramsTracerProvider,
+  diagLogLevel,
 }: {
   /**
    * The experiment to evaluate
@@ -502,6 +517,11 @@ export async function evaluateExperiment({
    * Intended as a pass-through from runExperiment
    */
   tracerProvider?: NodeTracerProvider | null;
+  /**
+   * Log level to set for the default DiagConsoleLogger when tracing.
+   * Omit to disable default diag logging, or to bring your own.
+   */
+  diagLogLevel?: DiagLogLevel;
 }): Promise<RanExperiment> {
   const isDryRun = typeof dryRun === "number" || dryRun === true;
   const client = _client ?? createClient();
@@ -521,6 +541,7 @@ export async function evaluateExperiment({
       baseUrl,
       headers: client.config.headers ?? {},
       useBatchSpanProcessor,
+      diagLogLevel,
     });
     if (setGlobalTracerProvider) {
       provider.register();

package/src/sessions/addSessionAnnotation.ts

@@ -0,0 +1,65 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { SessionAnnotation, toSessionAnnotationData } from "./types";
+
+/**
+ * Parameters to add a span annotation
+ */
+export interface AddSessionAnnotationParams extends ClientFn {
+  sessionAnnotation: SessionAnnotation;
+  /**
+   * If true, the request will be fulfilled synchronously and return the annotation ID.
+   * If false, the request will be processed asynchronously and return null.
+   * @default false
+   */
+  sync?: boolean;
+}
+
+/**
+ * Add an annotation to a session.
+ *
+ * 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 addSessionAnnotation({
+ *   sessionAnnotation: {
+ *     sessionId: "123abc",
+ *     name: "quality_score",
+ *     label: "good",
+ *     score: 0.95,
+ *     annotatorKind: "LLM",
+ *     identifier: "custom_id_123",
+ *     metadata: {
+ *       model: "gpt-4"
+ *     }
+ *   }
+ * });
+ * ```
+ */
+export async function addSessionAnnotation({
+  client: _client,
+  sessionAnnotation,
+  sync = false,
+}: AddSessionAnnotationParams): Promise<{ id: string } | null> {
+  const client = _client ?? createClient();
+
+  const { data, error } = await client.POST("/v1/session_annotations", {
+    params: {
+      query: { sync },
+    },
+    body: {
+      data: [toSessionAnnotationData(sessionAnnotation)],
+    },
+  });
+
+  if (error) {
+    throw new Error(`Failed to add session annotation: ${error}`);
+  }
+
+  return data?.data?.[0] || null;
+}

package/src/sessions/index.ts

@@ -0,0 +1,2 @@
+export * from "./addSessionAnnotation";
+export * from "./logSessionAnnotations";

package/src/sessions/logSessionAnnotations.ts

@@ -0,0 +1,77 @@
+import { createClient } from "../client";
+import { ClientFn } from "../types/core";
+import { SessionAnnotation, toSessionAnnotationData } from "./types";
+
+/**
+ * Parameters to log multiple session annotations
+ */
+export interface LogSessionAnnotationsParams extends ClientFn {
+  /**
+   * The session annotations to log
+   */
+  sessionAnnotations: SessionAnnotation[];
+  /**
+   * If true, the request will be fulfilled synchronously and return the annotation IDs.
+   * If false, the request will be processed asynchronously and return null.
+   * @default false
+   */
+  sync?: boolean;
+}
+
+/**
+ * Log multiple session 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 session annotations
+ * @returns The IDs of the created or updated annotations
+ *
+ * @example
+ * ```ts
+ * const results = await logSessionAnnotations({
+ *   sessionAnnotations: [
+ *     {
+ *       sessionId: "123abc",
+ *       name: "quality_score",
+ *       label: "good",
+ *       score: 0.95,
+ *       annotatorKind: "LLM",
+ *       identifier: "custom_id_123",
+ *       metadata: {
+ *         model: "gpt-4"
+ *       }
+ *     },
+ *     {
+ *       sessionId: "456def",
+ *       name: "sentiment",
+ *       label: "positive",
+ *       score: 0.8,
+ *       annotatorKind: "CODE"
+ *     }
+ *   ]
+ * });
+ * ```
+ */
+export async function logSessionAnnotations({
+  client: _client,
+  sessionAnnotations,
+  sync = false,
+}: LogSessionAnnotationsParams): Promise<{ id: string }[]> {
+  const client = _client ?? createClient();
+
+  const { data, error } = await client.POST("/v1/session_annotations", {
+    params: {
+      query: { sync },
+    },
+    body: {
+      data: sessionAnnotations.map(toSessionAnnotationData),
+    },
+  });
+
+  if (error) {
+    throw new Error(`Failed to log session annotations: ${error}`);
+  }
+
+  return data?.data || [];
+}

package/src/sessions/types.ts

@@ -0,0 +1,67 @@
+import { paths } from "../__generated__/api/v1";
+import { Annotation, AnnotationResult } from "../types/annotations";
+
+type SessionAnnotationData =
+  paths["/v1/session_annotations"]["post"]["requestBody"]["content"]["application/json"]["data"][0];
+
+/**
+ * Parameters for a single session annotation
+ */
+export interface SessionAnnotation extends Annotation {
+  /*
+   * The session ID used to track a conversation, thread, or session
+   */
+  sessionId: string;
+  /**
+   * The entity that performed the annotation
+   */
+  annotatorKind?: SessionAnnotationData["annotator_kind"];
+}
+
+/**
+ * Build and validate annotation result fields
+ */
+function buildSessionAnnotationResult(
+  annotation: Pick<SessionAnnotation, "label" | "score" | "explanation">
+): AnnotationResult {
+  const result: AnnotationResult = {};
+
+  // Build result with trimming for string fields
+  if (annotation.label !== undefined) {
+    result.label = annotation.label.trim() || null;
+  }
+  if (annotation.score !== undefined) {
+    result.score = annotation.score;
+  }
+  if (annotation.explanation !== undefined) {
+    result.explanation = annotation.explanation.trim() || null;
+  }
+
+  // Validate that at least one result field is provided
+  const hasValidResult =
+    result.label || result.score !== undefined || result.explanation;
+  if (!hasValidResult) {
+    throw new Error(
+      `At least one of label, score, or explanation must be provided for session annotation`
+    );
+  }
+  return result;
+}
+
+/**
+ * Convert a SessionAnnotation to the API format
+ */
+export function toSessionAnnotationData(
+  annotation: SessionAnnotation
+): SessionAnnotationData {
+  const result = buildSessionAnnotationResult(annotation);
+
+  return {
+    session_id: annotation.sessionId.trim(),
+    name: annotation.name.trim(),
+    annotator_kind: annotation.annotatorKind ?? "HUMAN",
+    result,
+    metadata: annotation.metadata ?? null,
+    identifier: annotation.identifier?.trim() ?? "",
+  };
+}

package/src/spans/types.ts

@@ -1,4 +1,5 @@
 import { paths } from "../__generated__/api/v1";
+import { Annotation } from "../types/annotations";
 
 type SpanAnnotationData =
   paths["/v1/span_annotations"]["post"]["requestBody"]["content"]["application/json"]["data"][0];
@@ -9,35 +10,11 @@ type SpanDocumentAnnotationData =
 /**
  * Parameters for a single span annotation
  */
-export interface SpanAnnotation {
+export interface SpanAnnotation extends Annotation {
   /**
    * 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;
-  /**
-   * Explanation of the annotation result
-   */
-  explanation?: string;
-  /**
-   * 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"
@@ -49,35 +26,11 @@ export interface SpanAnnotation {
 /**
  * Parameters for a single document annotation
  */
-export interface DocumentAnnotation {
-  /**
-   * The OpenTelemetry Span ID (hex format without 0x prefix)
-   */
-  spanId: string;
+export interface DocumentAnnotation extends SpanAnnotation {
   /**
    * The 0-based index of the document within the span
    */
   documentPosition: number;
-  /**
-   * The name of the annotation
-   */
-  name: string;
-  /**
-   * The label assigned by the annotation
-   */
-  label?: string;
-  /**
-   * The score assigned by the annotation
-   */
-  score?: number;
-  /**
-   * Explanation of the annotation result
-   */
-  explanation?: string;
-  /**
-   * Metadata for the annotation
-   */
-  metadata?: Record<string, unknown>;
   /**
    * The kind of annotator used for the annotation
    * Can be "HUMAN", "LLM", or "CODE"

package/src/types/annotations.ts

@@ -2,3 +2,42 @@ import { components } from "../__generated__/api/v1";
 
 export type AnnotatorKind =
   components["schemas"]["SpanAnnotationData"]["annotator_kind"];
+
+/**
+ * The result of an annotation from an author (e.x. an LLM or human)
+ */
+export type AnnotationResult = {
+  label?: string | null;
+  score?: number | null;
+  explanation?: string | null;
+};
+
+/**
+ * The base interface for all kinds of annotations (span, trace, session)
+ */
+export interface Annotation {
+  /**
+   * The name of the annotation
+   */
+  name: string;
+  /**
+   * The label assigned by the annotation
+   */
+  label?: string;
+  /**
+   * The score assigned by the annotation
+   */
+  score?: number;
+  /**
+   * Explanation of the annotation result
+   */
+  explanation?: string;
+  /**
+   * 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>;
+}