@arizeai/phoenix-client 7.6.0 → 7.7.1

package/src/experiments/runExperiment.ts

@@ -139,6 +139,199 @@ export type RunExperimentParams = ClientFn & {
   diagLogLevel?: DiagLogLevel;
 };
 
+type PreparedExperiment = {
+  experiment: ExperimentInfo;
+  projectName: string;
+  taskTracer: Tracer;
+  taskProvider?: NodeTracerProvider;
+  taskGlobalRegistration: GlobalTracerProviderRegistration | null;
+};
+
+function prepareDryRunExperiment({
+  dataset,
+  datasetSelector,
+  projectName,
+  experimentMetadata,
+  repetitions,
+  nExamples,
+}: {
+  dataset: Dataset;
+  datasetSelector: DatasetSelector;
+  projectName: string;
+  experimentMetadata: Record<string, unknown>;
+  repetitions: number;
+  nExamples: number;
+}): PreparedExperiment {
+  const now = new Date().toISOString();
+  return {
+    experiment: {
+      id: localId(),
+      datasetId: dataset.id,
+      datasetVersionId: dataset.versionId,
+      datasetSplits: datasetSelector.splits ?? [],
+      projectName,
+      metadata: experimentMetadata,
+      repetitions,
+      createdAt: now,
+      updatedAt: now,
+      exampleCount: nExamples,
+      successfulRunCount: 0,
+      failedRunCount: 0,
+      missingRunCount: nExamples * repetitions,
+    },
+    projectName,
+    taskTracer: createNoOpProvider().getTracer("no-op"),
+    taskGlobalRegistration: null,
+  };
+}
+
+async function prepareRecordedExperiment({
+  client,
+  dataset,
+  datasetSelector,
+  projectName: defaultProjectName,
+  experimentName,
+  experimentDescription,
+  experimentMetadata,
+  repetitions,
+  useBatchSpanProcessor,
+  diagLogLevel,
+  setGlobalTracerProvider,
+}: {
+  client: PhoenixClient;
+  dataset: Dataset;
+  datasetSelector: DatasetSelector;
+  projectName: string;
+  experimentName: string | undefined;
+  experimentDescription: string | undefined;
+  experimentMetadata: Record<string, unknown>;
+  repetitions: number;
+  useBatchSpanProcessor: boolean;
+  diagLogLevel: DiagLogLevel | undefined;
+  setGlobalTracerProvider: boolean;
+}): Promise<PreparedExperiment> {
+  const response = await client
+    .POST("/v1/datasets/{dataset_id}/experiments", {
+      params: { path: { dataset_id: dataset.id } },
+      body: {
+        name: experimentName,
+        description: experimentDescription,
+        metadata: experimentMetadata,
+        project_name: defaultProjectName,
+        repetitions,
+        ...(datasetSelector.splits ? { splits: datasetSelector.splits } : {}),
+        ...(dataset.versionId ? { version_id: dataset.versionId } : {}),
+      },
+    })
+    .then((result) => result.data?.data);
+  invariant(response, `Failed to create experiment`);
+  const projectName = response.project_name ?? defaultProjectName;
+  const experiment: ExperimentInfo = {
+    id: response.id,
+    datasetId: response.dataset_id,
+    datasetVersionId: response.dataset_version_id,
+    datasetSplits: datasetSelector.splits ?? [],
+    projectName,
+    repetitions: response.repetitions,
+    metadata: response.metadata || {},
+    createdAt: response.created_at,
+    updatedAt: response.updated_at,
+    exampleCount: response.example_count,
+    successfulRunCount: response.successful_run_count,
+    failedRunCount: response.failed_run_count,
+    missingRunCount: response.missing_run_count,
+  };
+  invariant(client.config.baseUrl, MISSING_BASE_URL_MESSAGE);
+  const taskProvider = register({
+    projectName,
+    url: getTraceExportUrl(client.config),
+    headers: client.config.headers
+      ? toObjectHeaders(client.config.headers)
+      : undefined,
+    batch: useBatchSpanProcessor,
+    diagLogLevel,
+    global: false,
+  });
+  return {
+    experiment,
+    projectName,
+    taskProvider,
+    taskGlobalRegistration: setGlobalTracerProvider
+      ? attachGlobalTracerProvider(taskProvider)
+      : null,
+    taskTracer: taskProvider.getTracer(projectName),
+  };
+}
+
+function buildExperimentLinks({
+  client,
+  datasetId,
+  experimentId,
+  isDryRun,
+}: {
+  client: PhoenixClient;
+  datasetId: string;
+  experimentId: string;
+  isDryRun: boolean;
+}): Array<{ label: string; url: string }> {
+  const baseUrl = client.config.baseUrl;
+  if (isDryRun || !baseUrl) return [];
+  return [
+    { label: "Dataset", url: getDatasetUrl({ baseUrl, datasetId }) },
+    {
+      label: "Experiments",
+      url: getDatasetExperimentsUrl({ baseUrl, datasetId }),
+    },
+    {
+      label: "Experiment",
+      url: getExperimentUrl({ baseUrl, datasetId, experimentId }),
+    },
+  ];
+}
+
+function logEvaluatorStart({
+  evaluators,
+  logger,
+}: {
+  evaluators: ExperimentEvaluatorLike[] | undefined;
+  logger: Logger;
+}): void {
+  if (!evaluators || evaluators.length === 0) return;
+  const names = getExperimentEvaluators(evaluators)
+    .map((evaluator) => evaluator.name)
+    .join(", ");
+  logger.info(`${PROGRESS_PREFIX.start}Evaluations (${names})`);
+}
+
+async function refreshRecordedExperiment({
+  isDryRun,
+  client,
+  experiment,
+  ranExperiment,
+}: {
+  isDryRun: boolean;
+  client: PhoenixClient;
+  experiment: ExperimentInfo;
+  ranExperiment: RanExperiment;
+}): Promise<void> {
+  if (isDryRun) return;
+  const updated = await getExperimentInfo({
+    client,
+    experimentId: experiment.id,
+  });
+  Object.assign(ranExperiment, updated);
+}
+
+function logEvaluationSummary({
+  logger,
+  evaluationRuns,
+}: {
+  logger: Logger;
+  evaluationRuns: ExperimentEvaluationRun[] | undefined;
+}): void {
+  if (evaluationRuns?.length) logEvalSummary(logger, evaluationRuns);
+}
+
 /**
  * Runs an experiment using a given set of dataset of examples.
  *
@@ -208,90 +401,32 @@ export async function runExperiment({
       ? 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: ExperimentInfo;
-  if (isDryRun) {
-    const now = new Date().toISOString();
-    const totalExamples = nExamples;
-    experiment = {
-      id: localId(),
-      datasetId: dataset.id,
-      datasetVersionId: dataset.versionId,
-      // @todo: the dataset should return splits in response body
-      datasetSplits: datasetSelector?.splits ?? [],
-      projectName,
-      metadata: experimentMetadata,
-      repetitions,
-      createdAt: now,
-      updatedAt: now,
-      exampleCount: totalExamples,
-      successfulRunCount: 0,
-      failedRunCount: 0,
-      missingRunCount: totalExamples * repetitions,
-    };
-    taskTracer = createNoOpProvider().getTracer("no-op");
-  } 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,
-          repetitions,
-          // @todo: the dataset should return splits in response body
-          ...(datasetSelector?.splits
-            ? { splits: datasetSelector.splits }
-            : {}),
-          ...(dataset?.versionId ? { version_id: dataset.versionId } : {}),
-        },
+  const defaultProjectName = `${dataset.name}-exp-${new Date().toISOString()}`;
+  const prepared = isDryRun
+    ? prepareDryRunExperiment({
+        dataset,
+        datasetSelector,
+        projectName: defaultProjectName,
+        experimentMetadata,
+        repetitions,
+        nExamples,
       })
-      .then((res) => res.data?.data);
-    invariant(experimentResponse, `Failed to create experiment`);
-    projectName = experimentResponse.project_name ?? projectName;
-    experiment = {
-      id: experimentResponse.id,
-      datasetId: experimentResponse.dataset_id,
-      datasetVersionId: experimentResponse.dataset_version_id,
-      // @todo: the dataset should return splits in response body
-      datasetSplits: datasetSelector?.splits ?? [],
-      projectName,
-      repetitions: experimentResponse.repetitions,
-      metadata: experimentResponse.metadata || {},
-      createdAt: experimentResponse.created_at,
-      updatedAt: experimentResponse.updated_at,
-      exampleCount: experimentResponse.example_count,
-      successfulRunCount: experimentResponse.successful_run_count,
-      failedRunCount: experimentResponse.failed_run_count,
-      missingRunCount: experimentResponse.missing_run_count,
-    };
-    // Initialize the tracer, now that we have a project name
-    const baseUrl = client.config.baseUrl;
-    invariant(baseUrl, MISSING_BASE_URL_MESSAGE);
-
-    taskProvider = register({
-      projectName,
-      url: getTraceExportUrl(client.config),
-      headers: client.config.headers
-        ? toObjectHeaders(client.config.headers)
-        : undefined,
-      batch: useBatchSpanProcessor,
-      diagLogLevel,
-      global: false,
-    });
-    taskGlobalRegistration = setGlobalTracerProvider
-      ? attachGlobalTracerProvider(taskProvider)
-      : null;
-
-    taskTracer = taskProvider.getTracer(projectName);
-  }
+    : await prepareRecordedExperiment({
+        client,
+        dataset,
+        datasetSelector,
+        projectName: defaultProjectName,
+        experimentName,
+        experimentDescription,
+        experimentMetadata,
+        repetitions,
+        useBatchSpanProcessor,
+        diagLogLevel,
+        setGlobalTracerProvider,
+      });
+  const { experiment, taskTracer } = prepared;
+  taskProvider = prepared.taskProvider;
+  taskGlobalRegistration = prepared.taskGlobalRegistration;
   try {
     if (!record) {
       logger.info(
@@ -299,31 +434,12 @@ export async function runExperiment({
       );
     }
 
-    const links: Array<{ label: string; url: string }> = [];
-    if (!isDryRun && client.config.baseUrl) {
-      links.push({
-        label: "Dataset",
-        url: getDatasetUrl({
-          baseUrl: client.config.baseUrl,
-          datasetId: dataset.id,
-        }),
-      });
-      links.push({
-        label: "Experiments",
-        url: getDatasetExperimentsUrl({
-          baseUrl: client.config.baseUrl,
-          datasetId: dataset.id,
-        }),
-      });
-      links.push({
-        label: "Experiment",
-        url: getExperimentUrl({
-          baseUrl: client.config.baseUrl,
-          datasetId: dataset.id,
-          experimentId: experiment.id,
-        }),
-      });
-    }
+    const links = buildExperimentLinks({
+      client,
+      datasetId: dataset.id,
+      experimentId: experiment.id,
+      isDryRun,
+    });
 
     const evCount = evaluators?.length ?? 0;
     logger.info(
@@ -367,12 +483,7 @@ export async function runExperiment({
     taskProvider = undefined;
     taskGlobalRegistration = null;
 
-    if (evaluators && evaluators.length > 0) {
-      const evNames = getExperimentEvaluators(evaluators)
-        .map((evaluator) => evaluator.name)
-        .join(", ");
-      logger.info(`${PROGRESS_PREFIX.start}Evaluations (${evNames})`);
-    }
+    logEvaluatorStart({ evaluators, logger });
 
     const { evaluationRuns } = await evaluateExperiment({
       experiment: ranExperiment,
@@ -387,15 +498,12 @@ export async function runExperiment({
     });
     ranExperiment.evaluationRuns = evaluationRuns;
 
-    // Refresh experiment info from server to get updated counts (non-dry-run only)
-    if (!isDryRun) {
-      const updatedExperiment = await getExperimentInfo({
-        client,
-        experimentId: experiment.id,
-      });
-      // Update the experiment info with the latest from the server
-      Object.assign(ranExperiment, updatedExperiment);
-    }
+    await refreshRecordedExperiment({
+      isDryRun,
+      client,
+      experiment,
+      ranExperiment,
+    });
 
     logTaskSummary(logger, {
       nExamples,
@@ -404,12 +512,10 @@ export async function runExperiment({
       nErrors: taskErrors,
     });
 
-    if (
-      ranExperiment.evaluationRuns &&
-      ranExperiment.evaluationRuns.length > 0
-    ) {
-      logEvalSummary(logger, ranExperiment.evaluationRuns);
-    }
+    logEvaluationSummary({
+      logger,
+      evaluationRuns: ranExperiment.evaluationRuns,
+    });
 
     logLinks(logger, links);
 

package/src/projects/index.ts

@@ -1 +1,2 @@
 export * from "./getProjects";
+export * from "./setProjectRetentionPolicy";

package/src/projects/setProjectRetentionPolicy.ts

@@ -0,0 +1,80 @@
+import invariant from "tiny-invariant";
+
+import type { components } from "../__generated__/api/v1";
+import { createClient } from "../client";
+import type { ClientFn } from "../types/core";
+import type { ProjectIdentifier } from "../types/projects";
+import { resolveProjectIdentifier } from "../types/projects";
+
+/**
+ * The retention-policy assignment returned for a project.
+ */
+export type ProjectRetentionPolicyAssignment =
+  components["schemas"]["ProjectRetentionPolicyData"];
+
+/**
+ * Parameters for assigning or resetting a project's retention policy.
+ */
+export type SetProjectRetentionPolicyParams = ClientFn &
+  ProjectIdentifier & {
+    /**
+     * The GlobalID of an existing trace retention policy, or `null` to reset
+     * the project to the default policy.
+     */
+    policyId: string | null;
+  };
+
+/**
+ * Assign an existing trace retention policy to a project, or reset the project
+ * to the default policy.
+ *
+ * This helper only changes the project's assignment. It does not create,
+ * retrieve, update, or delete retention policies.
+ *
+ * @param params - The project and policy assignment.
+ * @param params.project - A project name or GlobalID.
+ * @param params.projectId - A project GlobalID.
+ * @param params.projectName - A project name.
+ * @param params.policyId - An existing policy GlobalID, or `null` to reset.
+ * @param params.client - An optional Phoenix client instance.
+ * @returns The project's resulting retention-policy assignment.
+ *
+ * @example
+ * ```ts
+ * import { setProjectRetentionPolicy } from "@arizeai/phoenix-client/projects";
+ *
+ * await setProjectRetentionPolicy({
+ *   projectName: "support-bot",
+ *   policyId: "UHJvamVjdFRyYWNlUmV0ZW50aW9uUG9saWN5OjI=",
+ * });
+ *
+ * await setProjectRetentionPolicy({
+ *   projectId: "UHJvamVjdDox",
+ *   policyId: null,
+ * });
+ * ```
+ */
+export async function setProjectRetentionPolicy(
+  params: SetProjectRetentionPolicyParams
+): Promise<ProjectRetentionPolicyAssignment> {
+  const client = params.client ?? createClient();
+  const projectIdentifier = resolveProjectIdentifier(params);
+
+  const { data, error } = await client.PATCH(
+    "/v1/projects/{project_identifier}/retention",
+    {
+      params: {
+        path: {
+          project_identifier: projectIdentifier,
+        },
+      },
+      body: {
+        policy_id: params.policyId,
+      },
+    }
+  );
+
+  if (error) throw error;
+  invariant(data?.data, "Failed to set project retention policy");
+  return data.data;
+}

