Skip to content

Commit

Permalink
move guardrails check to server side
Browse files Browse the repository at this point in the history 
Signed-off-by: Joshua Li <[email protected]>
  • Loading branch information
@joshuali925
joshuali925 committed Mar 12, 2024
1 parent 9f78a1b commit 0b9d85b
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 36 deletions.
2 changes: 2 additions & 0 deletions common/constants/query_assist.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,5 @@ export const QUERY_ASSIST_API = {
};

export const ML_COMMONS_API_PREFIX = '/_plugins/_ml';

export const ERROR_DETAILS = { GUARDRAILS_TRIGGERED: 'guardrail triggered' };
24 changes: 22 additions & 2 deletions public/components/event_analytics/explorer/query_assist/__tests__/input.test.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import { fireEvent, render, waitFor } from '@testing-library/react';
import React, { ComponentProps } from 'react';
import { Provider } from 'react-redux';
import { coreMock } from '../../../../../../../../src/core/public/mocks';
import { QUERY_ASSIST_API } from '../../../../../../common/constants/query_assist';
import { ERROR_DETAILS, QUERY_ASSIST_API } from '../../../../../../common/constants/query_assist';
import * as coreServices from '../../../../../../common/utils/core_services';
import { coreRefs } from '../../../../../framework/core_refs';
import { rootReducer } from '../../../../../framework/redux/reducers';
Expand All @@ -29,6 +29,7 @@ const renderQueryAssistInput = (
selectedIndex: [{ label: 'selected-test-index' }],
nlqInput: 'test-input',
setNlqInput: jest.fn(),
handleTimePickerChange: jest.fn(),
},
overrideProps
);
Expand Down Expand Up @@ -74,6 +75,7 @@ describe('<QueryAssistInput /> spec', () => {

const { component } = renderQueryAssistInput();
await waitFor(() => {
// splitbutton dropdown buttons don't support custom test id in Oui 1.5
fireEvent.click(component.getByTestId('splitButton--dropdown'));
fireEvent.click(component.getByTestId('query-assist-generate-button'));
});
Expand Down Expand Up @@ -128,7 +130,25 @@ describe('<QueryAssistInput /> spec', () => {
});
expect(httpMock.post).toBeCalledWith(QUERY_ASSIST_API.SUMMARIZE, {
body:
'{"question":"test-input","index":"selected-test-index","isError":true,"query":"","response":"{\\"statusCode\\":429}"}',
'{"question":"test-input","index":"selected-test-index","isError":true,"query":"","response":"{\\"statusCode\\":429,\\"message\\":\\"Request is throttled. Try again later or contact your administrator\\"}"}',
});
});

