Add base FTR test coverage for inference APIs (#198000)

## Summary

Part of https://github.com/elastic/kibana-team/issues/1271

This PR introduces the first set of end to end integration test for the
inference APIs, and the tooling required to do so (see issue for more
context)

- Add a dedicated pipeline for ai-infra GenAI tests. pipeline is
triggered when:
  - genAI stack connectors, or ai-infra owned code is changed
  - when the `ci:all-gen-ai-suites` label is present on a PR
  - on merge
- adapt the `ftr_configs.sh` script to load GenAI connector
configuration from vault when a specific var env is set
- create the `@kbn/gen-ai-functional-testing` package, which for now
only contains utilities to load the GenAI connector configuration in FTR
tests
- Add FTR integration tests for the `chatComplete` API of the
`inference` plugin

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>

.buildkite/ftr_platform_stateful_configs.yml

@@ -42,6 +42,9 @@ disabled:
   # Default http2 config to use for performance journeys
   - x-pack/performance/configs/http2_config.ts
 
+  # Gen AI suites, running with their own pipeline
+  - x-pack/test/functional_gen_ai/inference/config.ts
+
 defaultQueue: 'n2-4-spot'
 enabled:
   - test/accessibility/config.ts

.buildkite/pipelines/on_merge.yml

@@ -169,6 +169,25 @@ steps:
         - exit_status: '*'
           limit: 1
 
+  - command: .buildkite/scripts/steps/test/ftr_configs.sh
+    env:
+      FTR_CONFIG: "x-pack/test/functional_gen_ai/inference/config.ts"
+      FTR_CONFIG_GROUP_KEY: 'ftr-ai-infra-gen-ai-inference-api'
+      FTR_GEN_AI: "1"
+    label: AppEx AI-Infra Inference APIs FTR tests
+    key: ai-infra-gen-ai-inference-api
+    timeout_in_minutes: 50
+    parallelism: 1
+    agents:
+      machineType: n2-standard-4
+      preemptible: true
+    retry:
+      automatic:
+        - exit_status: '-1'
+          limit: 3
+        - exit_status: '*'
+          limit: 1
+
   - command: .buildkite/scripts/steps/functional/security_serverless_entity_analytics.sh
     label: 'Serverless Entity Analytics - Security Cypress Tests'
     agents:

.buildkite/pipelines/pull_request/ai_infra_gen_ai.yml

@@ -0,0 +1,30 @@
+steps:
+  - group: AppEx AI-Infra genAI tests
+    key: ai-infra-gen-ai
+    depends_on:
+      - build
+      - quick_checks
+      - checks
+      - linting
+      - linting_with_types
+      - check_types
+      - check_oas_snapshot
+    steps:
+      - command: .buildkite/scripts/steps/test/ftr_configs.sh
+        env:
+          FTR_CONFIG: "x-pack/test/functional_gen_ai/inference/config.ts"
+          FTR_CONFIG_GROUP_KEY: 'ftr-ai-infra-gen-ai-inference-api'
+          FTR_GEN_AI: "1"
+        label: AppEx AI-Infra Inference APIs FTR tests
+        key: ai-infra-gen-ai-inference-api
+        timeout_in_minutes: 50
+        parallelism: 1
+        agents:
+          machineType: n2-standard-4
+          preemptible: true
+        retry:
+          automatic:
+            - exit_status: '-1'
+              limit: 3
+            - exit_status: '*'
+              limit: 1

.buildkite/scripts/common/setup_job_env.sh

@@ -132,6 +132,14 @@ EOF
   export ELASTIC_APM_API_KEY
 }
 
