github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-06-28 11:05:39 -04:00
Code Issues Projects Releases Packages Wiki Activity

[8.18] [Security Assistant] Fix telemetry tool reporting (#223832) (#224344)

Browse source 
# Backport

This will backport the following commits from `main` to `8.18`:
- [[Security Assistant] Fix telemetry tool reporting
(#223832)](https://github.com/elastic/kibana/pull/223832)

<!--- Backport version: 10.0.1 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sorenlouv/backport)

<!--BACKPORT [{"author":{"name":"Steph
Milovic","email":"stephanie.milovic@elastic.co"},"sourceCommit":{"committedDate":"2025-06-12T22:52:33Z","message":"[Security
Assistant] Fix telemetry tool reporting
(#223832)","sha":"602f96d46a3332471437e5e79faa5454293cbe14","branchLabelMapping":{"^v9.1.0$":"main","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","backport
missing","Team: SecuritySolution","Team:Security Generative
AI","backport:version","v9.1.0","v8.19.0","v9.0.3","v8.18.3"],"title":"[Security
Assistant] Fix telemetry tool
reporting","number":223832,"url":"https://github.com/elastic/kibana/pull/223832","mergeCommit":{"message":"[Security
Assistant] Fix telemetry tool reporting
(#223832)","sha":"602f96d46a3332471437e5e79faa5454293cbe14"}},"sourceBranch":"main","suggestedTargetBranches":["9.0","8.18"],"targetPullRequestStates":[{"branch":"main","label":"v9.1.0","branchLabelMappingKey":"^v9.1.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/223832","number":223832,"mergeCommit":{"message":"[Security
Assistant] Fix telemetry tool reporting
(#223832)","sha":"602f96d46a3332471437e5e79faa5454293cbe14"}},{"branch":"8.19","label":"v8.19.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"url":"https://github.com/elastic/kibana/pull/223845","number":223845,"state":"OPEN"},{"branch":"9.0","label":"v9.0.3","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"8.18","label":"v8.18.3","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->
This commit is contained in:
Steph Milovic 2025-06-18 02:20:50 -06:00 • committed by GitHub
parent ab128c2320
commit 7c47b7ef45
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 52 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

2
x-pack/platform/packages/shared/kbn-langchain/server/tracers/telemetry/telemetry_tracer.test.ts
View file

@ -162,7 +162,6 @@ describe('TelemetryTracer', () => {
elasticTools,
telemetry,
telemetryParams,
totalTools: 9,
},
logger
);
@ -173,7 +172,6 @@ describe('TelemetryTracer', () => {
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 () => {

3
x-pack/platform/packages/shared/kbn-langchain/server/tracers/telemetry/telemetry_tracer.ts
View file

@ -21,7 +21,6 @@ export interface LangChainTracerFields extends BaseCallbackHandlerInput {
elasticTools: string[];
telemetry: AnalyticsServiceSetup;
telemetryParams: TelemetryParams;
totalTools: number;
}
interface ToolRunStep {
action: {
@ -37,14 +36,12 @@ export class TelemetryTracer extends BaseTracer implements LangChainTracerFields
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> {

50
x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts
View file

@ -10,6 +10,7 @@ import { callAssistantGraph } from '.';
import { getDefaultAssistantGraph } from './graph';
import { invokeGraph, streamGraph } from './helpers';
import { loggerMock } from '@kbn/logging-mocks';
import { TelemetryTracer } from '@kbn/langchain/server/tracers/telemetry';
import { AgentExecutorParams, AssistantDataClients } from '../../executors/types';
import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks';
import { getPrompt, resolveProviderAndModel } from '@kbn/security-ai-prompts';
@ -30,6 +31,16 @@ jest.mock('@kbn/security-ai-prompts');
const getDefaultAssistantGraphMock = getDefaultAssistantGraph as jest.Mock;
const resolveProviderAndModelMock = resolveProviderAndModel as jest.Mock;
const getPromptMock = getPrompt as jest.Mock;
const telemetryTracerMock = TelemetryTracer as unknown as jest.Mock;
const getTool = jest.fn();
const mockTool: AssistantTool = {
id: 'id',
name: 'name',
description: 'description',
sourceRegister: 'sourceRegister',
isSupported: (params: AssistantToolParams) => true,
getTool: getTool.mockReturnValue({ name: 'name' }),
};
describe('callAssistantGraph', () => {
const mockDataClients = {
anonymizationFieldsDataClient: {
@ -37,7 +48,7 @@ describe('callAssistantGraph', () => {
},
kbDataClient: {
isInferenceEndpointExists: jest.fn(),
getAssistantTools: jest.fn(),
getAssistantTools: jest.fn().mockReturnValue([{ name: 'MyKBTool' }]),
},
} as unknown as AssistantDataClients;
@ -54,6 +65,7 @@ describe('callAssistantGraph', () => {
total: 0,
saved_objects: [],
});
const mockLogger = loggerMock.create();
const defaultParams = {
actionsClient: actionsClientMock.create(),
alertsIndexPattern: 'test-pattern',
@ -67,7 +79,7 @@ describe('callAssistantGraph', () => {
llmTasks: { retrieveDocumentationAvailable: jest.fn(), retrieveDocumentation: jest.fn() },
llmType: 'openai',
isOssModel: false,
logger: loggerMock.create(),
logger: mockLogger,
isStream: false,
onLlmResponse: jest.fn(),
onNewReplacements: jest.fn(),
@ -176,15 +188,6 @@ describe('callAssistantGraph', () => {
});
it('calls getPrompt for each tool and the default system prompt', async () => {
const getTool = jest.fn();
const mockTool: AssistantTool = {
id: 'id',
name: 'name',
description: 'description',
sourceRegister: 'sourceRegister',
isSupported: (params: AssistantToolParams) => true,
getTool,
};
const params = {
...defaultParams,
assistantTools: [
@ -227,6 +230,31 @@ describe('callAssistantGraph', () => {
);
});
it('Passes only Elastic tools, not custom, to Telemetry tracer', async () => {
await callAssistantGraph({
...defaultParams,
assistantTools: [
{ ...mockTool, name: 'test-tool', getTool: getTool.mockReturnValue({ name: 'test-tool' }) },
{
...mockTool,
name: 'test-tool2',
getTool: getTool.mockReturnValue({ name: 'test-tool2' }),
},
],
});
expect(telemetryTracerMock).toHaveBeenCalledWith(
{
elasticTools: ['test-tool2', 'test-tool2'],
telemetry: {},
telemetryParams: {},
},
{
...mockLogger,
context: ['defaultAssistantGraph'],
}
);
});
describe('agentRunnable', () => {
it('creates OpenAIToolsAgent for openai llmType', async () => {
const params = { ...defaultParams, llmType: 'openai' };

25
x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
View file

@ -156,6 +156,19 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
)
).filter((e) => e != null) as StructuredTool[];
const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger);
const telemetryTracer = telemetryParams
? new TelemetryTracer(
{
// this line MUST come before kbTools are added
elasticTools: tools.map(({ name }) => name),
telemetry,
telemetryParams,
},
logger
)
: undefined;
// If KB enabled, fetch for any KB IndexEntries and generate a tool for each
if (isEnabledKnowledgeBase) {
const kbTools = await dataClients?.kbDataClient?.getAssistantTools({
@ -191,18 +204,6 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
prompt: chatPromptTemplate,
});
const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger);
const telemetryTracer = telemetryParams
? new TelemetryTracer(
{
elasticTools: tools.map(({ name }) => name),
totalTools: tools.length,
telemetry,
telemetryParams,
},
logger
)
: undefined;
const { provider } =
!llmType || llmType === 'inference'
? await resolveProviderAndModel({