[9.0] [GenAI Connectors] Add support for telemetry metadata (elastic#208180) (elastic#210018)

# Backport

This will backport the following commits from `main` to `9.0`:
- [[GenAI Connectors] Add support for telemetry metadata
(elastic#208180)](elastic#208180)

<!--- Backport version: 9.4.3 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Marius
Iversen","email":"[email protected]"},"sourceCommit":{"committedDate":"2025-02-06T14:24:35Z","message":"[GenAI
Connectors] Add support for telemetry metadata (elastic#208180)\n\n##
Summary\n\nThis extends initial connector telemetry from PR
ref\nhttps://github.com/elastic/pull/186936.\n\nThe PR adds the
following optional fields when instantiating a new\nactionClient as part
of its `subActionParams`:\n\n```ts\n{\n telemetryMetadata : {\n
pluginId: \"your plugin name or unique identifier\",\n aggregateBy: \"ID
to aggregate on\" \n }\n}\n```\n\nThe support is added to all AI
connector models for both\nstream/non-stream/raw.\n\nThe PR also adds
token count usage for bedrock `InvokeAIRaw`, as that\nwas currently not
added correctly.\n\nPierre also helped with adding a new metadata
optional field for the `NL\nto ESQL functions`, so that users can pass
in similar metadata for LLM\nconversations using the
InfereceClient.\n\nPluginId is a field used to filter telemetry in the
way the team wants\nto implement it. It could be a team name, a plugin
name etc, all\ndepending on how the team wants to group and filter on
the telemetry\nevent.\n\nAggregateBy is intended to be used to group
multiple LLM calls for\naggregations and stats, for example a
conversationId that has multiple\nLLM calls.\n\nBoth fields are
optional, so when you do not want to aggregate the\noption can simply be
ignored.\n\n### Checklist\n\nCheck the PR satisfies following
conditions. \n\nReviewers should verify this PR satisfies this list as
well.\n\n- [x] Any text added follows [EUI's
writing\nguidelines](https://elastic.github.io/eui/#/guidelines/writing),
uses\nsentence case text and includes
[i18n\nsupport](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)\n-
[x] [Unit or
functional\ntests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)\nwere
updated or added to match the most common scenarios\n- [x] The PR
description includes the appropriate Release Notes section,\nand the
correct `release_note:*` label is applied per
the\n[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)\n\n---------\n\nCo-authored-by:
pgayvallet
<[email protected]>","sha":"3394b691b1582d504195182013f833ba727c5e7e","branchLabelMapping":{"^v9.1.0$":"main","^v8.19.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","v9.0.0","Team:
SecuritySolution","backport:version","v8.18.0","v9.1.0","v8.19.0"],"title":"[GenAI
Connectors] Add support for telemetry
metadata","number":208180,"url":"https://github.com/elastic/kibana/pull/208180","mergeCommit":{"message":"[GenAI
Connectors] Add support for telemetry metadata (elastic#208180)\n\n##
Summary\n\nThis extends initial connector telemetry from PR
ref\nhttps://github.com/elastic/pull/186936.\n\nThe PR adds the
following optional fields when instantiating a new\nactionClient as part
of its `subActionParams`:\n\n```ts\n{\n telemetryMetadata : {\n
pluginId: \"your plugin name or unique identifier\",\n aggregateBy: \"ID
to aggregate on\" \n }\n}\n```\n\nThe support is added to all AI
connector models for both\nstream/non-stream/raw.\n\nThe PR also adds
token count usage for bedrock `InvokeAIRaw`, as that\nwas currently not
added correctly.\n\nPierre also helped with adding a new metadata
optional field for the `NL\nto ESQL functions`, so that users can pass
in similar metadata for LLM\nconversations using the
InfereceClient.\n\nPluginId is a field used to filter telemetry in the
way the team wants\nto implement it. It could be a team name, a plugin
name etc, all\ndepending on how the team wants to group and filter on
the telemetry\nevent.\n\nAggregateBy is intended to be used to group
multiple LLM calls for\naggregations and stats, for example a
conversationId that has multiple\nLLM calls.\n\nBoth fields are
optional, so when you do not want to aggregate the\noption can simply be
ignored.\n\n### Checklist\n\nCheck the PR satisfies following
conditions. \n\nReviewers should verify this PR satisfies this list as
well.\n\n- [x] Any text added follows [EUI's
writing\nguidelines](https://elastic.github.io/eui/#/guidelines/writing),
uses\nsentence case text and includes
[i18n\nsupport](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)\n-
[x] [Unit or
functional\ntests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)\nwere
updated or added to match the most common scenarios\n- [x] The PR
description includes the appropriate Release Notes section,\nand the
correct `release_note:*` label is applied per
the\n[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)\n\n---------\n\nCo-authored-by:
pgayvallet
<[email protected]>","sha":"3394b691b1582d504195182013f833ba727c5e7e"}},"sourceBranch":"main","suggestedTargetBranches":["9.0","8.18","8.x"],"targetPullRequestStates":[{"branch":"9.0","label":"v9.0.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"8.18","label":"v8.18.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"main","label":"v9.1.0","branchLabelMappingKey":"^v9.1.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/208180","number":208180,"mergeCommit":{"message":"[GenAI
Connectors] Add support for telemetry metadata (elastic#208180)\n\n##
Summary\n\nThis extends initial connector telemetry from PR
ref\nhttps://github.com/elastic/pull/186936.\n\nThe PR adds the
following optional fields when instantiating a new\nactionClient as part
of its `subActionParams`:\n\n```ts\n{\n telemetryMetadata : {\n
pluginId: \"your plugin name or unique identifier\",\n aggregateBy: \"ID
to aggregate on\" \n }\n}\n```\n\nThe support is added to all AI
connector models for both\nstream/non-stream/raw.\n\nThe PR also adds
token count usage for bedrock `InvokeAIRaw`, as that\nwas currently not
added correctly.\n\nPierre also helped with adding a new metadata
optional field for the `NL\nto ESQL functions`, so that users can pass
in similar metadata for LLM\nconversations using the
InfereceClient.\n\nPluginId is a field used to filter telemetry in the
way the team wants\nto implement it. It could be a team name, a plugin
name etc, all\ndepending on how the team wants to group and filter on
the telemetry\nevent.\n\nAggregateBy is intended to be used to group
multiple LLM calls for\naggregations and stats, for example a
conversationId that has multiple\nLLM calls.\n\nBoth fields are
optional, so when you do not want to aggregate the\noption can simply be
ignored.\n\n### Checklist\n\nCheck the PR satisfies following
conditions. \n\nReviewers should verify this PR satisfies this list as
well.\n\n- [x] Any text added follows [EUI's
writing\nguidelines](https://elastic.github.io/eui/#/guidelines/writing),
uses\nsentence case text and includes
[i18n\nsupport](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)\n-
[x] [Unit or
functional\ntests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)\nwere
updated or added to match the most common scenarios\n- [x] The PR
description includes the appropriate Release Notes section,\nand the
correct `release_note:*` label is applied per
the\n[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)\n\n---------\n\nCo-authored-by:
pgayvallet
<[email protected]>","sha":"3394b691b1582d504195182013f833ba727c5e7e"}},{"branch":"8.x","label":"v8.19.0","branchLabelMappingKey":"^v8.19.0$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

Co-authored-by: Marius Iversen <[email protected]>

Author: kibanamachine

x-pack/platform/packages/shared/ai-infra/inference-common/index.ts

@@ -56,6 +56,8 @@ export {
   isToolValidationError,
   isTokenLimitReachedError,
   isToolNotFoundError,
+  type ChatCompleteMetadata,
+  type ConnectorTelemetryMetadata,
 } from './src/chat_complete';
 export {
   OutputEventType,

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts

@@ -9,6 +9,7 @@ import type { Observable } from 'rxjs';
 import type { ToolCallsOf, ToolOptions } from './tools';
 import type { Message } from './messages';
 import type { ChatCompletionEvent, ChatCompletionTokenCount } from './events';
+import type { ChatCompleteMetadata } from './metadata';
 
 /**
  * Request a completion from the LLM based on a prompt or conversation.
@@ -109,6 +110,10 @@ export type ChatCompleteOptions<
    * Optional signal that can be used to forcefully abort the request.
    */
   abortSignal?: AbortSignal;
+  /**
+   * Optional metadata related to call execution.
+   */
+  metadata?: ChatCompleteMetadata;
 } & TToolOptions;
 
 /**

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts

@@ -50,6 +50,7 @@ export {
   type UnvalidatedToolCall,
   type ToolChoice,
 } from './tools';
+export type { ChatCompleteMetadata, ConnectorTelemetryMetadata } from './metadata';
 export {
   isChatCompletionChunkEvent,
   isChatCompletionEvent,

x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/metadata.ts

@@ -0,0 +1,23 @@
+/*
+ * 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.
+ */
+
+/**
+ * Set of metadata that can be used then calling the inference APIs
+ *
+ * @public
+ */
+export interface ChatCompleteMetadata {
+  connectorTelemetry?: ConnectorTelemetryMetadata;
+}
+
+/**
+ * Pass through for the connector telemetry
+ */
+export interface ConnectorTelemetryMetadata {
+  pluginId?: string;
+  aggregateBy?: string;
+}

x-pack/platform/packages/shared/ai-infra/inference-common/src/output/api.ts

@@ -6,7 +6,13 @@
  */
 
 import type { Observable } from 'rxjs';
-import { Message, FunctionCallingMode, FromToolSchema, ToolSchema } from '../chat_complete';
+import {
+  Message,
+  FunctionCallingMode,
+  FromToolSchema,
+  ToolSchema,
+  ChatCompleteMetadata,
+} from '../chat_complete';
 import { Output, OutputEvent } from './events';
 
 /**
@@ -117,6 +123,10 @@ export interface OutputOptions<
      */
     onValidationError?: boolean | number;
   };
+  /**
+   * Optional metadata related to call execution.
+   */
+  metadata?: ChatCompleteMetadata;
 }
 
 /**

x-pack/platform/packages/shared/kbn-langchain/server/language_models/bedrock_chat.ts

@@ -10,6 +10,7 @@ import type { ActionsClient } from '@kbn/actions-plugin/server';
 import { BaseChatModelParams } from '@langchain/core/language_models/chat_models';
 import { Logger } from '@kbn/logging';
 import { PublicMethodsOf } from '@kbn/utility-types';
+import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib';
 import { prepareMessages, DEFAULT_BEDROCK_MODEL, DEFAULT_BEDROCK_REGION } from '../utils/bedrock';
 
 export interface CustomChatModelInput extends BaseChatModelParams {
@@ -20,6 +21,7 @@ export interface CustomChatModelInput extends BaseChatModelParams {
   signal?: AbortSignal;
   model?: string;
   maxTokens?: number;
+  telemetryMetadata?: TelemetryMetadata;
 }
 
 /**
@@ -49,6 +51,10 @@ export class ActionsClientBedrockChatModel extends _BedrockChat {
           params: {
             subAction: 'invokeAIRaw',
             subActionParams: {
+              telemetryMetadata: {
+                pluginId: params?.telemetryMetadata?.pluginId,
+                aggregateBy: params?.telemetryMetadata?.aggregateBy,
+              },
               messages: prepareMessages(inputBody.messages),
               temperature: params.temperature ?? inputBody.temperature,
               stopSequences: inputBody.stop_sequences,

x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/bedrock_runtime_client.ts

@@ -13,30 +13,33 @@ import {
   ConverseStreamCommand,
   ConverseStreamResponse,
 } from '@aws-sdk/client-bedrock-runtime';
+import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib';
 import { constructStack } from '@smithy/middleware-stack';
 import { HttpHandlerOptions } from '@smithy/types';
 import { PublicMethodsOf } from '@kbn/utility-types';
 import type { ActionsClient } from '@kbn/actions-plugin/server';
-
 import { prepareMessages } from '../../utils/bedrock';
 
 export interface CustomChatModelInput extends BedrockRuntimeClientConfig {
   actionsClient: PublicMethodsOf<ActionsClient>;
   connectorId: string;
   streaming?: boolean;
+  telemetryMetadata?: TelemetryMetadata;
 }
 
 export class BedrockRuntimeClient extends _BedrockRuntimeClient {
   middlewareStack: _BedrockRuntimeClient['middlewareStack'];
   streaming: boolean;
   actionsClient: PublicMethodsOf<ActionsClient>;
   connectorId: string;
+  telemetryMetadata?: TelemetryMetadata;
 
   constructor({ actionsClient, connectorId, ...fields }: CustomChatModelInput) {
     super(fields ?? {});
     this.streaming = fields.streaming ?? true;
     this.actionsClient = actionsClient;
     this.connectorId = connectorId;
+    this.telemetryMetadata = fields?.telemetryMetadata;
     // eliminate middleware steps that handle auth as Kibana connector handles auth
     this.middlewareStack = constructStack() as _BedrockRuntimeClient['middlewareStack'];
   }
@@ -56,6 +59,7 @@ export class BedrockRuntimeClient extends _BedrockRuntimeClient {
       params: {
         subAction: 'bedrockClientSend',
         subActionParams: {
+          telemetryMetadata: this.telemetryMetadata,
           command,
           signal: options?.abortSignal,
         },

x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/chat_bedrock_converse.ts

@@ -9,6 +9,7 @@ import type { ActionsClient } from '@kbn/actions-plugin/server';
 import { BaseChatModelParams } from '@langchain/core/language_models/chat_models';
 import { Logger } from '@kbn/logging';
 import { PublicMethodsOf } from '@kbn/utility-types';
+import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib';
 import { BedrockRuntimeClient } from './bedrock_runtime_client';
 import { DEFAULT_BEDROCK_MODEL, DEFAULT_BEDROCK_REGION } from '../../utils/bedrock';
 
@@ -18,6 +19,7 @@ export interface CustomChatModelInput extends BaseChatModelParams {
   logger: Logger;
   signal?: AbortSignal;
   model?: string;
+  telemetryMetadata?: TelemetryMetadata;
 }
 
 /**
@@ -45,6 +47,7 @@ export class ActionsClientChatBedrockConverse extends ChatBedrockConverse {
       connectorId,
       streaming: this.streaming,
       region: DEFAULT_BEDROCK_REGION,
+      telemetryMetadata: fields?.telemetryMetadata,
     });
   }
 }

x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_openai.ts

@@ -9,11 +9,12 @@ import { v4 as uuidv4 } from 'uuid';
 import { Logger } from '@kbn/core/server';
 import type { ActionsClient } from '@kbn/actions-plugin/server';
 import { get } from 'lodash/fp';
-
+import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib';
 import { ChatOpenAI } from '@langchain/openai';
 import { Stream } from 'openai/streaming';
 import type OpenAI from 'openai';
 import { PublicMethodsOf } from '@kbn/utility-types';
+
 import { DEFAULT_OPEN_AI_MODEL, DEFAULT_TIMEOUT } from './constants';
 import {
   InferenceChatCompleteParamsSchema,
@@ -36,6 +37,7 @@ export interface ActionsClientChatOpenAIParams {
   temperature?: number;
   signal?: AbortSignal;
   timeout?: number;
+  telemetryMetadata?: TelemetryMetadata;
 }
 
 /**
@@ -65,6 +67,7 @@ export class ActionsClientChatOpenAI extends ChatOpenAI {
   #traceId: string;
   #signal?: AbortSignal;
   #timeout?: number;
+  telemetryMetadata?: TelemetryMetadata;
 
   constructor({
     actionsClient,
@@ -79,6 +82,7 @@ export class ActionsClientChatOpenAI extends ChatOpenAI {
     temperature,
     timeout,
     maxTokens,
+    telemetryMetadata,
   }: ActionsClientChatOpenAIParams) {
     super({
       maxRetries,
@@ -109,6 +113,7 @@ export class ActionsClientChatOpenAI extends ChatOpenAI {
     // matters only for LangSmith logs (Metadata > Invocation Params)
     // the connector can be passed an undefined temperature through #temperature
     this.temperature = temperature ?? this.temperature;
+    this.telemetryMetadata = telemetryMetadata;
   }
 
   getActionResultData(): string {
@@ -237,6 +242,7 @@ export class ActionsClientChatOpenAI extends ChatOpenAI {
         : completionRequest.stream
         ? { ...body, timeout: this.#timeout ?? DEFAULT_TIMEOUT }
         : { body: JSON.stringify(body), timeout: this.#timeout ?? DEFAULT_TIMEOUT }),
+      telemetryMetadata: this.telemetryMetadata,
       signal: this.#signal,
     };
     return {

x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts

@@ -18,6 +18,7 @@ import { Logger } from '@kbn/logging';
 import { BaseChatModelParams } from '@langchain/core/language_models/chat_models';
 import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
 import { GeminiPartText } from '@langchain/google-common/dist/types';
+import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib';
 import {
   convertResponseBadFinishReasonToErrorMsg,
   convertResponseContentToChatGenerationChunk,
@@ -34,12 +35,14 @@ export interface CustomChatModelInput extends BaseChatModelParams {
   signal?: AbortSignal;
   model?: string;
   maxTokens?: number;
+  telemetryMetadata?: TelemetryMetadata;
 }
 
 export class ActionsClientChatVertexAI extends ChatVertexAI {
   #actionsClient: PublicMethodsOf<ActionsClient>;
   #connectorId: string;
   #model?: string;
+  telemetryMetadata?: TelemetryMetadata;
   constructor({ actionsClient, connectorId, ...props }: CustomChatModelInput) {
     super({
       ...props,
@@ -62,7 +65,8 @@ export class ActionsClientChatVertexAI extends ChatVertexAI {
       client,
       false,
       actionsClient,
-      connectorId
+      connectorId,
+      props?.telemetryMetadata
     );
   }
 
@@ -89,6 +93,7 @@ export class ActionsClientChatVertexAI extends ChatVertexAI {
           subAction: 'invokeStream',
           subActionParams: {
             model: this.#model,
+            telemetryMetadata: this.telemetryMetadata,
             messages: data?.contents,
             tools: data?.tools,
             temperature: this.temperature,