refactor to not modify map inside search request

Signed-off-by: Joshua Li <[email protected]>
joshuali925 · Jan 12, 2024 · a814add · a814add
1 parent 9617b7a
commit a814add
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 34 deletions.
diff --git 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
@@ -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
@@ -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<AgentResponse> =>
-  options.client.transport
+}): Promise<AgentResponse> => {
+  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<AgentResponse>;
+    .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<AgentResponse>;
+};