+# Set up GenAI keys
+{
+  if [[ "${FTR_GEN_AI:-}" =~ ^(1|true)$ ]]; then
+    echo "FTR_GEN_AI was set - exposing LLM connectors"
+    export KIBANA_TESTING_AI_CONNECTORS="$(vault_get ai-infra-ci-connectors connectors-config)"
+  fi
+}
+
 # Set up GCS Service Account for CDN
 {
   GCS_SA_CDN_KEY="$(vault_get gcs-sa-cdn-prod key)"

.buildkite/scripts/pipelines/pull_request/pipeline.ts

@@ -140,6 +140,20 @@ const getPipeline = (filename: string, removeSteps = true) => {
       pipeline.push(getPipeline('.buildkite/pipelines/pull_request/slo_plugin_e2e.yml'));
     }
 
+    if (
+      (await doAnyChangesMatch([
+        /^x-pack\/packages\/ai-infra/,
+        /^x-pack\/plugins\/ai_infra/,
+        /^x-pack\/plugins\/inference/,
+        /^x-pack\/plugins\/stack_connectors\/server\/connector_types\/bedrock/,
+        /^x-pack\/plugins\/stack_connectors\/server\/connector_types\/gemini/,
+        /^x-pack\/plugins\/stack_connectors\/server\/connector_types\/openai/,
+      ])) ||
+      GITHUB_PR_LABELS.includes('ci:all-gen-ai-suites')
+    ) {
+      pipeline.push(getPipeline('.buildkite/pipelines/pull_request/ai_infra_gen_ai.yml'));
+    }
+
     if (
       GITHUB_PR_LABELS.includes('ci:deploy-cloud') ||
       GITHUB_PR_LABELS.includes('ci:cloud-deploy') ||

.github/CODEOWNERS

@@ -370,6 +370,7 @@ packages/kbn-formatters @elastic/obs-ux-logs-team
 packages/kbn-ftr-common-functional-services @elastic/kibana-operations @elastic/appex-qa
 packages/kbn-ftr-common-functional-ui-services @elastic/appex-qa
 packages/kbn-ftr-screenshot-filename @elastic/kibana-operations @elastic/appex-qa
+packages/kbn-gen-ai-functional-testing @elastic/appex-ai-infra
 packages/kbn-generate @elastic/kibana-operations
 packages/kbn-generate-console-definitions @elastic/kibana-management
 packages/kbn-generate-csv @elastic/appex-sharedux
@@ -1819,6 +1820,7 @@ packages/kbn-monaco/src/esql @elastic/kibana-esql
 
 # AppEx AI Infra
 /x-pack/plugins/inference @elastic/appex-ai-infra @elastic/obs-ai-assistant @elastic/security-generative-ai
+/x-pack/test/functional_gen_ai/inference @elastic/appex-ai-infra
 
 # AppEx Platform Services Security
 //x-pack/test_serverless/api_integration/test_suites/common/security_response_headers.ts  @elastic/kibana-security
@@ -2104,7 +2106,7 @@ x-pack/test/api_integration/apis/management/index_management/inference_endpoints
 /x-pack/test/api_integration/services/security_solution_*.gen.ts @elastic/security-solution
 /x-pack/test/accessibility/apps/group3/security_solution.ts @elastic/security-solution
 /x-pack/test_serverless/functional/test_suites/security/config.ts @elastic/security-solution @elastic/appex-qa
-x-pack/test_serverless/functional/test_suites/security/config.mki_only.ts @elastic/security-solution @elastic/appex-qa 
+x-pack/test_serverless/functional/test_suites/security/config.mki_only.ts @elastic/security-solution @elastic/appex-qa
 x-pack/test_serverless/functional/test_suites/security/index.mki_only.ts @elastic/security-solution @elastic/appex-qa @elastic/kibana-cloud-security-posture
 /x-pack/test_serverless/functional/test_suites/security/config.feature_flags.ts @elastic/security-solution @elastic/kibana-cloud-security-posture
 /x-pack/test_serverless/api_integration/test_suites/observability/config.feature_flags.ts @elastic/security-solution

package.json

@@ -1454,6 +1454,7 @@
     "@kbn/ftr-common-functional-services": "link:packages/kbn-ftr-common-functional-services",
     "@kbn/ftr-common-functional-ui-services": "link:packages/kbn-ftr-common-functional-ui-services",
     "@kbn/ftr-screenshot-filename": "link:packages/kbn-ftr-screenshot-filename",
+    "@kbn/gen-ai-functional-testing": "link:packages/kbn-gen-ai-functional-testing",
     "@kbn/generate": "link:packages/kbn-generate",
     "@kbn/get-repo-files": "link:packages/kbn-get-repo-files",
     "@kbn/import-locator": "link:packages/kbn-import-locator",

packages/kbn-gen-ai-functional-testing/.gitignore

@@ -0,0 +1,2 @@
+## local version of the connector config
+connector_config.json

packages/kbn-gen-ai-functional-testing/README.md

@@ -0,0 +1,49 @@
+# @kbn/gen-ai-functional-testing
+
+Package exposing various utilities for GenAI/LLM related functional testing.
+
+## Features
+
+### LLM connectors
+
+Utilizing LLM connectors on FTR tests can be done via the `getPreconfiguredConnectorConfig` and `getAvailableConnectors` tools.
+
+`getPreconfiguredConnectorConfig` should be used to define the list of connectors when creating the FTR test's configuration.
+
+```ts
+import { getPreconfiguredConnectorConfig } from '@kbn/gen-ai-functional-testing'
+
+export default async function ({ readConfigFile }: FtrConfigProviderContext) {
+  const xpackFunctionalConfig = {...};
+  const preconfiguredConnectors = getPreconfiguredConnectorConfig();
+
+  return {
+    ...xpackFunctionalConfig.getAll(),
+    kbnTestServer: {
+      ...xpackFunctionalConfig.get('kbnTestServer'),
+      serverArgs: [
+        ...xpackFunctionalConfig.get('kbnTestServer.serverArgs'),
+        `--xpack.actions.preconfigured=${JSON.stringify(preconfiguredConnectors)}`,
+      ],
+    },
+  };
+}
+```
+
+then the `getAvailableConnectors` can be used during the test suite to retrieve the list of LLM connectors.
+
+For example to run some predefined test suite against all exposed LLM connectors:
+
+```ts
+import { getAvailableConnectors } from '@kbn/gen-ai-functional-testing';
+
+export default function (providerContext: FtrProviderContext) {
+  describe('Some GenAI FTR test suite', async () => {
+    getAvailableConnectors().forEach((connector) => {
+      describe(`Using connector ${connector.id}`, () => {
+        myTestSuite(connector, providerContext);
+      });
+    });
+  });
+}
+```

packages/kbn-gen-ai-functional-testing/index.ts

@@ -0,0 +1,16 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+export {
+  AI_CONNECTORS_VAR_ENV,
+  getPreconfiguredConnectorConfig,
+  getAvailableConnectors,
+  type AvailableConnector,
+  type AvailableConnectorWithId,
+} from './src/connectors';

packages/kbn-gen-ai-functional-testing/jest.config.js

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+module.exports = {
+  preset: '@kbn/test/jest_node',
+  rootDir: '../..',
+  roots: ['<rootDir>/packages/kbn-gen-ai-functional-testing'],
+};

packages/kbn-gen-ai-functional-testing/kibana.jsonc

@@ -0,0 +1,6 @@
+{
+  "type": "shared-common",
+  "id": "@kbn/gen-ai-functional-testing",
+  "owner": "@elastic/appex-ai-infra",
+  "devOnly": true
+}

packages/kbn-gen-ai-functional-testing/package.json

@@ -0,0 +1,6 @@
+{
+  "name": "@kbn/gen-ai-functional-testing",
+  "private": true,
+  "version": "1.0.0",
+  "license": "Elastic License 2.0 OR AGPL-3.0-only OR SSPL-1.0"
+}

packages/kbn-gen-ai-functional-testing/scripts/format_connector_config.js

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+require('@kbn/babel-register').install();
+require('../src/manage_connector_config').formatCurrentConfig();

packages/kbn-gen-ai-functional-testing/scripts/retrieve_connector_config.js

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+require('@kbn/babel-register').install();
+require('../src/manage_connector_config').retrieveFromVault();

packages/kbn-gen-ai-functional-testing/scripts/upload_connector_config.js

@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+require('@kbn/babel-register').install();
+require('../src/manage_connector_config').uploadToVault();

packages/kbn-gen-ai-functional-testing/src/connectors.ts

@@ -0,0 +1,91 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { schema } from '@kbn/config-schema';
+
+/**
+ * The environment variable that is used by the CI to load the connectors configuration
+ */
+export const AI_CONNECTORS_VAR_ENV = 'KIBANA_TESTING_AI_CONNECTORS';
+
+const connectorsSchema = schema.recordOf(
+  schema.string(),
+  schema.object({
+    name: schema.string(),
+    actionTypeId: schema.string(),
+    config: schema.recordOf(schema.string(), schema.any()),
+    secrets: schema.recordOf(schema.string(), schema.any()),
+  })
+);
+
+export interface AvailableConnector {
+  name: string;
+  actionTypeId: string;
+  config: Record<string, unknown>;
+  secrets: Record<string, unknown>;
+}
+
+export interface AvailableConnectorWithId extends AvailableConnector {
+  id: string;
+}
+
+const loadConnectors = (): Record<string, AvailableConnector> => {
+  const envValue = process.env[AI_CONNECTORS_VAR_ENV];
+  if (!envValue) {
+    return {};
+  }
+
+  let connectors: Record<string, AvailableConnector>;
+  try {
+    connectors = JSON.parse(Buffer.from(envValue, 'base64').toString('utf-8'));
+  } catch (e) {
+    throw new Error(
+      `Error trying to parse value from KIBANA_AI_CONNECTORS environment variable: ${e.message}`
+    );
+  }
+  return connectorsSchema.validate(connectors);
+};
+
+/**
+ * Retrieve the list of preconfigured connectors that should be used when defining the
+ * FTR configuration of suites using the connectors.
+ *
+ * @example
+ * ```ts
+ * import { getPreconfiguredConnectorConfig } from '@kbn/gen-ai-functional-testing'
+ *
+ * export default async function ({ readConfigFile }: FtrConfigProviderContext) {
+ *   const xpackFunctionalConfig = {...};
+ *   const preconfiguredConnectors = getPreconfiguredConnectorConfig();
+ *
+ *   return {
+ *     ...xpackFunctionalConfig.getAll(),
+ *     kbnTestServer: {
+ *       ...xpackFunctionalConfig.get('kbnTestServer'),
+ *       serverArgs: [
+ *         ...xpackFunctionalConfig.get('kbnTestServer.serverArgs'),
+ *         `--xpack.actions.preconfigured=${JSON.stringify(preconfiguredConnectors)}`,
+ *       ],
+ *     },
+ *   };
+ * }
+ * ```
+ */
+export const getPreconfiguredConnectorConfig = () => {
+  return loadConnectors();
+};
+
+export const getAvailableConnectors = (): AvailableConnectorWithId[] => {
+  return Object.entries(loadConnectors()).map(([id, connector]) => {
+    return {
+      id,
+      ...connector,
+    };
+  });
+};

packages/kbn-gen-ai-functional-testing/src/manage_connector_config.ts

@@ -0,0 +1,59 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import execa from 'execa';
+import Path from 'path';
+import { writeFile, readFile } from 'fs/promises';
+import { REPO_ROOT } from '@kbn/repo-info';
+
+const LOCAL_FILE = Path.join(
+  REPO_ROOT,
+  'packages',
+  'kbn-gen-ai-functional-testing',
+  'connector_config.json'
+);
+
+export const retrieveFromVault = async () => {
+  const { stdout } = await execa(
+    'vault',
+    ['read', '-field=connectors-config', 'secret/ci/elastic-kibana/ai-infra-ci-connectors'],
+    {
+      cwd: REPO_ROOT,
+      buffer: true,
+    }
+  );
+
+  const config = JSON.parse(Buffer.from(stdout, 'base64').toString('utf-8'));
+
+  await writeFile(LOCAL_FILE, JSON.stringify(config, undefined, 2));
+
+  // eslint-disable-next-line no-console
+  console.log(`Config dumped into ${LOCAL_FILE}`);
+};
+
+export const formatCurrentConfig = async () => {
+  const config = await readFile(LOCAL_FILE, 'utf-8');
+  const asB64 = Buffer.from(config).toString('base64');
+  // eslint-disable-next-line no-console
+  console.log(asB64);
+};
+
+export const uploadToVault = async () => {
+  const config = await readFile(LOCAL_FILE, 'utf-8');
+  const asB64 = Buffer.from(config).toString('base64');
+
+  await execa(
+    'vault',
+    ['write', 'secret/ci/elastic-kibana/ai-infra-ci-connectors', `connectors-config=${asB64}`],
+    {
+      cwd: REPO_ROOT,
+      buffer: true,
+    }
+  );
+};

packages/kbn-gen-ai-functional-testing/tsconfig.json

@@ -0,0 +1,20 @@
+{
+  "extends": "../../tsconfig.base.json",
+  "compilerOptions": {
+    "outDir": "target/types",
+    "types": [
+      "jest",
+      "node"
+    ]
+  },
+  "include": [
+    "**/*.ts",
+  ],
+  "exclude": [
+    "target/**/*"
+  ],
+  "kbn_references": [
+    "@kbn/config-schema",
+    "@kbn/repo-info",
+  ]
+}

packages/kbn-sse-utils-server/index.ts

@@ -8,3 +8,4 @@
  */
 
 export { observableIntoEventSourceStream } from './src/observable_into_event_source_stream';
+export { supertestToObservable } from './src/supertest_to_observable';

packages/kbn-sse-utils-server/src/supertest_to_observable.ts

@@ -0,0 +1,68 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type supertest from 'supertest';
+import { PassThrough } from 'stream';
+import { createParser } from 'eventsource-parser';
+import { Observable } from 'rxjs';
+
+/**
+ * Convert a supertest response to an SSE observable.
+ *
+ * Note: the supertest response should *NOT* be awaited when using that utility,
+ * or at least not before calling it.
+ *
+ * @example
+ * ```ts
+ * const response = supertest
+ *   .post(`/some/sse/endpoint`)
+ *   .set('kbn-xsrf', 'kibana')
+ *   .send({
+ *     some: 'thing'
+ *   });
+ * const events = supertestIntoObservable(response);
+ * ```
+ */
+export function supertestToObservable<T = any>(response: supertest.Test): Observable<T> {
+  const stream = new PassThrough();
+  response.pipe(stream);
+
+  return new Observable<T>((subscriber) => {
+    const parser = createParser({
+      onEvent: (event) => {
+        subscriber.next(JSON.parse(event.data));
+      },
+    });
+
+    const readStream = async () => {
+      return new Promise<void>((resolve, reject) => {
+        const decoder = new TextDecoder();
+
+        const processChunk = (value: BufferSource) => {
+          parser.feed(decoder.decode(value, { stream: true }));
+        };
+
+        stream.on('data', (chunk) => {
+          processChunk(chunk);
+        });
+
+        stream.on('end', () => resolve());
+        stream.on('error', (err) => reject(err));
+      });
+    };
+
+    readStream()
+      .then(() => {
+        subscriber.complete();
+      })
+      .catch((error) => {
+        subscriber.error(error);
+      });
+  });
+}

tsconfig.base.json

@@ -982,6 +982,8 @@
       "@kbn/ftr-screenshot-filename/*": ["packages/kbn-ftr-screenshot-filename/*"],
       "@kbn/functional-with-es-ssl-cases-test-plugin": ["x-pack/test/functional_with_es_ssl/plugins/cases"],
       "@kbn/functional-with-es-ssl-cases-test-plugin/*": ["x-pack/test/functional_with_es_ssl/plugins/cases/*"],
+      "@kbn/gen-ai-functional-testing": ["packages/kbn-gen-ai-functional-testing"],
+      "@kbn/gen-ai-functional-testing/*": ["packages/kbn-gen-ai-functional-testing/*"],
       "@kbn/gen-ai-streaming-response-example-plugin": ["x-pack/examples/gen_ai_streaming_response_example"],
       "@kbn/gen-ai-streaming-response-example-plugin/*": ["x-pack/examples/gen_ai_streaming_response_example/*"],
       "@kbn/generate": ["packages/kbn-generate"],

x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.test.ts

@@ -97,6 +97,11 @@ describe('processVertexStream', () => {
 
       expectObservable(processed$).toBe('--(ab)', {
         a: {
+          content: 'last chunk',
+          tool_calls: [],
+          type: ChatCompletionEventType.ChatCompletionChunk,
+        },
+        b: {
           tokens: {
             completion: 1,
             prompt: 2,
@@ -104,11 +109,6 @@ describe('processVertexStream', () => {
           },
           type: ChatCompletionEventType.ChatCompletionTokenCount,
         },
-        b: {
-          content: 'last chunk',
-          tool_calls: [],
-          type: ChatCompletionEventType.ChatCompletionChunk,
-        },
       });
     });
   });

x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts

@@ -18,18 +18,6 @@ export function processVertexStream() {
   return (source: Observable<GenerateContentResponseChunk>) =>
     new Observable<ChatCompletionChunkEvent | ChatCompletionTokenCountEvent>((subscriber) => {
       function handleNext(value: GenerateContentResponseChunk) {
-        // completion: only present on last chunk
-        if (value.usageMetadata) {
-          subscriber.next({
-            type: ChatCompletionEventType.ChatCompletionTokenCount,
-            tokens: {
-              prompt: value.usageMetadata.promptTokenCount,
-              completion: value.usageMetadata.candidatesTokenCount,
-              total: value.usageMetadata.totalTokenCount,
-            },
-          });
-        }
-
         const contentPart = value.candidates?.[0].content.parts[0];
         const completion = contentPart?.text;
         const toolCall = contentPart?.functionCall;
@@ -49,6 +37,18 @@ export function processVertexStream() {
               : [],
           });
         }
+
+        // completion: only present on last chunk
+        if (value.usageMetadata) {
+          subscriber.next({
+            type: ChatCompletionEventType.ChatCompletionTokenCount,
+            tokens: {
+              prompt: value.usageMetadata.promptTokenCount,
+              completion: value.usageMetadata.candidatesTokenCount,
+              total: value.usageMetadata.totalTokenCount,
+            },
+          });
+        }
       }
 
       source.subscribe({

x-pack/plugins/inference/server/routes/chat_complete.ts

@@ -13,7 +13,13 @@ import type {
   RequestHandlerContext,
   KibanaRequest,
 } from '@kbn/core/server';
-import { MessageRole, ToolCall, ToolChoiceType } from '@kbn/inference-common';
+import {
+  MessageRole,
+  ToolCall,
+  ToolChoiceType,
+  InferenceTaskEventType,
+  isInferenceError,
+} from '@kbn/inference-common';
 import type { ChatCompleteRequestBody } from '../../common/http_apis';
 import { createClient as createInferenceClient } from '../inference_client';
 import { InferenceServerStart, InferenceStartDependencies } from '../types';
@@ -130,10 +136,22 @@ export function registerChatCompleteRoute({
       },
     },
     async (context, request, response) => {
-      const chatCompleteResponse = await callChatComplete({ request, stream: false });
-      return response.ok({
-        body: chatCompleteResponse,
-      });
+      try {
+        const chatCompleteResponse = await callChatComplete({ request, stream: false });
+        return response.ok({
+          body: chatCompleteResponse,
+        });
+      } catch (e) {
+        return response.custom({
+          statusCode: isInferenceError(e) ? e.meta?.status ?? 500 : 500,
+          bypassErrorFormat: true,
+          body: {
+            type: InferenceTaskEventType.error,
+            code: e.code ?? 'unknown',
+            message: e.message,
+          },
+        });
+      }
     }
   );
 
@@ -145,9 +163,9 @@ export function registerChatCompleteRoute({
       },
     },
     async (context, request, response) => {
-      const chatCompleteResponse = await callChatComplete({ request, stream: true });
+      const chatCompleteEvents$ = await callChatComplete({ request, stream: true });
       return response.ok({
-        body: observableIntoEventSourceStream(chatCompleteResponse, logger),
+        body: observableIntoEventSourceStream(chatCompleteEvents$, logger),
       });
     }
   );

x-pack/test/functional_gen_ai/inference/config.ts

@@ -0,0 +1,32 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { FtrConfigProviderContext } from '@kbn/test';
+import { getPreconfiguredConnectorConfig } from '@kbn/gen-ai-functional-testing';
+import { services } from './ftr_provider_context';
+
+// eslint-disable-next-line import/no-default-export
+export default async function ({ readConfigFile }: FtrConfigProviderContext) {
+  const xpackFunctionalConfig = await readConfigFile(
+    require.resolve('../../functional/config.base.js')
+  );
+
+  const preconfiguredConnectors = getPreconfiguredConnectorConfig();
+
+  return {
+    ...xpackFunctionalConfig.getAll(),
+    services,
+    testFiles: [require.resolve('./tests')],
+    kbnTestServer: {
+      ...xpackFunctionalConfig.get('kbnTestServer'),
+      serverArgs: [
+        ...xpackFunctionalConfig.get('kbnTestServer.serverArgs'),
+        `--xpack.actions.preconfigured=${JSON.stringify(preconfiguredConnectors)}`,
+      ],
+    },
+  };
+}

x-pack/test/functional_gen_ai/inference/ftr_provider_context.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { GenericFtrProviderContext } from '@kbn/test';
+import { services } from '../../functional/services';
+import { pageObjects } from '../../functional/page_objects';
+
+export type FtrProviderContext = GenericFtrProviderContext<typeof services, typeof pageObjects>;
+
+export { services, pageObjects };

x-pack/test/functional_gen_ai/inference/tests/chat_complete.ts

@@ -0,0 +1,263 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { lastValueFrom, toArray } from 'rxjs';
+import expect from '@kbn/expect';
+import { supertestToObservable } from '@kbn/sse-utils-server';
+import type { AvailableConnectorWithId } from '@kbn/gen-ai-functional-testing';
+import type { FtrProviderContext } from '../ftr_provider_context';
+
+export const chatCompleteSuite = (
+  { id: connectorId, actionTypeId: connectorType }: AvailableConnectorWithId,
+  { getService }: FtrProviderContext
+) => {
+  const supertest = getService('supertest');
+
+  describe('chatComplete API', () => {
+    describe('streaming disabled', () => {
+      it('returns a chat completion message for a simple prompt', async () => {
+        const response = await supertest
+          .post(`/internal/inference/chat_complete`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId,
+            system: 'Please answer the user question',
+            messages: [{ role: 'user', content: '2+2 ?' }],
+          })
+          .expect(200);
+
+        const message = response.body;
+
+        expect(message.toolCalls.length).to.eql(0);
+        expect(message.content).to.contain('4');
+      });
+
+      it('executes a tool with native function calling', async () => {
+        const response = await supertest
+          .post(`/internal/inference/chat_complete`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId,
+            system:
+              'Please answer the user question. You can use the available tools if you think it can help',
+            messages: [{ role: 'user', content: 'What is the result of 2*4*6*8*10*123 ?' }],
+            toolChoice: 'required',
+            tools: {
+              calculator: {
+                description: 'The calculator can be used to resolve mathematical calculations',
+                schema: {
+                  type: 'object',
+                  properties: {
+                    formula: {
+                      type: 'string',
+                      description: `The input for the calculator, in plain text, e.g. "2+(4*8)"`,
+                    },
+                  },
+                },
+              },
+            },
+          })
+          .expect(200);
+
+        const message = response.body;
+
+        expect(message.toolCalls.length).to.eql(1);
+        expect(message.toolCalls[0].function.name).to.eql('calculator');
+        expect(message.toolCalls[0].function.arguments.formula).to.contain('123');
+      });
+
+      // simulated FC is only for openAI
+      if (connectorType === '.gen-ai') {
+        it('executes a tool with simulated function calling', async () => {
+          const response = await supertest
+            .post(`/internal/inference/chat_complete`)
+            .set('kbn-xsrf', 'kibana')
+            .send({
+              connectorId,
+              system:
+                'Please answer the user question. You can use the available tools if you think it can help',
+              messages: [{ role: 'user', content: 'What is the result of 2*4*6*8*10*123 ?' }],
+              functionCalling: 'simulated',
+              toolChoice: 'required',
+              tools: {
+                calculator: {
+                  description: 'The calculator can be used to resolve mathematical calculations',
+                  schema: {
+                    type: 'object',
+                    properties: {
+                      formula: {
+                        type: 'string',
+                        description: `The input for the calculator, in plain text, e.g. "2+(4*8)"`,
+                      },
+                    },
+                  },
+                },
+              },
+            })
+            .expect(200);
+
+          const message = response.body;
+
+          expect(message.toolCalls.length).to.eql(1);
+          expect(message.toolCalls[0].function.name).to.eql('calculator');
+          expect(message.toolCalls[0].function.arguments.formula).to.contain('123');
+        });
+      }
+
+      it('returns token counts', async () => {
+        const response = await supertest
+          .post(`/internal/inference/chat_complete`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId,
+            system: 'Please answer the user question',
+            messages: [{ role: 'user', content: '2+2 ?' }],
+          })
+          .expect(200);
+
+        const { tokens } = response.body;
+
+        expect(tokens.prompt).to.be.greaterThan(0);
+        expect(tokens.completion).to.be.greaterThan(0);
+        expect(tokens.total).eql(tokens.prompt + tokens.completion);
+      });
+
+      it('returns an error with the expected shape in case of error', async () => {
+        const response = await supertest
+          .post(`/internal/inference/chat_complete`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId: 'do-not-exist',
+            system: 'Please answer the user question',
+            messages: [{ role: 'user', content: '2+2 ?' }],
+          })
+          .expect(400);
+
+        const message = response.body;
+
+        expect(message).to.eql({
+          type: 'error',
+          code: 'requestError',
+          message: "No connector found for id 'do-not-exist'",
+        });
+      });
+    });
+
+    describe('streaming enabled', () => {
+      it('returns a chat completion message for a simple prompt', async () => {
+        const response = supertest
+          .post(`/internal/inference/chat_complete/stream`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId,
+            system: 'Please answer the user question',
+            messages: [{ role: 'user', content: '2+2 ?' }],
+          })
+          .expect(200);
+
+        const observable = supertestToObservable(response);
+
+        const message = await lastValueFrom(observable);
+
+        expect({
+          ...message,
+          content: '',
+        }).to.eql({ type: 'chatCompletionMessage', content: '', toolCalls: [] });
+        expect(message.content).to.contain('4');
+      });
+
+      it('executes a tool when explicitly requested', async () => {
+        const response = supertest
+          .post(`/internal/inference/chat_complete/stream`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId,
+            system:
+              'Please answer the user question. You can use the available tools if you think it can help',
+            messages: [{ role: 'user', content: 'What is the result of 2*4*6*8*10*123 ?' }],
+            toolChoice: 'required',
+            tools: {
+              calculator: {
+                description: 'The calculator can be used to resolve mathematical calculations',
+                schema: {
+                  type: 'object',
+                  properties: {
+                    formula: {
+                      type: 'string',
+                      description: `The input for the calculator, in plain text, e.g. "2+(4*8)"`,
+                    },
+                  },
+                },
+              },
+            },
+          })
+          .expect(200);
+
+        const observable = supertestToObservable(response);
+
+        const message = await lastValueFrom(observable);
+
+        expect(message.toolCalls.length).to.eql(1);
+        expect(message.toolCalls[0].function.name).to.eql('calculator');
+        expect(message.toolCalls[0].function.arguments.formula).to.contain('123');
+      });
+
+      it('returns a token count event', async () => {
+        const response = supertest
+          .post(`/internal/inference/chat_complete/stream`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId,
+            system: 'Please answer the user question',
+            messages: [{ role: 'user', content: '2+2 ?' }],
+          })
+          .expect(200);
+
+        const observable = supertestToObservable(response);
+
+        const events = await lastValueFrom(observable.pipe(toArray()));
+        const tokenEvent = events[events.length - 2];
+
+        expect(tokenEvent.type).to.eql('chatCompletionTokenCount');
+        expect(tokenEvent.tokens.prompt).to.be.greaterThan(0);
+        expect(tokenEvent.tokens.completion).to.be.greaterThan(0);
+        expect(tokenEvent.tokens.total).to.be(
+          tokenEvent.tokens.prompt + tokenEvent.tokens.completion
+        );
+      });
+
+      it('returns an error with the expected shape in case of error', async () => {
+        const response = supertest
+          .post(`/internal/inference/chat_complete/stream`)
+          .set('kbn-xsrf', 'kibana')
+          .send({
+            connectorId: 'do-not-exist',
+            system: 'Please answer the user question',
+            messages: [{ role: 'user', content: '2+2 ?' }],
+          })
+          .expect(200);
+
+        const observable = supertestToObservable(response);
+
+        const events = await lastValueFrom(observable.pipe(toArray()));
+
+        expect(events).to.eql([
+          {
+            type: 'error',
+            error: {
+              code: 'requestError',
+              message: "No connector found for id 'do-not-exist'",
+              meta: {
+                status: 400,
+              },
+            },
+          },
+        ]);
+      });
+    });
+  });
+};

