From 6d128ea795d4db8390e7178c4094991d3c621175 Mon Sep 17 00:00:00 2001
From: Joshua Li <joshuali925@gmail.com>
Date: Wed, 4 Sep 2024 09:40:59 -0700
Subject: [PATCH] [query assist] update api handler to accommodate new
 ml-commons config response (#2111)

* [query assist] update api handler to accommodate new ml-commons config response

Signed-off-by: Joshua Li <joshuali925@gmail.com>

* update `type` field in ml-commons response

Signed-off-by: Joshua Li <joshuali925@gmail.com>

---------

Signed-off-by: Joshua Li <joshuali925@gmail.com>
---
 .../utils/__tests__/agents.test.ts            | 19 +++++++++++++++++++
 server/routes/query_assist/utils/agents.ts    | 14 +++++++++++---
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/server/routes/query_assist/utils/__tests__/agents.test.ts b/server/routes/query_assist/utils/__tests__/agents.test.ts
index 1aa7628f82..2f92c4a480 100644
--- a/server/routes/query_assist/utils/__tests__/agents.test.ts
+++ b/server/routes/query_assist/utils/__tests__/agents.test.ts
@@ -40,6 +40,25 @@ describe('Agents helper functions', () => {
     `);
   });
 
+  it('acccepts ml_configuration key in get agent response', async () => {
+    mockedTransport.mockResolvedValueOnce({
+      body: {
+        config_type: 'agent',
+        ml_configuration: { agent_id: 'agentId' },
+      },
+    });
+    const id = await getAgentIdByConfig(client, 'test_agent');
+    expect(id).toEqual('agentId');
+    expect(mockedTransport.mock.calls[0]).toMatchInlineSnapshot(`
+      Array [
+        Object {
+          "method": "GET",
+          "path": "/_plugins/_ml/config/test_agent",
+        },
+      ]
+    `);
+  });
+
   it('handles not found errors', async () => {
     mockedTransport.mockRejectedValueOnce(
       new ResponseError(({
diff --git a/server/routes/query_assist/utils/agents.ts b/server/routes/query_assist/utils/agents.ts
index 5b660c5b8e..97a668bb88 100644
--- a/server/routes/query_assist/utils/agents.ts
+++ b/server/routes/query_assist/utils/agents.ts
@@ -34,12 +34,20 @@ export const getAgentIdByConfig = async (
     const response = (await opensearchClient.transport.request({
       method: 'GET',
       path: `${ML_COMMONS_API_PREFIX}/config/${configName}`,
-    })) as ApiResponse<{ type: string; configuration: { agent_id?: string } }>;
+    })) as ApiResponse<{
+      type?: string;
+      config_type?: string;
+      ml_configuration?: { agent_id?: 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);