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 b213958ae444a..656aca35a784b 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 88fa9a3a3c566..9e81152efb87b 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 @@ -34,6 +34,7 @@ export const bedrockClaudeAdapter: InferenceConnectorAdapter = { temperature = 0, modelName, abortSignal, + metadata, }) => { const noToolUsage = toolChoice === ToolChoiceType.none; @@ -46,6 +47,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 792c95d6bb91c..cbd1e0085cecf 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 @@ -32,6 +32,7 @@ export const geminiAdapter: InferenceConnectorAdapter = { temperature = 0, modelName, abortSignal, + metadata, }) => { return from( executor.invoke({ @@ -45,6 +46,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 5a00f8048c6a4..565e369fa70e8 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 83b1a47131bbd..9f1672f50e957 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 () => {