package/src/prompts/sdks/toOpenAI.ts

@@ -26,6 +26,61 @@ export type {
 export type ToOpenAIParams<PromptVariables extends Variables> =
   toSDKParamsBase<PromptVariables>;
 
+type PhoenixPrompt = ToOpenAIParams<Variables>["prompt"];
+
+function getInvocationParameters(
+  prompt: PhoenixPrompt
+): Partial<ChatCompletionCreateParams> {
+  const parameters = prompt.invocation_parameters;
+  switch (parameters.type) {
+    case "openai":
+      return parameters.openai;
+    case "azure_openai":
+      return parameters.azure_openai;
+    case "deepseek":
+      return parameters.deepseek;
+    case "xai":
+      return parameters.xai;
+    case "ollama":
+      return parameters.ollama;
+    case "cerebras":
+      return parameters.cerebras;
+    case "fireworks":
+      return parameters.fireworks;
+    case "groq":
+      return parameters.groq;
+    case "moonshot":
+      return parameters.moonshot;
+    case "perplexity":
+      return parameters.perplexity;
+    case "together":
+      return parameters.together;
+    default:
+      // eslint-disable-next-line no-console
+      console.warn(
+        "Prompt is not an OpenAI-family prompt, falling back to default OpenAI invocation parameters"
+      );
+      return {};
+  }
+}
+
+function getOpenAITools(
+  prompt: PhoenixPrompt
+): ChatCompletionCreateParams["tools"] {
+  const toolsList = prompt.tools?.tools ?? [];
+  if (toolsList.length === 0) return undefined;
+
+  return toolsList.map((tool) => {
+    if (isPromptToolRaw(tool)) return tool.raw;
+    const definition = safelyConvertToolDefinitionToProvider({
+      toolDefinition: tool,
+      targetProvider: "OPENAI",
+    });
+    invariant(definition, "Tool definition is not valid");
+    return definition;
+  }) as unknown as ChatCompletionCreateParams["tools"];
+}
+
 /**
  * Convert a Phoenix prompt to OpenAI client sdk's chat completion parameters
  *
@@ -36,49 +91,7 @@ export const toOpenAI = <PromptVariables extends Variables = Variables>({
   variables,
 }: ToOpenAIParams<PromptVariables>): ChatCompletionCreateParams | null => {
   try {
-    let invocationParameters: Partial<ChatCompletionCreateParams>;
-    switch (prompt.invocation_parameters.type) {
-      case "openai":
-        invocationParameters = prompt.invocation_parameters.openai;
-        break;
-      case "azure_openai":
-        invocationParameters = prompt.invocation_parameters.azure_openai;
-        break;
-      case "deepseek":
-        invocationParameters = prompt.invocation_parameters.deepseek;
-        break;
-      case "xai":
-        invocationParameters = prompt.invocation_parameters.xai;
-        break;
-      case "ollama":
-        invocationParameters = prompt.invocation_parameters.ollama;
-        break;
-      case "cerebras":
-        invocationParameters = prompt.invocation_parameters.cerebras;
-        break;
-      case "fireworks":
-        invocationParameters = prompt.invocation_parameters.fireworks;
-        break;
-      case "groq":
-        invocationParameters = prompt.invocation_parameters.groq;
-        break;
-      case "moonshot":
-        invocationParameters = prompt.invocation_parameters.moonshot;
-        break;
-      case "perplexity":
-        invocationParameters = prompt.invocation_parameters.perplexity;
-        break;
-      case "together":
-        invocationParameters = prompt.invocation_parameters.together;
-        break;
-      default:
-        // eslint-disable-next-line no-console
-        console.warn(
-          "Prompt is not an OpenAI-family prompt, falling back to default OpenAI invocation parameters"
-        );
-        invocationParameters = {};
-        break;
-    }
+    const invocationParameters = getInvocationParameters(prompt);
     // parts of the prompt that can be directly converted to OpenAI params
     const baseCompletionParams = {
       model: prompt.model_name,
@@ -109,24 +122,8 @@ export const toOpenAI = <PromptVariables extends Variables = Variables>({
       return openAIMessage;
     });
 
-    const toolsList = prompt.tools?.tools ?? [];
-    // Cast: raw tools are `Record<string, unknown>` straight from the prompt
-    // store. We trust the upstream caller to have stored a shape OpenAI's
-    // SDK accepts; no validation here.
-    const tools =
-      toolsList.length === 0
-        ? undefined
-        : (toolsList.map((tool) => {
-            if (isPromptToolRaw(tool)) {
-              return tool.raw;
-            }
-            const openAIToolDefinition = safelyConvertToolDefinitionToProvider({
-              toolDefinition: tool,
-              targetProvider: "OPENAI",
-            });
-            invariant(openAIToolDefinition, "Tool definition is not valid");
-            return openAIToolDefinition;
-          }) as unknown as ChatCompletionCreateParams["tools"]);
+    // Raw tools are trusted to already match the OpenAI SDK shape.
+    const tools = getOpenAITools(prompt);
 
     const tool_choice: OpenaiToolChoice | undefined = tools
       ? (safelyConvertToolChoiceToProvider({

package/src/sessions/sessionUtils.ts

@@ -13,6 +13,9 @@ export function toSession(data: SessionData): Session {
     projectId: data.project_id,
     startTime: data.start_time,
     endTime: data.end_time,
+    tokenCountPrompt: data.token_count_prompt,
+    tokenCountCompletion: data.token_count_completion,
+    tokenCountTotal: data.token_count_total,
     traces: data.traces.map((trace) => ({
       id: trace.id,
       traceId: trace.trace_id,