[Evals] Optimize tracing for Phoenix (elastic#238599)

## What

Optimizes inference instrumentation for use in Phoenix, by no longer
requiring a single root trace. Additionally, retry on missing scores in
the criteria evaluator.

## Why

We previously required a single root trace for its events to be
exported. This was done to not have to ingest all unrelated HTTP
requests that do not have any value for what we use Phoenix for:
debugging LLM-based workflows. The change we are making here is that we
use context to flag events as inference-related, instead of requiring it
to be a span. We don't need a single `RunExperiment` span anymore for
this reason, which messes up the token metrics in Phoenix for
evaluations.

I also added retries around scoring - in some cases the LLM did not
score all criteria. To that end, I've added `executeUntilValid` which
will ask the LLM to keep calling a tool until it no longer encounters
errors.

Some other minor changes:
- experiments can now contain metadata
- export types for `@kbn/evals` worker fixtures

I've targeted this at 9.2.0 as it is a dev-only change (mostly) and I
want to limit the chance of conflicts.

Author: dgieselaar

x-pack/platform/packages/shared/kbn-evals/index.ts

@@ -11,4 +11,4 @@ export { createPlaywrightEvalsConfig } from './src/config/create_playwright_eval
 export type { KibanaPhoenixClient } from './src/kibana_phoenix_client/client';
 export { createQuantitativeCorrectnessEvaluators } from './src/evaluators/correctness';
 export { createQuantitativeGroundednessEvaluator } from './src/evaluators/groundedness';
-export type { EvaluationDataset } from './src/types';
+export type { EvaluationDataset, EvaluationWorkerFixtures } from './src/types';

x-pack/platform/packages/shared/kbn-evals/src/evaluate.ts

@@ -6,22 +6,15 @@
  */
 
 import type { InferenceConnectorType, InferenceConnector, Model } from '@kbn/inference-common';
-import {
-  getConnectorModel,
-  type BoundInferenceClient,
-  getConnectorFamily,
-  getConnectorProvider,
-} from '@kbn/inference-common';
+import { getConnectorModel, getConnectorFamily, getConnectorProvider } from '@kbn/inference-common';
 import { createRestClient } from '@kbn/inference-plugin/common';
 import { test as base } from '@kbn/scout';
-import type { HttpHandler } from '@kbn/core/public';
-import type { AvailableConnectorWithId } from '@kbn/gen-ai-functional-testing';
 import { getPhoenixConfig } from './utils/get_phoenix_config';
 import { KibanaPhoenixClient } from './kibana_phoenix_client/client';
 import type { EvaluationTestOptions } from './config/create_playwright_eval_config';
 import { httpHandlerFromKbnClient } from './utils/http_handler_from_kbn_client';
 import { createCriteriaEvaluator } from './evaluators/criteria';
-import type { DefaultEvaluators } from './types';
+import type { DefaultEvaluators, EvaluationSpecificWorkerFixtures } from './types';
 import { reportModelScore } from './utils/report_model_score';
 import { createConnectorFixture } from './utils/create_connector_fixture';
 import { createCorrectnessAnalysisEvaluator } from './evaluators/correctness';
@@ -33,19 +26,7 @@ import { createGroundednessAnalysisEvaluator } from './evaluators/groundedness';
  * Test type for evaluations. Loads an inference client and a
  * (Kibana-flavored) Phoenix client.
  */
-export const evaluate = base.extend<
-  {},
-  {
-    inferenceClient: BoundInferenceClient;
-    phoenixClient: KibanaPhoenixClient;
-    evaluators: DefaultEvaluators;
-    fetch: HttpHandler;
-    connector: AvailableConnectorWithId;
-    evaluationConnector: AvailableConnectorWithId;
-    repetitions: number;
-    evaluationAnalysisService: EvaluationAnalysisService;
-  }
->({
+export const evaluate = base.extend<{}, EvaluationSpecificWorkerFixtures>({
   fetch: [
     async ({ kbnClient, log }, use) => {
       // add a HttpHandler as a fixture, so consumers can use

x-pack/platform/packages/shared/kbn-evals/src/evaluators/criteria/index.ts

@@ -8,8 +8,8 @@
 import type { BoundInferenceClient } from '@kbn/inference-common';
 import { ShortIdTable } from '@kbn/inference-common';
 import type { ToolingLog } from '@kbn/tooling-log';
-import { sumBy, uniqBy } from 'lodash';
-import pRetry from 'p-retry';
+import { difference, sumBy, uniqBy } from 'lodash';
+import { executeUntilValid } from '@kbn/inference-prompt-utils';
 import type { Evaluator } from '../../types';
 import { LlmCriteriaEvaluationPrompt } from './prompt';
 
@@ -55,22 +55,20 @@ export function createCriteriaEvaluator({
 
   return {
     evaluate: async ({ input, output }) => {
-      async function scoreTask() {
-        const response = await inferenceClient.prompt({
-          prompt: LlmCriteriaEvaluationPrompt,
-          input: {
-            input: JSON.stringify(input),
-            output: JSON.stringify(output),
-            criteria: structuredCriteria.map((criterion) => {
-              return `${criterion.id}: ${criterion.text}`;
-            }),
-          },
-        });
+      function toScores(
+        evaluatedCriteria: Array<{ id: string; result: 'PASS' | 'FAIL' | 'N/A'; reason?: string }>
+      ) {
+        const evaluations = uniqBy(evaluatedCriteria, (criterion) => criterion.id);
+
+        const evaluatedCriteriaIds = evaluations.map((evaluation) => evaluation.id);
+
+        const criteriaIds = Array.from(criteriaById.keys());
 
-        const evaluations = uniqBy(
-          response.toolCalls.flatMap((toolCall) => toolCall.function.arguments.criteria),
-          (criterion) => criterion.id
-        );
+        const unscored = difference(criteriaIds, evaluatedCriteriaIds);
+
+        if (unscored.length) {
+          throw new Error(`Missing scores for ${unscored.join(', ')}`);
+        }
 
         return evaluations.map((evaluation) => {
           const criterion = criteriaById.get(evaluation.id);
@@ -85,12 +83,40 @@ export function createCriteriaEvaluator({
         });
       }
 
-      const results = await pRetry(scoreTask, {
-        retries: 0,
-        onFailedAttempt: (error) => {
-          log.error(new Error(`Failed to score task`, { cause: error }));
-        },
-      });
+      async function scoreTask() {
+        const response = await executeUntilValid({
+          prompt: LlmCriteriaEvaluationPrompt,
+          inferenceClient,
+          input: {
+            input: JSON.stringify(input),
+            output: JSON.stringify(output),
+            criteria: structuredCriteria.map((criterion) => {
+              return `${criterion.id}: ${criterion.text}`;
+            }),
+          },
+          finalToolChoice: {
+            function: 'score',
+          },
+          maxRetries: 3,
+          toolCallbacks: {
+            score: async (toolCall) => {
+              return {
+                response: {
+                  scores: toScores(toolCall.function.arguments.criteria),
+                },
+              };
+            },
+          },
+        });
+
+        return response;
+      }
+
+      const { toolCalls } = await scoreTask();
+
+      const results = toolCalls.flatMap((toolCall) =>
+        toScores(toolCall.function.arguments.criteria)
+      );
 
       function normalize(val: number) {
         if (!isFinite(val)) {

x-pack/platform/packages/shared/kbn-evals/src/kibana_phoenix_client/client.ts

@@ -11,7 +11,7 @@ import type { RanExperiment, TaskOutput } from '@arizeai/phoenix-client/dist/esm
 import type { DatasetInfo, Example } from '@arizeai/phoenix-client/dist/esm/types/datasets';
 import type { SomeDevLog } from '@kbn/some-dev-log';
 import type { Model } from '@kbn/inference-common';
-import { withActiveInferenceSpan } from '@kbn/inference-tracing';
+import { withInferenceContext } from '@kbn/inference-tracing';
 import type { Evaluator, EvaluationDataset, ExperimentTask } from '../types';
 import { upsertDataset } from './upsert_dataset';
 import type { PhoenixConfig } from '../utils/get_phoenix_config';
@@ -86,9 +86,11 @@ export class KibanaPhoenixClient {
   async runExperiment<TEvaluationDataset extends EvaluationDataset, TTaskOutput extends TaskOutput>(
     {
       dataset,
+      metadata,
       task,
     }: {
       dataset: TEvaluationDataset;
+      metadata?: Record<string, unknown>;
       task: ExperimentTask<TEvaluationDataset['examples'][number], TTaskOutput>;
     },
     evaluators: Array<Evaluator<TEvaluationDataset['examples'][number], TTaskOutput>>
@@ -98,13 +100,15 @@ export class KibanaPhoenixClient {
     {
       dataset,
       task,
+      metadata: experimentMetadata,
     }: {
       dataset: EvaluationDataset;
       task: ExperimentTask<Example, TaskOutput>;
+      metadata?: Record<string, unknown>;
     },
     evaluators: Evaluator[]
   ): Promise<RanExperiment> {
-    return await withActiveInferenceSpan('RunExperiment', async (span) => {
+    return withInferenceContext(async () => {
       const { datasetId } = await this.syncDataSet(dataset);
 
       const experiments = await import('@arizeai/phoenix-client/experiments');
@@ -115,6 +119,7 @@ export class KibanaPhoenixClient {
         experimentName: `Run ID: ${this.options.runId} - Dataset: ${dataset.name}`,
         task,
         experimentMetadata: {
+          ...experimentMetadata,
           model: this.options.model,
           runId: this.options.runId,
         },

x-pack/platform/packages/shared/kbn-evals/src/kibana_phoenix_client/diff_examples.ts

@@ -7,13 +7,14 @@
 
 import type { Example } from '@arizeai/phoenix-client/dist/esm/types/datasets';
 import objectHash from 'object-hash';
+import { isEmpty, omitBy } from 'lodash';
 import type { ExampleWithId } from '../types';
 
 function normaliseExample(example: Example | ExampleWithId) {
   return {
     input: example.input,
     output: example.output,
-    metadata: example.metadata,
+    metadata: omitBy(example.metadata, isEmpty),
   };
 }
 

x-pack/platform/packages/shared/kbn-evals/src/types.ts

@@ -7,11 +7,17 @@
 
 import type { Example } from '@arizeai/phoenix-client/dist/esm/types/datasets';
 import type {
-  EvaluationResult,
+  EvaluationResult as PhoenixEvaluationResult,
   Evaluator as PhoenixEvaluator,
   TaskOutput,
 } from '@arizeai/phoenix-client/dist/esm/types/experiments';
+import type { BoundInferenceClient } from '@kbn/inference-common';
+import type { HttpHandler } from '@kbn/core/public';
+import type { AvailableConnectorWithId } from '@kbn/gen-ai-functional-testing';
+import type { ScoutWorkerFixtures } from '@kbn/scout';
+import type { KibanaPhoenixClient } from './kibana_phoenix_client/client';
 import type { EvaluationCriterion } from './evaluators/criteria';
+import type { EvaluationAnalysisService } from './utils/analysis';
 
 export interface EvaluationDataset {
   name: string;
@@ -31,6 +37,8 @@ export interface EvaluatorParams<TExample extends Example, TTaskOutput extends T
   metadata: TExample['metadata'];
 }
 
+export type EvaluationResult = PhoenixEvaluationResult;
+
 type EvaluatorCallback<TExample extends Example, TTaskOutput extends TaskOutput> = (
   params: EvaluatorParams<TExample, TTaskOutput>
 ) => Promise<EvaluationResult>;
@@ -54,3 +62,25 @@ export type ExperimentTask<TExample extends Example, TTaskOutput extends TaskOut
 
 // simple version of Phoenix's ExampleWithId
 export type ExampleWithId = Example & { id: string };
+
+export interface EvaluationSpecificWorkerFixtures {
+  inferenceClient: BoundInferenceClient;
+  phoenixClient: KibanaPhoenixClient;
+  evaluators: DefaultEvaluators;
+  fetch: HttpHandler;
+  connector: AvailableConnectorWithId;
+  evaluationConnector: AvailableConnectorWithId;
+  repetitions: number;
+  evaluationAnalysisService: EvaluationAnalysisService;
+}
+
+export interface EvaluationWorkerFixtures extends ScoutWorkerFixtures {
+  inferenceClient: BoundInferenceClient;
+  phoenixClient: KibanaPhoenixClient;
+  evaluators: DefaultEvaluators;
+  fetch: HttpHandler;
+  connector: AvailableConnectorWithId;
+  evaluationConnector: AvailableConnectorWithId;
+  repetitions: number;
+  evaluationAnalysisService: EvaluationAnalysisService;
+}

x-pack/platform/packages/shared/kbn-evals/tsconfig.json

@@ -29,5 +29,6 @@
     "@kbn/repo-info",
     "@kbn/std",
     "@kbn/test",
+    "@kbn/inference-prompt-utils",
   ]
 }

x-pack/platform/packages/shared/kbn-inference-prompt-utils/index.ts

@@ -10,3 +10,5 @@ export type {
   ReasoningPromptResponse,
   ReasoningPromptResponseOf,
 } from './src/flows/reasoning/types';
+
+export { executeUntilValid } from './src/flows/until_valid/execute_until_valid';