[Security solution] Add model parameter to token telemetry (elastic#1…

…87783)
cnasikas · Jul 8, 2024 · 2b5ff77 · 2b5ff77
1 parent 92099b2
commit 2b5ff77
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 1 deletion.
diff --git a/x-pack/plugins/actions/server/lib/action_executor.test.ts b/x-pack/plugins/actions/server/lib/action_executor.test.ts
@@ -1695,10 +1695,35 @@ describe('Event log', () => {
       { actionTypeId: '.gen-ai', completion_tokens: 9, prompt_tokens: 10, total_tokens: 19 }
     );
   });
+  test('reports telemetry for token count events with additional properties', async () => {
+    const executorMock = setupActionExecutorMock('.gen-ai', {} as ConnectorType['validate'], {
+      defaultModel: 'gpt-4',
+      apiProvider: 'OpenAI',
+    });
+    executorMock.mockResolvedValue({
+      actionId: '1',
+      status: 'ok',
+      // @ts-ignore
+      data: mockGenAi,
+    });
+    await actionExecutor.execute(executeParams);
+    expect(actionExecutorInitializationParams.analyticsService.reportEvent).toHaveBeenCalledWith(
+      GEN_AI_TOKEN_COUNT_EVENT.eventType,
+      {
+        actionTypeId: '.gen-ai',
+        completion_tokens: 9,
+        prompt_tokens: 10,
+        total_tokens: 19,
+        model: 'gpt-4',
+        provider: 'OpenAI',
+      }
+    );
+  });
 });
 function setupActionExecutorMock(
   actionTypeId = 'test',
-  validationOverride?: ConnectorType['validate']
+  validationOverride?: ConnectorType['validate'],
+  additionalConfig?: Record<string, unknown>
 ) {
   const thisConnectorType: jest.Mocked<ConnectorType> = {
     ...connectorType,
@@ -1712,6 +1737,7 @@ function setupActionExecutorMock(
       actionTypeId,
       config: {
         bar: true,
+        ...additionalConfig,
       },
       secrets: {
         baz: true,

diff --git a/x-pack/plugins/actions/server/lib/action_executor.ts b/x-pack/plugins/actions/server/lib/action_executor.ts
@@ -603,6 +603,7 @@ export class ActionExecutor {
                   ...(actionTypeId === '.gen-ai' && config?.apiProvider != null
                     ? { provider: config?.apiProvider }
                     : {}),
+                  ...(config?.defaultModel != null ? { model: config?.defaultModel } : {}),
                 });
               }
             })

diff --git a/x-pack/plugins/actions/server/lib/event_based_telemetry.ts b/x-pack/plugins/actions/server/lib/event_based_telemetry.ts
@@ -13,6 +13,7 @@ export const GEN_AI_TOKEN_COUNT_EVENT: EventTypeOpts<{
   prompt_tokens: number;
   completion_tokens: number;
   provider?: string;
+  model?: string;
 }> = {
   eventType: 'gen_ai_token_count',
   schema: {
@@ -51,6 +52,13 @@ export const GEN_AI_TOKEN_COUNT_EVENT: EventTypeOpts<{
         optional: true,
       },
     },
+    model: {
+      type: 'keyword',
+      _meta: {
+        description: 'LLM model',
+        optional: true,
+      },
+    },
   },
 };