send metadata to api

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

@@ -912,6 +912,7 @@ describe('AtlasAiService', function () {
         const mockAtlasService = new MockAtlasService();
         await preferences.savePreferences({
           enableChatbotEndpointForGenAI: true,
+          telemetryAtlasUserId: '1234',
         });
         atlasAiService = new AtlasAiService({
           apiURLPreset: 'cloud',
@@ -1047,7 +1048,11 @@ describe('AtlasAiService', function () {
             const requestBody = JSON.parse(args[1].body as string);
 
             expect(requestBody.model).to.equal('mongodb-chat-latest');
-            expect(requestBody.store).to.equal(false);
+            expect(requestBody.metadata).to.deep.equal({
+              userId: '1234',
+              store: 'true',
+              sensitiveStorage: 'sensitive',
+            });
             expect(requestBody.instructions).to.be.a('string');
             expect(requestBody.input).to.be.an('array');
 

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

@@ -259,6 +259,15 @@ export type MockDataSchemaResponse = z.infer<
   typeof MockDataSchemaResponseShape
 >;
 
+// TODO: Evaluate this
+function getActiveUserId(preferences: PreferencesAccess): string {
+  const { currentUserId, telemetryAnonymousId, telemetryAtlasUserId } =
+    preferences.getPreferences();
+  return (
+    currentUserId || telemetryAnonymousId || telemetryAtlasUserId || 'unknown'
+  );
+}
+
 /**
  * The type of resource from the natural language query REST API
  */
@@ -304,7 +313,13 @@ export class AtlasAiService {
           PLACEHOLDER_BASE_URL,
           this.atlasService.assistantApiEndpoint()
         );
-        return this.atlasService.authenticatedFetch(uri, init);
+        return this.atlasService.authenticatedFetch(uri, {
+          ...init,
+          headers: {
+            ...(init?.headers ?? {}),
+            entrypoint: 'natural-language-to-mql',
+          },
+        });
       },
       // TODO(COMPASS-10125): Switch the model to `mongodb-slim-latest` when
       // enabling this feature (to use edu-chatbot for GenAI).
@@ -445,7 +460,10 @@ export class AtlasAiService {
     connectionInfo: ConnectionInfo
   ) {
     if (this.preferences.getPreferences().enableChatbotEndpointForGenAI) {
-      const message = buildAggregateQueryPrompt(input);
+      const message = buildAggregateQueryPrompt({
+        ...input,
+        userId: getActiveUserId(this.preferences),
+      });
       return this.generateQueryUsingChatbot(
         message,
         validateAIAggregationResponse,
@@ -467,7 +485,10 @@ export class AtlasAiService {
     connectionInfo: ConnectionInfo
   ) {
     if (this.preferences.getPreferences().enableChatbotEndpointForGenAI) {
-      const message = buildFindQueryPrompt(input);
+      const message = buildFindQueryPrompt({
+        ...input,
+        userId: getActiveUserId(this.preferences),
+      });
       return this.generateQueryUsingChatbot(message, validateAIQueryResponse, {
         signal: input.signal,
       });

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

@@ -8,13 +8,14 @@ export async function getAiQueryResponse(
   message: AiQueryPrompt,
   abortSignal: AbortSignal
 ): Promise<string> {
+  const { instructions, ...restOfMetadata } = message.metadata;
   const response = streamText({
     model,
     messages: [{ role: 'user', content: message.prompt }],
     providerOptions: {
       openai: {
-        store: false,
-        instructions: message.metadata.instructions,
+        instructions,
+        metadata: restOfMetadata,
       },
     },
     abortSignal,