[query assist] update api handler to accommodate new ml-commons confi…

…g response Signed-off-by: Joshua Li <[email protected]>
joshuali925 · Sep 3, 2024 · 6ab3de6 · 6ab3de6
1 parent c655eb5
commit 6ab3de6
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git 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
@@ -34,12 +34,19 @@ 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;
+      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);