From 2fb31f9f6d98ddf16427f3a4f1eabf18990ea25b Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Mon, 22 Apr 2024 21:55:49 +0000 Subject: [PATCH] remove caching agent id (#1529) Signed-off-by: Joshua Li Signed-off-by: Paul Sebastian Co-authored-by: Paul Sebastian (cherry picked from commit 01bec056cb606810f57f1875d6651cceb0a609d3) Signed-off-by: github-actions[bot] --- server/routes/query_assist/routes.ts | 8 +-- .../utils/__tests__/agents.test.ts | 53 ++----------------- server/routes/query_assist/utils/agents.ts | 40 +++++--------- 3 files changed, 19 insertions(+), 82 deletions(-) diff --git a/server/routes/query_assist/routes.ts b/server/routes/query_assist/routes.ts index 2185045c31..59e5776542 100644 --- a/server/routes/query_assist/routes.ts +++ b/server/routes/query_assist/routes.ts @@ -12,7 +12,7 @@ import { import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors'; import { ERROR_DETAILS, QUERY_ASSIST_API } from '../../../common/constants/query_assist'; import { generateFieldContext } from '../../common/helpers/query_assist/generate_field_context'; -import { getAgentIdByConfig, requestWithRetryAgentSearch } from './utils/agents'; +import { getAgentIdByConfig, getAgentIdAndRequest } from './utils/agents'; import { AGENT_CONFIGS } from './utils/constants'; export function registerQueryAssistRoutes(router: IRouter) { @@ -57,7 +57,7 @@ export function registerQueryAssistRoutes(router: IRouter) { ): Promise> => { const client = context.core.opensearch.client.asCurrentUser; try { - const pplRequest = await requestWithRetryAgentSearch({ + const pplRequest = await getAgentIdAndRequest({ client, configName: AGENT_CONFIGS.PPL_AGENT, body: { @@ -118,7 +118,7 @@ export function registerQueryAssistRoutes(router: IRouter) { try { if (!isError) { - summaryRequest = await requestWithRetryAgentSearch({ + summaryRequest = await getAgentIdAndRequest({ client, configName: AGENT_CONFIGS.RESPONSE_SUMMARY_AGENT, body: { @@ -131,7 +131,7 @@ export function registerQueryAssistRoutes(router: IRouter) { client.search({ index, size: 1 }), ]); const fields = generateFieldContext(mappings, sampleDoc); - summaryRequest = await requestWithRetryAgentSearch({ + summaryRequest = await getAgentIdAndRequest({ client, configName: AGENT_CONFIGS.ERROR_SUMMARY_AGENT, body: { diff --git a/server/routes/query_assist/utils/__tests__/agents.test.ts b/server/routes/query_assist/utils/__tests__/agents.test.ts index ae6d43a4c5..1aa7628f82 100644 --- a/server/routes/query_assist/utils/__tests__/agents.test.ts +++ b/server/routes/query_assist/utils/__tests__/agents.test.ts @@ -7,7 +7,7 @@ import { ApiResponse } from '@opensearch-project/opensearch'; import { ResponseError } from '@opensearch-project/opensearch/lib/errors'; import { CoreRouteHandlerContext } from '../../../../../../../src/core/server/core_route_handler_context'; import { coreMock, httpServerMock } from '../../../../../../../src/core/server/mocks'; -import { agentIdMap, getAgentIdByConfig, requestWithRetryAgentSearch } from '../agents'; +import { getAgentIdByConfig, getAgentIdAndRequest } from '../agents'; describe('Agents helper functions', () => { const coreContext = new CoreRouteHandlerContext( @@ -75,27 +75,7 @@ describe('Agents helper functions', () => { ); }); - it('requests with valid agent id', async () => { - agentIdMap.test_agent = 'test-id'; - mockedTransport.mockResolvedValueOnce({ - body: { inference_results: [{ output: [{ result: 'test response' }] }] }, - }); - const response = await requestWithRetryAgentSearch({ - client, - configName: 'test_agent', - shouldRetryAgentSearch: true, - body: { parameters: { param1: 'value1' } }, - }); - expect(mockedTransport).toBeCalledWith( - expect.objectContaining({ - path: '/_plugins/_ml/agents/test-id/_execute', - }), - expect.anything() - ); - expect(response.body.inference_results[0].output[0].result).toEqual('test response'); - }); - - it('searches for agent id if id is undefined', async () => { + it('searches for agent id and sends request', async () => { mockedTransport .mockResolvedValueOnce({ body: { @@ -106,36 +86,9 @@ describe('Agents helper functions', () => { .mockResolvedValueOnce({ body: { inference_results: [{ output: [{ result: 'test response' }] }] }, }); - const response = await requestWithRetryAgentSearch({ + const response = await getAgentIdAndRequest({ client, configName: 'new_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'); - }); - - 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: { - type: 'agent', - configuration: { agent_id: 'new-id' }, - }, - }) - .mockResolvedValueOnce({ - body: { inference_results: [{ output: [{ result: 'test response' }] }] }, - }); - const response = await requestWithRetryAgentSearch({ - client, - configName: 'test_agent', - shouldRetryAgentSearch: true, body: { parameters: { param1: 'value1' } }, }); expect(mockedTransport).toBeCalledWith( diff --git a/server/routes/query_assist/utils/agents.ts b/server/routes/query_assist/utils/agents.ts index 84fd53fe09..5b660c5b8e 100644 --- a/server/routes/query_assist/utils/agents.ts +++ b/server/routes/query_assist/utils/agents.ts @@ -4,9 +4,8 @@ */ import { ApiResponse } from '@opensearch-project/opensearch/.'; -import { RequestBody } from '@opensearch-project/opensearch/lib/Transport'; +import { RequestBody, TransportRequestPromise } from '@opensearch-project/opensearch/lib/Transport'; import { OpenSearchClient } from '../../../../../../src/core/server'; -import { isResponseError } from '../../../../../../src/core/server/opensearch/client/errors'; import { ML_COMMONS_API_PREFIX } from '../../../../common/constants/query_assist'; const AGENT_REQUEST_OPTIONS = { @@ -27,8 +26,6 @@ type AgentResponse = ApiResponse<{ }>; }>; -export const agentIdMap: Record = {}; - export const getAgentIdByConfig = async ( opensearchClient: OpenSearchClient, configName: string @@ -49,32 +46,19 @@ export const getAgentIdByConfig = async ( } }; -export const requestWithRetryAgentSearch = async (options: { +export const getAgentIdAndRequest = async (options: { client: OpenSearchClient; configName: string; - shouldRetryAgentSearch?: boolean; body: RequestBody; }): Promise => { - const { client, configName, shouldRetryAgentSearch = true, body } = options; - let retry = shouldRetryAgentSearch; - if (!agentIdMap[configName]) { - agentIdMap[configName] = await getAgentIdByConfig(client, configName); - retry = false; - } - return client.transport - .request( - { - method: 'POST', - path: `${ML_COMMONS_API_PREFIX}/agents/${agentIdMap[configName]}/_execute`, - body, - }, - AGENT_REQUEST_OPTIONS - ) - .catch(async (error) => { - if (retry && isResponseError(error) && error.statusCode === 404) { - agentIdMap[configName] = await getAgentIdByConfig(client, configName); - return requestWithRetryAgentSearch({ ...options, shouldRetryAgentSearch: false }); - } - return Promise.reject(error); - }) as Promise; + const { client, configName, body } = options; + const agentId = await getAgentIdByConfig(client, configName); + return client.transport.request( + { + method: 'POST', + path: `${ML_COMMONS_API_PREFIX}/agents/${agentId}/_execute`, + body, + }, + AGENT_REQUEST_OPTIONS + ) as TransportRequestPromise; };