@arizeai/phoenix-client 4.0.2 → 4.1.0

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/runExperiment.ts

@@ -23,7 +23,7 @@ import { pluralize } from "../utils/pluralize";
 import { promisifyResult } from "../utils/promisifyResult";
 import { AnnotatorKind } from "../types/annotations";
 import { createProvider, createNoOpProvider } from "./instrumentation";
-import { SpanStatusCode, Tracer } from "@opentelemetry/api";
+import { SpanStatusCode, Tracer, trace } from "@opentelemetry/api";
 import {
   MimeType,
   OpenInferenceSpanKind,
@@ -290,11 +290,6 @@ export async function runExperiment({
     runs,
   };
 
-  // Shut down the provider so that the experiments run
-  if (provider) {
-    await provider.shutdown?.();
-  }
-
   const { evaluationRuns } = await evaluateExperiment({
     experiment: ranExperiment,
     evaluators: evaluators ?? [],
@@ -302,8 +297,7 @@ export async function runExperiment({
     logger,
     concurrency,
     dryRun,
-    setGlobalTracerProvider,
-    useBatchSpanProcessor,
+    tracerProvider: provider,
   });
   ranExperiment.evaluationRuns = evaluationRuns;
 
@@ -473,6 +467,7 @@ export async function evaluateExperiment({
   dryRun = false,
   setGlobalTracerProvider = true,
   useBatchSpanProcessor = true,
+  tracerProvider: paramsTracerProvider,
 }: {
   /**
    * The experiment to evaluate
@@ -502,6 +497,11 @@ export async function evaluateExperiment({
    * @default true
    */
   useBatchSpanProcessor?: boolean;
+  /**
+   * The tracer provider to use. If set, the other parameters will be ignored and the passed tracer provider will get used
+   * Intended as a pass-through from runExperiment
+   */
+  tracerProvider?: NodeTracerProvider | null;
 }): Promise<RanExperiment> {
   const isDryRun = typeof dryRun === "number" || dryRun === true;
   const client = _client ?? createClient();
@@ -511,7 +511,11 @@ export async function evaluateExperiment({
     "Phoenix base URL not found. Please set PHOENIX_HOST or set baseUrl on the client."
   );
   let provider: NodeTracerProvider;
-  if (!isDryRun) {
+
+  // Always allow changing of tracer providers
+  if (paramsTracerProvider) {
+    provider = paramsTracerProvider;
+  } else if (!isDryRun) {
     provider = createProvider({
       projectName: "evaluators",
       baseUrl,
@@ -668,7 +672,11 @@ export async function evaluateExperiment({
   logger.info(`✅ Evaluation runs completed`);
 
   if (provider) {
-    await provider.shutdown?.();
+    await provider.shutdown();
+    // Make sure it's not set globally anymore
+    if (setGlobalTracerProvider) {
+      trace.disable();
+    }
   }
 
   return {

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>;
+}

package/src/types/experiments.ts

@@ -51,7 +51,7 @@ export interface ExperimentRun extends Node {
   traceId: string | null;
 }
 
-export type EvaluatorParams = {
+export type EvaluatorParams<TaskOutputType = TaskOutput> = {
   /**
    * The input field of the Dataset Example
    */
@@ -59,7 +59,7 @@ export type EvaluatorParams = {
   /**
    * The output of the task
    */
-  output: TaskOutput;
+  output: TaskOutputType;
   /**
    * The expected or reference output of the Dataset Example
    */
@@ -79,10 +79,10 @@ export type Evaluator = {
 };
 
 export type EvaluationResult = {
-  score: number | null;
-  label: string | null;
-  metadata: Record<string, unknown>;
-  explanation: string | null;
+  score?: number | null;
+  label?: string | null;
+  metadata?: Record<string, unknown>;
+  explanation?: string | null;
 };
 
 export interface ExperimentEvaluationRun extends Node {