From 4cfd1713bf029b086083e43e6cbded430ce32f4d Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Sat, 20 Jan 2024 01:13:37 +0000 Subject: [PATCH] support disabling summarization Signed-off-by: Joshua Li --- .../query_assist/__tests__/input.test.tsx | 16 +++++++ .../explorer/query_assist/input.tsx | 48 +++++++++++-------- public/framework/core_refs.ts | 1 + public/plugin.ts | 4 ++ server/index.ts | 25 ++++++++-- server/routes/query_assist/routes.ts | 4 +- 6 files changed, 74 insertions(+), 24 deletions(-) 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 f6a9dd02de..7ced9cd288 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 @@ -85,7 +85,23 @@ describe(' spec', () => { ); }); + it('should not call summarize if disabled', async () => { + coreRefs.summarizeEnabled = false; + httpMock.post.mockRejectedValueOnce({ body: { statusCode: 429 } }); + + const { component } = renderQueryAssistInput(); + await waitFor(() => { + fireEvent.click(component.getByTestId('query-assist-generate-and-run-button')); + }); + + expect(httpMock.post).toBeCalledWith(QUERY_ASSIST_API.GENERATE_PPL, { + body: '{"question":"test-input","index":"selected-test-index"}', + }); + expect(httpMock.post).not.toBeCalledWith(QUERY_ASSIST_API.SUMMARIZE, expect.anything()); + }); + it('should call summarize for generate and run errors', async () => { + coreRefs.summarizeEnabled = true; httpMock.post.mockRejectedValueOnce({ body: { statusCode: 429 } }).mockResolvedValueOnce({ summary: 'too many requests', suggestedQuestions: ['1', '2'], diff --git a/public/components/event_analytics/explorer/query_assist/input.tsx b/public/components/event_analytics/explorer/query_assist/input.tsx index 4e1422c996..3ad8c6d29f 100644 --- a/public/components/event_analytics/explorer/query_assist/input.tsx +++ b/public/components/event_analytics/explorer/query_assist/input.tsx @@ -23,6 +23,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 { getOSDHttp } from '../../../../../common/utils'; import { coreRefs } from '../../../../framework/core_refs'; import chatLogo from '../../../datasources/icons/query-assistant-logo.svg'; @@ -34,7 +35,6 @@ import { } from '../../redux/slices/query_assistant_summarization_slice'; import { reset, selectQueryResult } from '../../redux/slices/query_result_slice'; import { changeQuery, selectQueries } from '../../redux/slices/query_slice'; -import { QUERY_ASSIST_API } from '../../../../../common/constants/query_assist'; interface SummarizationContext { question: string; @@ -87,22 +87,23 @@ export const QueryAssistInput: React.FC = (props) => { useEffect(() => { if ( - summaryData.responseForSummaryStatus === 'success' || - summaryData.responseForSummaryStatus === 'failure' - ) { - void (async () => { - await dispatch( - changeSummary({ - tabId: props.tabId, - data: { - summaryLoading: false, - }, - }) - ); - if (explorerData.total > 0 || summaryData.responseForSummaryStatus === 'failure') - generateSummary(); - })(); - } + props.nlqInput.trim().length === 0 || + (summaryData.responseForSummaryStatus !== 'success' && + summaryData.responseForSummaryStatus !== 'failure') + ) + return; + void (async () => { + await dispatch( + changeSummary({ + tabId: props.tabId, + data: { + summaryLoading: false, + }, + }) + ); + if (explorerData.total > 0 || summaryData.responseForSummaryStatus === 'failure') + generateSummary(); + })(); }, [summaryData.responseForSummaryStatus]); const [barSelected, setBarSelected] = useState(false); @@ -174,6 +175,11 @@ export const QueryAssistInput: React.FC = (props) => { } }; const generateSummary = async (context?: Partial) => { + if (!coreRefs.summarizeEnabled) { + if (summaryData.responseForSummaryStatus === 'failure') + coreRefs.toasts?.addError(explorerData.error.body, { title: 'Invalid query' }); + return; + } try { const isError = summaryData.responseForSummaryStatus === 'failure'; const summarizationContext: SummarizationContext = { @@ -334,7 +340,9 @@ export const QueryAssistInput: React.FC = (props) => { = (props) => { = { schema: observabilityConfig.schema, exposeToBrowser: { query_assist: true, + summarize: true, }, }; diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts index 68e1cd1bf8..e99f5f9729 100644 --- a/server/routes/query_assist/routes.ts +++ b/server/routes/query_assist/routes.ts @@ -16,11 +16,11 @@ import { generateFieldContext } from '../../common/helpers/query_assist/generate import { requestWithRetryAgentSearch } from './utils/agents'; export function registerQueryAssistRoutes(router: IRouter, config: ObservabilityConfig) { + const { ppl_agent_name: pplAgentName } = config.query_assist; const { - ppl_agent_name: pplAgentName, response_summary_agent_name: responseSummaryAgentName, error_summary_agent_name: errorSummaryAgentName, - } = config.query_assist; + } = config.summarize; router.post( {