feat(gen-ai): wrap prompt in tags for storage COMPASS-10083 (#7636)

* wrap prompt in tags

* send metadata to api

* assertions in e2e test

* clean up

* handle fle collections

* clean up and fix tests

* types fix

* pass request id in headers

* query-bar fixes

* aggregation fixes

* fix tests and check

* fix tests

* fix tests

* escape user input

* escape user_prompt tag from input

* allow headers for cors in tests

* hash user id

* fix test

Author: mabaasit

package-lock.json

@@ -1,6 +1,13 @@
 import type { DataService as OriginalDataService } from 'mongodb-data-service';
 
+type FetchCollectionMetadataDataServiceMethods =
+  | 'collectionStats'
+  | 'collectionInfo'
+  | 'listCollections'
+  | 'isListSearchIndexesSupported';
+
 export type RequiredDataServiceProps =
+  | FetchCollectionMetadataDataServiceMethods
   | 'isCancelError'
   | 'estimatedCount'
   | 'aggregate'

packages/compass-aggregations/src/modules/data-service.ts

@@ -49,6 +49,7 @@ import type {
   ConnectionScopedAppRegistry,
 } from '@mongodb-js/compass-connections/provider';
 import type { TrackFunction } from '@mongodb-js/compass-telemetry';
+import type Collection from 'mongodb-collection-model';
 /**
  * The main application reducer.
  *
@@ -110,6 +111,7 @@ export type PipelineBuilderExtraArgs = {
   connectionScopedAppRegistry: ConnectionScopedAppRegistry<
     'open-export' | 'view-edited' | 'agg-pipeline-out-executed'
   >;
+  collection: Collection;
 };
 
 export type PipelineBuilderThunkDispatch<A extends Action = AnyAction> =

packages/compass-aggregations/src/modules/index.ts

@@ -225,6 +225,7 @@ export const runAIPipelineGeneration = (
       logger: { log, mongoLogId },
       track,
       connectionInfoRef,
+      collection,
     }
   ) => {
     const {
@@ -286,6 +287,9 @@ export const runAIPipelineGeneration = (
           }
         )) || [];
       const schema = await getSimplifiedSchema(sampleDocuments);
+      const { isFLE } = await collection.fetchMetadata({
+        dataService: dataService!,
+      });
 
       const { collection: collectionName, database: databaseName } =
         toNS(namespace);
@@ -303,6 +307,7 @@ export const runAIPipelineGeneration = (
               }
             : undefined),
           requestId,
+          enableStorage: !isFLE,
         },
         connectionInfo
       );

packages/compass-aggregations/src/modules/pipeline-builder/pipeline-ai.ts

@@ -141,6 +141,11 @@ function createStore({
         dataService: {} as any,
         connectionInfoRef,
         connectionScopedAppRegistry,
+        collection: {
+          fetchMetadata() {
+            return Promise.resolve({ isFLE: false });
+          },
+        } as any,
       })
     )
   );

packages/compass-aggregations/src/modules/pipeline-builder/stage-editor.spec.ts

@@ -188,6 +188,7 @@ export function activateAggregationsPlugin(
         connectionInfoRef,
         connectionScopedAppRegistry,
         dataService,
+        collection: collectionModel,
       })
     )
   );

packages/compass-aggregations/src/stores/store.ts

@@ -32,6 +32,7 @@ function getMockedPluginArgs(
     CompassAggregationsPlugin.provider.withMockServices({
       atlasAiService,
       collection: {
+        fetchMetadata: () => ({}),
         toJSON: () => ({}),
         on: () => {},
         removeListener: () => {},

packages/compass-aggregations/test/configure-store.ts

@@ -18,11 +18,14 @@ export const MockQueryBarPlugin: typeof QueryBarPlugin =
       getConnectionString() {
         return { hosts: [] } as any;
       },
-    },
+    } as any,
     instance: { on() {}, removeListener() {} } as any,
     favoriteQueryStorageAccess: compassFavoriteQueryStorageAccess,
     recentQueryStorageAccess: compassRecentQueryStorageAccess,
     atlasAiService: {} as any,
+    collection: {
+      fetchMetadata: () => Promise.resolve({} as any),
+    } as any,
   });
 
 export const renderWithQueryBar = (

packages/compass-crud/test/render-with-query-bar.tsx

@@ -156,16 +156,16 @@ export async function startMockAssistantServer(
   getResponse: () => MockAssistantResponse;
   setResponse: (response: MockAssistantResponse) => void;
   getRequests: () => {
-    content: any;
-    req: any;
+    content: Record<string, any>;
+    req: http.IncomingMessage;
   }[];
   endpoint: string;
   server: http.Server;
   stop: () => Promise<void>;
 }> {
   let requests: {
-    content: any;
-    req: any;
+    content: Record<string, any>;
+    req: http.IncomingMessage;
   }[] = [];
   let response = _response;
   const server = http
@@ -174,7 +174,7 @@ export async function startMockAssistantServer(
       res.setHeader('Access-Control-Allow-Methods', 'POST, OPTIONS');
       res.setHeader(
         'Access-Control-Allow-Headers',
-        'Content-Type, Authorization, X-Request-Origin, User-Agent, X-CSRF-Token, X-CSRF-Time'
+        'Content-Type, Authorization, X-Request-Origin, User-Agent, X-CSRF-Token, X-CSRF-Time, Entrypoint, X-Client-Request-Id'
       );
       res.setHeader('Access-Control-Allow-Credentials', 'true');
 

packages/compass-e2e-tests/helpers/assistant-service.ts

@@ -265,10 +265,16 @@ describe('Collection ai query with chatbot (with mocked backend)', function () {
       expect(requests.length).to.equal(1);
 
       const queryRequest = requests[0];
+      expect(queryRequest.req.headers).to.have.property('x-client-request-id');
       // TODO(COMPASS-10125): Switch the model to `mongodb-slim-latest` when
       // enabling this feature.
       expect(queryRequest.content.model).to.equal('mongodb-chat-latest');
       expect(queryRequest.content.instructions).to.be.string;
+      expect(queryRequest.content.metadata).to.have.property('userId');
+      expect(queryRequest.content.metadata.store).to.have.equal('true');
+      expect(queryRequest.content.metadata.sensitiveStorage).to.have.equal(
+        'sensitive'
+      );
       expect(queryRequest.content.input).to.be.an('array').of.length(1);
 
       const message = queryRequest.content.input[0];