[Search] Playground sends telemetry headers with completion requests

x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/inference_chat_model.test.ts

@@ -93,6 +93,14 @@ const createChunkEvent = (input: ChunkEventInput): ChatCompletionChunkEvent => {
   }
 };
 
+const telemetryMetadata = {
+  pluginId: 'plugin-id',
+};
+
+const metadata = {
+  connectorTelemetry: telemetryMetadata,
+};
+
 describe('InferenceChatModel', () => {
   let chatComplete: ChatCompleteAPI & jest.MockedFn<ChatCompleteAPI>;
   let connector: InferenceConnector;
@@ -107,6 +115,7 @@ describe('InferenceChatModel', () => {
       const chatModel = new InferenceChatModel({
         chatComplete,
         connector,
+        telemetryMetadata,
       });
 
       const response = createResponse({ content: 'dummy' });
@@ -124,13 +133,15 @@ describe('InferenceChatModel', () => {
           },
         ],
         stream: false,
+        metadata,
       });
     });
 
     it('converts a complete conversation call', async () => {
       const chatModel = new InferenceChatModel({
         chatComplete,
         connector,
+        telemetryMetadata,
       });
 
       const response = createResponse({ content: 'dummy' });
@@ -170,13 +181,15 @@ describe('InferenceChatModel', () => {
           },
         ],
         stream: false,
+        metadata,
       });
     });
 
     it('converts a tool call conversation', async () => {
       const chatModel = new InferenceChatModel({
         chatComplete,
         connector,
+        telemetryMetadata,
       });
 
       const response = createResponse({ content: 'dummy' });
@@ -247,13 +260,15 @@ describe('InferenceChatModel', () => {
           },
         ],
         stream: false,
+        metadata,
       });
     });
 
     it('converts tools', async () => {
       const chatModel = new InferenceChatModel({
         chatComplete,
         connector,
+        telemetryMetadata,
       });
 
       const response = createResponse({ content: 'dummy' });
@@ -308,6 +323,7 @@ describe('InferenceChatModel', () => {
           },
         },
         stream: false,
+        metadata,
       });
     });
 
@@ -320,6 +336,7 @@ describe('InferenceChatModel', () => {
         model: 'super-duper-model',
         functionCallingMode: 'simulated',
         signal: abortCtrl.signal,
+        telemetryMetadata,
       });
 
       const response = createResponse({ content: 'dummy' });
@@ -336,6 +353,7 @@ describe('InferenceChatModel', () => {
         modelName: 'super-duper-model',
         abortSignal: abortCtrl.signal,
         stream: false,
+        metadata,
       });
     });
 
@@ -370,6 +388,9 @@ describe('InferenceChatModel', () => {
         modelName: 'some-other-model',
         abortSignal: abortCtrl.signal,
         stream: false,
+        metadata: {
+          connectorTelemetry: undefined,
+        },
       });
     });
   });
@@ -685,6 +706,7 @@ describe('InferenceChatModel', () => {
       const chatModel = new InferenceChatModel({
         chatComplete,
         connector,
+        telemetryMetadata,
       });
 
       const response = createResponse({ content: 'dummy' });
@@ -736,6 +758,7 @@ describe('InferenceChatModel', () => {
           },
         },
         stream: false,
+        metadata,
       });
     });
   });
@@ -746,6 +769,7 @@ describe('InferenceChatModel', () => {
         chatComplete,
         connector,
         model: 'my-super-model',
+        telemetryMetadata,
       });
 
       const identifyingParams = chatModel.identifyingParams();
@@ -754,6 +778,7 @@ describe('InferenceChatModel', () => {
         connectorId: 'connector-id',
         modelName: 'my-super-model',
         model_name: 'my-super-model',
+        metadata,
       });
     });
   });
@@ -792,6 +817,7 @@ describe('InferenceChatModel', () => {
       const chatModel = new InferenceChatModel({
         chatComplete,
         connector,
+        telemetryMetadata,
       });
 
       const structuredOutputModel = chatModel.withStructuredOutput(
@@ -858,6 +884,7 @@ describe('InferenceChatModel', () => {
           },
         },
         stream: false,
+        metadata,
       });
     });
 

x-pack/platform/packages/shared/ai-infra/inference-langchain/src/chat_model/inference_chat_model.ts

@@ -42,6 +42,7 @@ import {
   isToolValidationError,
   getConnectorDefaultModel,
   getConnectorProvider,
+  ConnectorTelemetryMetadata,
 } from '@kbn/inference-common';
 import type { ToolChoice } from './types';
 import { toAsyncIterator, wrapInferenceError } from './utils';
@@ -63,6 +64,7 @@ export interface InferenceChatModelParams extends BaseChatModelParams {
   temperature?: number;
   model?: string;
   signal?: AbortSignal;
+  telemetryMetadata?: ConnectorTelemetryMetadata;
 }
 
 export interface InferenceChatModelCallOptions extends BaseChatModelCallOptions {
@@ -94,6 +96,7 @@ export class InferenceChatModel extends BaseChatModel<InferenceChatModelCallOpti
   private readonly connector: InferenceConnector;
   // @ts-ignore unused for now
   private readonly logger: Logger;
+  private readonly telemetryMetadata?: ConnectorTelemetryMetadata;
 
   protected temperature?: number;
   protected functionCallingMode?: FunctionCallingMode;
@@ -104,6 +107,7 @@ export class InferenceChatModel extends BaseChatModel<InferenceChatModelCallOpti
     super(args);
     this.chatComplete = args.chatComplete;
     this.connector = args.connector;
+    this.telemetryMetadata = args.telemetryMetadata;
 
     this.temperature = args.temperature;
     this.functionCallingMode = args.functionCallingMode;
@@ -183,6 +187,7 @@ export class InferenceChatModel extends BaseChatModel<InferenceChatModelCallOpti
       tools: options.tools ? toolDefinitionToInference(options.tools) : undefined,
       toolChoice: options.tool_choice ? toolChoiceToInference(options.tool_choice) : undefined,
       abortSignal: options.signal ?? this.signal,
+      metadata: { connectorTelemetry: this.telemetryMetadata },
     };
   }
 

x-pack/solutions/search/plugins/search_playground/server/lib/get_chat_params.ts

@@ -86,6 +86,7 @@ export const getChatParams = async (
       // prevents the agent from retrying on failure
       // failure could be due to bad connector, we should deliver that result to the client asap
       maxRetries: 0,
+      telemetryMetadata: { pluginId: 'search_playground' }, // hard-coded because the pluginId is not snake cased and the telemetry expects snake case
     },
   });