[Stack connector] Forward telemetryMetadata.pluginId to EIS use case header (#214269)

2025-04-24 01:38:56 -04:00 · 2025-03-13 09:44:06 -06:00 · 2025-03-13 09:44:06 -06:00 · 13906cbb76
commit 13906cbb76
parent ebe7ef6da0
15 changed files with 70 additions and 3 deletions
--- a/x-pack/platform/packages/shared/kbn-langchain/server/language_models/llm.ts
+++ b/x-pack/platform/packages/shared/kbn-langchain/server/language_models/llm.ts
@ -11,6 +11,7 @@ import { LLM } from '@langchain/core/language_models/llms';
 import { get } from 'lodash/fp';
 import { v4 as uuidv4 } from 'uuid';
 import { PublicMethodsOf } from '@kbn/utility-types';
+import type { TelemetryMetadata } from '@kbn/actions-plugin/server/lib';
 import { DEFAULT_TIMEOUT, getDefaultArguments } from './constants';

 import { getMessageContentAndRole } from './helpers';
@ -28,6 +29,7 @@ interface ActionsClientLlmParams {
  timeout?: number;
  traceId?: string;
  traceOptions?: TraceOptions;
+  telemetryMetadata?: TelemetryMetadata;
 }

 export class ActionsClientLlm extends LLM {
@ -36,6 +38,7 @@ export class ActionsClientLlm extends LLM {
  #logger: Logger;
  #traceId: string;
  #timeout?: number;
+  telemetryMetadata?: TelemetryMetadata;

  // Local `llmType` as it can change and needs to be accessed by abstract `_llmType()` method
  // Not using getter as `this._llmType()` is called in the constructor via `super({})`
@ -54,6 +57,7 @@ export class ActionsClientLlm extends LLM {
    temperature,
    timeout,
    traceOptions,
+    telemetryMetadata,
  }: ActionsClientLlmParams) {
    super({
      callbacks: [...(traceOptions?.tracers ?? [])],
@ -67,6 +71,7 @@ export class ActionsClientLlm extends LLM {
    this.#timeout = timeout;
    this.model = model;
    this.temperature = temperature;
+    this.telemetryMetadata = telemetryMetadata;
  }

  _llmType() {
@ -102,6 +107,7 @@ export class ActionsClientLlm extends LLM {
                  model: this.model,
                  messages: [assistantMessage], // the assistant message
                },
+                telemetryMetadata: this.telemetryMetadata,
              },
            }
          : {
@ -113,6 +119,7 @@ export class ActionsClientLlm extends LLM {
                ...getDefaultArguments(this.llmType, this.temperature),
                // This timeout is large because LangChain prompts can be complicated and take a long time
                timeout: this.#timeout ?? DEFAULT_TIMEOUT,
+                telemetryMetadata: this.telemetryMetadata,
              },
            },
    };
--- a/x-pack/platform/plugins/shared/automatic_import/server/routes/analyze_api_route.ts
+++ b/x-pack/platform/plugins/shared/automatic_import/server/routes/analyze_api_route.ts
@ -77,6 +77,9 @@ export function registerApiAnalysisRoutes(router: IRouter<AutomaticImportRouteHa
            maxTokens: 4096,
            signal: abortSignal,
            streaming: false,
+            telemetryMetadata: {
+              pluginId: 'automatic_import',
+            },
          });

          const parameters = {
--- a/x-pack/platform/plugins/shared/automatic_import/server/routes/analyze_logs_routes.ts
+++ b/x-pack/platform/plugins/shared/automatic_import/server/routes/analyze_logs_routes.ts
@ -88,6 +88,9 @@ export function registerAnalyzeLogsRoutes(router: IRouter<AutomaticImportRouteHa
            maxTokens: 4096,
            signal: abortSignal,
            streaming: false,
+            telemetryMetadata: {
+              pluginId: 'automatic_import',
+            },
          });
          const options = {
            callbacks: [
--- a/x-pack/platform/plugins/shared/automatic_import/server/routes/categorization_routes.ts
+++ b/x-pack/platform/plugins/shared/automatic_import/server/routes/categorization_routes.ts
@ -93,6 +93,9 @@ export function registerCategorizationRoutes(router: IRouter<AutomaticImportRout
              maxTokens: 4096,
              signal: abortSignal,
              streaming: false,
+              telemetryMetadata: {
+                pluginId: 'automatic_import',
+              },
            });

            const parameters = {
--- a/x-pack/platform/plugins/shared/automatic_import/server/routes/cel_routes.ts
+++ b/x-pack/platform/plugins/shared/automatic_import/server/routes/cel_routes.ts
@ -77,6 +77,9 @@ export function registerCelInputRoutes(router: IRouter<AutomaticImportRouteHandl
            maxTokens: 4096,
            signal: abortSignal,
            streaming: false,
+            telemetryMetadata: {
+              pluginId: 'automatic_import',
+            },
          });

          const parameters = {
--- a/x-pack/platform/plugins/shared/automatic_import/server/routes/ecs_routes.ts
+++ b/x-pack/platform/plugins/shared/automatic_import/server/routes/ecs_routes.ts
@ -87,6 +87,9 @@ export function registerEcsRoutes(router: IRouter<AutomaticImportRouteHandlerCon
            maxTokens: 4096,
            signal: abortSignal,
            streaming: false,
+            telemetryMetadata: {
+              pluginId: 'automatic_import',
+            },
          });

          const parameters = {
--- a/x-pack/platform/plugins/shared/automatic_import/server/routes/related_routes.ts
+++ b/x-pack/platform/plugins/shared/automatic_import/server/routes/related_routes.ts
@ -87,6 +87,9 @@ export function registerRelatedRoutes(router: IRouter<AutomaticImportRouteHandle
            maxTokens: 4096,
            signal: abortSignal,
            streaming: false,
+            telemetryMetadata: {
+              pluginId: 'automatic_import',
+            },
          });

          const parameters = {
--- a/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/inference.test.ts
+++ b/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/inference.test.ts
@ -70,6 +70,7 @@ describe('InferenceConnector', () => {

      const response = await connector.performApiUnifiedCompletion({
        body: { messages: [{ content: 'What is Elastic?', role: 'user' }] },
+        telemetryMetadata: { pluginId: 'security_ai_assistant' },
      });
      expect(mockEsClient.transport.request).toBeCalledTimes(1);
      expect(mockEsClient.transport.request).toHaveBeenCalledWith(
@ -86,7 +87,13 @@ describe('InferenceConnector', () => {
          method: 'POST',
          path: '_inference/chat_completion/test/_stream',
        },
-        { asStream: true, meta: true }
+        {
+          asStream: true,
+          meta: true,
+          headers: {
+            'X-Elastic-Product-Use-Case': 'security_ai_assistant',
+          },
+        }
      );
      expect(response.choices[0].message.content).toEqual(' you');
    });
@ -290,7 +297,10 @@ describe('InferenceConnector', () => {
          method: 'POST',
          path: '_inference/chat_completion/test/_stream',
        },
-        { asStream: true, meta: true }
+        {
+          asStream: true,
+          meta: true,
+        }
      );
    });

@ -312,7 +322,11 @@ describe('InferenceConnector', () => {
          method: 'POST',
          path: '_inference/chat_completion/test/_stream',
        },
-        { asStream: true, meta: true, signal }
+        {
+          asStream: true,
+          meta: true,
+          signal,
+        }
      );
    });

--- a/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/inference.ts
+++ b/x-pack/platform/plugins/shared/stack_connectors/server/connector_types/inference/inference.ts
@ -196,6 +196,13 @@ export class InferenceConnector extends SubActionConnector<Config, Secrets> {
        asStream: true,
        meta: true,
        signal: params.signal,
+        ...(params.telemetryMetadata?.pluginId
+          ? {
+              headers: {
+                'X-Elastic-Product-Use-Case': params.telemetryMetadata?.pluginId,
+              },
+            }
+          : {}),
      }
    );
    // errors should be thrown as it will not be a stream response
--- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_evaluator_llm/index.ts
+++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/helpers/get_evaluator_llm/index.ts
@ -61,5 +61,8 @@ export const getEvaluatorLlm = async ({
    temperature: 0, // zero temperature for evaluation
    timeout: connectorTimeout,
    traceOptions,
+    telemetryMetadata: {
+      pluginId: 'security_attack_discovery',
+    },
  });
 };
--- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.ts
+++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.ts
@ -92,6 +92,9 @@ export const evaluateAttackDiscovery = async ({
        temperature: 0, // zero temperature for attack discovery, because we want structured JSON output
        timeout: connectorTimeout,
        traceOptions,
+        telemetryMetadata: {
+          pluginId: 'security_attack_discovery',
+        },
      });

      const graph = getDefaultAttackDiscoveryGraph({
--- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
+++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
@ -89,6 +89,9 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
      // failure could be due to bad connector, we should deliver that result to the client asap
      maxRetries: 0,
      convertSystemMessageToHumanContent: false,
+      telemetryMetadata: {
+        pluginId: 'security_ai_assistant',
+      },
    });

  const anonymizationFieldsRes =
--- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/attack_discovery/post/helpers/invoke_attack_discovery_graph/index.tsx
+++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/attack_discovery/post/helpers/invoke_attack_discovery_graph/index.tsx
@ -87,6 +87,9 @@ export const invokeAttackDiscoveryGraph = async ({
    temperature: 0, // zero temperature for attack discovery, because we want structured JSON output
    timeout: connectorTimeout,
    traceOptions,
+    telemetryMetadata: {
+      pluginId: 'security_attack_discovery',
+    },
  });

  if (llm == null) {
--- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts
+++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts
@ -156,6 +156,9 @@ export function getAssistantToolParams({
    temperature: 0, // zero temperature because we want structured JSON output
    timeout: connectorTimeout,
    traceOptions,
+    telemetryMetadata: {
+      pluginId: 'security_defend_insights',
+    },
  });

  return {
@ -443,6 +446,9 @@ export const invokeDefendInsightsGraph = async ({
    temperature: 0,
    timeout: connectorTimeout,
    traceOptions,
+    telemetryMetadata: {
+      pluginId: 'security_defend_insights',
+    },
  });

  if (llm == null) {
--- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts
+++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts
@ -250,6 +250,9 @@ export const postEvaluateRoute = (
                  streaming: false,
                  maxRetries: 0,
                  convertSystemMessageToHumanContent: false,
+                  telemetryMetadata: {
+                    pluginId: 'security_ai_assistant',
+                  },
                });
              const llm = createLlmInstance();
              const anonymizationFieldsRes =