From 0b9d85bf9be70b387c96f6f09a0c9de45b15f5f0 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Tue, 12 Mar 2024 21:15:05 +0000 Subject: [PATCH] move guardrails check to server side Signed-off-by: Joshua Li --- common/constants/query_assist.ts | 2 + .../query_assist/__tests__/input.test.tsx | 24 +++++++- .../explorer/query_assist/input.tsx | 56 +++++++++++-------- server/routes/query_assist/routes.ts | 27 +++++---- 4 files changed, 73 insertions(+), 36 deletions(-) diff --git a/common/constants/query_assist.ts b/common/constants/query_assist.ts index e0b3bb29e..8de10583d 100644 --- a/common/constants/query_assist.ts +++ b/common/constants/query_assist.ts @@ -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' }; diff --git a/public/components/event_analytics/explorer/query_assist/__tests__/input.test.tsx b/public/components/event_analytics/explorer/query_assist/__tests__/input.test.tsx index 9950d7ab4..31eb1c872 100644 --- a/public/components/event_analytics/explorer/query_assist/__tests__/input.test.tsx +++ b/public/components/event_analytics/explorer/query_assist/__tests__/input.test.tsx @@ -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'; @@ -29,6 +29,7 @@ const renderQueryAssistInput = ( selectedIndex: [{ label: 'selected-test-index' }], nlqInput: 'test-input', setNlqInput: jest.fn(), + handleTimePickerChange: jest.fn(), }, overrideProps ); @@ -74,6 +75,7 @@ describe(' 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')); }); @@ -128,7 +130,25 @@ describe(' 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(); + }); }); diff --git a/public/components/event_analytics/explorer/query_assist/input.tsx b/public/components/event_analytics/explorer/query_assist/input.tsx index 43284c9ca..a8bf8f3dc 100644 --- a/public/components/event_analytics/explorer/query_assist/input.tsx +++ b/public/components/event_analytics/explorer/query_assist/input.tsx @@ -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'; @@ -145,10 +145,6 @@ export const QueryAssistInput: React.FC> = (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({ @@ -160,13 +156,18 @@ export const QueryAssistInput: React.FC> = (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; @@ -178,12 +179,15 @@ export const QueryAssistInput: React.FC> = (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); } @@ -231,11 +235,13 @@ export const QueryAssistInput: React.FC> = (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({ @@ -261,17 +267,20 @@ export const QueryAssistInput: React.FC> = (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); @@ -341,6 +350,7 @@ export const QueryAssistInput: React.FC> = (props {prohibitedQuery ? (