add support for telemetry metadata to inference APIs / nlToEsql

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

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

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

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

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

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

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+/**
+ * Set of metadata that can be used then calling the inference APIs
+ *
+ * @public
+ */
+export interface ChatCompleteMetadata {
+  connectorTelemetry?: ConnectorTelemetryMetadata;
+}
+
+/**
+ * Pass through for the connector telemetry
+ */
+export interface ConnectorTelemetryMetadata {
+  pluginId?: string;
+  aggregateBy?: string;
+}

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

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

x-pack/platform/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<string, ToolSchema | undefined, boolean> {
     if (stream && retry !== undefined) {
@@ -56,6 +57,7 @@ export function createOutputApi(chatCompleteApi: ChatCompleteAPI) {
       modelName,
       functionCalling,
       abortSignal,
+      metadata,
       system,
       messages,
       ...(schema

x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts

@@ -35,7 +35,7 @@ export const bedrockClaudeAdapter: InferenceConnectorAdapter = {
     temperature = 0,
     modelName,
     abortSignal,
-    telemetryMetadata,
+    metadata,
   }) => {
     const noToolUsage = toolChoice === ToolChoiceType.none;
 
@@ -45,10 +45,10 @@ export const bedrockClaudeAdapter: InferenceConnectorAdapter = {
       tools: noToolUsage ? [] : toolsToBedrock(tools, messages),
       toolChoice: toolChoiceToBedrock(toolChoice),
       temperature,
-      telemetryMetadata,
       model: modelName,
       stopSequences: ['\n\nHuman:'],
       signal: abortSignal,
+      ...(metadata?.connectorTelemetry ? { telemetryMetadata: metadata.connectorTelemetry } : {}),
     };
 
     return from(

x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts

@@ -33,7 +33,7 @@ export const geminiAdapter: InferenceConnectorAdapter = {
     temperature = 0,
     modelName,
     abortSignal,
-    telemetryMetadata,
+    metadata,
   }) => {
     return from(
       executor.invoke({
@@ -47,7 +47,9 @@ export const geminiAdapter: InferenceConnectorAdapter = {
           model: modelName,
           signal: abortSignal,
           stopSequences: ['\n\nHuman:'],
-          telemetryMetadata,
+          ...(metadata?.connectorTelemetry
+            ? { telemetryMetadata: metadata.connectorTelemetry }
+            : {}),
         },
       })
     ).pipe(

x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/inference_adapter.ts

@@ -27,7 +27,7 @@ export const inferenceAdapter: InferenceConnectorAdapter = {
     modelName,
     logger,
     abortSignal,
-    telemetryMetadata,
+    metadata,
   }) => {
     const useSimulatedFunctionCalling =
       functionCalling === 'auto'
@@ -51,7 +51,9 @@ export const inferenceAdapter: InferenceConnectorAdapter = {
         subActionParams: {
           body: request,
           signal: abortSignal,
-          telemetryMetadata,
+          ...(metadata?.connectorTelemetry
+            ? { telemetryMetadata: metadata.connectorTelemetry }
+            : {}),
         },
       })
     ).pipe(

x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/openai/openai_adapter.ts

@@ -32,7 +32,7 @@ export const openAIAdapter: InferenceConnectorAdapter = {
     modelName,
     logger,
     abortSignal,
-    telemetryMetadata,
+    metadata,
   }) => {
     const useSimulatedFunctionCalling =
       functionCalling === 'auto'
@@ -71,7 +71,9 @@ export const openAIAdapter: InferenceConnectorAdapter = {
           body: JSON.stringify(request),
           signal: abortSignal,
           stream: true,
-          telemetryMetadata,
+          ...(metadata?.connectorTelemetry
+            ? { telemetryMetadata: metadata.connectorTelemetry }
+            : {}),
         },
       })
     ).pipe(

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<ToolOptions, boolean>): ChatCompleteCompositeResponse<
     ToolOptions,
     boolean
@@ -87,6 +88,7 @@ export function createChatCompleteApi({ request, actions, logger }: CreateChatCo
           functionCalling,
           modelName,
           abortSignal,
+          metadata,
         });
       }),
       chunksIntoMessage({

x-pack/platform/plugins/shared/inference/server/chat_complete/types.ts

@@ -5,13 +5,13 @@
  * 2.0.
  */
 
-import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib';
 import type {
   ChatCompletionChunkEvent,
   ChatCompletionTokenCountEvent,
   FunctionCallingMode,
   Message,
   ToolOptions,
+  ChatCompleteMetadata,
 } from '@kbn/inference-common';
 import type { Logger } from '@kbn/logging';
 import type { Observable } from 'rxjs';
@@ -43,7 +43,7 @@ export type InferenceAdapterChatCompleteOptions = {
   temperature?: number;
   modelName?: string;
   abortSignal?: AbortSignal;
-  telemetryMetadata?: TelemetryMetadata;
+  metadata?: ChatCompleteMetadata;
 } & ToolOptions;
 
 /**

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 = <TToolOptions extends ToolOptions>({
   functionCalling,
   logger,
   system,
+  metadata,
 }: {
   connectorId: string;
   systemMessage: string;
@@ -44,6 +46,7 @@ export const generateEsqlTask = <TToolOptions extends ToolOptions>({
   docBase: EsqlDocumentBase;
   functionCalling?: FunctionCallingMode;
   logger: Pick<Logger, 'debug'>;
+  metadata?: ChatCompleteMetadata;
   system?: string;
 }) => {
   return function askLlmToRespond({
@@ -73,6 +76,7 @@ export const generateEsqlTask = <TToolOptions extends ToolOptions>({
       chatCompleteApi({
         connectorId,
         functionCalling,
+        metadata,
         stream: true,
         system: `${systemMessage}
 

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,13 +23,15 @@ export const requestDocumentation = ({
   messages,
   connectorId,
   functionCalling,
+  metadata,
   toolOptions: { tools, toolChoice },
 }: {
   outputApi: InferenceClient['output'];
   system: string;
   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

x-pack/platform/plugins/shared/inference/server/tasks/nl_to_esql/task.ts

@@ -22,6 +22,7 @@ export function naturalLanguageToEsql<TToolOptions extends ToolOptions>({
   logger,
   functionCalling,
   system,
+  metadata,
   ...rest
 }: NlToEsqlTaskParams<TToolOptions>): Observable<NlToEsqlTaskEvent<TToolOptions>> {
   return from(loadDocBase()).pipe(
@@ -38,6 +39,7 @@ export function naturalLanguageToEsql<TToolOptions extends ToolOptions>({
         logger,
         systemMessage,
         functionCalling,
+        metadata,
         toolOptions: {
           tools,
           toolChoice,
@@ -51,6 +53,7 @@ export function naturalLanguageToEsql<TToolOptions extends ToolOptions>({
         outputApi: client.output,
         messages,
         system: systemMessage,
+        metadata,
         toolOptions: {
           tools,
           toolChoice,

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<TToolOptions extends ToolOptions> = {
   logger: Pick<Logger, 'debug'>;
   functionCalling?: FunctionCallingMode;
   system?: string;
+  metadata?: ChatCompleteMetadata;
 } & TToolOptions &
   ({ input: string } | { messages: Message[] });

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;