mirror of
https://github.com/elastic/kibana.git
synced 2025-04-24 09:48:58 -04:00
[Security solution] Knowledge base entry telemetry (#199225)
This commit is contained in:
Steph Milovic
GitHub
parent
f54b95179f
commit
1127bf491d
16 changed files with 578 additions and 29 deletions
|
|
@ -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';
|
||||
|
|
@ -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,
|
||||
},
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -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<void> {
|
||||
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<void> {}
|
||||
}
|
||||
|
|
@ -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<KnowledgeBaseEntryResponse | null> => {
|
||||
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;
|
||||
}
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<KnowledgeBaseEntryResponse | null> => {
|
||||
const authenticatedUser = this.options.currentUser;
|
||||
|
||||
|
|
@ -716,6 +720,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
|
|||
user: authenticatedUser,
|
||||
knowledgeBaseEntry,
|
||||
global,
|
||||
telemetry,
|
||||
isV2: this.options.v2KnowledgeBaseEnabled,
|
||||
});
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<T extends boolean> {
|
|||
response?: KibanaResponseFactory;
|
||||
size?: number;
|
||||
systemPrompt?: string;
|
||||
telemetry: AnalyticsServiceSetup;
|
||||
telemetryParams?: TelemetryParams;
|
||||
traceOptions?: TraceOptions;
|
||||
responseLanguage?: string;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<unknown, unknown, ExecuteConnectorRequestBody>;
|
||||
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<StreamResponseWithHeaders> => {
|
||||
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<InvokeGraphResponse> => {
|
||||
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 ?? [],
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<true | false> = async ({
|
|||
request,
|
||||
size,
|
||||
systemPrompt,
|
||||
telemetry,
|
||||
telemetryParams,
|
||||
traceOptions,
|
||||
responseLanguage = 'English',
|
||||
}) => {
|
||||
|
|
@ -107,6 +110,7 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
|
|||
replacements,
|
||||
request,
|
||||
size,
|
||||
telemetry,
|
||||
};
|
||||
|
||||
const tools: StructuredTool[] = assistantTools.flatMap(
|
||||
|
|
@ -150,7 +154,17 @@ export const callAssistantGraph: AgentExecutor<true | false> = 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<true | false> = async ({
|
|||
logger,
|
||||
onLlmResponse,
|
||||
request,
|
||||
telemetryTracer,
|
||||
traceOptions,
|
||||
});
|
||||
}
|
||||
|
|
@ -186,6 +201,7 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
|
|||
assistantGraph,
|
||||
inputs,
|
||||
onLlmResponse,
|
||||
telemetryTracer,
|
||||
traceOptions,
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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<EventTypeOpts<{ [key: string]: unknown }>> = [
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -282,6 +282,7 @@ export const postEvaluateRoute = (
|
|||
inference,
|
||||
connectorId: connector.id,
|
||||
size,
|
||||
telemetry: ctx.elasticAssistant.telemetry,
|
||||
};
|
||||
|
||||
const tools: StructuredTool[] = assistantTools.flatMap(
|
||||
|
|
|
|||
|
|
@ -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<boolean> = {
|
||||
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<StreamResponseWithHeaders['body'] | StaticReturnType['body']>(result);
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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<KnowledgeBaseEntryBulkCrudActionResponse> => {
|
||||
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<KnowledgeBaseEntryBulkCrudActionResponse>({
|
||||
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({
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -252,4 +252,5 @@ export interface AssistantToolParams {
|
|||
ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody
|
||||
>;
|
||||
size?: number;
|
||||
telemetry?: AnalyticsServiceSetup;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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/**/*",
|
||||
|
|
|
|||
|
|
@ -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.";
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue