Fix query assistant fetching agent bug (#7804)

* Fix query assistant fetching agent bug Signed-off-by: Liyun Xiu <[email protected]> * Add changelog Signed-off-by: Liyun Xiu <[email protected]> * fix style Signed-off-by: Liyun Xiu <[email protected]> * Add UT Signed-off-by: Liyun Xiu <[email protected]> --------- Signed-off-by: Liyun Xiu <[email protected]>
opensearch-project · Aug 23, 2024 · 3364d0c · 3364d0c
1 parent a59f8ab
commit 3364d0c
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 3 deletions.
diff --git a/changelogs/fragments/7804.yml b/changelogs/fragments/7804.yml
@@ -0,0 +1,2 @@
+fix:
+- Fix query assistant fetching agent bug ([#7804](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7804))
diff --git a/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts b/src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts
@@ -110,4 +110,27 @@ describe('Agents helper functions', () => {
     );
     expect(response.body.inference_results[0].output[0].result).toEqual('test response');
   });
+
+  it('searches for agent id and response contains ml_configuration', async () => {
+    mockedTransport
+      .mockResolvedValueOnce({
+        body: {
+          type: 'agent',
+          ml_configuration: { agent_id: 'new-id' },
+        },
+      })
+      .mockResolvedValueOnce({
+        body: { inference_results: [{ output: [{ result: 'test response' }] }] },
+      });
+    const response = await requestAgentByConfig({
+      context,
+      configName: 'new_agent',
+      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/src/plugins/query_enhancements/server/routes/query_assist/agents.ts b/src/plugins/query_enhancements/server/routes/query_assist/agents.ts
@@ -37,11 +37,13 @@ export const getAgentIdByConfig = async (
       method: 'GET',
       path: `${URI.ML}/config/${configName}`,
     })) as ApiResponse<{ type: string; configuration: { agent_id?: string } }>;
-
-    if (!response || response.body.configuration.agent_id === undefined) {
+    if (
+      !response ||
+      !(response.body.ml_configuration?.agent_id || response.body.configuration?.agent_id)
+    ) {
       throw new Error('cannot find any agent by configuration: ' + configName);
     }
-    return response.body.configuration.agent_id;
+    return response.body.ml_configuration?.agent_id || response.body.configuration.agent_id;
   } catch (error) {
     const errorMessage = JSON.stringify(error.meta?.body) || error;
     throw new Error(`Get agent '${configName}' failed, reason: ` + errorMessage);