diff --git a/x-pack/packages/kbn-langchain/server/tracers/telemetry/index.ts b/x-pack/packages/kbn-langchain/server/tracers/telemetry/index.ts new file mode 100644 index 0000000000000..079c0e9a33087 --- /dev/null +++ b/x-pack/packages/kbn-langchain/server/tracers/telemetry/index.ts @@ -0,0 +1,8 @@ +/* + * 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. + */ + +export { TelemetryTracer } from './telemetry_tracer'; diff --git a/x-pack/packages/kbn-langchain/server/tracers/telemetry/telemetry_tracer.test.ts b/x-pack/packages/kbn-langchain/server/tracers/telemetry/telemetry_tracer.test.ts new file mode 100644 index 0000000000000..bca293ba9957e --- /dev/null +++ b/x-pack/packages/kbn-langchain/server/tracers/telemetry/telemetry_tracer.test.ts @@ -0,0 +1,204 @@ +/* + * 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. + */ + +import { AnalyticsServiceSetup, Logger } from '@kbn/core/server'; +import { TelemetryTracer, TelemetryParams } from './telemetry_tracer'; +import { Run } from 'langsmith/schemas'; +import { loggerMock } from '@kbn/logging-mocks'; + +const mockRun = { + inputs: { + responseLanguage: 'English', + conversationId: 'db8f74c5-7dca-43a3-b592-d56f219dffab', + llmType: 'openai', + isStream: false, + isOssModel: false, + }, + outputs: { + input: + 'Generate an ESQL query to find documents with `host.name` that contains my favorite color', + lastNode: 'agent', + steps: [ + { + action: { + tool: 'KnowledgeBaseRetrievalTool', + toolInput: { + query: "user's favorite color", + }, + }, + observation: + '"[{\\"pageContent\\":\\"favorite color is blue\\",\\"metadata\\":{\\"source\\":\\"conversation\\",\\"required\\":false,\\"kbResource\\":\\"user\\"}},{\\"pageContent\\":\\"favorite food is pizza\\",\\"metadata\\":{\\"source\\":\\"conversation\\",\\"required\\":false,\\"kbResource\\":\\"user\\"}}]"', + }, + { + action: { + tool: 'NaturalLanguageESQLTool', + toolInput: { + question: 'Generate an ESQL query to find documents with host.name that contains blue', + }, + }, + observation: + '"To find documents with `host.name` that contains \\"blue\\", you can use the `LIKE` operator with wildcards. Here is the ES|QL query:\\n\\n```esql\\nFROM your_index\\n| WHERE host.name LIKE \\"*blue*\\"\\n```\\n\\nReplace `your_index` with the actual name of your index. This query will filter documents where the `host.name` field contains the substring \\"blue\\"."', + }, + { + action: { + tool: 'KnowledgeBaseRetrievalTool', + toolInput: { + query: "user's favorite food", + }, + }, + observation: + '"[{\\"pageContent\\":\\"favorite color is blue\\",\\"metadata\\":{\\"source\\":\\"conversation\\",\\"required\\":false,\\"kbResource\\":\\"user\\"}},{\\"pageContent\\":\\"favorite food is pizza\\",\\"metadata\\":{\\"source\\":\\"conversation\\",\\"required\\":false,\\"kbResource\\":\\"user\\"}}]"', + }, + { + action: { + tool: 'CustomIndexTool', + toolInput: { + query: 'query about index', + }, + }, + observation: '"Wow this is totally cool."', + }, + { + action: { + tool: 'CustomIndexTool', + toolInput: { + query: 'query about index', + }, + }, + observation: '"Wow this is totally cool."', + }, + { + action: { + tool: 'CustomIndexTool', + toolInput: { + query: 'query about index', + }, + }, + observation: '"Wow this is totally cool."', + }, + ], + hasRespondStep: false, + agentOutcome: { + returnValues: { + output: + 'To find documents with `host.name` that contains your favorite color "blue", you can use the `LIKE` operator with wildcards. Here is the ES|QL query:\n\n```esql\nFROM your_index\n| WHERE host.name LIKE "*blue*"\n```\n\nReplace `your_index` with the actual name of your index. This query will filter documents where the `host.name` field contains the substring "blue".', + }, + log: 'To find documents with `host.name` that contains your favorite color "blue", you can use the `LIKE` operator with wildcards. Here is the ES|QL query:\n\n```esql\nFROM your_index\n| WHERE host.name LIKE "*blue*"\n```\n\nReplace `your_index` with the actual name of your index. This query will filter documents where the `host.name` field contains the substring "blue".', + }, + messages: [], + chatTitle: 'Welcome', + llmType: 'openai', + isStream: false, + isOssModel: false, + conversation: { + timestamp: '2024-11-07T17:37:07.400Z', + createdAt: '2024-11-07T17:37:07.400Z', + users: [ + { + id: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', + name: 'elastic', + }, + ], + title: 'Welcome', + category: 'assistant', + apiConfig: { + connectorId: 'my-gpt4o-ai', + actionTypeId: '.gen-ai', + }, + isDefault: true, + messages: [ + { + timestamp: '2024-11-07T22:47:45.994Z', + content: + 'Generate an ESQL query to find documents with `host.name` that contains my favorite color', + role: 'user', + }, + ], + updatedAt: '2024-11-08T17:01:21.958Z', + replacements: {}, + namespace: 'default', + id: 'db8f74c5-7dca-43a3-b592-d56f219dffab', + }, + conversationId: 'db8f74c5-7dca-43a3-b592-d56f219dffab', + responseLanguage: 'English', + }, + end_time: 1731085297190, + start_time: 1731085289113, +} as unknown as Run; +const elasticTools = [ + 'AlertCountsTool', + 'NaturalLanguageESQLTool', + 'KnowledgeBaseRetrievalTool', + 'KnowledgeBaseWriteTool', + 'OpenAndAcknowledgedAlertsTool', + 'SecurityLabsKnowledgeBaseTool', +]; +const mockLogger = loggerMock.create(); + +describe('TelemetryTracer', () => { + let telemetry: AnalyticsServiceSetup; + let logger: Logger; + let telemetryParams: TelemetryParams; + let telemetryTracer: TelemetryTracer; + const reportEvent = jest.fn(); + beforeEach(() => { + telemetry = { + reportEvent, + } as unknown as AnalyticsServiceSetup; + logger = mockLogger; + telemetryParams = { + eventType: 'INVOKE_AI_SUCCESS', + assistantStreamingEnabled: true, + actionTypeId: '.gen-ai', + isEnabledKnowledgeBase: true, + model: 'test_model', + }; + telemetryTracer = new TelemetryTracer( + { + elasticTools, + telemetry, + telemetryParams, + totalTools: 9, + }, + logger + ); + }); + + it('should initialize correctly', () => { + expect(telemetryTracer.name).toBe('telemetry_tracer'); + expect(telemetryTracer.elasticTools).toEqual(elasticTools); + expect(telemetryTracer.telemetry).toBe(telemetry); + expect(telemetryTracer.telemetryParams).toBe(telemetryParams); + expect(telemetryTracer.totalTools).toBe(9); + }); + + it('should not log and report event on chain end if parent_run_id exists', async () => { + await telemetryTracer.onChainEnd({ ...mockRun, parent_run_id: '123' }); + + expect(logger.get().debug).not.toHaveBeenCalled(); + expect(telemetry.reportEvent).not.toHaveBeenCalled(); + }); + + it('should log and report event on chain end', async () => { + await telemetryTracer.onChainEnd(mockRun); + + expect(logger.get().debug).toHaveBeenCalledWith(expect.any(Function)); + expect(telemetry.reportEvent).toHaveBeenCalledWith('INVOKE_AI_SUCCESS', { + assistantStreamingEnabled: true, + actionTypeId: '.gen-ai', + isEnabledKnowledgeBase: true, + model: 'test_model', + isOssModel: false, + durationMs: 8077, + toolsInvoked: { + KnowledgeBaseRetrievalTool: 2, + NaturalLanguageESQLTool: 1, + CustomTool: 3, + }, + }); + }); +}); diff --git a/x-pack/packages/kbn-langchain/server/tracers/telemetry/telemetry_tracer.ts b/x-pack/packages/kbn-langchain/server/tracers/telemetry/telemetry_tracer.ts new file mode 100644 index 0000000000000..7031e638c1fa4 --- /dev/null +++ b/x-pack/packages/kbn-langchain/server/tracers/telemetry/telemetry_tracer.ts @@ -0,0 +1,94 @@ +/* + * 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. + */ + +import { BaseCallbackHandlerInput } from '@langchain/core/callbacks/base'; +import type { Run } from 'langsmith/schemas'; +import { BaseTracer } from '@langchain/core/tracers/base'; +import { AnalyticsServiceSetup, Logger } from '@kbn/core/server'; + +export interface TelemetryParams { + assistantStreamingEnabled: boolean; + actionTypeId: string; + isEnabledKnowledgeBase: boolean; + eventType: string; + model?: string; +} +export interface LangChainTracerFields extends BaseCallbackHandlerInput { + elasticTools: string[]; + telemetry: AnalyticsServiceSetup; + telemetryParams: TelemetryParams; + totalTools: number; +} +interface ToolRunStep { + action: { + tool: string; + }; +} +/** + * TelemetryTracer is a tracer that uses event based telemetry to track LangChain events. + */ +export class TelemetryTracer extends BaseTracer implements LangChainTracerFields { + name = 'telemetry_tracer'; + logger: Logger; + elasticTools: string[]; + telemetry: AnalyticsServiceSetup; + telemetryParams: TelemetryParams; + totalTools: number; + constructor(fields: LangChainTracerFields, logger: Logger) { + super(fields); + this.logger = logger.get('telemetryTracer'); + this.elasticTools = fields.elasticTools; + this.telemetry = fields.telemetry; + this.telemetryParams = fields.telemetryParams; + this.totalTools = fields.totalTools; + } + + async onChainEnd(run: Run): Promise { + if (!run.parent_run_id) { + const { eventType, ...telemetryParams } = this.telemetryParams; + const toolsInvoked = + run?.outputs && run?.outputs.steps.length + ? run.outputs.steps.reduce((acc: { [k: string]: number }, event: ToolRunStep | never) => { + if ('action' in event && event?.action?.tool) { + if (this.elasticTools.includes(event.action.tool)) { + return { + ...acc, + ...(event.action.tool in acc + ? { [event.action.tool]: acc[event.action.tool] + 1 } + : { [event.action.tool]: 1 }), + }; + } else { + // Custom tool names are user data, so we strip them out + return { + ...acc, + ...('CustomTool' in acc + ? { CustomTool: acc.CustomTool + 1 } + : { CustomTool: 1 }), + }; + } + } + return acc; + }, {}) + : {}; + const telemetryValue = { + ...telemetryParams, + durationMs: (run.end_time ?? 0) - (run.start_time ?? 0), + toolsInvoked, + ...(telemetryParams.actionTypeId === '.gen-ai' + ? { isOssModel: run.inputs.isOssModel } + : {}), + }; + this.logger.debug( + () => `Invoke ${eventType} telemetry:\n${JSON.stringify(telemetryValue, null, 2)}` + ); + this.telemetry.reportEvent(eventType, telemetryValue); + } + } + + // everything below is required for type only + protected async persistRun(_run: Run): Promise {} +} diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts index 09bb5b291ef9a..77a1e37df965f 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts @@ -6,7 +6,12 @@ */ import { v4 as uuidv4 } from 'uuid'; -import { AuthenticatedUser, ElasticsearchClient, Logger } from '@kbn/core/server'; +import { + AnalyticsServiceSetup, + AuthenticatedUser, + ElasticsearchClient, + Logger, +} from '@kbn/core/server'; import { DocumentEntryCreateFields, @@ -15,6 +20,10 @@ import { KnowledgeBaseEntryUpdateProps, Metadata, } from '@kbn/elastic-assistant-common'; +import { + CREATE_KNOWLEDGE_BASE_ENTRY_ERROR_EVENT, + CREATE_KNOWLEDGE_BASE_ENTRY_SUCCESS_EVENT, +} from '../../lib/telemetry/event_based_telemetry'; import { getKnowledgeBaseEntry } from './get_knowledge_base_entry'; import { CreateKnowledgeBaseEntrySchema, UpdateKnowledgeBaseEntrySchema } from './types'; @@ -27,6 +36,7 @@ export interface CreateKnowledgeBaseEntryParams { knowledgeBaseEntry: KnowledgeBaseEntryCreateProps | LegacyKnowledgeBaseEntryCreateProps; global?: boolean; isV2?: boolean; + telemetry: AnalyticsServiceSetup; } export const createKnowledgeBaseEntry = async ({ @@ -38,6 +48,7 @@ export const createKnowledgeBaseEntry = async ({ logger, global = false, isV2 = false, + telemetry, }: CreateKnowledgeBaseEntryParams): Promise => { const createdAt = new Date().toISOString(); const body = isV2 @@ -55,6 +66,12 @@ export const createKnowledgeBaseEntry = async ({ entry: knowledgeBaseEntry as unknown as TransformToLegacyCreateSchemaProps['entry'], global, }); + const telemetryPayload = { + entryType: body.type, + required: body.required ?? false, + sharing: body.users.length ? 'private' : 'global', + ...(body.type === 'document' ? { source: body.source } : {}), + }; try { const response = await esClient.create({ body, @@ -63,17 +80,24 @@ export const createKnowledgeBaseEntry = async ({ refresh: 'wait_for', }); - return await getKnowledgeBaseEntry({ + const newKnowledgeBaseEntry = await getKnowledgeBaseEntry({ esClient, knowledgeBaseIndex, id: response._id, logger, user, }); + + telemetry.reportEvent(CREATE_KNOWLEDGE_BASE_ENTRY_SUCCESS_EVENT.eventType, telemetryPayload); + return newKnowledgeBaseEntry; } catch (err) { logger.error( `Error creating Knowledge Base Entry: ${err} with kbResource: ${knowledgeBaseEntry.name}` ); + telemetry.reportEvent(CREATE_KNOWLEDGE_BASE_ENTRY_ERROR_EVENT.eventType, { + ...telemetryPayload, + errorMessage: err.message ?? 'Unknown error', + }); throw err; } }; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 333fbb796ddd9..50e124321fe6c 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -25,7 +25,7 @@ import { import pRetry from 'p-retry'; import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey'; import { StructuredTool } from '@langchain/core/tools'; -import { ElasticsearchClient } from '@kbn/core/server'; +import { AnalyticsServiceSetup, ElasticsearchClient } from '@kbn/core/server'; import { IndexPatternsFetcher } from '@kbn/data-views-plugin/server'; import { map } from 'lodash'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; @@ -459,6 +459,8 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { filter: docsCreated.map((c) => `_id:${c}`).join(' OR '), }) : undefined; + // Intentionally no telemetry here - this path only used to install security docs + // Plans to make this function private in a different PR so no user entry ever is created in this path this.options.logger.debug(`created: ${created?.data.hits.hits.length ?? '0'}`); this.options.logger.debug(() => `errors: ${JSON.stringify(errors, null, 2)}`); @@ -686,10 +688,12 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { */ public createKnowledgeBaseEntry = async ({ knowledgeBaseEntry, + telemetry, global = false, }: { knowledgeBaseEntry: KnowledgeBaseEntryCreateProps | LegacyKnowledgeBaseEntryCreateProps; global?: boolean; + telemetry: AnalyticsServiceSetup; }): Promise => { const authenticatedUser = this.options.currentUser; @@ -716,6 +720,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { user: authenticatedUser, knowledgeBaseEntry, global, + telemetry, isV2: this.options.v2KnowledgeBaseEnabled, }); }; diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index 3e573aff2f4c8..da560dfae72dd 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -15,6 +15,8 @@ import { ExecuteConnectorRequestBody, Message, Replacements } from '@kbn/elastic import { StreamResponseWithHeaders } from '@kbn/ml-response-stream/server'; import { PublicMethodsOf } from '@kbn/utility-types'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; +import { AnalyticsServiceSetup } from '@kbn/core-analytics-server'; +import { TelemetryParams } from '@kbn/langchain/server/tracers/telemetry/telemetry_tracer'; import { ResponseBody } from '../types'; import type { AssistantTool } from '../../../types'; import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base'; @@ -55,6 +57,8 @@ export interface AgentExecutorParams { response?: KibanaResponseFactory; size?: number; systemPrompt?: string; + telemetry: AnalyticsServiceSetup; + telemetryParams?: TelemetryParams; traceOptions?: TraceOptions; responseLanguage?: string; } diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index d1b3514b15b73..0126692b5b6a5 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts @@ -7,6 +7,7 @@ import agent, { Span } from 'elastic-apm-node'; import type { Logger } from '@kbn/logging'; +import { TelemetryTracer } from '@kbn/langchain/server/tracers/telemetry'; import { streamFactory, StreamResponseWithHeaders } from '@kbn/ml-response-stream/server'; import { transformError } from '@kbn/securitysolution-es-utils'; import type { KibanaRequest } from '@kbn/core-http-server'; @@ -26,6 +27,7 @@ interface StreamGraphParams { logger: Logger; onLlmResponse?: OnLlmResponse; request: KibanaRequest; + telemetryTracer?: TelemetryTracer; traceOptions?: TraceOptions; } @@ -38,6 +40,7 @@ interface StreamGraphParams { * @param logger * @param onLlmResponse * @param request + * @param telemetryTracer * @param traceOptions */ export const streamGraph = async ({ @@ -47,6 +50,7 @@ export const streamGraph = async ({ logger, onLlmResponse, request, + telemetryTracer, traceOptions, }: StreamGraphParams): Promise => { let streamingSpan: Span | undefined; @@ -84,7 +88,11 @@ export const streamGraph = async ({ const stream = await assistantGraph.streamEvents( inputs, { - callbacks: [apmTracer, ...(traceOptions?.tracers ?? [])], + callbacks: [ + apmTracer, + ...(traceOptions?.tracers ?? []), + ...(telemetryTracer ? [telemetryTracer] : []), + ], runName: DEFAULT_ASSISTANT_GRAPH_ID, tags: traceOptions?.tags ?? [], version: 'v2', @@ -120,7 +128,11 @@ export const streamGraph = async ({ let finalMessage = ''; let conversationId: string | undefined; const stream = assistantGraph.streamEvents(inputs, { - callbacks: [apmTracer, ...(traceOptions?.tracers ?? [])], + callbacks: [ + apmTracer, + ...(traceOptions?.tracers ?? []), + ...(telemetryTracer ? [telemetryTracer] : []), + ], runName: DEFAULT_ASSISTANT_GRAPH_ID, streamMode: 'values', tags: traceOptions?.tags ?? [], @@ -187,6 +199,7 @@ interface InvokeGraphParams { assistantGraph: DefaultAssistantGraph; inputs: GraphInputs; onLlmResponse?: OnLlmResponse; + telemetryTracer?: TelemetryTracer; traceOptions?: TraceOptions; } interface InvokeGraphResponse { @@ -202,6 +215,7 @@ interface InvokeGraphResponse { * @param assistantGraph * @param inputs * @param onLlmResponse + * @param telemetryTracer * @param traceOptions */ export const invokeGraph = async ({ @@ -209,6 +223,7 @@ export const invokeGraph = async ({ assistantGraph, inputs, onLlmResponse, + telemetryTracer, traceOptions, }: InvokeGraphParams): Promise => { return withAssistantSpan(DEFAULT_ASSISTANT_GRAPH_ID, async (span) => { @@ -222,7 +237,11 @@ export const invokeGraph = async ({ span.addLabels({ evaluationId: traceOptions?.evaluationId }); } const r = await assistantGraph.invoke(inputs, { - callbacks: [apmTracer, ...(traceOptions?.tracers ?? [])], + callbacks: [ + apmTracer, + ...(traceOptions?.tracers ?? []), + ...(telemetryTracer ? [telemetryTracer] : []), + ], runName: DEFAULT_ASSISTANT_GRAPH_ID, tags: traceOptions?.tags ?? [], }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index 4f043c681f8df..f55006e452cd0 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts @@ -13,6 +13,7 @@ import { createToolCallingAgent, } from 'langchain/agents'; import { APMTracer } from '@kbn/langchain/server/tracers/apm'; +import { TelemetryTracer } from '@kbn/langchain/server/tracers/telemetry'; import { getLlmClass } from '../../../../routes/utils'; import { EsAnonymizationFieldsSchema } from '../../../../ai_assistant_data_clients/anonymization_fields/types'; import { AssistantToolParams } from '../../../../types'; @@ -44,6 +45,8 @@ export const callAssistantGraph: AgentExecutor = async ({ request, size, systemPrompt, + telemetry, + telemetryParams, traceOptions, responseLanguage = 'English', }) => { @@ -107,6 +110,7 @@ export const callAssistantGraph: AgentExecutor = async ({ replacements, request, size, + telemetry, }; const tools: StructuredTool[] = assistantTools.flatMap( @@ -150,7 +154,17 @@ export const callAssistantGraph: AgentExecutor = async ({ }); const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger); - + const telemetryTracer = telemetryParams + ? new TelemetryTracer( + { + elasticTools: assistantTools.map(({ name }) => name), + totalTools: tools.length, + telemetry, + telemetryParams, + }, + logger + ) + : undefined; const assistantGraph = getDefaultAssistantGraph({ agentRunnable, dataClients, @@ -177,6 +191,7 @@ export const callAssistantGraph: AgentExecutor = async ({ logger, onLlmResponse, request, + telemetryTracer, traceOptions, }); } @@ -186,6 +201,7 @@ export const callAssistantGraph: AgentExecutor = async ({ assistantGraph, inputs, onLlmResponse, + telemetryTracer, traceOptions, }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts b/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts index 5ff5ff894dffe..1087703ba13a4 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts @@ -76,7 +76,16 @@ export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{ assistantStreamingEnabled: boolean; actionTypeId: string; isEnabledKnowledgeBase: boolean; + durationMs: number; + ['toolsInvoked.AlertCountsTool']?: number; + ['toolsInvoked.NaturalLanguageESQLTool']?: number; + ['toolsInvoked.KnowledgeBaseRetrievalTool']?: number; + ['toolsInvoked.KnowledgeBaseWriteTool']?: number; + ['toolsInvoked.OpenAndAcknowledgedAlertsTool']?: number; + ['toolsInvoked.SecurityLabsKnowledgeBaseTool']?: number; + ['toolsInvoked.CustomTool']?: number; model?: string; + isOssModel?: boolean; }> = { eventType: 'invoke_assistant_success', schema: { @@ -105,6 +114,68 @@ export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{ description: 'Is knowledge base enabled', }, }, + isOssModel: { + type: 'boolean', + _meta: { + description: 'Is OSS model used on the request', + optional: true, + }, + }, + durationMs: { + type: 'integer', + _meta: { + description: 'The duration of the request.', + }, + }, + 'toolsInvoked.AlertCountsTool': { + type: 'long', + _meta: { + description: 'Number of times tool was invoked.', + optional: true, + }, + }, + 'toolsInvoked.NaturalLanguageESQLTool': { + type: 'long', + _meta: { + description: 'Number of times tool was invoked.', + optional: true, + }, + }, + 'toolsInvoked.KnowledgeBaseRetrievalTool': { + type: 'long', + _meta: { + description: 'Number of times tool was invoked.', + optional: true, + }, + }, + 'toolsInvoked.KnowledgeBaseWriteTool': { + type: 'long', + _meta: { + description: 'Number of times tool was invoked.', + optional: true, + }, + }, + 'toolsInvoked.OpenAndAcknowledgedAlertsTool': { + type: 'long', + _meta: { + description: 'Number of times tool was invoked.', + optional: true, + }, + }, + 'toolsInvoked.SecurityLabsKnowledgeBaseTool': { + type: 'long', + _meta: { + description: 'Number of times tool was invoked.', + optional: true, + }, + }, + 'toolsInvoked.CustomTool': { + type: 'long', + _meta: { + description: 'Number of times tool was invoked.', + optional: true, + }, + }, }, }; @@ -261,9 +332,90 @@ export const ATTACK_DISCOVERY_ERROR_EVENT: EventTypeOpts<{ }, }; +export const CREATE_KNOWLEDGE_BASE_ENTRY_SUCCESS_EVENT: EventTypeOpts<{ + entryType: 'index' | 'document'; + required: boolean; + sharing: 'private' | 'global'; + source?: string; +}> = { + eventType: 'create_knowledge_base_entry_success', + schema: { + entryType: { + type: 'keyword', + _meta: { + description: 'Index entry or document entry', + }, + }, + sharing: { + type: 'keyword', + _meta: { + description: 'Sharing setting: private or global', + }, + }, + required: { + type: 'boolean', + _meta: { + description: 'Whether this resource should always be included', + }, + }, + source: { + type: 'keyword', + _meta: { + description: 'Where the knowledge base document entry was created', + optional: true, + }, + }, + }, +}; + +export const CREATE_KNOWLEDGE_BASE_ENTRY_ERROR_EVENT: EventTypeOpts<{ + entryType: 'index' | 'document'; + required: boolean; + sharing: 'private' | 'global'; + source?: string; + errorMessage: string; +}> = { + eventType: 'create_knowledge_base_entry_error', + schema: { + entryType: { + type: 'keyword', + _meta: { + description: 'Index entry or document entry', + }, + }, + sharing: { + type: 'keyword', + _meta: { + description: 'Sharing setting: private or global', + }, + }, + required: { + type: 'boolean', + _meta: { + description: 'Whether this resource should always be included', + }, + }, + source: { + type: 'keyword', + _meta: { + description: 'Where the knowledge base document entry was created', + optional: true, + }, + }, + errorMessage: { + type: 'keyword', + _meta: { + description: 'Error message', + }, + }, + }, +}; + export const events: Array> = [ KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT, KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT, + CREATE_KNOWLEDGE_BASE_ENTRY_SUCCESS_EVENT, + CREATE_KNOWLEDGE_BASE_ENTRY_ERROR_EVENT, INVOKE_ASSISTANT_SUCCESS_EVENT, INVOKE_ASSISTANT_ERROR_EVENT, ATTACK_DISCOVERY_SUCCESS_EVENT, diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index d5db24d44f3e4..e4f520b190b5a 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -282,6 +282,7 @@ export const postEvaluateRoute = ( inference, connectorId: connector.id, size, + telemetry: ctx.elasticAssistant.telemetry, }; const tools: StructuredTool[] = assistantTools.flatMap( diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts index d25ed5fc77f10..0c5c39f77d692 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -30,6 +30,7 @@ import { ActionsClient } from '@kbn/actions-plugin/server'; import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabilities'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import type { InferenceServerStart } from '@kbn/inference-plugin/server'; +import { INVOKE_ASSISTANT_SUCCESS_EVENT } from '../lib/telemetry/event_based_telemetry'; import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; import { FindResponse } from '../ai_assistant_data_clients/find'; import { EsPromptsSchema } from '../ai_assistant_data_clients/prompts/types'; @@ -46,7 +47,6 @@ import { executeAction, StaticResponse } from '../lib/executor'; import { getLangChainMessages } from '../lib/langchain/helpers'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; -import { INVOKE_ASSISTANT_SUCCESS_EVENT } from '../lib/telemetry/event_based_telemetry'; import { ElasticAssistantRequestHandlerContext, GetElser } from '../types'; import { callAssistantGraph } from '../lib/langchain/graphs/default_assistant_graph'; @@ -399,6 +399,7 @@ export const langChainExecute = async ({ kbDataClient, }; + const isKnowledgeBaseInstalled = await getIsKnowledgeBaseInstalled(kbDataClient); // Shared executor params const executorParams: AgentExecutorParams = { abortSignal, @@ -422,6 +423,14 @@ export const langChainExecute = async ({ responseLanguage, size: request.body.size, systemPrompt, + telemetry, + telemetryParams: { + actionTypeId, + model: request.body.model, + assistantStreamingEnabled: isStream, + isEnabledKnowledgeBase: isKnowledgeBaseInstalled, + eventType: INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, + }, traceOptions: { projectName: request.body.langSmithProject, tracers: getLangSmithTracer({ @@ -436,14 +445,6 @@ export const langChainExecute = async ({ executorParams ); - const isKnowledgeBaseInstalled = await getIsKnowledgeBaseInstalled(kbDataClient); - - telemetry.reportEvent(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, { - actionTypeId, - model: request.body.model, - assistantStreamingEnabled: isStream, - isEnabledKnowledgeBase: isKnowledgeBaseInstalled, - }); return response.ok(result); }; diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts index fbe73525578b0..fc49068a09cc9 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts @@ -6,7 +6,7 @@ */ import moment from 'moment'; -import type { IKibanaResponse, KibanaResponseFactory } from '@kbn/core/server'; +import { AnalyticsServiceSetup, IKibanaResponse, KibanaResponseFactory } from '@kbn/core/server'; import { transformError } from '@kbn/securitysolution-es-utils'; import { @@ -20,6 +20,7 @@ import { } from '@kbn/elastic-assistant-common'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; +import { CREATE_KNOWLEDGE_BASE_ENTRY_SUCCESS_EVENT } from '../../../lib/telemetry/event_based_telemetry'; import { performChecks } from '../../helpers'; import { KNOWLEDGE_BASE_ENTRIES_TABLE_MAX_PAGE_SIZE } from '../../../../common/constants'; import { @@ -62,7 +63,8 @@ const buildBulkResponse = ( created = [], deleted = [], skipped = [], - }: KnowledgeBaseEntryBulkCrudActionResults & { errors: BulkOperationError[] } + }: KnowledgeBaseEntryBulkCrudActionResults & { errors: BulkOperationError[] }, + telemetry: AnalyticsServiceSetup ): IKibanaResponse => { const numSucceeded = updated.length + created.length + deleted.length; const numSkipped = skipped.length; @@ -82,6 +84,16 @@ const buildBulkResponse = ( skipped, }; + if (created.length) { + created.forEach((entry) => { + telemetry.reportEvent(CREATE_KNOWLEDGE_BASE_ENTRY_SUCCESS_EVENT.eventType, { + entryType: entry.type, + required: 'required' in entry ? entry.required ?? false : false, + sharing: entry.users.length ? 'private' : 'global', + ...(entry.type === 'document' ? { source: entry.source } : {}), + }); + }); + } if (numFailed > 0) { return response.custom({ headers: { 'content-type': 'application/json' }, @@ -289,14 +301,18 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug }) : undefined; - return buildBulkResponse(response, { - // @ts-ignore-next-line TS2322 - updated: transformESToKnowledgeBase(docsUpdated), - created: created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : [], - deleted: docsDeleted ?? [], - skipped: [], - errors, - }); + return buildBulkResponse( + response, + { + // @ts-ignore-next-line TS2322 + updated: transformESToKnowledgeBase(docsUpdated), + created: created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : [], + deleted: docsDeleted ?? [], + skipped: [], + errors, + }, + ctx.elasticAssistant.telemetry + ); } catch (err) { const error = transformError(err); return assistantResponse.error({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts index 0bfe9de269f7c..d5df2d02055fd 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts @@ -65,6 +65,7 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout const createResponse = await kbDataClient?.createKnowledgeBaseEntry({ knowledgeBaseEntry: request.body, global: request.body.users != null && request.body.users.length === 0, + telemetry: ctx.elasticAssistant.telemetry, }); if (createResponse == null) { diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index e84b97ab43d7a..00fec0dcabc6d 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -252,4 +252,5 @@ export interface AssistantToolParams { ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody >; size?: number; + telemetry?: AnalyticsServiceSetup; } diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json index d3436f28a1d3e..52ed30dde67f8 100644 --- a/x-pack/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/plugins/elastic_assistant/tsconfig.json @@ -49,7 +49,8 @@ "@kbn/std", "@kbn/zod", "@kbn/inference-plugin", - "@kbn/data-views-plugin" + "@kbn/data-views-plugin", + "@kbn/core-analytics-server" ], "exclude": [ "target/**/*", diff --git a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts index 4069eeeef5b97..3ae47afbf05bf 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts @@ -12,10 +12,12 @@ import type { AIAssistantKnowledgeBaseDataClient } from '@kbn/elastic-assistant- import { DocumentEntryType } from '@kbn/elastic-assistant-common'; import type { KnowledgeBaseEntryCreateProps } from '@kbn/elastic-assistant-common'; import type { LegacyKnowledgeBaseEntryCreateProps } from '@kbn/elastic-assistant-plugin/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry'; +import type { AnalyticsServiceSetup } from '@kbn/core-analytics-server'; import { APP_UI_ID } from '../../../../common'; export interface KnowledgeBaseWriteToolParams extends AssistantToolParams { kbDataClient: AIAssistantKnowledgeBaseDataClient; + telemetry: AnalyticsServiceSetup; } const toolDetails = { @@ -34,7 +36,7 @@ export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { getTool(params: AssistantToolParams) { if (!this.isSupported(params)) return null; - const { kbDataClient, logger } = params as KnowledgeBaseWriteToolParams; + const { telemetry, kbDataClient, logger } = params as KnowledgeBaseWriteToolParams; if (kbDataClient == null) return null; return new DynamicStructuredTool({ @@ -77,7 +79,7 @@ export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = { }; logger.debug(() => `knowledgeBaseEntry\n ${JSON.stringify(knowledgeBaseEntry, null, 2)}`); - const resp = await kbDataClient.createKnowledgeBaseEntry({ knowledgeBaseEntry }); + const resp = await kbDataClient.createKnowledgeBaseEntry({ knowledgeBaseEntry, telemetry }); if (resp == null) { return "I'm sorry, but I was unable to add this entry to your knowledge base.";