From c2b7b9ef41a9c858da4ba2b674e3ca1f423f0992 Mon Sep 17 00:00:00 2001 From: Marius Iversen Date: Thu, 6 Feb 2025 15:24:35 +0100 Subject: [PATCH] [GenAI Connectors] Add support for telemetry metadata (#208180) ## Summary This extends initial connector telemetry from PR ref https://github.com/elastic/kibana/pull/186936. The PR adds the following optional fields when instantiating a new actionClient as part of its `subActionParams`: ```ts { telemetryMetadata : { pluginId: "your plugin name or unique identifier", aggregateBy: "ID to aggregate on" } } ``` The support is added to all AI connector models for both stream/non-stream/raw. The PR also adds token count usage for bedrock `InvokeAIRaw`, as that was currently not added correctly. Pierre also helped with adding a new metadata optional field for the `NL to ESQL functions`, so that users can pass in similar metadata for LLM conversations using the InfereceClient. PluginId is a field used to filter telemetry in the way the team wants to implement it. It could be a team name, a plugin name etc, all depending on how the team wants to group and filter on the telemetry event. AggregateBy is intended to be used to group multiple LLM calls for aggregations and stats, for example a conversationId that has multiple LLM calls. Both fields are optional, so when you do not want to aggregate the option can simply be ignored. ### Checklist Check the PR satisfies following conditions. Reviewers should verify this PR satisfies this list as well. - [x] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md) - [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [x] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) --------- Co-authored-by: pgayvallet (cherry picked from commit 3394b691b1582d504195182013f833ba727c5e7e) --- .../shared/ai-infra/inference-common/index.ts | 2 + .../inference-common/src/chat_complete/api.ts | 5 + .../src/chat_complete/index.ts | 1 + .../src/chat_complete/metadata.ts | 23 + .../inference-common/src/output/api.ts | 12 +- .../server/language_models/bedrock_chat.ts | 6 + .../bedrock_runtime_client.ts | 6 +- .../chat_bedrock_converse.ts | 3 + .../server/language_models/chat_openai.ts | 8 +- .../chat_vertex/chat_vertex.ts | 7 +- .../language_models/chat_vertex/connection.ts | 7 +- .../server/language_models/gemini_chat.ts | 6 + .../language_models/simple_chat_model.test.ts | 10 + .../language_models/simple_chat_model.ts | 13 + .../server/language_models/types.ts | 2 + .../connector_types.test.ts.snap | 513 ++++++++++++++++++ .../actions/server/lib/action_executor.ts | 5 +- .../server/lib/event_based_telemetry.ts | 17 + .../server/lib/gen_ai_token_tracking.ts | 44 ++ .../shared/actions/server/lib/index.ts | 1 + .../common/output/create_output_api.ts | 2 + .../bedrock/bedrock_claude_adapter.ts | 2 + .../adapters/gemini/gemini_adapter.ts | 4 + .../adapters/inference/inference_adapter.ts | 4 + .../adapters/openai/openai_adapter.ts | 4 + .../inference/server/chat_complete/api.ts | 2 + .../inference/server/chat_complete/types.ts | 4 +- .../tasks/nl_to_esql/actions/generate_esql.ts | 4 + .../actions/request_documentation.ts | 4 + .../inference/server/tasks/nl_to_esql/task.ts | 3 + .../server/tasks/nl_to_esql/types.ts | 2 + .../stack_connectors/common/bedrock/schema.ts | 15 + .../stack_connectors/common/gemini/schema.ts | 8 + .../common/inference/schema.ts | 6 + .../stack_connectors/common/openai/schema.ts | 7 + .../server/connector_types/bedrock/bedrock.ts | 2 +- .../task/util/esql_knowledge_base_caller.ts | 6 +- .../tests/actions/connector_types/bedrock.ts | 8 +- 38 files changed, 764 insertions(+), 14 deletions(-) create mode 100644 x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/metadata.ts diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts index f1835ab1ab52c..622a09bfbc71e 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts +++ b/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, diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts index 155de9b286c9b..0315d30033a23 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/api.ts +++ b/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; /** diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts index 227e72d93ca92..810f1d7a870eb 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/index.ts +++ b/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, diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/metadata.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/chat_complete/metadata.ts new file mode 100644 index 0000000000000..527cb0d346a8b --- /dev/null +++ b/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; +} diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/output/api.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/output/api.ts index a2d68883638c3..79c30cb22a689 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/output/api.ts +++ b/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; } /** diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/bedrock_chat.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/bedrock_chat.ts index 70395298d3c98..7b5781205cf7f 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/bedrock_chat.ts +++ b/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, diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/bedrock_runtime_client.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/bedrock_runtime_client.ts index 7f20591bd51a4..465b70a2cbc8f 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/bedrock_runtime_client.ts +++ b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/bedrock_runtime_client.ts @@ -13,17 +13,18 @@ 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; connectorId: string; streaming?: boolean; + telemetryMetadata?: TelemetryMetadata; } export class BedrockRuntimeClient extends _BedrockRuntimeClient { @@ -31,12 +32,14 @@ export class BedrockRuntimeClient extends _BedrockRuntimeClient { streaming: boolean; actionsClient: PublicMethodsOf; 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, }, diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/chat_bedrock_converse.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/chat_bedrock_converse.ts index bdc84130925d6..c2bac841a376a 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_bedrock_converse/chat_bedrock_converse.ts +++ b/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, }); } } diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_openai.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_openai.ts index 2933695a94bbf..0bd84a26b8b41 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_openai.ts +++ b/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 { diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts index 5c7a9ef918da3..af5e3eda5c7fb 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/chat_vertex.ts +++ b/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; #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, diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/connection.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/connection.ts index 442e6b079db9b..3f28500e81898 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/connection.ts +++ b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/chat_vertex/connection.ts @@ -16,6 +16,7 @@ import { ActionsClient } from '@kbn/actions-plugin/server'; import { PublicMethodsOf } from '@kbn/utility-types'; import { EnhancedGenerateContentResponse } from '@google/generative-ai'; import { AsyncCaller } from '@langchain/core/utils/async_caller'; +import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib'; import { convertResponseBadFinishReasonToErrorMsg } from '../../utils/gemini'; // only implements non-streaming requests @@ -26,17 +27,20 @@ export class ActionsClientChatConnection extends ChatConnection { #model?: string; temperature: number; caller: AsyncCaller; + telemetryMetadata?: TelemetryMetadata; constructor( fields: GoogleAIBaseLLMInput, caller: AsyncCaller, client: GoogleAbstractedClient, _streaming: boolean, // defaulting to false in the super actionsClient: PublicMethodsOf, - connectorId: string + connectorId: string, + telemetryMetadata?: TelemetryMetadata ) { super(fields, caller, client, false); this.actionsClient = actionsClient; this.connectorId = connectorId; + this.telemetryMetadata = telemetryMetadata; this.caller = caller; this.#model = fields.model; this.temperature = fields.temperature ?? 0; @@ -77,6 +81,7 @@ export class ActionsClientChatConnection extends ChatConnection { params: { subAction: 'invokeAIRaw', subActionParams: { + telemetryMetadata: this.telemetryMetadata, model: this.#model, messages: data?.contents, tools: data?.tools, diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/gemini_chat.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/gemini_chat.ts index f8755af19d78c..a226e16736f6d 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/gemini_chat.ts +++ b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/gemini_chat.ts @@ -21,6 +21,7 @@ import { Logger } from '@kbn/logging'; import { BaseChatModelParams } from '@langchain/core/language_models/chat_models'; import { get } from 'lodash/fp'; import { Readable } from 'stream'; +import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib'; import { convertBaseMessagesToContent, convertResponseBadFinishReasonToErrorMsg, @@ -36,6 +37,7 @@ export interface CustomChatModelInput extends BaseChatModelParams { signal?: AbortSignal; model?: string; maxTokens?: number; + telemetryMetadata?: TelemetryMetadata; } export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI { @@ -43,6 +45,7 @@ export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI { #connectorId: string; #temperature: number; #model?: string; + telemetryMetadata?: TelemetryMetadata; constructor({ actionsClient, connectorId, ...props }: CustomChatModelInput) { super({ @@ -50,6 +53,7 @@ export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI { apiKey: 'asda', maxOutputTokens: props.maxTokens ?? 2048, }); + this.telemetryMetadata = props.telemetryMetadata; // LangChain needs model to be defined for logging purposes this.model = props.model ?? this.model; // If model is not specified by consumer, the connector will defin eit so do not pass @@ -71,6 +75,7 @@ export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI { params: { subAction: 'invokeAIRaw', subActionParams: { + telemetryMetadata: this.telemetryMetadata, model: this.#model, messages: request.contents, tools: request.tools, @@ -159,6 +164,7 @@ export class ActionsClientGeminiChatModel extends ChatGoogleGenerativeAI { }, []), temperature: this.#temperature, tools: request.tools, + telemetryMetadata: this.telemetryMetadata, }, }, }; diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.test.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.test.ts index 634d8260cc3b8..89e32d24857c8 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.test.ts +++ b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.test.ts @@ -227,6 +227,11 @@ describe('ActionsClientSimpleChatModel', () => { temperature: 0, stopSequences: ['\n'], maxTokens: 333, + model: undefined, + telemetryMetadata: { + aggregateBy: undefined, + pluginId: undefined, + }, }); expect(result).toEqual(mockActionResponse.message); @@ -252,6 +257,11 @@ describe('ActionsClientSimpleChatModel', () => { expect(rest).toEqual({ temperature: 0, + model: undefined, + telemetryMetadata: { + aggregateBy: undefined, + pluginId: undefined, + }, }); expect(result).toEqual(mockActionResponse.message); diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.ts index 787aed559e285..1215ae435d49e 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.ts +++ b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/simple_chat_model.ts @@ -18,6 +18,7 @@ import { get } from 'lodash/fp'; import { ChatGenerationChunk } from '@langchain/core/outputs'; import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager'; import { PublicMethodsOf } from '@kbn/utility-types'; +import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib'; import { parseGeminiStreamAsAsyncIterator, parseGeminiStream } from '../utils/gemini'; import { parseBedrockStreamAsAsyncIterator, parseBedrockStream } from '../utils/bedrock'; import { getDefaultArguments } from './constants'; @@ -37,6 +38,7 @@ export interface CustomChatModelInput extends BaseChatModelParams { temperature?: number; streaming: boolean; maxTokens?: number; + telemetryMetadata?: TelemetryMetadata; } function _formatMessages(messages: BaseMessage[]) { @@ -62,6 +64,7 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel { streaming: boolean; model?: string; temperature?: number; + telemetryMetadata?: TelemetryMetadata; constructor({ actionsClient, @@ -73,6 +76,7 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel { signal, streaming, maxTokens, + telemetryMetadata, }: CustomChatModelInput) { super({}); @@ -86,6 +90,7 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel { this.model = model; this.temperature = temperature; this.streaming = streaming; + this.telemetryMetadata = telemetryMetadata; } _llmType() { @@ -119,6 +124,10 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel { subActionParams: { model: this.model, messages: formattedMessages, + telemetryMetadata: { + pluginId: this.telemetryMetadata?.pluginId, + aggregateBy: this.telemetryMetadata?.aggregateBy, + }, ...getDefaultArguments(this.llmType, this.temperature, options.stop, this.#maxTokens), }, }, @@ -214,6 +223,10 @@ export class ActionsClientSimpleChatModel extends SimpleChatModel { subActionParams: { model: this.model, messages: formattedMessages, + telemetryMetadata: { + pluginId: this.telemetryMetadata?.pluginId, + aggregateBy: this.telemetryMetadata?.aggregateBy, + }, ...getDefaultArguments(this.llmType, this.temperature, options.stop, this.#maxTokens), }, }, diff --git a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/types.ts b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/types.ts index 69d18d4f1b2a0..30534225a1407 100644 --- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/types.ts +++ b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/types.ts @@ -5,6 +5,7 @@ * 2.0. */ +import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib'; import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; import type OpenAI from 'openai'; @@ -39,6 +40,7 @@ export interface InvokeAIActionParamsSchema { functions?: OpenAI.ChatCompletionCreateParamsNonStreaming['functions']; signal?: AbortSignal; timeout?: number; + telemetryMetadata?: TelemetryMetadata; } export interface RunActionParamsSchema { body: string; diff --git a/x-pack/platform/plugins/shared/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap b/x-pack/platform/plugins/shared/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap index b6abe08b9bd9b..23a7e09e2cb01 100644 --- a/x-pack/platform/plugins/shared/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap +++ b/x-pack/platform/plugins/shared/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap @@ -435,6 +435,63 @@ Object { ], "type": "string", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], @@ -840,6 +897,63 @@ Object { ], "type": "string", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], @@ -1200,6 +1314,63 @@ Object { ], "type": "string", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], @@ -1420,6 +1591,63 @@ Object { ], "type": "any", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, }, "type": "object", } @@ -4570,6 +4798,63 @@ Object { ], "type": "array", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], @@ -4731,6 +5016,63 @@ Object { ], "type": "array", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], @@ -4871,6 +5213,63 @@ Object { ], "type": "string", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], @@ -5127,6 +5526,63 @@ Object { ], "type": "string", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], @@ -5294,6 +5750,63 @@ Object { ], "type": "string", }, + "telemetryMetadata": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "keys": Object { + "aggregateBy": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + "pluginId": Object { + "flags": Object { + "default": [Function], + "error": [Function], + "presence": "optional", + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "rules": Array [ + Object { + "args": Object { + "method": [Function], + }, + "name": "custom", + }, + ], + "type": "string", + }, + }, + "metas": Array [ + Object { + "x-oas-optional": true, + }, + ], + "type": "object", + }, "temperature": Object { "flags": Object { "default": [Function], diff --git a/x-pack/platform/plugins/shared/actions/server/lib/action_executor.ts b/x-pack/platform/plugins/shared/actions/server/lib/action_executor.ts index 232c82fc36c4c..b84e5843ee195 100644 --- a/x-pack/platform/plugins/shared/actions/server/lib/action_executor.ts +++ b/x-pack/platform/plugins/shared/actions/server/lib/action_executor.ts @@ -422,7 +422,7 @@ export class ActionExecutor { const actionType = actionTypeRegistry.get(actionTypeId); const configurationUtilities = actionTypeRegistry.getUtils(); - let validatedParams; + let validatedParams: Record; let validatedConfig; let validatedSecrets; try { @@ -610,11 +610,14 @@ export class ActionExecutor { prompt_tokens: tokenTracking.prompt_tokens ?? 0, completion_tokens: tokenTracking.completion_tokens ?? 0, }); + analyticsService.reportEvent(GEN_AI_TOKEN_COUNT_EVENT.eventType, { actionTypeId, total_tokens: tokenTracking.total_tokens ?? 0, prompt_tokens: tokenTracking.prompt_tokens ?? 0, completion_tokens: tokenTracking.completion_tokens ?? 0, + aggregateBy: tokenTracking?.telemetry_metadata?.aggregateBy, + pluginId: tokenTracking?.telemetry_metadata?.pluginId, ...(actionTypeId === '.gen-ai' && config?.apiProvider != null ? { provider: config?.apiProvider } : {}), diff --git a/x-pack/platform/plugins/shared/actions/server/lib/event_based_telemetry.ts b/x-pack/platform/plugins/shared/actions/server/lib/event_based_telemetry.ts index 2a8006758489f..7aa1762ca9fce 100644 --- a/x-pack/platform/plugins/shared/actions/server/lib/event_based_telemetry.ts +++ b/x-pack/platform/plugins/shared/actions/server/lib/event_based_telemetry.ts @@ -14,6 +14,8 @@ export const GEN_AI_TOKEN_COUNT_EVENT: EventTypeOpts<{ completion_tokens: number; provider?: string; model?: string; + pluginId?: string; + aggregateBy?: string; }> = { eventType: 'gen_ai_token_count', schema: { @@ -59,6 +61,21 @@ export const GEN_AI_TOKEN_COUNT_EVENT: EventTypeOpts<{ optional: true, }, }, + pluginId: { + type: 'keyword', + _meta: { + description: 'Optional Kibana plugin ID that can be used to filter/aggregate telemetry', + optional: true, + }, + }, + aggregateBy: { + type: 'keyword', + _meta: { + description: + 'Optional field used to group telemetry data by a specific field that is important to the consumer, like a task or conversation ID', + optional: true, + }, + }, }, }; diff --git a/x-pack/platform/plugins/shared/actions/server/lib/gen_ai_token_tracking.ts b/x-pack/platform/plugins/shared/actions/server/lib/gen_ai_token_tracking.ts index ff73095ac2427..5eb452bd1388b 100644 --- a/x-pack/platform/plugins/shared/actions/server/lib/gen_ai_token_tracking.ts +++ b/x-pack/platform/plugins/shared/actions/server/lib/gen_ai_token_tracking.ts @@ -26,6 +26,11 @@ import { parseGeminiStreamForUsageMetadata, } from './get_token_count_from_invoke_stream'; +export interface TelemetryMetadata { + pluginId?: string; + aggregateBy?: string; +} + interface OwnProps { actionTypeId: string; logger: Logger; @@ -50,8 +55,13 @@ export const getGenAiTokenTracking = async ({ total_tokens: number; prompt_tokens: number; completion_tokens: number; + telemetry_metadata?: TelemetryMetadata; } | null> => { // this is an async iterator from the OpenAI sdk + let telemetryMetadata: TelemetryMetadata | undefined; + if (hasTelemetryMetadata(validatedParams.subActionParams)) { + telemetryMetadata = validatedParams.subActionParams.telemetryMetadata; + } if (validatedParams.subAction === 'invokeAsyncIterator' && actionTypeId === '.gen-ai') { try { const data = result.data as { @@ -69,6 +79,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: total, prompt_tokens: prompt, completion_tokens: completion, + telemetry_metadata: telemetryMetadata, }; } logger.error( @@ -78,6 +89,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: 0, prompt_tokens: 0, completion_tokens: 0, + telemetry_metadata: telemetryMetadata, }; } catch (e) { logger.error( @@ -105,6 +117,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: totalTokenCount, prompt_tokens: promptTokenCount, completion_tokens: candidatesTokenCount, + telemetry_metadata: telemetryMetadata, }; } catch (e) { logger.error('Failed to calculate tokens from Invoke Stream subaction streaming response'); @@ -130,6 +143,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: total, prompt_tokens: prompt, completion_tokens: completion, + telemetry_metadata: telemetryMetadata, }; } catch (e) { logger.error('Failed to calculate tokens from Invoke Stream subaction streaming response'); @@ -150,6 +164,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: total, prompt_tokens: prompt, completion_tokens: completion, + telemetry_metadata: telemetryMetadata, }; } catch (e) { logger.error('Failed to calculate tokens from streaming response'); @@ -171,6 +186,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: data.usage?.total_tokens ?? 0, prompt_tokens: data.usage?.prompt_tokens ?? 0, completion_tokens: data.usage?.completion_tokens ?? 0, + telemetry_metadata: telemetryMetadata, }; } @@ -195,6 +211,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: total, prompt_tokens: prompt, completion_tokens: completion, + telemetry_metadata: telemetryMetadata, }; } else { logger.error('Response from Bedrock run response did not contain completion string'); @@ -202,6 +219,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: 0, prompt_tokens: 0, completion_tokens: 0, + telemetry_metadata: telemetryMetadata, }; } } catch (e) { @@ -227,6 +245,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: data.usageMetadata?.totalTokenCount ?? 0, prompt_tokens: data.usageMetadata?.promptTokenCount ?? 0, completion_tokens: data.usageMetadata?.candidatesTokenCount ?? 0, + telemetry_metadata: telemetryMetadata, }; } @@ -253,6 +272,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: total, prompt_tokens: prompt, completion_tokens: completion, + telemetry_metadata: telemetryMetadata, }; } else { logger.error('Response from Bedrock invoke response did not contain message string'); @@ -260,6 +280,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: 0, prompt_tokens: 0, completion_tokens: 0, + telemetry_metadata: telemetryMetadata, }; } } catch (e) { @@ -284,6 +305,7 @@ export const getGenAiTokenTracking = async ({ total_tokens: usage.totalTokens, prompt_tokens: usage.inputTokens, completion_tokens: usage.outputTokens, + telemetry_metadata: telemetryMetadata, }; } else { logger.error('Response from Bedrock converse API did not contain usage object'); @@ -291,6 +313,24 @@ export const getGenAiTokenTracking = async ({ } } + if (actionTypeId === '.bedrock' && validatedParams.subAction === 'invokeAIRaw') { + const results = result.data as unknown as { + content: Array<{ type: string; text: string }>; + usage?: { input_tokens: number; output_tokens: number }; + }; + if (results?.usage) { + const { input_tokens: inputTokens = 0, output_tokens: outputTokens = 0 } = results.usage; + return { + total_tokens: inputTokens + outputTokens, + prompt_tokens: inputTokens, + completion_tokens: outputTokens, + telemetry_metadata: telemetryMetadata, + }; + } else { + logger.error('Response from Bedrock converse API did not contain usage object'); + return null; + } + } return null; }; @@ -299,3 +339,7 @@ export const shouldTrackGenAiToken = (actionTypeId: string) => actionTypeId === '.bedrock' || actionTypeId === '.gemini' || actionTypeId === '.inference'; + +function hasTelemetryMetadata(obj: unknown): obj is { telemetryMetadata: TelemetryMetadata } { + return obj !== null && typeof obj === 'object' && 'telemetryMetadata' in obj; +} diff --git a/x-pack/platform/plugins/shared/actions/server/lib/index.ts b/x-pack/platform/plugins/shared/actions/server/lib/index.ts index e13fb85008a84..9afa038d79711 100644 --- a/x-pack/platform/plugins/shared/actions/server/lib/index.ts +++ b/x-pack/platform/plugins/shared/actions/server/lib/index.ts @@ -39,3 +39,4 @@ export { parseDate } from './parse_date'; export type { RelatedSavedObjects } from './related_saved_objects'; export { getBasicAuthHeader, combineHeadersWithBasicAuthHeader } from './get_basic_auth_header'; export { tryCatch } from './try_catch'; +export type { TelemetryMetadata } from './gen_ai_token_tracking'; diff --git a/x-pack/platform/plugins/shared/inference/common/output/create_output_api.ts b/x-pack/platform/plugins/shared/inference/common/output/create_output_api.ts index 94899add465ff..cedaaa38dc4ce 100644 --- a/x-pack/platform/plugins/shared/inference/common/output/create_output_api.ts +++ b/x-pack/platform/plugins/shared/inference/common/output/create_output_api.ts @@ -36,6 +36,7 @@ export function createOutputApi(chatCompleteApi: ChatCompleteAPI) { functionCalling, stream, abortSignal, + metadata, retry, }: DefaultOutputOptions): OutputCompositeResponse { if (stream && retry !== undefined) { @@ -56,6 +57,7 @@ export function createOutputApi(chatCompleteApi: ChatCompleteAPI) { modelName, functionCalling, abortSignal, + metadata, system, messages, ...(schema diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts index e4b91d493db1b..cef0f870824b9 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts @@ -35,6 +35,7 @@ export const bedrockClaudeAdapter: InferenceConnectorAdapter = { temperature = 0, modelName, abortSignal, + metadata, }) => { const noToolUsage = toolChoice === ToolChoiceType.none; @@ -47,6 +48,7 @@ export const bedrockClaudeAdapter: InferenceConnectorAdapter = { model: modelName, stopSequences: ['\n\nHuman:'], signal: abortSignal, + ...(metadata?.connectorTelemetry ? { telemetryMetadata: metadata.connectorTelemetry } : {}), }; return from( diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts index a96b35f2bec84..4a877e2ada860 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts @@ -33,6 +33,7 @@ export const geminiAdapter: InferenceConnectorAdapter = { temperature = 0, modelName, abortSignal, + metadata, }) => { return from( executor.invoke({ @@ -46,6 +47,9 @@ export const geminiAdapter: InferenceConnectorAdapter = { model: modelName, signal: abortSignal, stopSequences: ['\n\nHuman:'], + ...(metadata?.connectorTelemetry + ? { telemetryMetadata: metadata.connectorTelemetry } + : {}), }, }) ).pipe( diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/inference_adapter.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/inference_adapter.ts index 5e2a92ad90ace..62cc506624c07 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/inference_adapter.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/inference_adapter.ts @@ -27,6 +27,7 @@ export const inferenceAdapter: InferenceConnectorAdapter = { modelName, logger, abortSignal, + metadata, }) => { const useSimulatedFunctionCalling = functionCalling === 'auto' @@ -50,6 +51,9 @@ export const inferenceAdapter: InferenceConnectorAdapter = { subActionParams: { body: request, signal: abortSignal, + ...(metadata?.connectorTelemetry + ? { telemetryMetadata: metadata.connectorTelemetry } + : {}), }, }) ).pipe( diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/openai_adapter.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/openai_adapter.ts index 4cb4cf7d72cdc..083cf88586a6f 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/openai_adapter.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/openai_adapter.ts @@ -32,6 +32,7 @@ export const openAIAdapter: InferenceConnectorAdapter = { modelName, logger, abortSignal, + metadata, }) => { const useSimulatedFunctionCalling = functionCalling === 'auto' @@ -70,6 +71,9 @@ export const openAIAdapter: InferenceConnectorAdapter = { body: JSON.stringify(request), signal: abortSignal, stream: true, + ...(metadata?.connectorTelemetry + ? { telemetryMetadata: metadata.connectorTelemetry } + : {}), }, }) ).pipe( diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/api.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/api.ts index e69d09b97a671..4b74860ae9be5 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/api.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/api.ts @@ -44,6 +44,7 @@ export function createChatCompleteApi({ request, actions, logger }: CreateChatCo modelName, stream, abortSignal, + metadata, }: ChatCompleteOptions): ChatCompleteCompositeResponse< ToolOptions, boolean @@ -87,6 +88,7 @@ export function createChatCompleteApi({ request, actions, logger }: CreateChatCo functionCalling, modelName, abortSignal, + metadata, }); }), chunksIntoMessage({ diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/types.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/types.ts index abafe70b619fb..9fdd0b19ff64c 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/types.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/types.ts @@ -13,6 +13,7 @@ import type { FunctionCallingMode, Message, ToolOptions, + ChatCompleteMetadata, } from '@kbn/inference-common'; import type { InferenceExecutor } from './utils'; @@ -36,12 +37,13 @@ export interface InferenceConnectorAdapter { export type InferenceAdapterChatCompleteOptions = { executor: InferenceExecutor; messages: Message[]; + logger: Logger; system?: string; functionCalling?: FunctionCallingMode; temperature?: number; modelName?: string; abortSignal?: AbortSignal; - logger: Logger; + metadata?: ChatCompleteMetadata; } & ToolOptions; /** diff --git a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index 5c2612aa0a4d4..54fa27b2bdafa 100644 --- a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -17,6 +17,7 @@ import { OutputCompleteEvent, OutputEventType, FunctionCallingMode, + ChatCompleteMetadata, } from '@kbn/inference-common'; import { correctCommonEsqlMistakes, generateFakeToolCallId } from '../../../../common'; import { InferenceClient } from '../../..'; @@ -35,6 +36,7 @@ export const generateEsqlTask = ({ functionCalling, logger, system, + metadata, }: { connectorId: string; systemMessage: string; @@ -44,6 +46,7 @@ export const generateEsqlTask = ({ docBase: EsqlDocumentBase; functionCalling?: FunctionCallingMode; logger: Pick; + metadata?: ChatCompleteMetadata; system?: string; }) => { return function askLlmToRespond({ @@ -73,6 +76,7 @@ export const generateEsqlTask = ({ chatCompleteApi({ connectorId, functionCalling, + metadata, stream: true, system: `${systemMessage} diff --git a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/request_documentation.ts b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/request_documentation.ts index 06e75db09bdc9..02f91a1d86c75 100644 --- a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/request_documentation.ts +++ b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/actions/request_documentation.ts @@ -12,6 +12,7 @@ import { Message, withoutOutputUpdateEvents, FunctionCallingMode, + ChatCompleteMetadata, } from '@kbn/inference-common'; import { InferenceClient } from '../../..'; import { requestDocumentationSchema } from './shared'; @@ -22,6 +23,7 @@ export const requestDocumentation = ({ messages, connectorId, functionCalling, + metadata, toolOptions: { tools, toolChoice }, }: { outputApi: InferenceClient['output']; @@ -29,6 +31,7 @@ export const requestDocumentation = ({ messages: Message[]; connectorId: string; functionCalling?: FunctionCallingMode; + metadata?: ChatCompleteMetadata; toolOptions: ToolOptions; }) => { const hasTools = !isEmpty(tools) && toolChoice !== ToolChoiceType.none; @@ -38,6 +41,7 @@ export const requestDocumentation = ({ connectorId, stream: true, functionCalling, + metadata, system, previousMessages: messages, input: `Based on the previous conversation, request documentation diff --git a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/task.ts b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/task.ts index 801d80a30174e..eba2a8bbb3903 100644 --- a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/task.ts +++ b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/task.ts @@ -22,6 +22,7 @@ export function naturalLanguageToEsql({ logger, functionCalling, system, + metadata, ...rest }: NlToEsqlTaskParams): Observable> { return from(loadDocBase()).pipe( @@ -38,6 +39,7 @@ export function naturalLanguageToEsql({ logger, systemMessage, functionCalling, + metadata, toolOptions: { tools, toolChoice, @@ -51,6 +53,7 @@ export function naturalLanguageToEsql({ outputApi: client.output, messages, system: systemMessage, + metadata, toolOptions: { tools, toolChoice, diff --git a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/types.ts b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/types.ts index 5a1477524dbd4..3303efbe438e9 100644 --- a/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/types.ts +++ b/x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/types.ts @@ -13,6 +13,7 @@ import type { Message, ToolOptions, OutputCompleteEvent, + ChatCompleteMetadata, } from '@kbn/inference-common'; import type { InferenceClient } from '../../inference_client'; @@ -30,5 +31,6 @@ export type NlToEsqlTaskParams = { logger: Pick; functionCalling?: FunctionCallingMode; system?: string; + metadata?: ChatCompleteMetadata; } & TToolOptions & ({ input: string } | { messages: Message[] }); diff --git a/x-pack/platform/plugins/shared/stack_connectors/common/bedrock/schema.ts b/x-pack/platform/plugins/shared/stack_connectors/common/bedrock/schema.ts index e9194a752300c..1b40efe223e1a 100644 --- a/x-pack/platform/plugins/shared/stack_connectors/common/bedrock/schema.ts +++ b/x-pack/platform/plugins/shared/stack_connectors/common/bedrock/schema.ts @@ -8,6 +8,11 @@ import { schema } from '@kbn/config-schema'; import { DEFAULT_BEDROCK_MODEL } from './constants'; +export const TelemtryMetadataSchema = schema.object({ + pluginId: schema.maybe(schema.string()), + aggregateBy: schema.maybe(schema.string()), +}); + // Connector schema export const ConfigSchema = schema.object({ apiUrl: schema.string(), @@ -71,10 +76,17 @@ export const InvokeAIActionParamsSchema = schema.object({ ) ), toolChoice: schema.maybe(BedrockToolChoiceSchema), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const InvokeAIActionResponseSchema = schema.object({ message: schema.string(), + usage: schema.maybe( + schema.object({ + input_tokens: schema.number(), + output_tokens: schema.number(), + }) + ), }); export const InvokeAIRawActionParamsSchema = schema.object({ @@ -85,6 +97,7 @@ export const InvokeAIRawActionParamsSchema = schema.object({ }) ), model: schema.maybe(schema.string()), + temperature: schema.maybe(schema.number()), stopSequences: schema.maybe(schema.arrayOf(schema.string())), system: schema.maybe(schema.string()), @@ -103,6 +116,7 @@ export const InvokeAIRawActionParamsSchema = schema.object({ ) ), toolChoice: schema.maybe(BedrockToolChoiceSchema), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const InvokeAIRawActionResponseSchema = schema.object({}, { unknowns: 'allow' }); @@ -154,6 +168,7 @@ export const BedrockClientSendParamsSchema = schema.object({ command: schema.any(), // Kibana related properties signal: schema.maybe(schema.any()), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const BedrockClientSendResponseSchema = schema.object({}, { unknowns: 'allow' }); diff --git a/x-pack/platform/plugins/shared/stack_connectors/common/gemini/schema.ts b/x-pack/platform/plugins/shared/stack_connectors/common/gemini/schema.ts index 1171cdc2a037a..c12f3f113d044 100644 --- a/x-pack/platform/plugins/shared/stack_connectors/common/gemini/schema.ts +++ b/x-pack/platform/plugins/shared/stack_connectors/common/gemini/schema.ts @@ -8,6 +8,11 @@ import { schema } from '@kbn/config-schema'; import { DEFAULT_GEMINI_MODEL } from './constants'; +export const TelemtryMetadataSchema = schema.object({ + pluginId: schema.maybe(schema.string()), + aggregateBy: schema.maybe(schema.string()), +}); + export const ConfigSchema = schema.object({ apiUrl: schema.string(), defaultModel: schema.string({ defaultValue: DEFAULT_GEMINI_MODEL }), @@ -27,6 +32,7 @@ export const RunActionParamsSchema = schema.object({ temperature: schema.maybe(schema.number()), stopSequences: schema.maybe(schema.arrayOf(schema.string())), raw: schema.maybe(schema.boolean()), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const RunApiResponseSchema = schema.object( @@ -73,6 +79,7 @@ export const InvokeAIActionParamsSchema = schema.object({ allowedFunctionNames: schema.maybe(schema.arrayOf(schema.string())), }) ), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const InvokeAIRawActionParamsSchema = schema.object({ @@ -84,6 +91,7 @@ export const InvokeAIRawActionParamsSchema = schema.object({ signal: schema.maybe(schema.any()), timeout: schema.maybe(schema.number()), tools: schema.maybe(schema.arrayOf(schema.any())), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const InvokeAIActionResponseSchema = schema.object({ diff --git a/x-pack/platform/plugins/shared/stack_connectors/common/inference/schema.ts b/x-pack/platform/plugins/shared/stack_connectors/common/inference/schema.ts index 2213efef1d6e8..2650e7b0bb3c1 100644 --- a/x-pack/platform/plugins/shared/stack_connectors/common/inference/schema.ts +++ b/x-pack/platform/plugins/shared/stack_connectors/common/inference/schema.ts @@ -7,6 +7,11 @@ import { schema } from '@kbn/config-schema'; +export const TelemtryMetadataSchema = schema.object({ + pluginId: schema.maybe(schema.string()), + aggregateBy: schema.maybe(schema.string()), +}); + export const ConfigSchema = schema.object({ provider: schema.string(), taskType: schema.string(), @@ -137,6 +142,7 @@ export const UnifiedChatCompleteParamsSchema = schema.object({ }), // abort signal from client signal: schema.maybe(schema.any()), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const UnifiedChatCompleteResponseSchema = schema.object({ diff --git a/x-pack/platform/plugins/shared/stack_connectors/common/openai/schema.ts b/x-pack/platform/plugins/shared/stack_connectors/common/openai/schema.ts index 7c3d4afcb8d1e..029972332fb14 100644 --- a/x-pack/platform/plugins/shared/stack_connectors/common/openai/schema.ts +++ b/x-pack/platform/plugins/shared/stack_connectors/common/openai/schema.ts @@ -8,6 +8,11 @@ import { schema } from '@kbn/config-schema'; import { DEFAULT_OPENAI_MODEL, OpenAiProviderType } from './constants'; +export const TelemtryMetadataSchema = schema.object({ + pluginId: schema.maybe(schema.string()), + aggregateBy: schema.maybe(schema.string()), +}); + // Connector schema export const ConfigSchema = schema.oneOf([ schema.object({ @@ -37,6 +42,7 @@ export const RunActionParamsSchema = schema.object({ // abort signal from client signal: schema.maybe(schema.any()), timeout: schema.maybe(schema.number()), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); const AIMessage = schema.object({ @@ -145,6 +151,7 @@ export const InvokeAIActionParamsSchema = schema.object({ // abort signal from client signal: schema.maybe(schema.any()), timeout: schema.maybe(schema.number()), + telemetryMetadata: schema.maybe(TelemtryMetadataSchema), }); export const InvokeAIActionResponseSchema = schema.object({ diff --git a/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/bedrock/bedrock.ts b/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/bedrock/bedrock.ts index 339efa49f69bf..c97d341a12e04 100644 --- a/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/bedrock/bedrock.ts +++ b/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/bedrock/bedrock.ts @@ -392,7 +392,7 @@ The Kibana Connector in use may need to be reconfigured with an updated Amazon B }, connectorUsageCollector )) as RunActionResponse; - return { message: res.completion.trim() }; + return { message: res.completion.trim(), usage: res?.usage }; } public async invokeAIRaw( diff --git a/x-pack/solutions/security/plugins/security_solution/server/lib/siem_migrations/rules/task/util/esql_knowledge_base_caller.ts b/x-pack/solutions/security/plugins/security_solution/server/lib/siem_migrations/rules/task/util/esql_knowledge_base_caller.ts index 2277f2fae41a9..73db3e8f72a95 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/lib/siem_migrations/rules/task/util/esql_knowledge_base_caller.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/lib/siem_migrations/rules/task/util/esql_knowledge_base_caller.ts @@ -25,11 +25,7 @@ export const getEsqlKnowledgeBase: GetEsqlTranslatorToolParams = client, connectorId, input, - logger: { - debug: (source) => { - logger.debug(typeof source === 'function' ? source() : source); - }, - }, + logger, }) ); return content; diff --git a/x-pack/test/alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts b/x-pack/test/alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts index ca8ae72ee06ce..2cf1d774aa24a 100644 --- a/x-pack/test/alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts +++ b/x-pack/test/alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts @@ -456,7 +456,13 @@ export default function bedrockTest({ getService }: FtrProviderContext) { expect(body).to.eql({ status: 'ok', connector_id: bedrockActionId, - data: { message: bedrockClaude2SuccessResponse.completion }, + data: { + message: bedrockClaude2SuccessResponse.completion, + usage: { + input_tokens: 41, + output_tokens: 64, + }, + }, }); const events: IValidatedEvent[] = await retry.try(async () => {