From a814add319c9436871dfaf0f4a87b377e20a770a Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Fri, 12 Jan 2024 21:30:52 +0000 Subject: [PATCH] refactor to not modify map inside search request Signed-off-by: Joshua Li --- server/routes/query_assist/routes.ts | 20 +----------- .../utils/__tests__/agents.test.ts | 24 ++++++++++++-- server/routes/query_assist/utils/agents.ts | 32 +++++++++++-------- 3 files changed, 42 insertions(+), 34 deletions(-) diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts index 6690c5ad4..a9bd90348 100644 --- a/server/routes/query_assist/routes.ts +++ b/server/routes/query_assist/routes.ts @@ -12,7 +12,7 @@ import { } from '../../../../../src/core/server'; import { QUERY_ASSIST_API } from '../../../common/constants/query_assist'; import { generateFieldContext } from '../../common/helpers/query_assist/generate_field_context'; -import { agentIdMap, searchAgentIdByName, requestWithRetryAgentSearch } from './utils/agents'; +import { requestWithRetryAgentSearch } from './utils/agents'; export function registerQueryAssistRoutes(router: IRouter, config: ObservabilityConfig) { const { @@ -44,17 +44,10 @@ export function registerQueryAssistRoutes(router: IRouter, config: Observability }); const client = context.core.opensearch.client.asCurrentUser; - let shouldRetryAgentSearch = true; - if (!agentIdMap[pplAgentName]) { - await searchAgentIdByName(client, pplAgentName); - shouldRetryAgentSearch = false; - } - try { const pplRequest = await requestWithRetryAgentSearch({ client, agentName: pplAgentName, - shouldRetryAgentSearch, body: { parameters: { index: request.body.index, @@ -113,27 +106,17 @@ export function registerQueryAssistRoutes(router: IRouter, config: Observability const { index, question, query, response: _response, isError } = request.body; const queryResponse = JSON.stringify(_response); let summaryRequest; - let shouldRetryAgentSearch = true; try { if (!isError) { - if (!agentIdMap[responseSummaryAgentName]) { - await searchAgentIdByName(client, responseSummaryAgentName); - shouldRetryAgentSearch = false; - } summaryRequest = await requestWithRetryAgentSearch({ client, agentName: responseSummaryAgentName, - shouldRetryAgentSearch, body: { parameters: { index, question, query, response: queryResponse }, }, }); } else { - if (!agentIdMap[errorSummaryAgentName]) { - await searchAgentIdByName(client, errorSummaryAgentName); - shouldRetryAgentSearch = false; - } const [mappings, sampleDoc] = await Promise.all([ client.indices.getMapping({ index }), client.search({ index, size: 1 }), @@ -142,7 +125,6 @@ export function registerQueryAssistRoutes(router: IRouter, config: Observability summaryRequest = await requestWithRetryAgentSearch({ client, agentName: errorSummaryAgentName, - shouldRetryAgentSearch, body: { parameters: { index, question, query, response: queryResponse, fields }, }, diff --git a/server/routes/query_assist/utils/__tests__/agents.test.ts b/server/routes/query_assist/utils/__tests__/agents.test.ts index 90ec2725c..713299fa3 100644 --- a/server/routes/query_assist/utils/__tests__/agents.test.ts +++ b/server/routes/query_assist/utils/__tests__/agents.test.ts @@ -83,9 +83,8 @@ describe('Agents helper functions', () => { expect(response.body.inference_results[0].output[0].result).toEqual('test response'); }); - it('searches for agent id if not found', async () => { + it('searches for agent id if id is undefined', async () => { mockedTransport - .mockRejectedValueOnce({ statusCode: 404, body: {}, headers: {} }) .mockResolvedValueOnce({ body: { hits: { total: { value: 1 }, hits: [{ _id: 'new-id' }] } } }) .mockResolvedValueOnce({ body: { inference_results: [{ output: [{ result: 'test response' }] }] }, @@ -102,4 +101,25 @@ describe('Agents helper functions', () => { ); expect(response.body.inference_results[0].output[0].result).toEqual('test response'); }); + + it('searches for agent id if id is not found', async () => { + agentIdMap['test agent'] = 'non-exist-agent'; + mockedTransport + .mockRejectedValueOnce({ statusCode: 404, body: {}, headers: {} }) + .mockResolvedValueOnce({ body: { hits: { total: { value: 1 }, hits: [{ _id: 'new-id' }] } } }) + .mockResolvedValueOnce({ + body: { inference_results: [{ output: [{ result: 'test response' }] }] }, + }); + const response = await requestWithRetryAgentSearch({ + client, + agentName: 'test agent', + shouldRetryAgentSearch: true, + body: { parameters: { param1: 'value1' } }, + }); + expect(mockedTransport).toBeCalledWith( + expect.objectContaining({ path: '/_plugins/_ml/agents/new-id/_execute' }), + expect.anything() + ); + expect(response.body.inference_results[0].output[0].result).toEqual('test response'); + }); }); diff --git a/server/routes/query_assist/utils/agents.ts b/server/routes/query_assist/utils/agents.ts index 64ff1791b..ae909790b 100644 --- a/server/routes/query_assist/utils/agents.ts +++ b/server/routes/query_assist/utils/agents.ts @@ -58,7 +58,6 @@ export const searchAgentIdByName = async ( throw new Error('cannot find any agent by name: ' + name); } const id = response.body.hits.hits[0]._id; - agentIdMap[name] = id; return id; } catch (error) { const errorMessage = JSON.stringify(error.meta?.body) || error; @@ -69,22 +68,29 @@ export const searchAgentIdByName = async ( export const requestWithRetryAgentSearch = async (options: { client: OpenSearchClient; agentName: string; - shouldRetryAgentSearch: boolean; + shouldRetryAgentSearch?: boolean; body: RequestBody; -}): Promise => - options.client.transport +}): Promise => { + const { client, agentName, shouldRetryAgentSearch = true, body } = options; + let retry = shouldRetryAgentSearch; + if (!agentIdMap[agentName]) { + agentIdMap[agentName] = await searchAgentIdByName(client, agentName); + retry = false; + } + return client.transport .request( { method: 'POST', - path: `${ML_COMMONS_API_PREFIX}/agents/${agentIdMap[options.agentName]}/_execute`, - body: options.body, + path: `${ML_COMMONS_API_PREFIX}/agents/${agentIdMap[agentName]}/_execute`, + body, }, AGENT_REQUEST_OPTIONS ) - .catch((error) => - options.shouldRetryAgentSearch && isResponseError(error) && error.statusCode === 404 - ? searchAgentIdByName(options.client, options.agentName).then(() => - requestWithRetryAgentSearch({ ...options, shouldRetryAgentSearch: false }) - ) - : Promise.reject(error) - ) as Promise; + .catch(async (error) => { + if (retry && isResponseError(error) && error.statusCode === 404) { + agentIdMap[agentName] = await searchAgentIdByName(client, agentName); + return requestWithRetryAgentSearch({ ...options, shouldRetryAgentSearch: false }); + } + return Promise.reject(error); + }) as Promise; +};