x-pack/test/functional_gen_ai/inference/tests/index.ts

@@ -0,0 +1,21 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getAvailableConnectors } from '@kbn/gen-ai-functional-testing';
+import { FtrProviderContext } from '../ftr_provider_context';
+import { chatCompleteSuite } from './chat_complete';
+
+// eslint-disable-next-line import/no-default-export
+export default function (providerContext: FtrProviderContext) {
+  describe('Inference plugin - API integration tests', async () => {
+    getAvailableConnectors().forEach((connector) => {
+      describe(`Connector ${connector.id}`, () => {
+        chatCompleteSuite(connector, providerContext);
+      });
+    });
+  });
+}

x-pack/test/tsconfig.json

@@ -188,6 +188,8 @@
     "@kbn/ai-assistant-common",
     "@kbn/core-deprecations-common",
     "@kbn/usage-collection-plugin",
+    "@kbn/sse-utils-server",
+    "@kbn/gen-ai-functional-testing",
     "@kbn/integration-assistant-plugin"
   ]
 }

yarn.lock

@@ -5781,6 +5781,10 @@
   version "0.0.0"
   uid ""
 
+"@kbn/gen-ai-functional-testing@link:packages/kbn-gen-ai-functional-testing":
+  version "0.0.0"
+  uid ""
+
 "@kbn/gen-ai-streaming-response-example-plugin@link:x-pack/examples/gen_ai_streaming_response_example":
   version "0.0.0"
   uid ""