chore(query-bar): open query options when applied from ai, better error reporting (#4626)

Author: Anemy

packages/compass-query-bar/src/modules/ai-query-request.spec.ts

@@ -136,7 +136,7 @@ describe('#runFetchAIQuery', function () {
         ...mockUserPrompt,
         sampleDocuments: [
           {
-            test: '4'.repeat(6000),
+            test: '4'.repeat(60000),
           },
         ],
         signal: new AbortController().signal,
@@ -161,7 +161,7 @@ describe('#runFetchAIQuery', function () {
             a: ['3'],
           },
           {
-            a: ['4'.repeat(5000)],
+            a: ['4'.repeat(50000)],
           },
         ],
         signal: new AbortController().signal,
@@ -190,15 +190,17 @@ describe('#runFetchAIQuery', function () {
     });
   });
 
-  describe('when the endpoint is set and the server errors', function () {
+  describe('when the server errors', function () {
     let stopServer: () => Promise<void>;
 
     beforeEach(async function () {
       // Start a mock server to pass an ai response.
       // Set the server endpoint in the env.
       const { endpoint, stop } = await startMockAIServer({
-        response: {},
-        sendError: true,
+        response: {
+          status: 500,
+          body: 'error',
+        },
       });
 
       stopServer = stop;
@@ -225,4 +227,44 @@ describe('#runFetchAIQuery', function () {
       );
     });
   });
+
+  describe('when the server errors with an AIError', function () {
+    let stopServer: () => Promise<void>;
+
+    beforeEach(async function () {
+      // Start a mock server to pass an ai response.
+      // Set the server endpoint in the env.
+      const { endpoint, stop } = await startMockAIServer({
+        response: {
+          status: 500,
+          body: {
+            name: 'AIError',
+            errorMessage: 'tortillas',
+            codeName: 'ExampleCode',
+          },
+        },
+      });
+
+      stopServer = stop;
+      process.env.DEV_AI_QUERY_ENDPOINT = endpoint;
+      process.env.DEV_AI_USERNAME = TEST_AUTH_USERNAME;
+      process.env.DEV_AI_PASSWORD = TEST_AUTH_PASSWORD;
+    });
+
+    afterEach(async function () {
+      await stopServer();
+      delete process.env.DEV_AI_QUERY_ENDPOINT;
+      delete process.env.DEV_AI_USERNAME;
+      delete process.env.DEV_AI_PASSWORD;
+    });
+
+    it('throws the error', async function () {
+      const promise = runFetchAIQuery({
+        ...mockUserPrompt,
+        signal: new AbortController().signal,
+      });
+
+      await expect(promise).to.be.rejectedWith('Error: ExampleCode: tortillas');
+    });
+  });
 });

packages/compass-query-bar/src/modules/ai-query-request.ts

@@ -5,6 +5,8 @@ import type { AbortSignal as NodeFetchAbortSignal } from 'node-fetch/externals';
 import type { SimplifiedSchema } from 'mongodb-schema';
 import type { Document } from 'mongodb';
 
+const serverErrorMessageName = 'AIError';
+
 function getAIQueryEndpoint(): string {
   if (!process.env.DEV_AI_QUERY_ENDPOINT) {
     throw new Error(
@@ -81,7 +83,20 @@ export async function runFetchAIQuery({
   });
 
   if (!res.ok) {
-    throw new Error(`Error: ${res.status} ${res.statusText}`);
+    // We try to parse the response to see if the server returned any
+    // information we can show a user.
+    let serverErrorMessage = `${res.status} ${res.statusText}`;
+    try {
+      const messageJSON = await res.json();
+      if (messageJSON.name === serverErrorMessageName) {
+        serverErrorMessage = `${messageJSON.codeName as string}: ${
+          messageJSON.errorMessage as string
+        }`;
+      }
+    } catch (err) {
+      // no-op, use the default status and statusText in the message.
+    }
+    throw new Error(`Error: ${serverErrorMessage}`);
   }
 
   const jsonResponse = await res.json();

packages/compass-query-bar/src/stores/ai-query-reducer.spec.ts

@@ -169,8 +169,10 @@ describe('aiQueryReducer', function () {
 
       beforeEach(async function () {
         const { endpoint, stop } = await startMockAIServer({
-          response: {},
-          sendError: true,
+          response: {
+            status: 500,
+            body: 'test',
+          },
         });
 
         stopServer = stop;

packages/compass-query-bar/src/stores/ai-query-reducer.ts

@@ -7,6 +7,9 @@ import preferences from 'compass-preferences-model';
 import type { QueryBarThunkAction } from './query-bar-store';
 import { isAction } from '../utils';
 import { runFetchAIQuery } from '../modules/ai-query-request';
+import { mapQueryToFormFields } from '../utils/query';
+import type { QueryFormFields } from '../constants/query-properties';
+import { DEFAULT_FIELD_VALUES } from '../constants/query-bar-store';
 
 const { log, mongoLogId } = createLoggerAndTelemetry('AI-QUERY-UI');
 
@@ -79,7 +82,7 @@ type AIQueryFailedAction = {
 
 export type AIQuerySucceededAction = {
   type: AIQueryActionTypes.AIQuerySucceeded;
-  query: unknown;
+  fields: QueryFormFields;
 };
 
 function logFailed(errorMessage: string) {
@@ -166,15 +169,19 @@ export const runAIQuery = (
       return;
     }
 
-    let query;
+    let fields;
     try {
       if (!jsonResponse?.content?.query) {
         throw new Error(
           'No query returned. Please try again with a different prompt.'
         );
       }
 
-      query = jsonResponse?.content?.query;
+      const query = jsonResponse?.content?.query;
+      fields = mapQueryToFormFields({
+        ...DEFAULT_FIELD_VALUES,
+        ...(query ?? {}),
+      });
     } catch (err: any) {
       logFailed(err?.message);
       dispatch({
@@ -184,9 +191,8 @@ export const runAIQuery = (
       return;
     }
 
-    // Error if the response is empty. TODO: We'll want to also parse if no
-    // applicable query fields are detected.
-    if (!query || Object.keys(query).length === 0) {
+    // Error when the response is empty or there is nothing to map.
+    if (!fields || Object.keys(fields).length === 0) {
       const msg =
         'No query was returned from the ai. Consider re-wording your prompt.';
       logFailed(msg);
@@ -202,13 +208,15 @@ export const runAIQuery = (
       'AIQuery',
       'AI query request succeeded',
       {
-        query,
+        query: {
+          ...fields,
+        },
       }
     );
 
     dispatch({
       type: AIQueryActionTypes.AIQuerySucceeded,
-      query,
+      fields,
     });
   };
 };

packages/compass-query-bar/src/stores/query-bar-reducer.ts

@@ -23,6 +23,7 @@ import {
   isQueryFieldsValid,
   validateField,
   isEqualDefaultQuery,
+  doesQueryHaveExtraOptionsSet,
 } from '../utils/query';
 import type { ChangeFilterEvent } from '../modules/change-filter';
 import { changeFilter } from '../modules/change-filter';
@@ -473,14 +474,7 @@ export const queryBarReducer: Reducer<QueryBarState> = (
   }
 
   if (
-    isAction<ApplyFromHistoryAction>(
-      action,
-      QueryBarActions.ApplyFromHistory
-    ) ||
-    isAction<AIQuerySucceededAction>(
-      action,
-      AIQueryActionTypes.AIQuerySucceeded
-    )
+    isAction<ApplyFromHistoryAction>(action, QueryBarActions.ApplyFromHistory)
   ) {
     return {
       ...state,
@@ -491,6 +485,19 @@ export const queryBarReducer: Reducer<QueryBarState> = (
     };
   }
 
+  if (
+    isAction<AIQuerySucceededAction>(
+      action,
+      AIQueryActionTypes.AIQuerySucceeded
+    )
+  ) {
+    return {
+      ...state,
+      expanded: state.expanded || doesQueryHaveExtraOptionsSet(action.fields),
+      fields: action.fields,
+    };
+  }
+
   if (
     isAction<RecentQueriesFetchedAction>(
       action,

packages/compass-query-bar/src/utils/query.spec.ts

@@ -0,0 +1,60 @@
+import { expect } from 'chai';
+
+import { doesQueryHaveExtraOptionsSet, mapQueryToFormFields } from './query';
+import { DEFAULT_FIELD_VALUES } from '../constants/query-bar-store';
+
+describe('#doesQueryHaveExtraOptionsSet', function () {
+  it('returns true when there is a non filter option, false otherwise', function () {
+    const defaultFields = mapQueryToFormFields(DEFAULT_FIELD_VALUES);
+
+    expect(
+      doesQueryHaveExtraOptionsSet({
+        ...defaultFields,
+      })
+    ).to.be.false;
+    expect(
+      doesQueryHaveExtraOptionsSet({
+        ...defaultFields,
+        filter: {
+          valid: true,
+          string: '{test: 2}',
+          value: {
+            test: 2,
+          },
+        },
+      })
+    ).to.be.false;
+    expect(
+      doesQueryHaveExtraOptionsSet({
+        ...defaultFields,
+        sort: {
+          valid: true,
+          string: '[["a", -1]]',
+          value: [['a', -1]],
+        },
+      })
+    ).to.be.true;
+    expect(
+      doesQueryHaveExtraOptionsSet({
+        ...defaultFields,
+        skip: {
+          valid: true,
+          string: '25',
+          value: 25,
+        },
+      })
+    ).to.be.true;
+    expect(
+      doesQueryHaveExtraOptionsSet({
+        ...defaultFields,
+        project: {
+          valid: true,
+          string: '{test: 1}',
+          value: {
+            test: 1,
+          },
+        },
+      })
+    ).to.be.true;
+  });
+});

packages/compass-query-bar/src/utils/query.ts

@@ -30,10 +30,34 @@ export function mapFormFieldsToQuery(fields: QueryFormFields): BaseQuery {
   };
 }
 
+// Returns true if any fields that aren't filter have non-default values.
+export function doesQueryHaveExtraOptionsSet(fields?: QueryFormFields) {
+  if (!fields) {
+    return false;
+  }
+
+  for (const property of QUERY_PROPERTIES) {
+    if (property === 'filter') {
+      continue;
+    }
+
+    if (
+      !isEqual(fields[property].value, DEFAULT_QUERY_VALUES[property]) &&
+      !isEqual(fields[property].value, DEFAULT_FIELD_VALUES[property])
+    ) {
+      return true;
+    }
+  }
+  return false;
+}
+
 /**
  * Map query document to the query fields state only preserving valid values
  */
-export function mapQueryToFormFields(query?: BaseQuery, onlyValid = true) {
+export function mapQueryToFormFields(
+  query?: BaseQuery,
+  onlyValid = true
+): QueryFormFields {
   return Object.fromEntries(
     Object.entries(query ?? {})
       .map(([key, _value]) => {

packages/compass-query-bar/test/create-mock-ai-endpoint.ts

@@ -20,16 +20,20 @@ function checkReqAuth(req: http.IncomingMessage) {
 export async function startMockAIServer(
   {
     response,
-    sendError,
   }: {
-    response: any;
-    sendError?: boolean;
+    response: {
+      status: number;
+      body: any;
+    };
   } = {
     response: {
-      content: {
-        query: {
-          find: {
-            test: 'pineapple',
+      status: 200,
+      body: {
+        content: {
+          query: {
+            find: {
+              test: 'pineapple',
+            },
           },
         },
       },
@@ -71,14 +75,11 @@ export async function startMockAIServer(
             content: jsonObject,
           });
 
-          if (sendError) {
-            res.writeHead(500);
-            res.end('Error occurred.');
-            return;
-          }
-
           res.setHeader('Content-Type', 'application/json');
-          return res.end(JSON.stringify(response));
+          if (response.status !== 200) {
+            res.writeHead(response.status);
+          }
+          return res.end(JSON.stringify(response.body));
         });
     })
     .listen(0);