@arizeai/phoenix-client 7.7.0 → 7.8.0

package/src/spans/getSpans.ts

@@ -155,28 +155,23 @@ export type GetSpansResult = {
  * } while (cursor);
  * ```
  */
-export async function getSpans({
-  client: _client,
-  project,
-  cursor,
-  limit = 100,
-  startTime,
-  endTime,
+type SpansQuery = NonNullable<operations["getSpans"]["parameters"]["query"]>;
+
+async function ensureSpanFilterCapabilities({
+  client,
   traceIds,
   spanIds,
-  parentId,
   name,
   spanKind,
   statusCode,
-  attributes,
-}: GetSpansParams): Promise<GetSpansResult> {
-  const client = _client ?? createClient();
-  const serializedAttributes =
-    attributes != null ? serializeAttributes(attributes) : undefined;
-  const attributeFilters =
-    serializedAttributes != null && serializedAttributes.length > 0
-      ? serializedAttributes
-      : undefined;
+  attributeFilters,
+}: Pick<
+  GetSpansParams,
+  "traceIds" | "spanIds" | "name" | "spanKind" | "statusCode"
+> & {
+  client: NonNullable<GetSpansParams["client"]>;
+  attributeFilters: string[] | undefined;
+}): Promise<void> {
   if (traceIds) {
     await ensureServerCapability({ client, requirement: GET_SPANS_TRACE_IDS });
   }
@@ -186,58 +181,99 @@ export async function getSpans({
   if (name != null || spanKind != null || statusCode != null) {
     await ensureServerCapability({ client, requirement: GET_SPANS_FILTERS });
   }
-  if (attributeFilters != null) {
+  if (attributeFilters) {
     await ensureServerCapability({
       client,
       requirement: GET_SPANS_BY_ATTRIBUTE,
     });
   }
-  const projectIdentifier = resolveProjectIdentifier(project);
-
-  const params: NonNullable<operations["getSpans"]["parameters"]["query"]> = {
-    limit,
-  };
-
-  if (cursor) {
-    params.cursor = cursor;
-  }
+}
 
+function buildSpansQuery({
+  cursor,
+  limit,
+  startTime,
+  endTime,
+  traceIds,
+  spanIds,
+  parentId,
+  name,
+  spanKind,
+  statusCode,
+  attributeFilters,
+}: Omit<GetSpansParams, "client" | "project" | "attributes"> & {
+  limit: number;
+  attributeFilters: string[] | undefined;
+}): SpansQuery {
+  const params: SpansQuery = { limit };
+  if (cursor) params.cursor = cursor;
   if (startTime) {
     params.start_time =
       startTime instanceof Date ? startTime.toISOString() : startTime;
   }
-
   if (endTime) {
     params.end_time = endTime instanceof Date ? endTime.toISOString() : endTime;
   }
-
-  if (traceIds) {
-    params.trace_id = traceIds;
-  }
-
-  if (spanIds) {
-    params.span_id = spanIds;
-  }
-
+  if (traceIds) params.trace_id = traceIds;
+  if (spanIds) params.span_id = spanIds;
   if (parentId !== undefined) {
     params.parent_id = parentId === null ? "null" : parentId;
   }
-
-  if (name) {
-    params.name = Array.isArray(name) ? name : [name];
-  }
-
-  if (spanKind) {
+  if (name) params.name = Array.isArray(name) ? name : [name];
+  if (spanKind)
     params.span_kind = Array.isArray(spanKind) ? spanKind : [spanKind];
-  }
-
   if (statusCode) {
     params.status_code = Array.isArray(statusCode) ? statusCode : [statusCode];
   }
+  if (attributeFilters) params.attribute = attributeFilters;
+  return params;
+}
 
-  if (attributeFilters != null) {
-    params.attribute = attributeFilters;
-  }
+export async function getSpans({
+  client: _client,
+  project,
+  cursor,
+  limit = 100,
+  startTime,
+  endTime,
+  traceIds,
+  spanIds,
+  parentId,
+  name,
+  spanKind,
+  statusCode,
+  attributes,
+}: GetSpansParams): Promise<GetSpansResult> {
+  const client = _client ?? createClient();
+  const serializedAttributes =
+    attributes != null ? serializeAttributes(attributes) : undefined;
+  const attributeFilters =
+    serializedAttributes != null && serializedAttributes.length > 0
+      ? serializedAttributes
+      : undefined;
+  await ensureSpanFilterCapabilities({
+    client,
+    traceIds,
+    spanIds,
+    name,
+    spanKind,
+    statusCode,
+    attributeFilters,
+  });
+  const projectIdentifier = resolveProjectIdentifier(project);
+  const params = buildSpansQuery({
+    cursor,
+    limit,
+    startTime,
+    endTime,
+    traceIds,
+    spanIds,
+    parentId,
+    name,
+    spanKind,
+    statusCode,
+    attributeFilters,
+  });
 
   const { data, error } = await client.GET(
     "/v1/projects/{project_identifier}/spans",

package/src/testing/phoenix-test-tracking.ts

@@ -220,31 +220,33 @@ function buildLinks(
   ];
 }
 
-/**
- * Initialize the suite: upload the dataset, create the experiment, and
- * register the OpenInference tracer.
- *
- * If tracing is disabled (no Phoenix env vars, or PHOENIX_TEST_TRACKING=false),
- * this populates a no-op tracer and exits without making any network calls.
- */
-export async function initializeSuite(suite: SuiteState): Promise<void> {
-  const tracking = isTrackingEnabled(suite);
-  if (!tracking.enabled) {
-    suite.trackingDisabled = true;
-    suite.trackingDisabledReason = tracking.reason;
-    suite.tracer = createNoOpProvider().getTracer("no-op");
-    suite.evaluatorTracer = suite.tracer;
-    return;
+function disableSuiteTracking({
+  suite,
+  error,
+}: {
+  suite: SuiteState;
+  error?: unknown;
+}): void {
+  suite.trackingDisabled = true;
+  if (error !== undefined) {
+    suite.setupError =
+      error instanceof Error ? error : new Error(String(error));
   }
+  suite.tracer = createNoOpProvider().getTracer("no-op");
+  suite.evaluatorTracer = suite.tracer;
+}
 
-  const client = suite.config.client ?? createClient();
-  suite.client = client;
-
-  const datasetName = suite.config.datasetName ?? suite.name;
-  const description =
-    suite.config.description ??
-    `Phoenix test dataset auto-generated from ${suite.name}`;
-
+async function createSuiteDataset({
+  suite,
+  client,
+  datasetName,
+  description,
+}: {
+  suite: SuiteState;
+  client: PhoenixClient;
+  datasetName: string;
+  description: string;
+}): Promise<string | null> {
   const examples = Array.from(suite.registeredExamples.values()).map(
     (registered) => ({
       id: registered.params.id ?? null,
@@ -254,8 +256,6 @@ export async function initializeSuite(suite: SuiteState): Promise<void> {
       splits: registered.params.splits,
     })
   );
-
-  let datasetId: string;
   try {
     const created = await createDataset({
       client,
@@ -263,67 +263,81 @@ export async function initializeSuite(suite: SuiteState): Promise<void> {
       description,
       examples,
     });
-    datasetId = created.datasetId;
-  } catch (err) {
-    suite.trackingDisabled = true;
-    suite.setupError = err instanceof Error ? err : new Error(String(err));
-    suite.tracer = createNoOpProvider().getTracer("no-op");
-    suite.evaluatorTracer = suite.tracer;
-    return;
+    return created.datasetId;
+  } catch (error) {
+    disableSuiteTracking({ suite, error });
+    return null;
   }
-  suite.datasetId = datasetId;
+}
 
-  // Map test names to server-side example ids by re-fetching the dataset.
-  // The server doesn't promise that the GET response order matches the
-  // upload order, so we match by user-supplied `id` first, then by
-  // `JSON.stringify(input)` deep-equality with FIFO-on-collision.
+async function resolveSuiteExampleIds({
+  suite,
+  client,
+  datasetId,
+}: {
+  suite: SuiteState;
+  client: PhoenixClient;
+  datasetId: string;
+}): Promise<void> {
   try {
     const { data: response } = await client.GET("/v1/datasets/{id}/examples", {
       params: { path: { id: datasetId } },
     });
     const fetched = response?.data?.examples ?? [];
-
     const idToTestName = new Map<string, string>();
     const inputKeyToTestNames = new Map<string, string[]>();
     for (const [testName, registered] of suite.registeredExamples.entries()) {
       if (registered.params.id) {
         idToTestName.set(registered.params.id, testName);
-        continue;
+      } else {
+        const key = stableKey(registered.params.input);
+        inputKeyToTestNames.set(key, [
+          ...(inputKeyToTestNames.get(key) ?? []),
+          testName,
+        ]);
       }
-      const key = stableKey(registered.params.input);
-      const arr = inputKeyToTestNames.get(key) ?? [];
-      arr.push(testName);
-      inputKeyToTestNames.set(key, arr);
     }
-
-    for (const ex of fetched) {
-      const byId = idToTestName.get(ex.id);
-      if (byId) {
-        suite.exampleIdsByTest.set(byId, {
-          exampleId: ex.id,
-          nodeId: ex.node_id,
+    for (const example of fetched) {
+      const testName = idToTestName.get(example.id);
+      if (testName) {
+        suite.exampleIdsByTest.set(testName, {
+          exampleId: example.id,
+          nodeId: example.node_id,
         });
-        idToTestName.delete(ex.id);
+        idToTestName.delete(example.id);
         continue;
       }
-      const queue = inputKeyToTestNames.get(stableKey(ex.input));
-      if (queue && queue.length) {
-        const testName = queue.shift() as string;
-        suite.exampleIdsByTest.set(testName, {
-          exampleId: ex.id,
-          nodeId: ex.node_id,
+      const queue = inputKeyToTestNames.get(stableKey(example.input));
+      const matchedTestName = queue?.shift();
+      if (matchedTestName) {
+        suite.exampleIdsByTest.set(matchedTestName, {
+          exampleId: example.id,
+          nodeId: example.node_id,
         });
       }
     }
   } catch {
-    // If we cannot resolve example ids, runs will be logged without one.
+    // Runs can still be logged without resolved example IDs.
   }
+}
 
+async function createSuiteExperiment({
+  suite,
+  client,
+  datasetId,
+  datasetName,
+  description,
+}: {
+  suite: SuiteState;
+  client: PhoenixClient;
+  datasetId: string;
+  datasetName: string;
+  description: string;
+}): Promise<boolean> {
   const projectName = `${datasetName}-${new Date().toISOString()}`;
   suite.projectName = projectName;
-
   try {
-    const experimentResponse = await client
+    const response = await client
       .POST("/v1/datasets/{dataset_id}/experiments", {
         params: { path: { dataset_id: datasetId } },
         body: {
@@ -334,37 +348,46 @@ export async function initializeSuite(suite: SuiteState): Promise<void> {
           repetitions: Math.max(1, suite.maxRepetitions ?? 1),
         },
       })
-      .then((res) => res.data?.data);
-    if (!experimentResponse) {
-      throw new Error("Failed to create experiment");
-    }
-    suite.experimentId = experimentResponse.id;
-    suite.projectName = experimentResponse.project_name ?? projectName;
-  } catch (err) {
-    suite.trackingDisabled = true;
-    suite.setupError = err instanceof Error ? err : new Error(String(err));
-    suite.tracer = createNoOpProvider().getTracer("no-op");
-    suite.evaluatorTracer = suite.tracer;
-    return;
+      .then((result) => result.data?.data);
+    if (!response) throw new Error("Failed to create experiment");
+    suite.experimentId = response.id;
+    suite.projectName = response.project_name ?? projectName;
+    return true;
+  } catch (error) {
+    disableSuiteTracking({ suite, error });
+    return false;
   }
+}
 
+function setupSuiteTracer({
+  suite,
+  client,
+}: {
+  suite: SuiteState;
+  client: PhoenixClient;
+}): boolean {
   const baseUrl = client.config.baseUrl;
   if (!baseUrl) {
-    suite.trackingDisabled = true;
-    suite.setupError = new Error(
-      "Phoenix base URL not found. Set PHOENIX_ENDPOINT (or PHOENIX_COLLECTOR_ENDPOINT) or pass baseUrl on the client."
-    );
-    suite.tracer = createNoOpProvider().getTracer("no-op");
-    suite.evaluatorTracer = suite.tracer;
-    return;
+    disableSuiteTracking({
+      suite,
+      error: new Error(
+        "Phoenix base URL not found. Set PHOENIX_ENDPOINT (or PHOENIX_COLLECTOR_ENDPOINT) or pass baseUrl on the client."
+      ),
+    });
+    return false;
   }
-
   maybeWarnHttpScheme(baseUrl, client.config.headers);
-
-  let provider: NodeTracerProvider;
+  const projectName = suite.projectName;
+  if (!projectName) {
+    disableSuiteTracking({
+      suite,
+      error: new Error("Experiment project name is missing."),
+    });
+    return false;
+  }
   try {
-    provider = register({
-      projectName: suite.projectName,
+    const provider = register({
+      projectName,
       url: getTraceExportUrl(client.config),
       headers: client.config.headers
         ? toObjectHeaders(client.config.headers)
@@ -374,18 +397,59 @@ export async function initializeSuite(suite: SuiteState): Promise<void> {
     });
     suite.tracerProvider = provider;
     suite.globalRegistration = attachGlobalTracerProvider(provider);
-  } catch (err) {
+    suite.tracer = provider.getTracer(projectName);
+    suite.evaluatorTracer = provider.getTracer(`${projectName}-evaluators`);
+    return true;
+  } catch (error) {
+    disableSuiteTracking({ suite, error });
+    return false;
+  }
+}
+
+/**
+ * Initialize the suite: upload the dataset, create the experiment, and
+ * register the OpenInference tracer.
+ *
+ * If tracing is disabled (no Phoenix env vars, or PHOENIX_TEST_TRACKING=false),
+ * this populates a no-op tracer and exits without making any network calls.
+ */
+export async function initializeSuite(suite: SuiteState): Promise<void> {
+  const tracking = isTrackingEnabled(suite);
+  if (!tracking.enabled) {
     suite.trackingDisabled = true;
-    suite.setupError = err instanceof Error ? err : new Error(String(err));
+    suite.trackingDisabledReason = tracking.reason;
     suite.tracer = createNoOpProvider().getTracer("no-op");
     suite.evaluatorTracer = suite.tracer;
     return;
   }
-  suite.tracer = provider.getTracer(suite.projectName);
-  suite.evaluatorTracer = provider.getTracer(`${suite.projectName}-evaluators`);
 
-  if (suite.datasetId && suite.experimentId) {
-    suite.links = buildLinks(client, suite.datasetId, suite.experimentId);
+  const client = suite.config.client ?? createClient();
+  suite.client = client;
+
+  const datasetName = suite.config.datasetName ?? suite.name;
+  const description =
+    suite.config.description ??
+    `Phoenix test dataset auto-generated from ${suite.name}`;
+
+  const datasetId = await createSuiteDataset({
+    suite,
+    client,
+    datasetName,
+    description,
+  });
+  if (!datasetId) return;
+  suite.datasetId = datasetId;
+  await resolveSuiteExampleIds({ suite, client, datasetId });
+  const hasExperiment = await createSuiteExperiment({
+    suite,
+    client,
+    datasetId,
+    datasetName,
+    description,
+  });
+  if (!hasExperiment || !setupSuiteTracer({ suite, client })) return;
+  if (suite.experimentId) {
+    suite.links = buildLinks(client, datasetId, suite.experimentId);
   }
 }
 

package/src/types/prompts.ts

@@ -16,7 +16,8 @@ export type PromptModelProvider =
   | "FIREWORKS"
   | "GROQ"
   | "MOONSHOT"
-  | "PERPLEXITY";
+  | "PERPLEXITY"
+  | "ZAI";
 
 /**
  * Supported prompt provider SDKs

package/src/types/sessions.ts

@@ -24,6 +24,12 @@ export interface Session extends Node {
   startTime: string;
   /** ISO 8601 timestamp of when the last trace in the session ended */
   endTime: string;
+  /** Cumulative prompt token count across all spans in the session */
+  tokenCountPrompt?: number;
+  /** Cumulative completion token count across all spans in the session */
+  tokenCountCompletion?: number;
+  /** Cumulative total token count across all spans in the session */
+  tokenCountTotal?: number;
   /** The traces that belong to this session */
   traces: SessionTrace[];
 }