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

[Security solution] Bedrock token tracking - dashboard link added to connector UI (#172115)

Browse source
This commit is contained in:
Steph Milovic 2023-12-01 11:53:45 -07:00 • committed by GitHub
parent 33c74aeb03
commit fed4af1f31
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
31 changed files with 536 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

42
x-pack/plugins/actions/server/lib/gen_ai_token_tracking.test.ts
View file

@ -67,7 +67,7 @@ describe('getGenAiTokenTracking', () => {
expect(logger.error).not.toHaveBeenCalled();
});
it('should return the total, prompt, and completion token counts when given a valid Bedrock response', async () => {
it('should return the total, prompt, and completion token counts when given a valid Bedrock response for run/test subactions', async () => {
const actionTypeId = '.bedrock';
const result = {
@ -103,6 +103,46 @@ describe('getGenAiTokenTracking', () => {
});
});
it('should return the total, prompt, and completion token counts when given a valid Bedrock response for invokeAI subaction', async () => {
const actionTypeId = '.bedrock';
const result = {
actionId: '123',
status: 'ok' as const,
data: {
message: 'Sample completion',
},
};
const validatedParams = {
subAction: 'invokeAI',
subActionParams: {
messages: [
{
role: 'user',
content: 'Sample message',
},
],
},
};
const tokenTracking = await getGenAiTokenTracking({
actionTypeId,
logger,
result,
validatedParams,
});
expect(tokenTracking).toEqual({
total_tokens: 100,
prompt_tokens: 50,
completion_tokens: 50,
});
expect(logger.error).not.toHaveBeenCalled();
expect(mockGetTokenCountFromBedrockInvoke).toHaveBeenCalledWith({
response: 'Sample completion',
body: '{"prompt":"Sample message"}',
});
});
it('should return the total, prompt, and completion token counts when given a valid OpenAI streamed response', async () => {
const mockReader = new IncomingMessage(new Socket());
const actionTypeId = '.gen-ai';

26
x-pack/plugins/actions/server/lib/gen_ai_token_tracking.ts
View file

@ -117,6 +117,32 @@ export const getGenAiTokenTracking = async ({
logger.error(e);
}
}
// this is a non-streamed Bedrock response used by security solution
if (actionTypeId === '.bedrock' && validatedParams.subAction === 'invokeAI') {
try {
const { total, prompt, completion } = await getTokenCountFromBedrockInvoke({
response: (
result.data as unknown as {
message: string;
}
).message,
body: JSON.stringify({
prompt: (validatedParams as { subActionParams: { messages: Array<{ content: string }> } })
.subActionParams.messages[0].content,
}),
});
return {
total_tokens: total,
prompt_tokens: prompt,
completion_tokens: completion,
};
} catch (e) {
logger.error('Failed to calculate tokens from Bedrock invoke response');
logger.error(e);
}
}
return null;
};

1
x-pack/plugins/stack_connectors/common/bedrock/constants.ts
View file

@ -18,6 +18,7 @@ export enum SUB_ACTION {
RUN = 'run',
INVOKE_AI = 'invokeAI',
INVOKE_STREAM = 'invokeStream',
DASHBOARD = 'getDashboard',
TEST = 'test',
}

9
x-pack/plugins/stack_connectors/common/bedrock/schema.ts
View file

@ -52,3 +52,12 @@ export const RunActionResponseSchema = schema.object(
);
export const StreamingResponseSchema = schema.any();
// Run action schema
export const DashboardActionParamsSchema = schema.object({
dashboardId: schema.string(),
});
export const DashboardActionResponseSchema = schema.object({
available: schema.boolean(),
});

4
x-pack/plugins/stack_connectors/common/bedrock/types.ts
View file

@ -8,6 +8,8 @@
import { TypeOf } from '@kbn/config-schema';
import {
ConfigSchema,
DashboardActionParamsSchema,
DashboardActionResponseSchema,
SecretsSchema,
RunActionParamsSchema,
RunActionResponseSchema,
@ -25,3 +27,5 @@ export type InvokeAIActionResponse = TypeOf<typeof InvokeAIActionResponseSchema>
export type RunActionResponse = TypeOf<typeof RunActionResponseSchema>;
export type StreamActionParams = TypeOf<typeof StreamActionParamsSchema>;
export type StreamingResponse = TypeOf<typeof StreamingResponseSchema>;
export type DashboardActionParams = TypeOf<typeof DashboardActionParamsSchema>;
export type DashboardActionResponse = TypeOf<typeof DashboardActionResponseSchema>;

1
x-pack/plugins/stack_connectors/public/connector_types/bedrock/bedrock.tsx
View file

@ -57,5 +57,6 @@ export function getConnectorType(): BedrockConnector {
},
actionConnectorFields: lazy(() => import('./connector')),
actionParamsFields: lazy(() => import('./params')),
actionReadOnlyExtraComponent: lazy(() => import('./dashboard_link')),
};
}

52
x-pack/plugins/stack_connectors/public/connector_types/bedrock/connector.test.tsx
View file

@ -8,13 +8,17 @@
import React from 'react';
import BedrockConnectorFields from './connector';
import { ConnectorFormTestProvider } from '../lib/test_utils';
import { act, render, waitFor } from '@testing-library/react';
import { act, fireEvent, render, waitFor } from '@testing-library/react';
import userEvent from '@testing-library/user-event';
import { useKibana } from '@kbn/triggers-actions-ui-plugin/public';
import { DEFAULT_BEDROCK_MODEL } from '../../../common/bedrock/constants';
import { useGetDashboard } from '../lib/gen_ai/use_get_dashboard';
jest.mock('@kbn/triggers-actions-ui-plugin/public/common/lib/kibana');
jest.mock('../lib/gen_ai/use_get_dashboard');
const useKibanaMock = useKibana as jest.Mocked<typeof useKibana>;
const mockDashboard = useGetDashboard as jest.Mock;
const bedrockConnector = {
actionTypeId: '.bedrock',
name: 'bedrock',
@ -36,6 +40,9 @@ describe('BedrockConnectorFields renders', () => {
beforeEach(() => {
jest.clearAllMocks();
useKibanaMock().services.application.navigateToUrl = navigateToUrl;
mockDashboard.mockImplementation(({ connectorId }) => ({
dashboardUrl: `https://dashboardurl.com/${connectorId}`,
}));
});
test('Bedrock connector fields are rendered', async () => {
const { getAllByTestId } = render(
@ -58,6 +65,49 @@ describe('BedrockConnectorFields renders', () => {
expect(getAllByTestId('bedrock-api-model-doc')[0]).toBeInTheDocument();
});
describe('Dashboard link', () => {
it('Does not render if isEdit is false and dashboardUrl is defined', async () => {
const { queryByTestId } = render(
<ConnectorFormTestProvider connector={bedrockConnector}>
<BedrockConnectorFields
readOnly={false}
isEdit={false}
registerPreSubmitValidator={() => {}}
/>
</ConnectorFormTestProvider>
);
expect(queryByTestId('link-gen-ai-token-dashboard')).not.toBeInTheDocument();
});
it('Does not render if isEdit is true and dashboardUrl is null', async () => {
mockDashboard.mockImplementation((id: string) => ({
dashboardUrl: null,
}));
const { queryByTestId } = render(
<ConnectorFormTestProvider connector={bedrockConnector}>
<BedrockConnectorFields readOnly={false} isEdit registerPreSubmitValidator={() => {}} />
</ConnectorFormTestProvider>
);
expect(queryByTestId('link-gen-ai-token-dashboard')).not.toBeInTheDocument();
});
it('Renders if isEdit is true and dashboardUrl is defined', async () => {
const { getByTestId } = render(
<ConnectorFormTestProvider connector={bedrockConnector}>
<BedrockConnectorFields readOnly={false} isEdit registerPreSubmitValidator={() => {}} />
</ConnectorFormTestProvider>
);
expect(getByTestId('link-gen-ai-token-dashboard')).toBeInTheDocument();
});
it('On click triggers redirect with correct saved object id', async () => {
const { getByTestId } = render(
<ConnectorFormTestProvider connector={bedrockConnector}>
<BedrockConnectorFields readOnly={false} isEdit registerPreSubmitValidator={() => {}} />
</ConnectorFormTestProvider>
);
fireEvent.click(getByTestId('link-gen-ai-token-dashboard'));
expect(navigateToUrl).toHaveBeenCalledWith(`https://dashboardurl.com/123`);
});
});
describe('Validation', () => {
const onSubmit = jest.fn();

19
x-pack/plugins/stack_connectors/public/connector_types/bedrock/connector.tsx
View file

@ -10,16 +10,23 @@ import {
ActionConnectorFieldsProps,
SimpleConnectorForm,
} from '@kbn/triggers-actions-ui-plugin/public';
import { useFormData } from '@kbn/es-ui-shared-plugin/static/forms/hook_form_lib';
import DashboardLink from './dashboard_link';
import { BEDROCK } from './translations';
import { bedrockConfig, bedrockSecrets } from './constants';
const BedrockConnectorFields: React.FC<ActionConnectorFieldsProps> = ({ readOnly, isEdit }) => {
const [{ id, name }] = useFormData();
return (
<SimpleConnectorForm
isEdit={isEdit}
readOnly={readOnly}
configFormSchema={bedrockConfig}
secretsFormSchema={bedrockSecrets}
/>
<>
<SimpleConnectorForm
isEdit={isEdit}
readOnly={readOnly}
configFormSchema={bedrockConfig}
secretsFormSchema={bedrockSecrets}
/>
{isEdit && <DashboardLink connectorId={id} connectorName={name} selectedProvider={BEDROCK} />}
</>
);
};

51
x-pack/plugins/stack_connectors/public/connector_types/bedrock/dashboard_link.tsx Normal file
View file

@ -0,0 +1,51 @@
/*
* 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 React, { useCallback } from 'react';
import { EuiLink } from '@elastic/eui';
import { useKibana } from '@kbn/triggers-actions-ui-plugin/public';
import * as i18n from './translations';
import { useGetDashboard } from '../lib/gen_ai/use_get_dashboard';
interface Props {
connectorId: string;
connectorName: string;
selectedProvider?: string;
}
// tested from ./connector.test.tsx
export const DashboardLink: React.FC<Props> = ({
connectorId,
connectorName,
selectedProvider = 'Bedrock',
}) => {
const { dashboardUrl } = useGetDashboard({ connectorId, selectedProvider });
const {
services: {
application: { navigateToUrl },
},
} = useKibana();
const onClick = useCallback(
(e) => {
e.preventDefault();
if (dashboardUrl) {
navigateToUrl(dashboardUrl);
}
},
[dashboardUrl, navigateToUrl]
);
return dashboardUrl != null ? (
// href gives us right click -> open in new tab
// onclick prevents page reload
// eslint-disable-next-line @elastic/eui/href-or-on-click
<EuiLink data-test-subj="link-gen-ai-token-dashboard" onClick={onClick} href={dashboardUrl}>
{i18n.USAGE_DASHBOARD_LINK(selectedProvider, connectorName)}
</EuiLink>
) : null;
};
// eslint-disable-next-line import/no-default-export
export { DashboardLink as default };

6
x-pack/plugins/stack_connectors/public/connector_types/bedrock/translations.ts
View file

@ -97,3 +97,9 @@ export const BODY_DESCRIPTION = i18n.translate(
export const MODEL = i18n.translate('xpack.stackConnectors.components.bedrock.model', {
defaultMessage: 'Model',
});
export const USAGE_DASHBOARD_LINK = (apiProvider: string, connectorName: string) =>
i18n.translate('xpack.stackConnectors.components.genAi.dashboardLink', {
values: { apiProvider, connectorName },
defaultMessage: 'View {apiProvider} Usage Dashboard for "{ connectorName }" Connector',
});

2
x-pack/plugins/stack_connectors/public/connector_types/openai/api.test.ts → x-pack/plugins/stack_connectors/public/connector_types/lib/gen_ai/api.test.ts
View file

@ -7,7 +7,7 @@
import { httpServiceMock } from '@kbn/core-http-browser-mocks';
import { getDashboard } from './api';
import { SUB_ACTION } from '../../../common/openai/constants';
import { SUB_ACTION } from '../../../../common/openai/constants';
const response = {
available: true,
};

4
x-pack/plugins/stack_connectors/public/connector_types/openai/api.ts → x-pack/plugins/stack_connectors/public/connector_types/lib/gen_ai/api.ts
View file

@ -7,8 +7,8 @@
import { HttpSetup } from '@kbn/core-http-browser';
import { ActionTypeExecutorResult, BASE_ACTION_API_PATH } from '@kbn/actions-plugin/common';
import { SUB_ACTION } from '../../../common/openai/constants';
import { ConnectorExecutorResult, rewriteResponseToCamelCase } from '../lib/rewrite_response_body';
import { SUB_ACTION } from '../../../../common/openai/constants';
import { ConnectorExecutorResult, rewriteResponseToCamelCase } from '../rewrite_response_body';
export async function getDashboard({
http,

14
x-pack/plugins/stack_connectors/public/connector_types/lib/gen_ai/translations.ts Normal file
View file

@ -0,0 +1,14 @@
/*
* 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 { i18n } from '@kbn/i18n';
export const GET_DASHBOARD_API_ERROR = (apiProvider: string) =>
i18n.translate('xpack.stackConnectors.components.genAi.error.dashboardApiError', {
values: { apiProvider },
defaultMessage: 'Error finding {apiProvider} Token Usage Dashboard.',
});

69
x-pack/plugins/stack_connectors/public/connector_types/openai/use_get_dashboard.test.ts → x-pack/plugins/stack_connectors/public/connector_types/lib/gen_ai/use_get_dashboard.test.ts
View file

@ -39,7 +39,9 @@ const mockServices = {
const mockDashboard = getDashboard as jest.Mock;
const mockKibana = useKibana as jest.Mock;
describe('useGetDashboard_function', () => {
const defaultArgs = { connectorId, selectedProvider: 'OpenAI' };
describe('useGetDashboard', () => {
beforeEach(() => {
jest.clearAllMocks();
mockDashboard.mockResolvedValue({ data: { available: true } });
@ -48,36 +50,45 @@ describe('useGetDashboard_function', () => {
});
});
it('fetches the dashboard and sets the dashboard URL', async () => {
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard({ connectorId }));
await waitForNextUpdate();
expect(mockDashboard).toHaveBeenCalledWith(
expect.objectContaining({
connectorId,
dashboardId: 'generative-ai-token-usage-space',
})
);
expect(mockGetRedirectUrl).toHaveBeenCalledWith({
query: {
language: 'kuery',
query: `kibana.saved_objects: { id : ${connectorId} }`,
},
dashboardId: 'generative-ai-token-usage-space',
});
expect(result.current.isLoading).toBe(false);
expect(result.current.dashboardUrl).toBe(
'http://localhost:5601/app/dashboards#/view/generative-ai-token-usage-space'
);
});
it.each([
['Azure OpenAI', 'openai'],
['OpenAI', 'openai'],
['Bedrock', 'bedrock'],
])(
'fetches the %p dashboard and sets the dashboard URL with %p',
async (selectedProvider, urlKey) => {
const { result, waitForNextUpdate } = renderHook(() =>
useGetDashboard({ ...defaultArgs, selectedProvider })
);
await waitForNextUpdate();
expect(mockDashboard).toHaveBeenCalledWith(
expect.objectContaining({
connectorId,
dashboardId: `generative-ai-token-usage-${urlKey}-space`,
})
);
expect(mockGetRedirectUrl).toHaveBeenCalledWith({
query: {
language: 'kuery',
query: `kibana.saved_objects: { id : ${connectorId} }`,
},
dashboardId: `generative-ai-token-usage-${urlKey}-space`,
});
expect(result.current.isLoading).toBe(false);
expect(result.current.dashboardUrl).toBe(
`http://localhost:5601/app/dashboards#/view/generative-ai-token-usage-${urlKey}-space`
);
}
);
it('handles the case where the dashboard is not available.', async () => {
mockDashboard.mockResolvedValue({ data: { available: false } });
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard({ connectorId }));
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard(defaultArgs));
await waitForNextUpdate();
expect(mockDashboard).toHaveBeenCalledWith(
expect.objectContaining({
connectorId,
dashboardId: 'generative-ai-token-usage-space',
dashboardId: 'generative-ai-token-usage-openai-space',
})
);
expect(mockGetRedirectUrl).not.toHaveBeenCalled();
@ -91,7 +102,7 @@ describe('useGetDashboard_function', () => {
services: { ...mockServices, spaces: null },
});
const { result } = renderHook(() => useGetDashboard({ connectorId }));
const { result } = renderHook(() => useGetDashboard(defaultArgs));
expect(mockDashboard).not.toHaveBeenCalled();
expect(mockGetRedirectUrl).not.toHaveBeenCalled();
expect(result.current.isLoading).toBe(false);
@ -99,7 +110,9 @@ describe('useGetDashboard_function', () => {
});
it('handles the case where connectorId is empty string', async () => {
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard({ connectorId: '' }));
const { result, waitForNextUpdate } = renderHook(() =>
useGetDashboard({ ...defaultArgs, connectorId: '' })
);
await waitForNextUpdate();
expect(mockDashboard).not.toHaveBeenCalled();
expect(mockGetRedirectUrl).not.toHaveBeenCalled();
@ -111,7 +124,7 @@ describe('useGetDashboard_function', () => {
mockKibana.mockReturnValue({
services: { ...mockServices, dashboard: {} },
});
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard({ connectorId }));
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard(defaultArgs));
await waitForNextUpdate();
expect(result.current.isLoading).toBe(false);
expect(result.current.dashboardUrl).toBe(null);
@ -119,7 +132,7 @@ describe('useGetDashboard_function', () => {
it('correctly handles errors and displays the appropriate toast messages.', async () => {
mockDashboard.mockRejectedValue(new Error('Error fetching dashboard'));
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard({ connectorId }));
const { result, waitForNextUpdate } = renderHook(() => useGetDashboard(defaultArgs));
await waitForNextUpdate();
expect(result.current.isLoading).toBe(false);
expect(mockToasts.addDanger).toHaveBeenCalledWith({

28
x-pack/plugins/stack_connectors/public/connector_types/openai/use_get_dashboard.ts → x-pack/plugins/stack_connectors/public/connector_types/lib/gen_ai/use_get_dashboard.ts
View file

@ -7,20 +7,20 @@
import { useState, useEffect, useRef, useCallback } from 'react';
import { useKibana } from '@kbn/triggers-actions-ui-plugin/public';
import { getDashboardId } from './constants';
import { getDashboard } from './api';
import * as i18n from './translations';
interface Props {
connectorId: string;
selectedProvider: string;
}
export interface UseGetDashboard {
dashboardUrl: string | null;
isLoading: boolean;
}
export const useGetDashboard = ({ connectorId }: Props): UseGetDashboard => {
export const useGetDashboard = ({ connectorId, selectedProvider }: Props): UseGetDashboard => {
const {
dashboard,
http,
@ -84,7 +84,7 @@ export const useGetDashboard = ({ connectorId }: Props): UseGetDashboard => {
if (res.status && res.status === 'error') {
toasts.addDanger({
title: i18n.GET_DASHBOARD_API_ERROR,
title: i18n.GET_DASHBOARD_API_ERROR(selectedProvider),
text: `${res.serviceMessage ?? res.message}`,
});
}
@ -94,7 +94,7 @@ export const useGetDashboard = ({ connectorId }: Props): UseGetDashboard => {
setDashboardCheckComplete(true);
setIsLoading(false);
toasts.addDanger({
title: i18n.GET_DASHBOARD_API_ERROR,
title: i18n.GET_DASHBOARD_API_ERROR(selectedProvider),
text: error.message,
});
}
@ -103,7 +103,7 @@ export const useGetDashboard = ({ connectorId }: Props): UseGetDashboard => {
if (spaceId != null && connectorId.length > 0 && !dashboardCheckComplete) {
abortCtrl.current.abort();
fetchData(getDashboardId(spaceId));
fetchData(getDashboardId(selectedProvider, spaceId));
}
return () => {
@ -111,10 +111,24 @@ export const useGetDashboard = ({ connectorId }: Props): UseGetDashboard => {
setIsLoading(false);
abortCtrl.current.abort();
};
}, [connectorId, dashboardCheckComplete, dashboardUrl, http, setUrl, spaceId, toasts]);
}, [
connectorId,
dashboardCheckComplete,
dashboardUrl,
http,
selectedProvider,
setUrl,
spaceId,
toasts,
]);
return {
isLoading,
dashboardUrl,
};
};
const getDashboardId = (selectedProvider: string, spaceId: string): string =>
`generative-ai-token-usage-${
selectedProvider.toLowerCase().includes('openai') ? 'openai' : 'bedrock'
}-${spaceId}`;

6
x-pack/plugins/stack_connectors/public/connector_types/openai/connector.test.tsx
View file

@ -12,10 +12,10 @@ import { act, fireEvent, render, waitFor } from '@testing-library/react';
import userEvent from '@testing-library/user-event';
import { DEFAULT_OPENAI_MODEL, OpenAiProviderType } from '../../../common/openai/constants';
import { useKibana } from '@kbn/triggers-actions-ui-plugin/public';
import { useGetDashboard } from './use_get_dashboard';
import { useGetDashboard } from '../lib/gen_ai/use_get_dashboard';
jest.mock('@kbn/triggers-actions-ui-plugin/public/common/lib/kibana');
jest.mock('./use_get_dashboard');
jest.mock('../lib/gen_ai/use_get_dashboard');
const useKibanaMock = useKibana as jest.Mocked<typeof useKibana>;
const mockDashboard = useGetDashboard as jest.Mock;
@ -101,7 +101,7 @@ describe('ConnectorFields renders', () => {
}));
const { queryByTestId } = render(
<ConnectorFormTestProvider connector={openAiConnector}>
<ConnectorFields readOnly={false} isEdit={false} registerPreSubmitValidator={() => {}} />
<ConnectorFields readOnly={false} isEdit registerPreSubmitValidator={() => {}} />
</ConnectorFormTestProvider>
);
expect(queryByTestId('link-gen-ai-token-dashboard')).not.toBeInTheDocument();

2
x-pack/plugins/stack_connectors/public/connector_types/openai/constants.tsx
View file

@ -154,5 +154,3 @@ export const providerOptions = [
label: i18n.AZURE_AI,
},
];
export const getDashboardId = (spaceId: string): string => `generative-ai-token-usage-${spaceId}`;

11
x-pack/plugins/stack_connectors/public/connector_types/openai/dashboard_link.tsx
View file

@ -9,7 +9,7 @@ import React, { useCallback } from 'react';
import { EuiLink } from '@elastic/eui';
import { useKibana } from '@kbn/triggers-actions-ui-plugin/public';
import * as i18n from './translations';
import { useGetDashboard } from './use_get_dashboard';
import { useGetDashboard } from '../lib/gen_ai/use_get_dashboard';
interface Props {
connectorId: string;
@ -20,9 +20,9 @@ interface Props {
export const DashboardLink: React.FC<Props> = ({
connectorId,
connectorName,
selectedProvider = '',
selectedProvider = 'OpenAI',
}) => {
const { dashboardUrl } = useGetDashboard({ connectorId });
const { dashboardUrl } = useGetDashboard({ connectorId, selectedProvider });
const {
services: {
application: { navigateToUrl },
@ -38,7 +38,10 @@ export const DashboardLink: React.FC<Props> = ({
[dashboardUrl, navigateToUrl]
);
return dashboardUrl != null ? (
<EuiLink data-test-subj="link-gen-ai-token-dashboard" onClick={onClick}>
// href gives us right click -> open in new tab
// onclick prevents page reload
// eslint-disable-next-line @elastic/eui/href-or-on-click
<EuiLink data-test-subj="link-gen-ai-token-dashboard" onClick={onClick} href={dashboardUrl}>
{i18n.USAGE_DASHBOARD_LINK(selectedProvider, connectorName)}
</EuiLink>
) : null;

7
x-pack/plugins/stack_connectors/public/connector_types/openai/translations.ts
View file

@ -100,10 +100,3 @@ export const USAGE_DASHBOARD_LINK = (apiProvider: string, connectorName: string)
values: { apiProvider, connectorName },
defaultMessage: 'View {apiProvider} Usage Dashboard for "{ connectorName }" Connector',
});
export const GET_DASHBOARD_API_ERROR = i18n.translate(
'xpack.stackConnectors.components.genAi.error.dashboardApiError',
{
defaultMessage: 'Error finding OpenAI Token Usage Dashboard.',
}
);

93
x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.test.ts
View file

@ -18,7 +18,9 @@ import {
DEFAULT_TOKEN_LIMIT,
} from '../../../common/bedrock/constants';
import { DEFAULT_BODY } from '../../../public/connector_types/bedrock/constants';
import { initDashboard } from '../lib/gen_ai/create_gen_ai_dashboard';
import { AxiosError } from 'axios';
jest.mock('../lib/gen_ai/create_gen_ai_dashboard');
// @ts-ignore
const mockSigner = jest.spyOn(aws, 'sign').mockReturnValue({ signed: true });
@ -41,18 +43,19 @@ describe('BedrockConnector', () => {
});
});
const connector = new BedrockConnector({
configurationUtilities: actionsConfigMock.create(),
connector: { id: '1', type: BEDROCK_CONNECTOR_ID },
config: {
apiUrl: DEFAULT_BEDROCK_URL,
defaultModel: DEFAULT_BEDROCK_MODEL,
},
secrets: { accessKey: '123', secret: 'secret' },
logger: loggingSystemMock.createLogger(),
services: actionsMock.createServices(),
});
describe('Bedrock', () => {
const connector = new BedrockConnector({
configurationUtilities: actionsConfigMock.create(),
connector: { id: '1', type: BEDROCK_CONNECTOR_ID },
config: {
apiUrl: DEFAULT_BEDROCK_URL,
defaultModel: DEFAULT_BEDROCK_MODEL,
},
secrets: { accessKey: '123', secret: 'secret' },
logger: loggingSystemMock.createLogger(),
services: actionsMock.createServices(),
});
beforeEach(() => {
// @ts-ignore
connector.request = mockRequest;
@ -335,6 +338,74 @@ describe('BedrockConnector', () => {
});
});
});
describe('Token dashboard', () => {
const mockGenAi = initDashboard as jest.Mock;
beforeEach(() => {
// @ts-ignore
connector.esClient.transport.request = mockRequest;
mockRequest.mockResolvedValue({ has_all_requested: true });
mockGenAi.mockResolvedValue({ success: true });
jest.clearAllMocks();
});
it('the create dashboard API call returns available: true when user has correct permissions', async () => {
const response = await connector.getDashboard({ dashboardId: '123' });
expect(mockRequest).toBeCalledTimes(1);
expect(mockRequest).toHaveBeenCalledWith({
path: '/_security/user/_has_privileges',
method: 'POST',
body: {
index: [
{
names: ['.kibana-event-log-*'],
allow_restricted_indices: true,
privileges: ['read'],
},
],
},
});
expect(response).toEqual({ available: true });
});
it('the create dashboard API call returns available: false when user has correct permissions', async () => {
mockRequest.mockResolvedValue({ has_all_requested: false });
const response = await connector.getDashboard({ dashboardId: '123' });
expect(mockRequest).toBeCalledTimes(1);
expect(mockRequest).toHaveBeenCalledWith({
path: '/_security/user/_has_privileges',
method: 'POST',
body: {
index: [
{
names: ['.kibana-event-log-*'],
allow_restricted_indices: true,
privileges: ['read'],
},
],
},
});
expect(response).toEqual({ available: false });
});
it('the create dashboard API call returns available: false when init dashboard fails', async () => {
mockGenAi.mockResolvedValue({ success: false });
const response = await connector.getDashboard({ dashboardId: '123' });
expect(mockRequest).toBeCalledTimes(1);
expect(mockRequest).toHaveBeenCalledWith({
path: '/_security/user/_has_privileges',
method: 'POST',
body: {
index: [
{
names: ['.kibana-event-log-*'],
allow_restricted_indices: true,
privileges: ['read'],
},
],
},
});
expect(response).toEqual({ available: false });
});
});
});
function createStreamMock() {

53
x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
View file

@ -10,6 +10,7 @@ import aws from 'aws4';
import type { AxiosError } from 'axios';
import { IncomingMessage } from 'http';
import { PassThrough } from 'stream';
import { initDashboard } from '../lib/gen_ai/create_gen_ai_dashboard';
import {
RunActionParamsSchema,
RunActionResponseSchema,
@ -26,7 +27,12 @@ import type {
StreamActionParams,
} from '../../../common/bedrock/types';
import { SUB_ACTION, DEFAULT_TOKEN_LIMIT } from '../../../common/bedrock/constants';
import { StreamingResponse } from '../../../common/bedrock/types';
import {
DashboardActionParams,
DashboardActionResponse,
StreamingResponse,
} from '../../../common/bedrock/types';
import { DashboardActionParamsSchema } from '../../../common/bedrock/schema';
interface SignedRequest {
host: string;
@ -55,6 +61,12 @@ export class BedrockConnector extends SubActionConnector<Config, Secrets> {
schema: RunActionParamsSchema,
});
this.registerSubAction({
name: SUB_ACTION.DASHBOARD,
method: 'getDashboard',
schema: DashboardActionParamsSchema,
});
this.registerSubAction({
name: SUB_ACTION.TEST,
method: 'runApi',
@ -119,6 +131,42 @@ export class BedrockConnector extends SubActionConnector<Config, Secrets> {
) as SignedRequest;
}
/**
* retrieves a dashboard from the Kibana server and checks if the
* user has the necessary privileges to access it.
* @param dashboardId The ID of the dashboard to retrieve.
*/
public async getDashboard({
dashboardId,
}: DashboardActionParams): Promise<DashboardActionResponse> {
const privilege = (await this.esClient.transport.request({
path: '/_security/user/_has_privileges',
method: 'POST',
body: {
index: [
{
names: ['.kibana-event-log-*'],
allow_restricted_indices: true,
privileges: ['read'],
},
],
},
})) as { has_all_requested: boolean };
if (!privilege?.has_all_requested) {
return { available: false };
}
const response = await initDashboard({
logger: this.logger,
savedObjectsClient: this.savedObjectsClient,
dashboardId,
genAIProvider: 'Bedrock',
});
return { available: response.success };
}
/**
* responsible for making a POST request to the external API endpoint and returning the response data
* @param body The stringified request body to be sent in the POST request.
@ -186,9 +234,8 @@ export class BedrockConnector extends SubActionConnector<Config, Secrets> {
/**
* Deprecated. Use invokeStream instead.
* TODO: remove before 8.12 FF in part 3 of streaming work for security solution
* TODO: remove once streaming work is implemented in langchain mode for security solution
* tracked here: https://github.com/elastic/security-team/issues/7363
* No token tracking implemented for this method
*/
public async invokeAI({
messages,

17
x-pack/plugins/stack_connectors/server/connector_types/openai/create_dashboard.test.ts → x-pack/plugins/stack_connectors/server/connector_types/lib/gen_ai/create_gen_ai_dashboard.test.ts
View file

@ -5,8 +5,8 @@
* 2.0.
*/
import { initDashboard } from './create_dashboard';
import { getDashboard } from './dashboard';
import { initDashboard } from './create_gen_ai_dashboard';
import { getDashboard } from './gen_ai_dashboard';
import { savedObjectsClientMock } from '@kbn/core-saved-objects-api-server-mocks';
import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
import { Logger } from '@kbn/logging';
@ -16,22 +16,22 @@ jest.mock('uuid', () => ({
}));
const logger = loggingSystemMock.create().get() as jest.Mocked<Logger>;
const dashboardId = 'test-dashboard-id';
const savedObjectsClient = savedObjectsClientMock.create();
const defaultArgs = { logger, savedObjectsClient, dashboardId, genAIProvider: 'OpenAI' as const };
describe('createDashboard', () => {
beforeEach(() => {
jest.clearAllMocks();
});
it('fetches the Gen Ai Dashboard saved object', async () => {
const dashboardId = 'test-dashboard-id';
const result = await initDashboard({ logger, savedObjectsClient, dashboardId });
const result = await initDashboard(defaultArgs);
expect(result.success).toBe(true);
expect(logger.error).not.toHaveBeenCalled();
expect(savedObjectsClient.get).toHaveBeenCalledWith('dashboard', dashboardId);
});
it('creates the Gen Ai Dashboard saved object when the dashboard saved object does not exist', async () => {
const dashboardId = 'test-dashboard-id';
const soClient = {
...savedObjectsClient,
get: jest.fn().mockRejectedValue({
@ -46,12 +46,12 @@ describe('createDashboard', () => {
},
}),
};
const result = await initDashboard({ logger, savedObjectsClient: soClient, dashboardId });
const result = await initDashboard({ ...defaultArgs, savedObjectsClient: soClient });
expect(soClient.get).toHaveBeenCalledWith('dashboard', dashboardId);
expect(soClient.create).toHaveBeenCalledWith(
'dashboard',
getDashboard(dashboardId).attributes,
getDashboard(defaultArgs.genAIProvider, dashboardId).attributes,
{ overwrite: true, id: dashboardId }
);
expect(result.success).toBe(true);
@ -72,8 +72,7 @@ describe('createDashboard', () => {
},
}),
};
const dashboardId = 'test-dashboard-id';
const result = await initDashboard({ logger, savedObjectsClient: soClient, dashboardId });
const result = await initDashboard({ ...defaultArgs, savedObjectsClient: soClient });
expect(result.success).toBe(false);
expect(result.error?.message).toBe('Internal Server Error: Error happened');
expect(result.error?.statusCode).toBe(500);

8
x-pack/plugins/stack_connectors/server/connector_types/openai/create_dashboard.ts → x-pack/plugins/stack_connectors/server/connector_types/lib/gen_ai/create_gen_ai_dashboard.ts
View file

@ -8,7 +8,7 @@ import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-ser
import { DashboardAttributes } from '@kbn/dashboard-plugin/common';
import { Logger } from '@kbn/logging';
import { getDashboard } from './dashboard';
import { getDashboard } from './gen_ai_dashboard';
export interface OutputError {
message: string;
@ -19,10 +19,12 @@ export const initDashboard = async ({
logger,
savedObjectsClient,
dashboardId,
genAIProvider,
}: {
logger: Logger;
savedObjectsClient: SavedObjectsClientContract;
dashboardId: string;
genAIProvider: 'OpenAI' | 'Bedrock';
}): Promise<{
success: boolean;
error?: OutputError;
@ -50,13 +52,13 @@ export const initDashboard = async ({
try {
await savedObjectsClient.create<DashboardAttributes>(
'dashboard',
getDashboard(dashboardId).attributes,
getDashboard(genAIProvider, dashboardId).attributes,
{
overwrite: true,
id: dashboardId,
}
);
logger.info(`Successfully created Gen Ai Dashboard ${dashboardId}`);
logger.info(`Successfully created Generative AI Dashboard ${dashboardId}`);
return { success: true };
} catch (error) {
return {

41
x-pack/plugins/stack_connectors/server/connector_types/openai/dashboard.ts → x-pack/plugins/stack_connectors/server/connector_types/lib/gen_ai/gen_ai_dashboard.ts
View file

@ -8,10 +8,28 @@
import { DashboardAttributes } from '@kbn/dashboard-plugin/common';
import { v4 as uuidv4 } from 'uuid';
import { SavedObject } from '@kbn/core-saved-objects-common/src/server_types';
import { OPENAI_TITLE, OPENAI_CONNECTOR_ID } from '../../../../common/openai/constants';
import { BEDROCK_TITLE, BEDROCK_CONNECTOR_ID } from '../../../../common/bedrock/constants';
export const dashboardTitle = `OpenAI Token Usage`;
const getDashboardTitle = (title: string) => `${title} Token Usage`;
export const getDashboard = (
genAIProvider: 'OpenAI' | 'Bedrock',
dashboardId: string
): SavedObject<DashboardAttributes> => {
const attributes =
genAIProvider === 'OpenAI'
? {
provider: OPENAI_TITLE,
dashboardTitle: getDashboardTitle(OPENAI_TITLE),
actionTypeId: OPENAI_CONNECTOR_ID,
}
: {
provider: BEDROCK_TITLE,
dashboardTitle: getDashboardTitle(BEDROCK_TITLE),
actionTypeId: BEDROCK_CONNECTOR_ID,
};
export const getDashboard = (dashboardId: string): SavedObject<DashboardAttributes> => {
const ids: Record<string, string> = {
genAiSavedObjectId: dashboardId,
tokens: uuidv4(),
@ -20,10 +38,9 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
};
return {
attributes: {
description: 'Displays OpenAI token consumption per Kibana user',
description: `Displays ${attributes.provider} token consumption per Kibana user`,
kibanaSavedObjectMeta: {
searchSourceJSON:
'{"query":{"query":"kibana.saved_objects: { type_id : \\".gen-ai\\" } ","language":"kuery"},"filter":[]}',
searchSourceJSON: `{"query":{"query":"kibana.saved_objects: { type_id : \\"${attributes.actionTypeId}\\" } ","language":"kuery"},"filter":[]}`,
},
optionsJSON:
'{"useMargins":true,"syncColors":false,"syncCursor":true,"syncTooltips":false,"hidePanelTitles":false}',
@ -125,7 +142,7 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
yLeft: 0,
yRight: 0,
},
yTitle: 'Sum of OpenAI Completion + Prompt Tokens',
yTitle: `Sum of ${attributes.provider} Completion + Prompt Tokens`,
axisTitlesVisibilitySettings: {
x: true,
yLeft: true,
@ -133,7 +150,7 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
},
},
query: {
query: 'kibana.saved_objects:{ type_id: ".gen-ai" }',
query: `kibana.saved_objects:{ type_id: "${attributes.actionTypeId}" }`,
language: 'kuery',
},
filters: [],
@ -143,7 +160,7 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
'475e8ca0-e78e-454a-8597-a5492f70dce3': {
columns: {
'0f9814ec-0964-4efa-93a3-c7f173df2483': {
label: 'OpenAI Completion Tokens',
label: `${attributes.provider} Completion Tokens`,
dataType: 'number',
operationType: 'sum',
sourceField: 'kibana.action.execution.gen_ai.usage.completion_tokens',
@ -192,7 +209,7 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
customLabel: true,
},
'b0e390e4-d754-4eb4-9fcc-4347dadda394': {
label: 'OpenAI Prompt Tokens',
label: `${attributes.provider} Prompt Tokens`,
dataType: 'number',
operationType: 'sum',
sourceField: 'kibana.action.execution.gen_ai.usage.prompt_tokens',
@ -298,7 +315,7 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
],
},
query: {
query: 'kibana.saved_objects: { type_id : ".gen-ai" } ',
query: `kibana.saved_objects: { type_id : "${attributes.actionTypeId}" } `,
language: 'kuery',
},
filters: [],
@ -334,7 +351,7 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
customLabel: true,
},
'b0e390e4-d754-4eb4-9fcc-4347dadda394': {
label: 'Sum of OpenAI Total Tokens',
label: `Sum of ${attributes.provider} Total Tokens`,
dataType: 'number',
operationType: 'sum',
sourceField: 'kibana.action.execution.gen_ai.usage.total_tokens',
@ -392,7 +409,7 @@ export const getDashboard = (dashboardId: string): SavedObject<DashboardAttribut
},
]),
timeRestore: false,
title: dashboardTitle,
title: attributes.dashboardTitle,
version: 1,
},
coreMigrationVersion: '8.8.0',

4
x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts
View file

@ -16,9 +16,9 @@ import {
import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
import { actionsMock } from '@kbn/actions-plugin/server/mocks';
import { RunActionResponseSchema, StreamingResponseSchema } from '../../../common/openai/schema';
import { initDashboard } from './create_dashboard';
import { initDashboard } from '../lib/gen_ai/create_gen_ai_dashboard';
import { PassThrough, Transform } from 'stream';
jest.mock('./create_dashboard');
jest.mock('../lib/gen_ai/create_gen_ai_dashboard');
describe('OpenAIConnector', () => {
let mockRequest: jest.Mock;

5
x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts
View file

@ -31,7 +31,7 @@ import {
InvokeAIActionParams,
InvokeAIActionResponse,
} from '../../../common/openai/types';
import { initDashboard } from './create_dashboard';
import { initDashboard } from '../lib/gen_ai/create_gen_ai_dashboard';
import {
getAxiosOptions,
getRequestWithStreamOption,
@ -187,6 +187,7 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
logger: this.logger,
savedObjectsClient: this.savedObjectsClient,
dashboardId,
genAIProvider: 'OpenAI',
});
return { available: response.success };
@ -209,7 +210,7 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
/**
* Deprecated. Use invokeStream instead.
* TODO: remove before 8.12 FF in part 3 of streaming work for security solution
* TODO: remove once streaming work is implemented in langchain mode for security solution
* tracked here: https://github.com/elastic/security-team/issues/7363
*/
public async invokeAI(body: InvokeAIActionParams): Promise<InvokeAIActionResponse> {

1
x-pack/plugins/translations/translations/fr-FR.json
View file

@ -38619,7 +38619,6 @@
"xpack.stackConnectors.components.genAi.defaultModelTextFieldLabel": "Modèle par défaut",
"xpack.stackConnectors.components.genAi.defaultModelTooltipContent": "Si une requête ne comprend pas de modèle, le modèle par défaut est utilisé.",
"xpack.stackConnectors.components.genAi.documentation": "documentation",
"xpack.stackConnectors.components.genAi.error.dashboardApiError": "Une erreur s'est produite lors de la recherche du tableau de bord de l'utilisation des tokens d'OpenAI.",
"xpack.stackConnectors.components.genAi.error.requiredApiProviderText": "Un fournisseur d’API est nécessaire.",
"xpack.stackConnectors.components.genAi.error.requiredGenerativeAiBodyText": "Le corps est requis.",
"xpack.stackConnectors.components.genAi.openAi": "OpenAI",

1
x-pack/plugins/translations/translations/ja-JP.json
View file

@ -38618,7 +38618,6 @@
"xpack.stackConnectors.components.genAi.defaultModelTextFieldLabel": "デフォルトモデル",
"xpack.stackConnectors.components.genAi.defaultModelTooltipContent": "リクエストにモデルが含まれていない場合、デフォルトが使われます。",
"xpack.stackConnectors.components.genAi.documentation": "ドキュメンテーション",
"xpack.stackConnectors.components.genAi.error.dashboardApiError": "OpenAIトークン使用状況ダッシュボードの検索エラー。",
"xpack.stackConnectors.components.genAi.error.requiredApiProviderText": "APIプロバイダーは必須です。",
"xpack.stackConnectors.components.genAi.error.requiredGenerativeAiBodyText": "本文が必要です。",
"xpack.stackConnectors.components.genAi.openAi": "OpenAI",

1
x-pack/plugins/translations/translations/zh-CN.json
View file

@ -38611,7 +38611,6 @@
"xpack.stackConnectors.components.genAi.defaultModelTextFieldLabel": "默认模型",
"xpack.stackConnectors.components.genAi.defaultModelTooltipContent": "如果请求不包含模型，它将使用默认值。",
"xpack.stackConnectors.components.genAi.documentation": "文档",
"xpack.stackConnectors.components.genAi.error.dashboardApiError": "查找 OpenAI 令牌使用情况仪表板时出错。",
"xpack.stackConnectors.components.genAi.error.requiredApiProviderText": "“API 提供商”必填。",
"xpack.stackConnectors.components.genAi.error.requiredGenerativeAiBodyText": "“正文”必填。",
"xpack.stackConnectors.components.genAi.openAi": "OpenAI",

64
x-pack/test/alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts
View file

@ -33,7 +33,9 @@ const defaultConfig = {
export default function bedrockTest({ getService }: FtrProviderContext) {
const supertest = getService('supertest');
const objectRemover = new ObjectRemover(supertest);
const supertestWithoutAuth = getService('supertestWithoutAuth');
const configService = getService('config');
const retry = getService('retry');
const createConnector = async (apiUrl: string, spaceId?: string) => {
const result = await supertest
.post(`${getUrlPrefix(spaceId ?? 'default')}/api/actions/connector`)
@ -446,6 +448,68 @@ export default function bedrockTest({ getService }: FtrProviderContext) {
});
});
});
describe('Token tracking dashboard', () => {
const dashboardId = 'specific-dashboard-id-default';
it('should not create a dashboard when user does not have kibana event log permissions', async () => {
const { body } = await supertestWithoutAuth
.post(`/api/actions/connector/${bedrockActionId}/_execute`)
.auth('global_read', 'global_read-password')
.set('kbn-xsrf', 'foo')
.send({
params: {
subAction: 'getDashboard',
subActionParams: {
dashboardId,
},
},
})
.expect(200);
// check dashboard has not been created
await supertest
.get(`/api/saved_objects/dashboard/${dashboardId}`)
.set('kbn-xsrf', 'foo')
.expect(404);
expect(body).to.eql({
status: 'ok',
connector_id: bedrockActionId,
data: { available: false },
});
});
it('should create a dashboard when user has correct permissions', async () => {
const { body } = await supertest
.post(`/api/actions/connector/${bedrockActionId}/_execute`)
.set('kbn-xsrf', 'foo')
.send({
params: {
subAction: 'getDashboard',
subActionParams: {
dashboardId,
},
},
})
.expect(200);
// check dashboard has been created
await retry.try(async () =>
supertest
.get(`/api/saved_objects/dashboard/${dashboardId}`)
.set('kbn-xsrf', 'foo')
.expect(200)
);
objectRemover.add('default', dashboardId, 'dashboard', 'saved_objects');
expect(body).to.eql({
status: 'ok',
connector_id: bedrockActionId,
data: { available: true },
});
});
});
});
});

2
x-pack/test/alerting_api_integration/security_and_spaces/group2/tests/actions/connector_types/openai.ts
View file

@ -313,7 +313,7 @@ export default function genAiTest({ getService }: FtrProviderContext) {
data: genAiSuccessResponse,
});
});
describe('OpenAI dashboard', () => {
describe('Token tracking dashboard', () => {
const dashboardId = 'specific-dashboard-id-default';
it('should not create a dashboard when user does not have kibana event log permissions', async () => {