handle fle collections

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

@@ -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/index.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/pipeline-builder/pipeline-ai.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/stores/store.ts

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

packages/compass-generative-ai/src/atlas-ai-service.ts

@@ -34,6 +34,7 @@ type GenerativeAiInput = {
   sampleDocuments?: Document[];
   signal: AbortSignal;
   requestId: string;
+  enableStorage: boolean;
 };
 
 // The size/token validation happens on the server, however, we do

packages/compass-generative-ai/src/utils/gen-ai-prompt.ts

@@ -62,6 +62,7 @@ export type PromptContextOptions = {
   databaseName: string;
   collectionName: string;
   userId: string;
+  enableStorage: boolean;
   schema?: unknown;
   sampleDocuments?: unknown[];
 };
@@ -82,7 +83,9 @@ function buildUserPromptForQuery({
   collectionName,
   schema,
   sampleDocuments,
-}: PromptContextOptions & { type: 'find' | 'aggregate' }): string {
+}: Omit<PromptContextOptions, 'userId' | 'enableStorage'> & {
+  type: 'find' | 'aggregate';
+}): string {
   const messages = [];
 
   const queryPrompt = [
@@ -156,50 +159,72 @@ export type AiQueryPrompt = {
   metadata: {
     instructions: string;
     userId: string;
-    store: 'true' | 'false';
-    sensitiveStorage: 'sensitive';
-  };
+  } & (
+    | {
+        store: 'true';
+        sensitiveStorage: 'sensitive';
+      }
+    | {
+        store: 'false';
+      }
+  );
 };
 
 function buildMetadata({
   instructions,
   userId,
+  enableStorage,
 }: {
   instructions: string;
   userId: string;
+  enableStorage: boolean;
 }): AiQueryPrompt['metadata'] {
   return {
     instructions,
     userId,
-    sensitiveStorage: 'sensitive',
-    store: 'true',
+    ...(enableStorage
+      ? {
+          sensitiveStorage: 'sensitive',
+          store: 'true',
+        }
+      : {
+          store: 'false',
+        }),
   };
 }
 
-export function buildFindQueryPrompt(
-  options: PromptContextOptions
-): AiQueryPrompt {
+export function buildFindQueryPrompt({
+  userId,
+  enableStorage,
+  ...restOfTheOptions
+}: PromptContextOptions): AiQueryPrompt {
   const prompt = buildUserPromptForQuery({
     type: 'find',
-    ...options,
+    ...restOfTheOptions,
   });
   const instructions = buildInstructionsForFindQuery();
   return {
     prompt,
-    metadata: buildMetadata({ instructions, userId: options.userId }),
+    metadata: buildMetadata({
+      instructions,
+      userId,
+      enableStorage,
+    }),
   };
 }
 
-export function buildAggregateQueryPrompt(
-  options: PromptContextOptions
-): AiQueryPrompt {
+export function buildAggregateQueryPrompt({
+  userId,
+  enableStorage,
+  ...restOfTheOptions
+}: PromptContextOptions): AiQueryPrompt {
   const prompt = buildUserPromptForQuery({
     type: 'aggregate',
-    ...options,
+    ...restOfTheOptions,
   });
   const instructions = buildInstructionsForAggregateQuery();
   return {
     prompt,
-    metadata: buildMetadata({ instructions, userId: options.userId }),
+    metadata: buildMetadata({ instructions, userId, enableStorage }),
   };
 }

packages/compass-query-bar/src/index.tsx

@@ -6,7 +6,10 @@ import {
   dataServiceLocator,
   type DataServiceLocator,
 } from '@mongodb-js/compass-connections/provider';
-import { mongoDBInstanceLocator } from '@mongodb-js/compass-app-stores/provider';
+import {
+  mongoDBInstanceLocator,
+  collectionModelLocator,
+} from '@mongodb-js/compass-app-stores/provider';
 import {
   QueryBarComponentProvider,
   useQueryBarComponent,
@@ -56,6 +59,7 @@ const QueryBarPlugin = registerCompassPlugin(
     atlasAiService: atlasAiServiceLocator,
     favoriteQueryStorageAccess: favoriteQueryStorageAccessLocator,
     recentQueryStorageAccess: recentQueryStorageAccessLocator,
+    collection: collectionModelLocator,
   }
 );
 

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

@@ -166,6 +166,7 @@ export const runAIQuery = (
       logger: { log },
       connectionInfoRef,
       track,
+      collection,
     }
   ) => {
     const provideSampleDocuments =
@@ -218,6 +219,7 @@ export const runAIQuery = (
         }
       );
       const schema = await getSimplifiedSchema(sampleDocuments);
+      const { isFLE } = await collection.fetchMetadata({ dataService });
 
       const { collection: collectionName, database: databaseName } =
         toNS(namespace);
@@ -235,6 +237,7 @@ export const runAIQuery = (
               }
             : undefined),
           requestId,
+          enableStorage: !isFLE,
         },
         connectionInfo
       );

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

@@ -35,9 +35,18 @@ import type {
   RecentQueryStorage,
 } from '@mongodb-js/my-queries-storage/provider';
 import type { TrackFunction } from '@mongodb-js/compass-telemetry';
+import type Collection from 'mongodb-collection-model';
 
 // Partial of DataService that mms shares with Compass.
-type QueryBarDataService = Pick<DataService, 'sample' | 'getConnectionString'>;
+type FetchCollectionMetadataDataServiceMethods =
+  | 'collectionStats'
+  | 'collectionInfo'
+  | 'listCollections'
+  | 'isListSearchIndexesSupported';
+type QueryBarDataService = Pick<
+  DataService,
+  'sample' | 'getConnectionString' | FetchCollectionMetadataDataServiceMethods
+>;
 
 type QueryBarServices = {
   instance: MongoDBInstance;
@@ -51,6 +60,7 @@ type QueryBarServices = {
   atlasAiService: AtlasAiService;
   favoriteQueryStorageAccess?: FavoriteQueryStorageAccess;
   recentQueryStorageAccess?: RecentQueryStorageAccess;
+  collection: Collection;
 };
 
 // TODO(COMPASS-7412): this doesn't have service injector
@@ -73,14 +83,18 @@ export type RootState = ReturnType<typeof rootQueryBarReducer>;
 export type QueryBarExtraArgs = {
   globalAppRegistry: AppRegistry;
   localAppRegistry: AppRegistry;
-  dataService: Pick<QueryBarDataService, 'sample'>;
+  dataService: Pick<
+    QueryBarDataService,
+    'sample' | FetchCollectionMetadataDataServiceMethods
+  >;
   preferences: PreferencesAccess;
   favoriteQueryStorage?: FavoriteQueryStorage;
   recentQueryStorage?: RecentQueryStorage;
   logger: Logger;
   track: TrackFunction;
   connectionInfoRef: ConnectionInfoRef;
   atlasAiService: AtlasAiService;
+  collection: Collection;
 };
 
 export type QueryBarThunkDispatch<A extends AnyAction = AnyAction> =
@@ -126,6 +140,7 @@ export function activatePlugin(
     atlasAiService,
     favoriteQueryStorageAccess,
     recentQueryStorageAccess,
+    collection,
   } = services;
 
   const favoriteQueryStorage = favoriteQueryStorageAccess?.getStorage();
@@ -158,6 +173,7 @@ export function activatePlugin(
       track,
       connectionInfoRef,
       atlasAiService,
+      collection,
     }
   );