[Playground] Propagate Error message into FE (#182201)

Samiul-TheSoccerFan · kibanamachine · web-flow · commit 0288bb466403 · 2024-05-01T10:22:41.000-04:00
## Summary - Fix error not being propagated into FE - Added tests ## UI ### Rate limit error message: ![Screenshot 2024-04-30 at 1 05 40 PM](https://github.com/elastic/kibana/assets/150824886/2734d27b-ef2b-469b-9344-a7c62cd502bc) ### BAD LLM ![Screenshot 2024-04-30 at 1 05 47 PM](https://github.com/elastic/kibana/assets/150824886/fb3025e5-5f5e-49e6-8277-7dba52cc76cd) ### Invalid API key ![Screenshot 2024-04-30 at 1 06 26 PM](https://github.com/elastic/kibana/assets/150824886/67c1bdf6-6785-48ed-870a-ae24e8ac7e9f) ### Checklist Delete any items that are not applicable to this PR. - [ ] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md) - [ ] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [ ] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [ ] [Flaky Test Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was used on any tests changed - [ ] Any UI touched in this PR is usable by keyboard only (learn more about [keyboard accessibility](https://webaim.org/techniques/keyboard/)) - [ ] Any UI touched in this PR does not create any new axe failures (run axe in browser: [FF](https://addons.mozilla.org/en-US/firefox/addon/axe-devtools/), [Chrome](https://chrome.google.com/webstore/detail/axe-web-accessibility-tes/lhdoppojpmngadmnindnejefpokejbdd?hl=en-US)) - [ ] If a plugin configuration key changed, check if it needs to be allowlisted in the cloud and added to the [docker list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker) - [ ] This renders correctly on smaller devices using a responsive layout. (You can test this [in your browser](https://www.browserstack.com/guide/responsive-testing-on-local-server)) - [ ] This was checked for [cross-browser compatibility](https://www.elastic.co/support/matrix#matrix_browsers) ### For maintainers - [ ] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
diff --git a/x-pack/plugins/search_playground/__mocks__/router.mock.ts b/x-pack/plugins/search_playground/__mocks__/router.mock.ts
@@ -0,0 +1,95 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import {
+  IRouter,
+  KibanaRequest,
+  RequestHandlerContext,
+  RouteValidatorConfig,
+} from '@kbn/core/server';
+import { httpServiceMock, httpServerMock } from '@kbn/core/server/mocks';
+
+/**
+ * Test helper that mocks Kibana's router and DRYs out various helper (callRoute, schema validation)
+ */
+
+type MethodType = 'get' | 'post' | 'put' | 'patch' | 'delete';
+type PayloadType = 'params' | 'query' | 'body';
+
+interface IMockRouter {
+  method: MethodType;
+  path: string;
+  context?: jest.Mocked<RequestHandlerContext>;
+}
+interface IMockRouterRequest {
+  body?: object;
+  query?: object;
+  params?: object;
+}
+type MockRouterRequest = KibanaRequest | IMockRouterRequest;
+
+export class MockRouter {
+  public router!: jest.Mocked<IRouter>;
+  public method: MethodType;
+  public path: string;
+  public context: jest.Mocked<RequestHandlerContext>;
+  public payload?: PayloadType;
+  public response = httpServerMock.createResponseFactory();
+
+  constructor({ method, path, context = {} as jest.Mocked<RequestHandlerContext> }: IMockRouter) {
+    this.createRouter();
+    this.method = method;
+    this.path = path;
+    this.context = context;
+  }
+
+  public createRouter = () => {
+    this.router = httpServiceMock.createRouter();
+  };
+
+  public callRoute = async (request: MockRouterRequest) => {
+    const route = this.findRouteRegistration();
+    const [, handler] = route;
+    await handler(this.context, httpServerMock.createKibanaRequest(request as any), this.response);
+  };
+
+  /**
+   * Schema validation helpers
+   */
+
+  public validateRoute = (request: MockRouterRequest) => {
+    const route = this.findRouteRegistration();
+    const [config] = route;
+    const validate = config.validate as RouteValidatorConfig<{}, {}, {}>;
+    const payloads = Object.keys(request) as PayloadType[];
+
+    payloads.forEach((payload: PayloadType) => {
+      const payloadValidation = validate[payload] as { validate(request: KibanaRequest): void };
+      const payloadRequest = request[payload] as KibanaRequest;
+
+      payloadValidation.validate(payloadRequest);
+    });
+  };
+
+  public shouldValidate = (request: MockRouterRequest) => {
+    expect(() => this.validateRoute(request)).not.toThrow();
+  };
+
+  public shouldThrow = (request: MockRouterRequest) => {
+    expect(() => this.validateRoute(request)).toThrow();
+  };
+
+  private findRouteRegistration = () => {
+    const routerCalls = this.router[this.method].mock.calls as any[];
+    if (!routerCalls.length) throw new Error('No routes registered.');
+
+    const route = routerCalls.find(([router]: any) => router.path === this.path);
+    if (!route) throw new Error('No matching registered routes found - check method/path keys');
+
+    return route;
+  };
+}
diff --git a/x-pack/plugins/search_playground/server/routes.test.ts b/x-pack/plugins/search_playground/server/routes.test.ts
@@ -5,7 +5,19 @@
  * 2.0.
  */
 
-import { createRetriever } from './routes';
+import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
+import { RequestHandlerContext } from '@kbn/core/server';
+import { coreMock } from '@kbn/core/server/mocks';
+import { MockRouter } from '../__mocks__/router.mock';
+import { ConversationalChain } from './lib/conversational_chain';
+import { getChatParams } from './lib/get_chat_params';
+import { createRetriever, defineRoutes } from './routes';
+
+jest.mock('./lib/get_chat_params', () => ({
+  getChatParams: jest.fn(),
+}));
+
+jest.mock('./lib/conversational_chain');
 
 describe('createRetriever', () => {
   test('works when the question has quotes', () => {
@@ -18,3 +30,75 @@ describe('createRetriever', () => {
     expect(result).toEqual({ query: { match: { text: 'How can I "do something" with quotes?' } } });
   });
 });
+
+describe('Search Playground routes', () => {
+  let mockRouter: MockRouter;
+  const mockClient = {
+    asCurrentUser: {},
+  };
+
+  const mockCore = {
+    elasticsearch: { client: mockClient },
+  };
+  const mockLogger = loggingSystemMock.createLogger().get();
+
+  describe('POST - Chat Messages', () => {
+    const mockData = {
+      connector_id: 'open-ai',
+      indices: 'my-index',
+      prompt: 'You are an assistant',
+      citations: true,
+      elasticsearch_query: {},
+      summarization_model: 'GPT-4',
+      doc_size: 3,
+      source_fields: '{}',
+    };
+
+    const mockRequestBody = {
+      data: mockData,
+    };
+
+    beforeEach(() => {
+      jest.clearAllMocks();
+
+      const coreStart = coreMock.createStart();
+
+      const context = {
+        core: Promise.resolve(mockCore),
+      } as unknown as jest.Mocked<RequestHandlerContext>;
+
+      mockRouter = new MockRouter({
+        context,
+        method: 'post',
+        path: '/internal/search_playground/chat',
+      });
+
+      defineRoutes({
+        logger: mockLogger,
+        router: mockRouter.router,
+        getStartServices: jest.fn().mockResolvedValue([coreStart, {}, {}]),
+      });
+    });
+
+    it('responds with error message if stream throws an error', async () => {
+      (getChatParams as jest.Mock).mockResolvedValue({ model: 'open-ai' });
+      (ConversationalChain as jest.Mock).mockImplementation(() => {
+        return {
+          stream: jest
+            .fn()
+            .mockRejectedValue(new Error('Unexpected API error - Some Open AI error message')),
+        };
+      });
+
+      await mockRouter.callRoute({
+        body: mockRequestBody,
+      });
+
+      expect(mockRouter.response.badRequest).toHaveBeenCalledWith({
+        body: {
+          message: 'Unexpected API error - Some Open AI error message',
+        },
+      });
+    });
+  });
+});
diff --git a/x-pack/plugins/search_playground/server/routes.ts b/x-pack/plugins/search_playground/server/routes.ts
@@ -131,10 +131,10 @@ export function defineRoutes({
       } catch (e) {
         logger.error('Failed to create the chat stream', e);
 
-        if (typeof e === 'string') {
+        if (typeof e === 'object') {
           return response.badRequest({
             body: {
-              message: e,
+              message: e.message,
             },
           });
         }
diff --git a/x-pack/plugins/search_playground/tsconfig.json b/x-pack/plugins/search_playground/tsconfig.json
@@ -35,7 +35,8 @@
     "@kbn/elastic-assistant-common",
     "@kbn/logging",
     "@kbn/react-kibana-context-render",
-    "@kbn/doc-links"
+    "@kbn/doc-links",
+    "@kbn/core-logging-server-mocks"
   ],
   "exclude": [
     "target/**/*",
