github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-04-23 01:13:23 -04:00
Code Issues Projects Releases Packages Wiki Activity

[Security solution] Generative AI Connector - track token usage per user (#158508)

Browse source
This commit is contained in:
Steph Milovic 2023-06-05 17:23:34 -06:00 • committed by GitHub
parent 8f6469a48b
commit 531129b945
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 199 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

110
x-pack/plugins/actions/server/lib/action_executor.test.ts
View file

@ -19,6 +19,7 @@ import {
asHttpRequestExecutionSource,
asSavedObjectExecutionSource,
} from './action_execution_source';
import { securityMock } from '@kbn/security-plugin/server/mocks';
const actionExecutor = new ActionExecutor({ isESOCanEncrypt: true });
const services = actionsMock.createServices();
@ -40,10 +41,12 @@ const executeParams = {
const spacesMock = spacesServiceMock.createStartContract();
const loggerMock: ReturnType<typeof loggingSystemMock.createLogger> =
loggingSystemMock.createLogger();
const securityMockStart = securityMock.createStart();
actionExecutor.initialize({
logger: loggerMock,
spaces: spacesMock,
security: securityMockStart,
getServices: () => services,
actionTypeRegistry,
encryptedSavedObjectsClient,
@ -69,6 +72,18 @@ beforeEach(() => {
jest.resetAllMocks();
spacesMock.getSpaceId.mockReturnValue('some-namespace');
loggerMock.get.mockImplementation(() => loggerMock);
const mockRealm = { name: 'default_native', type: 'native' };
securityMockStart.authc.getCurrentUser.mockImplementation(() => ({
authentication_realm: mockRealm,
authentication_provider: mockRealm,
authentication_type: 'realm',
lookup_realm: mockRealm,
elastic_cloud_user: true,
enabled: true,
profile_uid: '123',
roles: ['superuser'],
username: 'coolguy',
}));
});
test('successfully executes', async () => {
@ -203,6 +218,10 @@ test('successfully executes', async () => {
],
},
"message": "action executed: test:1: 1",
"user": Object {
"id": "123",
"name": "coolguy",
},
},
],
]
@ -346,6 +365,10 @@ test('successfully executes when http_request source is specified', async () =>
],
},
"message": "action executed: test:1: 1",
"user": Object {
"id": "123",
"name": "coolguy",
},
},
],
]
@ -492,6 +515,10 @@ test('successfully executes when saved_object source is specified', async () =>
],
},
"message": "action executed: test:1: 1",
"user": Object {
"id": "123",
"name": "coolguy",
},
},
],
]
@ -613,6 +640,10 @@ test('successfully executes with preconfigured connector', async () => {
],
},
"message": "action executed: test:preconfigured: Preconfigured",
"user": Object {
"id": "123",
"name": "coolguy",
},
},
],
]
@ -1063,6 +1094,10 @@ test('should not throw error if action is preconfigured and isESOCanEncrypt is f
],
},
"message": "action executed: test:preconfigured: Preconfigured",
"user": Object {
"id": "123",
"name": "coolguy",
},
},
],
]
@ -1279,7 +1314,78 @@ test('writes to event log for execute and execute start when consumer and relate
});
});
function setupActionExecutorMock() {
test('writes usage data to event log for gen ai events', async () => {
const executorMock = setupActionExecutorMock('.gen-ai');
const mockGenAi = {
id: 'chatcmpl-7LztF5xsJl2z5jcNpJKvaPm4uWt8x',
object: 'chat.completion',
created: 1685477149,
model: 'gpt-3.5-turbo-0301',
usage: {
prompt_tokens: 10,
completion_tokens: 9,
total_tokens: 19,
},
choices: [
{
message: {
role: 'assistant',
content: 'Hello! How can I assist you today?',
},
finish_reason: 'stop',
index: 0,
},
],
};
executorMock.mockResolvedValue({
actionId: '1',
status: 'ok',
// @ts-ignore
data: mockGenAi,
});
await actionExecutor.execute(executeParams);
expect(eventLogger.logEvent).toHaveBeenCalledTimes(2);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(2, {
event: {
action: 'execute',
kind: 'action',
outcome: 'success',
},
kibana: {
action: {
execution: {
uuid: '2',
gen_ai: {
usage: mockGenAi.usage,
},
},
name: 'action-1',
id: '1',
},
alert: {
rule: {
execution: {
uuid: '123abc',
},
},
},
saved_objects: [
{
id: '1',
namespace: 'some-namespace',
rel: 'primary',
type: 'action',
type_id: '.gen-ai',
},
],
space_ids: ['some-namespace'],
},
message: 'action executed: .gen-ai:1: action-1',
user: { name: 'coolguy', id: '123' },
});
});
function setupActionExecutorMock(actionTypeId = 'test') {
const actionType: jest.Mocked<ActionType> = {
id: 'test',
name: 'Test',
@ -1297,7 +1403,7 @@ function setupActionExecutorMock() {
type: 'action',
attributes: {
name: 'action-1',
actionTypeId: 'test',
actionTypeId,
config: {
bar: true,
},

37
x-pack/plugins/actions/server/lib/action_executor.ts
View file

@ -12,6 +12,7 @@ import { withSpan } from '@kbn/apm-utils';
import { EncryptedSavedObjectsClient } from '@kbn/encrypted-saved-objects-plugin/server';
import { SpacesServiceStart } from '@kbn/spaces-plugin/server';
import { IEventLogger, SAVED_OBJECT_REL_PRIMARY } from '@kbn/event-log-plugin/server';
import { SecurityPluginStart } from '@kbn/security-plugin/server';
import {
validateParams,
validateConfig,
@ -40,6 +41,7 @@ const Millis2Nanos = 1000 * 1000;
export interface ActionExecutorContext {
logger: Logger;
spaces?: SpacesServiceStart;
security?: SecurityPluginStart;
getServices: GetServicesFunction;
encryptedSavedObjectsClient: EncryptedSavedObjectsClient;
actionTypeRegistry: ActionTypeRegistryContract;
@ -101,7 +103,6 @@ export class ActionExecutor {
if (!this.isInitialized) {
throw new Error('ActionExecutor not initialized');
}
return withSpan(
{
name: `execute_action`,
@ -118,6 +119,7 @@ export class ActionExecutor {
actionTypeRegistry,
eventLogger,
preconfiguredActions,
security,
} = this.actionExecutorContext!;
const services = getServices(request);
@ -251,6 +253,39 @@ export class ActionExecutor {
event.event = event.event || {};
// start gen_ai extension
// add event.kibana.action.execution.gen_ai to event log when GenerativeAi Connector is executed
if (result.status === 'ok' && actionTypeId === '.gen-ai') {
const data = result.data as unknown as {
usage: { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number };
};
event.kibana = event.kibana || {};
event.kibana.action = event.kibana.action || {};
event.kibana = {
...event.kibana,
action: {
...event.kibana.action,
execution: {
...event.kibana.action.execution,
gen_ai: {
usage: {
total_tokens: data.usage?.total_tokens,
prompt_tokens: data.usage?.prompt_tokens,
completion_tokens: data.usage?.completion_tokens,
},
},
},
},
};
}
// end gen_ai extension
const currentUser = security?.authc.getCurrentUser(request);
event.user = event.user || {};
event.user.name = currentUser?.username;
event.user.id = currentUser?.profile_uid;
if (result.status === 'ok') {
span?.setOutcome('success');
event.event.outcome = 'success';

4
x-pack/plugins/actions/server/plugin.ts
View file

@ -32,7 +32,7 @@ import {
import { LicensingPluginSetup, LicensingPluginStart } from '@kbn/licensing-plugin/server';
import { SpacesPluginStart, SpacesPluginSetup } from '@kbn/spaces-plugin/server';
import { PluginSetupContract as FeaturesPluginSetup } from '@kbn/features-plugin/server';
import { SecurityPluginSetup } from '@kbn/security-plugin/server';
import { SecurityPluginSetup, SecurityPluginStart } from '@kbn/security-plugin/server';
import {
IEventLogClientService,
IEventLogger,
@ -176,6 +176,7 @@ export interface ActionsPluginsStart {
licensing: LicensingPluginStart;
eventLog: IEventLogClientService;
spaces?: SpacesPluginStart;
security?: SecurityPluginStart;
}
const includedHiddenTypes = [
@ -488,6 +489,7 @@ export class ActionsPlugin implements Plugin<PluginSetupContract, PluginStartCon
logger,
eventLogger: this.eventLogger!,
spaces: plugins.spaces?.spacesService,
security: plugins.security,
getServices: this.getServicesFactory(
getScopedSavedObjectsClientWithoutAccessToActions,
core.elasticsearch,

2
x-pack/plugins/event_log/generated/README.md
View file

@ -1,7 +1,7 @@
# Generating event schema
The files in this directory were generated by manually running the script
`../scripts/create-schemas.js` from the root directory of the repository.
`../scripts/create_schemas.js` from the root directory of the repository.
**These files should not be edited by hand.**

21
x-pack/plugins/event_log/generated/mappings.json
View file

@ -228,6 +228,10 @@
},
"ignore_above": 1024,
"type": "keyword"
},
"id": {
"ignore_above": 1024,
"type": "keyword"
}
}
},
@ -466,6 +470,23 @@
"uuid": {
"ignore_above": 1024,
"type": "keyword"
},
"gen_ai": {
"properties": {
"usage": {
"properties": {
"prompt_tokens": {
"type": "long"
},
"completion_tokens": {
"type": "long"
},
"total_tokens": {
"type": "long"
}
}
}
}
}
}
}

12
x-pack/plugins/event_log/generated/schemas.ts
View file

@ -97,6 +97,7 @@ export const EventSchema = schema.maybe(
user: schema.maybe(
schema.object({
name: ecsString(),
id: ecsString(),
})
),
kibana: schema.maybe(
@ -206,6 +207,17 @@ export const EventSchema = schema.maybe(
schema.object({
source: ecsString(),
uuid: ecsString(),
gen_ai: schema.maybe(
schema.object({
usage: schema.maybe(
schema.object({
prompt_tokens: ecsStringOrNumber(),
completion_tokens: ecsStringOrNumber(),
total_tokens: ecsStringOrNumber(),
})
),
})
),
})
),
})

18
x-pack/plugins/event_log/scripts/mappings.js
View file

@ -246,6 +246,23 @@ exports.EcsCustomPropertyMappings = {
ignore_above: 1024,
type: 'keyword',
},
gen_ai: {
properties: {
usage: {
properties: {
prompt_tokens: {
type: 'long',
},
completion_tokens: {
type: 'long',
},
total_tokens: {
type: 'long',
},
},
},
},
},
},
},
},
@ -269,6 +286,7 @@ exports.EcsPropertiesToGenerate = [
'log.logger',
'rule',
'user.name',
'user.id',
'kibana',
];