Fix returns 500 error for AI agent APIs when OpenSearch returns 4xx e…

…rror with json format error message (#380) (#381) * Fix returns 500 error for AI agent APIs when OpenSearch returns 4xx code with json format error message Signed-off-by: gaobinlong <[email protected]> * Modify change log Signed-off-by: gaobinlong <[email protected]> --------- Signed-off-by: gaobinlong <[email protected]> (cherry picked from commit ffa509b) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> # Conflicts: # CHANGELOG.md Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Dec 27, 2024 · cfc975c · cfc975c
1 parent 3456f3f
commit cfc975c
Show file tree

Hide file tree

Showing 6 changed files with 701 additions and 19 deletions.
diff --git a/server/routes/agent_routes.test.ts b/server/routes/agent_routes.test.ts
@@ -0,0 +1,135 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Boom } from '@hapi/boom';
+import { OpenSearchDashboardsRequest, Router } from '../../../../src/core/server/http/router';
+import { enhanceWithContext, triggerHandler } from './router.mock';
+import { resetMocks } from '../services/storage/agent_framework_storage_service.mock';
+import { httpServerMock } from '../../../../src/core/server/http/http_server.mocks';
+import { loggerMock } from '../../../../src/core/server/logging/logger.mock';
+import { AGENT_API } from '../../common/constants/llm';
+import { AssistantClient } from '../services/assistant_client';
+import { RequestHandlerContext } from '../../../../src/core/server';
+import { registerAgentRoutes } from './agent_routes';
+const mockedLogger = loggerMock.create();
+
+export const createMockedAssistantClient = (
+  request: OpenSearchDashboardsRequest
+): AssistantClient => {
+  return new AssistantClient(request, {} as RequestHandlerContext);
+};
+
+const mockedAssistantClient = createMockedAssistantClient({} as OpenSearchDashboardsRequest);
+
+describe('test execute agent route', () => {
+  const router = new Router(
+    '',
+    mockedLogger,
+    enhanceWithContext({
+      assistant_plugin: {
+        logger: mockedLogger,
+      },
+    })
+  );
+  registerAgentRoutes(router, {
+    getScopedClient: jest.fn(
+      (request: OpenSearchDashboardsRequest, context: RequestHandlerContext) => {
+        return mockedAssistantClient;
+      }
+    ),
+  });
+  const executeAgentRequest = (payload: {}, query: {}) =>
+    triggerHandler(router, {
+      method: 'post',
+      path: AGENT_API.EXECUTE,
+      req: httpServerMock.createRawRequest({
+        payload: JSON.stringify(payload),
+        query,
+      }),
+    });
+  beforeEach(() => {
+    loggerMock.clear(mockedLogger);
+    resetMocks();
+  });
+
+  it('return 4xx when execute agent throws 4xx error', async () => {
+    mockedAssistantClient.executeAgentByConfigName = jest.fn().mockRejectedValue({
+      statusCode: 429,
+      body: {
+        status: 429,
+        error: {
+          type: 'OpenSearchStatusException',
+          reason: 'System Error',
+          details: 'Request is throttled at model level.',
+        },
+      },
+    });
+    const result = (await executeAgentRequest(
+      {},
+      {
+        agentConfigName: 'os_insight',
+      }
+    )) as Boom;
+    expect(result.output).toMatchInlineSnapshot(`
+      Object {
+        "headers": Object {},
+        "payload": Object {
+          "error": "Too Many Requests",
+          "message": "{\\"status\\":429,\\"error\\":{\\"type\\":\\"OpenSearchStatusException\\",\\"reason\\":\\"System Error\\",\\"details\\":\\"Request is throttled at model level.\\"}}",
+          "statusCode": 429,
+        },
+        "statusCode": 429,
+      }
+    `);
+  });
+
+  it('return 4xx when executeAgent throws 4xx error in string format', async () => {
+    mockedAssistantClient.executeAgentByConfigName = jest.fn().mockRejectedValue({
+      statusCode: 429,
+      body: 'Request is throttled at model level',
+    });
+    const result = (await executeAgentRequest(
+      {},
+      {
+        agentConfigName: 'os_insight',
+      }
+    )) as Boom;
+    expect(result.output).toMatchInlineSnapshot(`
+      Object {
+        "headers": Object {},
+        "payload": Object {
+          "error": "Too Many Requests",
+          "message": "Request is throttled at model level",
+          "statusCode": 429,
+        },
+        "statusCode": 429,
+      }
+    `);
+  });
+
+  it('return 5xx when executeAgent throws 5xx error', async () => {
+    mockedAssistantClient.executeAgentByConfigName = jest.fn().mockRejectedValue({
+      statusCode: 500,
+      body: 'Server error',
+    });
+    const result = (await executeAgentRequest(
+      {},
+      {
+        agentConfigName: 'os_insight',
+      }
+    )) as Boom;
+    expect(result.output).toMatchInlineSnapshot(`
+      Object {
+        "headers": Object {},
+        "payload": Object {
+          "error": "Internal Server Error",
+          "message": "Execute agent failed!",
+          "statusCode": 500,
+        },
+        "statusCode": 500,
+      }
+    `);
+  });
+});
diff --git a/server/routes/agent_routes.ts b/server/routes/agent_routes.ts
@@ -42,15 +42,13 @@ export function registerAgentRoutes(router: IRouter, assistantService: Assistant
         context.assistant_plugin.logger.error('Execute agent failed!', e);
         if (e.statusCode >= 400 && e.statusCode <= 499) {
           return res.customError({
-            body: e.body,
+            body: { message: typeof e.body === 'string' ? e.body : JSON.stringify(e.body) },
             statusCode: e.statusCode,
-            headers: e.headers,
           });
         } else {
           return res.customError({
             body: 'Execute agent failed!',
             statusCode: 500,
-            headers: e.headers,
           });
         }
       }