it('should display callout when response returned 400 with guardrails denied', async () => {
coreRefs.summarizeEnabled = true;
httpMock.post.mockRejectedValueOnce({
body: { statusCode: 400, message: ERROR_DETAILS.GUARDRAILS_TRIGGERED },
});

const { component } = renderQueryAssistInput();
await waitFor(() => {
// splitbutton data-test-subj doesn't work in Oui 1.5, this should be query-assist-generate-and-run-button
fireEvent.click(component.getByText('Generate and run'));
});

expect(httpMock.post).toBeCalledWith(QUERY_ASSIST_API.GENERATE_PPL, {
body: '{"question":"test-input","index":"selected-test-index"}',
});
expect(component.getByTestId('query-assist-guard-callout')).toBeInTheDocument();
});
});
56 changes: 33 additions & 23 deletions public/components/event_analytics/explorer/query_assist/input.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import { ResponseError } from '@opensearch-project/opensearch/lib/errors';
import React, { useEffect, useState } from 'react';
import { useDispatch, useSelector } from 'react-redux';
import { RAW_QUERY } from '../../../../../common/constants/explorer';
import { QUERY_ASSIST_API } from '../../../../../common/constants/query_assist';
import { ERROR_DETAILS, QUERY_ASSIST_API } from '../../../../../common/constants/query_assist';
import { QUERY_ASSIST_START_TIME } from '../../../../../common/constants/shared';
import { getOSDHttp } from '../../../../../common/utils';
import { coreRefs } from '../../../../framework/core_refs';
Expand Down Expand Up @@ -145,10 +145,6 @@ export const QueryAssistInput: React.FC<React.PropsWithChildren<Props>> = (props
index: props.selectedIndex[0].label,
}),
});
// TODO remove
const mockProhibited = Math.random() > 0.5;
setProhibitedQuery(mockProhibited);
if (mockProhibited) throw new ProhibitedQueryError();
await props.handleQueryChange(generatedPPL);
await dispatch(
changeQuery({
Expand All @@ -160,13 +156,18 @@ export const QueryAssistInput: React.FC<React.PropsWithChildren<Props>> = (props
);
return generatedPPL;
};
const formatError = (error: ResponseError): Error => {
if (error.body) {
const formatError = (error: ResponseError | Error): Error => {
if ('body' in error) {
if (error.body.statusCode === 429)
return {
...error.body,
message: 'Request is throttled. Try again later or contact your administrator',
} as Error;
if (
error.body.statusCode === 400 &&
error.body.message === ERROR_DETAILS.GUARDRAILS_TRIGGERED
)
return new ProhibitedQueryError(error.body.message);
return error.body as Error;
}
return error;
Expand All @@ -178,12 +179,15 @@ export const QueryAssistInput: React.FC<React.PropsWithChildren<Props>> = (props
if (!props.selectedIndex.length) return;
try {
setGenerating(true);
setProhibitedQuery(false);
await request();
} catch (error) {
if (error instanceof ProhibitedQueryError) return;
coreRefs.toasts?.addError(formatError(error as ResponseError), {
title: 'Failed to generate results',
});
} catch (err) {
const error = formatError(err);
if (error instanceof ProhibitedQueryError) {
setProhibitedQuery(true);
return;
}
coreRefs.toasts?.addError(error, { title: 'Failed to generate results' });
} finally {
setGenerating(false);
}
Expand Down Expand Up @@ -231,11 +235,13 @@ export const QueryAssistInput: React.FC<React.PropsWithChildren<Props>> = (props
},
})
);
} catch (error) {
if (error instanceof ProhibitedQueryError) return;
coreRefs.toasts?.addError(error, {
title: 'Failed to summarize results',
});
} catch (err) {
const error = formatError(err);
if (error instanceof ProhibitedQueryError) {
setProhibitedQuery(true);
return;
}
coreRefs.toasts?.addError(error, { title: 'Failed to summarize results' });
} finally {
await dispatch(
changeSummary({
Expand All @@ -261,17 +267,20 @@ export const QueryAssistInput: React.FC<React.PropsWithChildren<Props>> = (props
try {
setGenerating(true);
setGeneratingOrRunning(true);
setProhibitedQuery(false);
await request();
await props.handleTimePickerChange([QUERY_ASSIST_START_TIME, 'now']);
await props.handleTimeRangePickerRefresh(undefined, true);
} catch (error) {
if (error instanceof ProhibitedQueryError) return;
} catch (err) {
const error = formatError(err);
if (error instanceof ProhibitedQueryError) {
setProhibitedQuery(true);
return;
}
if (coreRefs.summarizeEnabled) {
generateSummary({ isError: true, response: JSON.stringify((error as ResponseError).body) });
generateSummary({ isError: true, response: JSON.stringify(error) });
} else {
coreRefs.toasts?.addError(formatError(error as ResponseError), {
title: 'Failed to generate results',
});
coreRefs.toasts?.addError(error, { title: 'Failed to generate results' });
}
} finally {
setGenerating(false);
Expand Down Expand Up @@ -341,6 +350,7 @@ export const QueryAssistInput: React.FC<React.PropsWithChildren<Props>> = (props
</EuiFlexGroup>
{prohibitedQuery ? (
<EuiCallOut
data-test-subj="query-assist-guard-callout"
title="I am unable to respond to this query. Try another question."
size="s"
color="danger"
Expand Down
27 changes: 16 additions & 11 deletions server/routes/query_assist/routes.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@ import {
ResponseError,
} from '../../../../../src/core/server';
import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors';
import { QUERY_ASSIST_API } from '../../../common/constants/query_assist';
import { ERROR_DETAILS, QUERY_ASSIST_API } from '../../../common/constants/query_assist';
import { generateFieldContext } from '../../common/helpers/query_assist/generate_field_context';
import { requestWithRetryAgentSearch, getAgentIdByConfig } from './utils/agents';
import { getAgentIdByConfig, requestWithRetryAgentSearch } from './utils/agents';
import { AGENT_CONFIGS } from './utils/constants';

export function registerQueryAssistRoutes(router: IRouter) {
Expand Down Expand Up @@ -81,11 +81,12 @@ export function registerQueryAssistRoutes(router: IRouter) {
.replace(/\bSPAN\(/g, 'span('); // https://github.com/opensearch-project/dashboards-observability/issues/759
return response.ok({ body: ppl });
} catch (error) {
// parse PPL query from error response if exists
// TODO remove after https://github.com/opensearch-project/skills/issues/138
if (isResponseError(error) && error.body.error?.reason) {
const pplMatch = error.body.error.reason.match(/execute ppl:(.+), get error:/);
if (pplMatch[1]) return response.ok({ body: pplMatch[1] });
if (
isResponseError(error) &&
error.statusCode === 400 &&
error.body.error.details === ERROR_DETAILS.GUARDRAILS_TRIGGERED
) {
return response.badRequest({ body: ERROR_DETAILS.GUARDRAILS_TRIGGERED });
}
return response.custom({ statusCode: error.statusCode || 500, body: error.message });
}
Expand Down Expand Up @@ -152,10 +153,14 @@ export function registerQueryAssistRoutes(router: IRouter) {
},
});
} catch (error) {
return response.custom({
statusCode: error.statusCode || 500,
body: error.message,
});
if (
isResponseError(error) &&
error.statusCode === 400 &&
error.body.error.details === ERROR_DETAILS.GUARDRAILS_TRIGGERED
) {
return response.badRequest({ body: ERROR_DETAILS.GUARDRAILS_TRIGGERED });
}
return response.custom({ statusCode: error.statusCode || 500, body: error.message });
}
}
);
Expand Down

0 comments on commit 0b9d85b

Please sign in to comment.