From 1295d519972f48ef933196dd57fdb7cef8bd0c18 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Tue, 25 Nov 2025 12:16:13 +0100 Subject: [PATCH 01/23] add feature flag --- packages/compass-preferences-model/src/feature-flags.ts | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/packages/compass-preferences-model/src/feature-flags.ts b/packages/compass-preferences-model/src/feature-flags.ts index b2517c6aa51..d3d1ebfcb92 100644 --- a/packages/compass-preferences-model/src/feature-flags.ts +++ b/packages/compass-preferences-model/src/feature-flags.ts @@ -224,6 +224,14 @@ export const FEATURE_FLAG_DEFINITIONS = [ 'Enable automatic relationship inference during data model generation', }, }, + { + name: 'enableChatbotEndpointForGenAI', + stage: 'development', + atlasCloudFeatureFlagName: null, + description: { + short: 'Enable Chatbot API for Generative AI', + }, + }, ] as const satisfies ReadonlyArray; type FeatureFlagDefinitions = typeof FEATURE_FLAG_DEFINITIONS; From 7a44de4464413bed1fb958473e9c7042e2fb4aef Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Tue, 25 Nov 2025 12:16:22 +0100 Subject: [PATCH 02/23] migrate prompts to compass --- .../src/utils/gen-api-prompt.ts | 130 ++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 packages/compass-generative-ai/src/utils/gen-api-prompt.ts diff --git a/packages/compass-generative-ai/src/utils/gen-api-prompt.ts b/packages/compass-generative-ai/src/utils/gen-api-prompt.ts new file mode 100644 index 00000000000..92b0995b509 --- /dev/null +++ b/packages/compass-generative-ai/src/utils/gen-api-prompt.ts @@ -0,0 +1,130 @@ +const MAX_TOTAL_PROMPT_LENGTH = 512000; +const MIN_SAMPLE_DOCUMENTS = 1; + +function getCurrentTimeString() { + const dateTime = new Date(); + const options: Intl.DateTimeFormatOptions = { + weekday: 'short', + year: 'numeric', + month: 'short', + day: '2-digit', + hour: '2-digit', + minute: '2-digit', + second: '2-digit', + timeZoneName: 'short', + hour12: false, + }; + // Tue, Nov 25, 2025, 12:06:20 GMT+1 + return dateTime.toLocaleString('en-US', options); +} + +function buildInstructionsForFindQuery() { + return [ + 'Reduce prose to the minimum, your output will be parsed by a machine. ' + + 'You generate MongoDB find query arguments. Provide filter, project, sort, skip, ' + + 'limit and aggregation in shell syntax, wrap each argument with XML delimiters as follows:', + '{}', + '{}', + '{}', + '0', + '0', + '[]', + 'Additional instructions:', + '- Only use the aggregation field when the request cannot be represented with the other fields.', + '- Do not use the aggregation field if a find query fulfills the objective.', + '- If specifying latitude and longitude coordinates, list the longitude first, and then latitude.', + `- The current date is ${getCurrentTimeString()}`, + ].join('\n'); +} + +function buildInstructionsForAggregateQuery() { + return [ + 'Reduce prose to the minimum, your output will be parsed by a machine. ' + + 'You generate MongoDB aggregation pipelines. Provide only the aggregation ' + + 'pipeline contents in an array in shell syntax, wrapped with XML delimiters as follows:', + '[]', + 'Additional instructions:', + '- If specifying latitude and longitude coordinates, list the longitude first, and then latitude.', + '- Only pass the contents of the aggregation, no surrounding syntax.', + `- The current date is ${getCurrentTimeString()}`, + ].join('\n'); +} + +type QueryType = 'find' | 'aggregate'; + +export type UserPromptForQueryOptions = { + type: QueryType; + userPrompt: string; + databaseName?: string; + collectionName?: string; + schema?: unknown; + sampleDocuments?: unknown[]; +}; + +export function buildInstructionsForQuery({ type }: { type: QueryType }) { + return type === 'find' + ? buildInstructionsForFindQuery() + : buildInstructionsForAggregateQuery(); +} + +export function buildUserPromptForQuery({ + type, + userPrompt, + databaseName, + collectionName, + schema, + sampleDocuments, +}: UserPromptForQueryOptions): string { + const messages = []; + + const queryPrompt = [ + type === 'find' ? 'Write a query' : 'Generate an aggregation', + 'that does the following:', + `"${userPrompt}"`, + ].join(' '); + + if (databaseName) { + messages.push(`Database name: "${databaseName}"`); + } + if (collectionName) { + messages.push(`Collection name: "${collectionName}"`); + } + if (schema) { + messages.push( + 'Schema from a sample of documents from the collection: ```' + + JSON.stringify(schema) + + '```' + ); + } + if (sampleDocuments) { + // When attaching the sample documents, we want to ensure that we do not + // exceed the token limit. So we try following: + // 1. If attaching all the sample documents exceeds then limit, we attach only 1 document. + // 2. If attaching 1 document still exceeds the limit, we do not attach any sample documents. + const sampleDocumentsStr = JSON.stringify(sampleDocuments); + const singleDocumentStr = JSON.stringify( + sampleDocuments.slice(0, MIN_SAMPLE_DOCUMENTS) + ); + if ( + sampleDocumentsStr.length + + messages.join('\n').length + + queryPrompt.length <= + MAX_TOTAL_PROMPT_LENGTH + ) { + messages.push( + 'Sample documents from the collection: ```' + sampleDocumentsStr + '```' + ); + } else if ( + singleDocumentStr.length + + messages.join('\n').length + + queryPrompt.length <= + MAX_TOTAL_PROMPT_LENGTH + ) { + messages.push( + 'Sample document from the collection: ```' + singleDocumentStr + '```' + ); + } + } + messages.push(queryPrompt); + return messages.join('\n'); +} From d52468238fe08a6dce87bfa23584e57023f0990e Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Tue, 25 Nov 2025 13:41:11 +0100 Subject: [PATCH 03/23] co-pilot feedback --- .../src/utils/gen-api-prompt.ts | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/packages/compass-generative-ai/src/utils/gen-api-prompt.ts b/packages/compass-generative-ai/src/utils/gen-api-prompt.ts index 92b0995b509..a67f0e22b6c 100644 --- a/packages/compass-generative-ai/src/utils/gen-api-prompt.ts +++ b/packages/compass-generative-ai/src/utils/gen-api-prompt.ts @@ -1,3 +1,7 @@ +// When including sample documents, we want to ensure that we do not +// attach large documents and exceed the limit. OpenAI roughly estimates +// 4 characters = 1 token and we should not exceed context window limits. +// This roughly translates to 128k tokens. const MAX_TOTAL_PROMPT_LENGTH = 512000; const MIN_SAMPLE_DOCUMENTS = 1; @@ -14,7 +18,7 @@ function getCurrentTimeString() { timeZoneName: 'short', hour12: false, }; - // Tue, Nov 25, 2025, 12:06:20 GMT+1 + // e.g. Tue, Nov 25, 2025, 12:00:00 GMT+1 return dateTime.toLocaleString('en-US', options); } @@ -105,19 +109,17 @@ export function buildUserPromptForQuery({ const singleDocumentStr = JSON.stringify( sampleDocuments.slice(0, MIN_SAMPLE_DOCUMENTS) ); + const promptLengthWithoutSampleDocs = + messages.join('\n').length + queryPrompt.length; if ( - sampleDocumentsStr.length + - messages.join('\n').length + - queryPrompt.length <= + sampleDocumentsStr.length + promptLengthWithoutSampleDocs <= MAX_TOTAL_PROMPT_LENGTH ) { messages.push( 'Sample documents from the collection: ```' + sampleDocumentsStr + '```' ); } else if ( - singleDocumentStr.length + - messages.join('\n').length + - queryPrompt.length <= + singleDocumentStr.length + promptLengthWithoutSampleDocs <= MAX_TOTAL_PROMPT_LENGTH ) { messages.push( From c96161be71aff540bb886be766bb0ba1aaf12263 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Mon, 1 Dec 2025 10:13:31 +0300 Subject: [PATCH 04/23] clean up --- .../{gen-api-prompt.ts => gen-ai-prompt.ts} | 68 ++++++++++++++++--- 1 file changed, 57 insertions(+), 11 deletions(-) rename packages/compass-generative-ai/src/utils/{gen-api-prompt.ts => gen-ai-prompt.ts} (79%) diff --git a/packages/compass-generative-ai/src/utils/gen-api-prompt.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts similarity index 79% rename from packages/compass-generative-ai/src/utils/gen-api-prompt.ts rename to packages/compass-generative-ai/src/utils/gen-ai-prompt.ts index a67f0e22b6c..caa8d7f6639 100644 --- a/packages/compass-generative-ai/src/utils/gen-api-prompt.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts @@ -54,10 +54,7 @@ function buildInstructionsForAggregateQuery() { ].join('\n'); } -type QueryType = 'find' | 'aggregate'; - export type UserPromptForQueryOptions = { - type: QueryType; userPrompt: string; databaseName?: string; collectionName?: string; @@ -65,20 +62,14 @@ export type UserPromptForQueryOptions = { sampleDocuments?: unknown[]; }; -export function buildInstructionsForQuery({ type }: { type: QueryType }) { - return type === 'find' - ? buildInstructionsForFindQuery() - : buildInstructionsForAggregateQuery(); -} - -export function buildUserPromptForQuery({ +function buildUserPromptForQuery({ type, userPrompt, databaseName, collectionName, schema, sampleDocuments, -}: UserPromptForQueryOptions): string { +}: UserPromptForQueryOptions & { type: 'find' | 'aggregate' }): string { const messages = []; const queryPrompt = [ @@ -130,3 +121,58 @@ export function buildUserPromptForQuery({ messages.push(queryPrompt); return messages.join('\n'); } + +export type AiQueryPrompt = { + prompt: string; + metadata: { + instructions: string; + }; +}; + +export function buildFindQueryPrompt({ + userPrompt, + databaseName, + collectionName, + schema, + sampleDocuments, +}: UserPromptForQueryOptions): AiQueryPrompt { + const prompt = buildUserPromptForQuery({ + type: 'find', + userPrompt, + databaseName, + collectionName, + schema, + sampleDocuments, + }); + const instructions = buildInstructionsForFindQuery(); + return { + prompt, + metadata: { + instructions, + }, + }; +} + +export function buildAggregateQueryPrompt({ + userPrompt, + databaseName, + collectionName, + schema, + sampleDocuments, +}: UserPromptForQueryOptions): AiQueryPrompt { + const prompt = buildUserPromptForQuery({ + type: 'aggregate', + userPrompt, + databaseName, + collectionName, + schema, + sampleDocuments, + }); + const instructions = buildInstructionsForAggregateQuery(); + return { + prompt, + metadata: { + instructions, + }, + }; +} From f4d05a773a9a265bd74707dc58b9f7fe3e3bf1b9 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Mon, 1 Dec 2025 17:27:39 +0300 Subject: [PATCH 05/23] use edu api for gen ai --- package-lock.json | 140 ++++++++++++++ packages/compass-generative-ai/package.json | 4 +- .../src/atlas-ai-errors.ts | 13 +- .../src/atlas-ai-service.ts | 172 ++++++++++++++++-- .../src/utils/ai-chat-transport.ts | 81 +++++++++ .../src/utils/gen-ai-prompt.ts | 38 ++-- packages/compass-web/src/entrypoint.tsx | 8 +- .../compass/src/app/components/entrypoint.tsx | 1 + 8 files changed, 412 insertions(+), 45 deletions(-) create mode 100644 packages/compass-generative-ai/src/utils/ai-chat-transport.ts diff --git a/package-lock.json b/package-lock.json index c60ba80e5c7..3f105fcfda1 100644 --- a/package-lock.json +++ b/package-lock.json @@ -16785,6 +16785,15 @@ } } }, + "node_modules/@vercel/oidc": { + "version": "3.0.5", + "resolved": "https://registry.npmjs.org/@vercel/oidc/-/oidc-3.0.5.tgz", + "integrity": "sha512-fnYhv671l+eTTp48gB4zEsTW/YtRgRPnkI2nT7x6qw5rkI1Lq2hTmQIpHPgyThI0znLK+vX2n9XxKdXZ7BUbbw==", + "license": "Apache-2.0", + "engines": { + "node": ">= 20" + } + }, "node_modules/@vue/compiler-core": { "version": "3.5.6", "resolved": "https://registry.npmjs.org/@vue/compiler-core/-/compiler-core-3.5.6.tgz", @@ -25641,6 +25650,15 @@ "node": ">=0.4.x" } }, + "node_modules/eventsource-parser": { + "version": "3.0.6", + "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.0.6.tgz", + "integrity": "sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg==", + "license": "MIT", + "engines": { + "node": ">=18.0.0" + } + }, "node_modules/evp_bytestokey": { "version": "1.0.3", "resolved": "https://registry.npmjs.org/evp_bytestokey/-/evp_bytestokey-1.0.3.tgz", @@ -49683,6 +49701,7 @@ "version": "0.68.0", "license": "SSPL", "dependencies": { + "@ai-sdk/openai": "^2.0.4", "@mongodb-js/atlas-service": "^0.73.0", "@mongodb-js/compass-app-registry": "^9.4.29", "@mongodb-js/compass-components": "^1.59.2", @@ -49691,6 +49710,7 @@ "@mongodb-js/compass-telemetry": "^1.19.5", "@mongodb-js/compass-utils": "^0.9.23", "@mongodb-js/connection-info": "^0.24.0", + "ai": "^5.0.26", "bson": "^6.10.4", "compass-preferences-model": "^2.66.3", "mongodb": "^6.19.0", @@ -49723,6 +49743,74 @@ "xvfb-maybe": "^0.2.1" } }, + "packages/compass-generative-ai/node_modules/@ai-sdk/gateway": { + "version": "2.0.17", + "resolved": "https://registry.npmjs.org/@ai-sdk/gateway/-/gateway-2.0.17.tgz", + "integrity": "sha512-oVAG6q72KsjKlrYdLhWjRO7rcqAR8CjokAbYuyVZoCO4Uh2PH/VzZoxZav71w2ipwlXhHCNaInGYWNs889MMDA==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "2.0.0", + "@ai-sdk/provider-utils": "3.0.18", + "@vercel/oidc": "3.0.5" + }, + "engines": { + "node": ">=18" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "packages/compass-generative-ai/node_modules/@ai-sdk/openai": { + "version": "2.0.75", + "resolved": "https://registry.npmjs.org/@ai-sdk/openai/-/openai-2.0.75.tgz", + "integrity": "sha512-ThDHg1+Jes7S0AOXa01EyLBSzZiZwzB5do9vAlufNkoiRHGTH1BmoShrCyci/TUsg4ky1HwbK4hPK+Z0isiE6g==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "2.0.0", + "@ai-sdk/provider-utils": "3.0.18" + }, + "engines": { + "node": ">=18" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "packages/compass-generative-ai/node_modules/@ai-sdk/provider-utils": { + "version": "3.0.18", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.18.tgz", + "integrity": "sha512-ypv1xXMsgGcNKUP+hglKqtdDuMg68nWHucPPAhIENrbFAI+xCHiqPVN8Zllxyv1TNZwGWUghPxJXU+Mqps0YRQ==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "2.0.0", + "@standard-schema/spec": "^1.0.0", + "eventsource-parser": "^3.0.6" + }, + "engines": { + "node": ">=18" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "packages/compass-generative-ai/node_modules/ai": { + "version": "5.0.104", + "resolved": "https://registry.npmjs.org/ai/-/ai-5.0.104.tgz", + "integrity": "sha512-MZOkL9++nY5PfkpWKBR3Rv+Oygxpb9S16ctv8h91GvrSif7UnNEdPMVZe3bUyMd2djxf0AtBk/csBixP0WwWZQ==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/gateway": "2.0.17", + "@ai-sdk/provider": "2.0.0", + "@ai-sdk/provider-utils": "3.0.18", + "@opentelemetry/api": "1.9.0" + }, + "engines": { + "node": ">=18" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, "packages/compass-generative-ai/node_modules/diff": { "version": "4.0.2", "resolved": "https://registry.npmjs.org/diff/-/diff-4.0.2.tgz", @@ -62394,6 +62482,7 @@ "@mongodb-js/compass-generative-ai": { "version": "file:packages/compass-generative-ai", "requires": { + "@ai-sdk/openai": "^2.0.4", "@mongodb-js/atlas-service": "^0.73.0", "@mongodb-js/compass-app-registry": "^9.4.29", "@mongodb-js/compass-components": "^1.59.2", @@ -62412,6 +62501,7 @@ "@types/mocha": "^9.0.0", "@types/react": "^17.0.5", "@types/sinon-chai": "^3.2.5", + "ai": "^5.0.26", "bson": "^6.10.4", "chai": "^4.3.6", "compass-preferences-model": "^2.66.3", @@ -62432,6 +62522,46 @@ "zod": "^3.25.76" }, "dependencies": { + "@ai-sdk/gateway": { + "version": "2.0.17", + "resolved": "https://registry.npmjs.org/@ai-sdk/gateway/-/gateway-2.0.17.tgz", + "integrity": "sha512-oVAG6q72KsjKlrYdLhWjRO7rcqAR8CjokAbYuyVZoCO4Uh2PH/VzZoxZav71w2ipwlXhHCNaInGYWNs889MMDA==", + "requires": { + "@ai-sdk/provider": "2.0.0", + "@ai-sdk/provider-utils": "3.0.18", + "@vercel/oidc": "3.0.5" + } + }, + "@ai-sdk/openai": { + "version": "2.0.75", + "resolved": "https://registry.npmjs.org/@ai-sdk/openai/-/openai-2.0.75.tgz", + "integrity": "sha512-ThDHg1+Jes7S0AOXa01EyLBSzZiZwzB5do9vAlufNkoiRHGTH1BmoShrCyci/TUsg4ky1HwbK4hPK+Z0isiE6g==", + "requires": { + "@ai-sdk/provider": "2.0.0", + "@ai-sdk/provider-utils": "3.0.18" + } + }, + "@ai-sdk/provider-utils": { + "version": "3.0.18", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.18.tgz", + "integrity": "sha512-ypv1xXMsgGcNKUP+hglKqtdDuMg68nWHucPPAhIENrbFAI+xCHiqPVN8Zllxyv1TNZwGWUghPxJXU+Mqps0YRQ==", + "requires": { + "@ai-sdk/provider": "2.0.0", + "@standard-schema/spec": "^1.0.0", + "eventsource-parser": "^3.0.6" + } + }, + "ai": { + "version": "5.0.104", + "resolved": "https://registry.npmjs.org/ai/-/ai-5.0.104.tgz", + "integrity": "sha512-MZOkL9++nY5PfkpWKBR3Rv+Oygxpb9S16ctv8h91GvrSif7UnNEdPMVZe3bUyMd2djxf0AtBk/csBixP0WwWZQ==", + "requires": { + "@ai-sdk/gateway": "2.0.17", + "@ai-sdk/provider": "2.0.0", + "@ai-sdk/provider-utils": "3.0.18", + "@opentelemetry/api": "1.9.0" + } + }, "diff": { "version": "4.0.2", "resolved": "https://registry.npmjs.org/diff/-/diff-4.0.2.tgz", @@ -70995,6 +71125,11 @@ "dev": true, "requires": {} }, + "@vercel/oidc": { + "version": "3.0.5", + "resolved": "https://registry.npmjs.org/@vercel/oidc/-/oidc-3.0.5.tgz", + "integrity": "sha512-fnYhv671l+eTTp48gB4zEsTW/YtRgRPnkI2nT7x6qw5rkI1Lq2hTmQIpHPgyThI0znLK+vX2n9XxKdXZ7BUbbw==" + }, "@vue/compiler-core": { "version": "3.5.6", "resolved": "https://registry.npmjs.org/@vue/compiler-core/-/compiler-core-3.5.6.tgz", @@ -78150,6 +78285,11 @@ "resolved": "https://registry.npmjs.org/events/-/events-1.1.1.tgz", "integrity": "sha512-kEcvvCBByWXGnZy6JUlgAp2gBIUjfCAV6P6TgT1/aaQKcmuAEC4OZTV1I4EWQLz2gxZw76atuVyvHhTxvi0Flw==" }, + "eventsource-parser": { + "version": "3.0.6", + "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.0.6.tgz", + "integrity": "sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg==" + }, "evp_bytestokey": { "version": "1.0.3", "resolved": "https://registry.npmjs.org/evp_bytestokey/-/evp_bytestokey-1.0.3.tgz", diff --git a/packages/compass-generative-ai/package.json b/packages/compass-generative-ai/package.json index 5b9b01abb92..5450a68ea6e 100644 --- a/packages/compass-generative-ai/package.json +++ b/packages/compass-generative-ai/package.json @@ -68,7 +68,9 @@ "react-redux": "^8.1.3", "redux": "^4.2.1", "redux-thunk": "^2.4.2", - "zod": "^3.25.76" + "zod": "^3.25.76", + "@ai-sdk/openai": "^2.0.4", + "ai": "^5.0.26" }, "devDependencies": { "@mongodb-js/eslint-config-compass": "^1.4.12", diff --git a/packages/compass-generative-ai/src/atlas-ai-errors.ts b/packages/compass-generative-ai/src/atlas-ai-errors.ts index 992377cf18b..e8160b88324 100644 --- a/packages/compass-generative-ai/src/atlas-ai-errors.ts +++ b/packages/compass-generative-ai/src/atlas-ai-errors.ts @@ -18,4 +18,15 @@ class AtlasAiServiceApiResponseParseError extends Error { } } -export { AtlasAiServiceInvalidInputError, AtlasAiServiceApiResponseParseError }; +class AtlasAiServiceGenAiResponseError extends Error { + constructor(message: string) { + super(message); + this.name = 'AtlasAiServiceGenAiResponseError'; + } +} + +export { + AtlasAiServiceInvalidInputError, + AtlasAiServiceApiResponseParseError, + AtlasAiServiceGenAiResponseError, +}; diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 04f75a90161..a6ceb53a68e 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -15,7 +15,17 @@ import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer'; import { AtlasAiServiceInvalidInputError, AtlasAiServiceApiResponseParseError, + AtlasAiServiceGenAiResponseError, } from './atlas-ai-errors'; +import { createOpenAI } from '@ai-sdk/openai'; +import { + AiChatTransport, + type SendMessagesPayload, +} from './utils/ai-chat-transport'; +import { + buildAggregateQueryPrompt, + buildFindQueryPrompt, +} from './utils/gen-ai-prompt'; type GenerativeAiInput = { userInput: string; @@ -40,14 +50,6 @@ type AIAggregation = { }; }; -type AIFeatureEnablement = { - features: { - [featureName: string]: { - enabled: boolean; - }; - }; -}; - type AIQuery = { content: { query: Record< @@ -100,6 +102,40 @@ function hasExtraneousKeys(obj: any, expectedKeys: string[]) { return Object.keys(obj).some((key) => !expectedKeys.includes(key)); } +/** + * For generating queries with the AI chatbot, currently we don't + * have a chat interface and only need to send a single message. + * This function builds the message so it can be sent to the AI model. + */ +function buildChatMessageForAiModel( + { signal, ...restOfInput }: Omit, + { type }: { type: 'find' | 'aggregate' } +): SendMessagesPayload { + const { prompt, metadata } = + type === 'aggregate' + ? buildAggregateQueryPrompt(restOfInput) + : buildFindQueryPrompt(restOfInput); + return { + trigger: 'submit-message', + chatId: crypto.randomUUID(), + messageId: undefined, + messages: [ + { + id: crypto.randomUUID(), + role: 'user', + parts: [ + { + type: 'text', + text: prompt, + }, + ], + metadata, + }, + ], + abortSignal: signal, + }; +} + export function validateAIQueryResponse( response: any ): asserts response is AIQuery { @@ -271,6 +307,8 @@ export class AtlasAiService { private preferences: PreferencesAccess; private logger: Logger; + private chatTransport: AiChatTransport; + constructor({ apiURLPreset, atlasService, @@ -286,8 +324,16 @@ export class AtlasAiService { this.atlasService = atlasService; this.preferences = preferences; this.logger = logger; - this.initPromise = this.setupAIAccess(); + + const model = createOpenAI({ + apiKey: '', + baseURL: this.atlasService.assistantApiEndpoint(), + fetch: this.atlasService.authenticatedFetch.bind(this.atlasService), + // TODO(COMPASS-10125): Switch the model to `mongodb-slim-latest` when + // enabling this feature (to use edu-chatbot for GenAI). + }).responses('mongodb-chat-latest'); + this.chatTransport = new AiChatTransport({ model }); } /** @@ -423,6 +469,13 @@ export class AtlasAiService { input: GenerativeAiInput, connectionInfo: ConnectionInfo ) { + if (this.preferences.getPreferences().enableChatbotEndpointForGenAI) { + return this.generateQueryUsingChatbot( + input, + validateAIAggregationResponse, + { type: 'aggregate' } + ); + } return this.getQueryOrAggregationFromUserInput( { connectionInfo, @@ -437,6 +490,11 @@ export class AtlasAiService { input: GenerativeAiInput, connectionInfo: ConnectionInfo ) { + if (this.preferences.getPreferences().enableChatbotEndpointForGenAI) { + return this.generateQueryUsingChatbot(input, validateAIQueryResponse, { + type: 'find', + }); + } return this.getQueryOrAggregationFromUserInput( { urlId: 'query', @@ -527,12 +585,96 @@ export class AtlasAiService { }); } - private validateAIFeatureEnablementResponse( - response: any - ): asserts response is AIFeatureEnablement { - const { features } = response; - if (typeof features !== 'object') { - throw new Error('Unexpected response: expected features to be an object'); + private parseXmlResponseAndMapToMmsJsonResponse(responseStr: string): any { + const expectedTags = [ + 'filter', + 'project', + 'sort', + 'skip', + 'limit', + 'aggregation', + ]; + + // Currently the prompt forces LLM to return xml-styled data + const result: any = {}; + for (const tag of expectedTags) { + const regex = new RegExp(`<${tag}>([\\s\\S]*?)<\\/${tag}>`, 'i'); + const match = responseStr.match(regex); + if (match && match[1]) { + const value = match[1].trim(); + try { + // Here the value is valid js string, but not valid json, so we use eval to parse it. + const tagValue = eval(`(${value})`); + if ( + !tagValue || + (typeof tagValue === 'object' && Object.keys(tagValue).length === 0) + ) { + result[tag] = null; + } else { + result[tag] = JSON.stringify(tagValue); + } + } catch (e) { + this.logger.log.warn( + this.logger.mongoLogId(1_001_000_309), + 'AtlasAiService', + `Failed to parse value for tag <${tag}>: ${value}`, + { error: e } + ); + result[tag] = null; + } + } } + + // Keep the response same as we have from mms api + return { + content: { + aggregation: { + pipeline: result.aggregation, + }, + query: { + filter: result.filter, + project: result.project, + sort: result.sort, + skip: result.skip, + limit: result.limit, + }, + }, + }; + } + + private async generateQueryUsingChatbot( + // TODO(COMPASS-10083): When storing data, we want have requestId as well. + input: Omit, + validateFn: (res: any) => asserts res is T, + options: { type: 'find' | 'aggregate' } + ): Promise { + this.throwIfAINotEnabled(); + const response = await this.chatTransport.sendMessages( + buildChatMessageForAiModel(input, options) + ); + + const chunks: string[] = []; + let done = false; + const reader = response.getReader(); + while (!done) { + const { done: _done, value } = await reader.read(); + if (_done) { + done = true; + break; + } + if (value.type === 'text-delta') { + chunks.push(value.delta); + } + if (value.type === 'error') { + throw new AtlasAiServiceGenAiResponseError(value.errorText); + } + } + + const responseText = chunks.join(''); + const parsedResponse = + this.parseXmlResponseAndMapToMmsJsonResponse(responseText); + validateFn(parsedResponse); + + return parsedResponse; } } diff --git a/packages/compass-generative-ai/src/utils/ai-chat-transport.ts b/packages/compass-generative-ai/src/utils/ai-chat-transport.ts new file mode 100644 index 00000000000..bc51640b627 --- /dev/null +++ b/packages/compass-generative-ai/src/utils/ai-chat-transport.ts @@ -0,0 +1,81 @@ +import { + type ChatTransport, + type LanguageModel, + type UIMessageChunk, + type UIMessage, + convertToModelMessages, + streamText, +} from 'ai'; + +type ChatMessageMetadata = { + confirmation?: boolean; + instructions?: string; + sendWithoutHistory?: boolean; +}; +type ChatMessage = UIMessage; + +export type SendMessagesPayload = Parameters< + ChatTransport['sendMessages'] +>[0]; + +/** Returns true if the message should be excluded from being sent to the API. */ +export function shouldExcludeMessage({ metadata }: ChatMessage) { + if (metadata?.confirmation) { + return true; + } + return false; +} + +export class AiChatTransport implements ChatTransport { + private model: LanguageModel; + + constructor({ model }: { model: LanguageModel }) { + this.model = model; + } + + static emptyStream = new ReadableStream({ + start(controller) { + controller.close(); + }, + }); + + sendMessages({ messages, abortSignal }: SendMessagesPayload) { + // If the most recent message is a message that is meant to be excluded + // then we do not need to send this request to the assistant API as it's likely + // redundant otherwise. + if (shouldExcludeMessage(messages[messages.length - 1])) { + return Promise.resolve(AiChatTransport.emptyStream); + } + + const filteredMessages = messages.filter( + (message) => !shouldExcludeMessage(message) + ); + + // If no messages remain after filtering, return an empty stream + if (filteredMessages.length === 0) { + return Promise.resolve(AiChatTransport.emptyStream); + } + + const lastMessage = filteredMessages[filteredMessages.length - 1]; + const result = streamText({ + model: this.model, + messages: lastMessage.metadata?.sendWithoutHistory + ? convertToModelMessages([lastMessage]) + : convertToModelMessages(filteredMessages), + abortSignal: abortSignal, + providerOptions: { + openai: { + store: false, + instructions: lastMessage.metadata?.instructions ?? '', + }, + }, + }); + + return Promise.resolve(result.toUIMessageStream({ sendSources: true })); + } + + reconnectToStream(): Promise | null> { + // For this implementation, we don't support reconnecting to streams + return Promise.resolve(null); + } +} diff --git a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts index caa8d7f6639..eb8d61459c1 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts @@ -55,7 +55,7 @@ function buildInstructionsForAggregateQuery() { } export type UserPromptForQueryOptions = { - userPrompt: string; + userInput: string; databaseName?: string; collectionName?: string; schema?: unknown; @@ -64,7 +64,7 @@ export type UserPromptForQueryOptions = { function buildUserPromptForQuery({ type, - userPrompt, + userInput, databaseName, collectionName, schema, @@ -75,7 +75,7 @@ function buildUserPromptForQuery({ const queryPrompt = [ type === 'find' ? 'Write a query' : 'Generate an aggregation', 'that does the following:', - `"${userPrompt}"`, + `"${userInput}"`, ].join(' '); if (databaseName) { @@ -129,20 +129,12 @@ export type AiQueryPrompt = { }; }; -export function buildFindQueryPrompt({ - userPrompt, - databaseName, - collectionName, - schema, - sampleDocuments, -}: UserPromptForQueryOptions): AiQueryPrompt { +export function buildFindQueryPrompt( + options: UserPromptForQueryOptions +): AiQueryPrompt { const prompt = buildUserPromptForQuery({ type: 'find', - userPrompt, - databaseName, - collectionName, - schema, - sampleDocuments, + ...options, }); const instructions = buildInstructionsForFindQuery(); return { @@ -153,20 +145,12 @@ export function buildFindQueryPrompt({ }; } -export function buildAggregateQueryPrompt({ - userPrompt, - databaseName, - collectionName, - schema, - sampleDocuments, -}: UserPromptForQueryOptions): AiQueryPrompt { +export function buildAggregateQueryPrompt( + options: UserPromptForQueryOptions +): AiQueryPrompt { const prompt = buildUserPromptForQuery({ type: 'aggregate', - userPrompt, - databaseName, - collectionName, - schema, - sampleDocuments, + ...options, }); const instructions = buildInstructionsForAggregateQuery(); return { diff --git a/packages/compass-web/src/entrypoint.tsx b/packages/compass-web/src/entrypoint.tsx index 830c297ed22..a6efbce17d0 100644 --- a/packages/compass-web/src/entrypoint.tsx +++ b/packages/compass-web/src/entrypoint.tsx @@ -105,7 +105,13 @@ const WithAtlasProviders: React.FC<{ children: React.ReactNode }> = ({ return ( - + {children} diff --git a/packages/compass/src/app/components/entrypoint.tsx b/packages/compass/src/app/components/entrypoint.tsx index ec6911bdd65..2cbf3e97d99 100644 --- a/packages/compass/src/app/components/entrypoint.tsx +++ b/packages/compass/src/app/components/entrypoint.tsx @@ -61,6 +61,7 @@ export const WithAtlasProviders: React.FC = ({ children }) => { options={{ defaultHeaders: { 'User-Agent': `${getAppName()}/${getAppVersion()}`, + 'X-Request-Origin': 'mongodb-compass', }, }} > From 17635315c80460b5bcaef72934000f5bb7eef63f Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Mon, 1 Dec 2025 18:15:18 +0300 Subject: [PATCH 06/23] clean up a bit --- .../src/atlas-ai-service.spec.ts | 1 + .../src/atlas-ai-service.ts | 65 ++--------------- .../src/utils/xml-to-mms-response.spec.ts | 69 +++++++++++++++++++ .../src/utils/xml-to-mms-response.ts | 58 ++++++++++++++++ 4 files changed, 133 insertions(+), 60 deletions(-) create mode 100644 packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts create mode 100644 packages/compass-generative-ai/src/utils/xml-to-mms-response.ts diff --git a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts index 1767fef6dd5..a9dc4016f0f 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -51,6 +51,7 @@ class MockAtlasService { getCurrentUser = () => Promise.resolve(ATLAS_USER); cloudEndpoint = (url: string) => `${['/cloud', url].join('/')}`; adminApiEndpoint = (url: string) => `${[BASE_URL, url].join('/')}`; + assistantApiEndpoint = (url: string) => `${[BASE_URL, url].join('/')}`; authenticatedFetch = (url: string, init: RequestInit) => { return fetch(url, init); }; diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index a6ceb53a68e..f8b40576252 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -26,6 +26,7 @@ import { buildAggregateQueryPrompt, buildFindQueryPrompt, } from './utils/gen-ai-prompt'; +import { parseXmlToMmsJsonResponse } from './utils/xml-to-mms-response'; type GenerativeAiInput = { userInput: string; @@ -585,63 +586,6 @@ export class AtlasAiService { }); } - private parseXmlResponseAndMapToMmsJsonResponse(responseStr: string): any { - const expectedTags = [ - 'filter', - 'project', - 'sort', - 'skip', - 'limit', - 'aggregation', - ]; - - // Currently the prompt forces LLM to return xml-styled data - const result: any = {}; - for (const tag of expectedTags) { - const regex = new RegExp(`<${tag}>([\\s\\S]*?)<\\/${tag}>`, 'i'); - const match = responseStr.match(regex); - if (match && match[1]) { - const value = match[1].trim(); - try { - // Here the value is valid js string, but not valid json, so we use eval to parse it. - const tagValue = eval(`(${value})`); - if ( - !tagValue || - (typeof tagValue === 'object' && Object.keys(tagValue).length === 0) - ) { - result[tag] = null; - } else { - result[tag] = JSON.stringify(tagValue); - } - } catch (e) { - this.logger.log.warn( - this.logger.mongoLogId(1_001_000_309), - 'AtlasAiService', - `Failed to parse value for tag <${tag}>: ${value}`, - { error: e } - ); - result[tag] = null; - } - } - } - - // Keep the response same as we have from mms api - return { - content: { - aggregation: { - pipeline: result.aggregation, - }, - query: { - filter: result.filter, - project: result.project, - sort: result.sort, - skip: result.skip, - limit: result.limit, - }, - }, - }; - } - private async generateQueryUsingChatbot( // TODO(COMPASS-10083): When storing data, we want have requestId as well. input: Omit, @@ -670,9 +614,10 @@ export class AtlasAiService { } } - const responseText = chunks.join(''); - const parsedResponse = - this.parseXmlResponseAndMapToMmsJsonResponse(responseText); + const parsedResponse = parseXmlToMmsJsonResponse( + chunks.join(''), + this.logger + ); validateFn(parsedResponse); return parsedResponse; diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts new file mode 100644 index 00000000000..ca02709890e --- /dev/null +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts @@ -0,0 +1,69 @@ +import { expect } from 'chai'; +import { parseXmlToMmsJsonResponse } from './xml-to-mms-response'; +import type { Logger } from '@mongodb-js/compass-logging'; + +const loggerMock = { + log: { + warn: () => { + /* noop */ + }, + }, + mongoLogId: (id: number) => id, +} as unknown as Logger; +describe('parseXmlToMmsJsonResponse', function () { + it('should parse valid XML string to MMS JSON response', function () { + const xmlString = ` + { age: { $gt: 25 } } + { name: 1, age: 1 } + { age: -1 } + 5 + 10 + [{ $match: { status: "A" } }] + `; + + const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + + expect(result).to.deep.equal({ + content: { + aggregation: { + pipeline: '[{"$match":{"status":"A"}}]', + }, + query: { + filter: '{"age":{"$gt":25}}', + project: '{"name":1,"age":1}', + sort: '{"age":-1}', + skip: '5', + limit: '10', + }, + }, + }); + }); + + context('it should return null values for invalid data', function () { + it('invalid json', function () { + const result = parseXmlToMmsJsonResponse( + `{ age: { $gt: 25 `, + loggerMock + ); + expect(result.content.query.filter).to.equal(null); + }); + it('empty object', function () { + const result = parseXmlToMmsJsonResponse( + `{}`, + loggerMock + ); + expect(result.content.query.filter).to.equal(null); + }); + it('empty array', function () { + const result = parseXmlToMmsJsonResponse( + `[]`, + loggerMock + ); + expect(result.content.aggregation.pipeline).to.equal(null); + }); + it('zero value', function () { + const result = parseXmlToMmsJsonResponse(`0`, loggerMock); + expect(result.content.query.limit).to.equal(null); + }); + }); +}); diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts new file mode 100644 index 00000000000..e62c122d1f3 --- /dev/null +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -0,0 +1,58 @@ +import type { Logger } from '@mongodb-js/compass-logging'; + +export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { + const expectedTags = [ + 'filter', + 'project', + 'sort', + 'skip', + 'limit', + 'aggregation', + ]; + + // Currently the prompt forces LLM to return xml-styled data + const result: any = {}; + for (const tag of expectedTags) { + const regex = new RegExp(`<${tag}>([\\s\\S]*?)<\\/${tag}>`, 'i'); + const match = xmlString.match(regex); + if (match && match[1]) { + const value = match[1].trim(); + try { + // Here the value is valid js string, but not valid json, so we use eval to parse it. + const tagValue = eval(`(${value})`); + if ( + !tagValue || + (typeof tagValue === 'object' && Object.keys(tagValue).length === 0) + ) { + result[tag] = null; + } else { + result[tag] = JSON.stringify(tagValue); + } + } catch (e) { + logger.log.warn( + logger.mongoLogId(1_001_000_309), + 'AtlasAiService', + `Failed to parse value for tag <${tag}>: ${value}`, + { error: e } + ); + result[tag] = null; + } + } + } + + // Keep the response same as we have from mms api + return { + content: { + aggregation: { + pipeline: result.aggregation, + }, + query: { + filter: result.filter, + project: result.project, + sort: result.sort, + skip: result.skip, + limit: result.limit, + }, + }, + }; +} From d5e2c8454ba1048a39adfa450a59c64df9c10ad5 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 01:37:44 +0300 Subject: [PATCH 07/23] fix url for test and ensure aggregations have content --- .../compass-generative-ai/src/atlas-ai-service.ts | 14 ++++++++++++-- .../src/utils/xml-to-mms-response.ts | 4 +--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index f8b40576252..521563118e2 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -327,10 +327,20 @@ export class AtlasAiService { this.logger = logger; this.initPromise = this.setupAIAccess(); + const initialBaseUrl = 'http://PLACEHOLDER_BASE_URL_TO_BE_REPLACED.invalid'; const model = createOpenAI({ apiKey: '', - baseURL: this.atlasService.assistantApiEndpoint(), - fetch: this.atlasService.authenticatedFetch.bind(this.atlasService), + baseURL: initialBaseUrl, + fetch: (url, init) => { + // The `baseUrl` can be dynamically changed, but `createOpenAI` + // doesn't allow us to change it after initial call. Instead + // we're going to update it every time the fetch call happens + const uri = String(url).replace( + initialBaseUrl, + this.atlasService.assistantApiEndpoint() + ); + return this.atlasService.authenticatedFetch(uri, init); + }, // TODO(COMPASS-10125): Switch the model to `mongodb-slim-latest` when // enabling this feature (to use edu-chatbot for GenAI). }).responses('mongodb-chat-latest'); diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts index e62c122d1f3..bd543aa1444 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -43,9 +43,7 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { // Keep the response same as we have from mms api return { content: { - aggregation: { - pipeline: result.aggregation, - }, + ...(result.aggregation ? { aggregation: result.aggregation } : {}), query: { filter: result.filter, project: result.project, From 07815800e112776264c939af1c3e18d44c78e4e6 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 10:47:36 +0300 Subject: [PATCH 08/23] tests --- .../helpers/assistant-service.ts | 2 +- .../tests/collection-ai-query.test.ts | 145 ++++++++++++++++++ 2 files changed, 146 insertions(+), 1 deletion(-) diff --git a/packages/compass-e2e-tests/helpers/assistant-service.ts b/packages/compass-e2e-tests/helpers/assistant-service.ts index c36a3d9ae3d..5abca2c7d5f 100644 --- a/packages/compass-e2e-tests/helpers/assistant-service.ts +++ b/packages/compass-e2e-tests/helpers/assistant-service.ts @@ -212,8 +212,8 @@ export async function startMockAssistantServer( }); if (response.status !== 200) { - res.writeHead(response.status); res.setHeader('Content-Type', 'application/json'); + res.writeHead(response.status); return res.end(JSON.stringify({ error: response.body })); } diff --git a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts index 919aa54490b..e84159481f1 100644 --- a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts +++ b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts @@ -8,12 +8,14 @@ import { cleanup, screenshotIfFailed, DEFAULT_CONNECTION_NAME_1, + screenshotPathName, } from '../helpers/compass'; import type { Compass } from '../helpers/compass'; import * as Selectors from '../helpers/selectors'; import { createNumbersCollection } from '../helpers/insert-data'; import { startMockAtlasServiceServer } from '../helpers/mock-atlas-service'; import type { MockAtlasServerResponse } from '../helpers/mock-atlas-service'; +import { startMockAssistantServer } from '../helpers/assistant-service'; describe('Collection ai query (with mocked backend)', function () { let compass: Compass; @@ -171,3 +173,146 @@ describe('Collection ai query (with mocked backend)', function () { }); }); }); + +async function setup( + browser: CompassBrowser, + dbName: string, + collName: string +) { + await createNumbersCollection(); + await browser.setupDefaultConnections(); + await browser.connectToDefaults(); + await browser.navigateToCollectionTab( + DEFAULT_CONNECTION_NAME_1, + dbName, + collName, + 'Documents' + ); + + await browser.setFeature('enableChatbotEndpointForGenAI', true); + await browser.setFeature('enableGenAIFeatures', true); + await browser.setFeature('enableGenAISampleDocumentPassing', true); + await browser.setFeature('optInGenAIFeatures', true); +} + +describe.only('Collection ai query with chatbot (with mocked backend)', function () { + const dbName = 'test'; + const collName = 'numbers'; + let compass: Compass; + let browser: CompassBrowser; + + let mockAssistantServer: Awaited>; + + before(async function () { + mockAssistantServer = await startMockAssistantServer(); + compass = await init(this.test?.fullTitle()); + browser = compass.browser; + + await browser.setEnv( + 'COMPASS_ASSISTANT_BASE_URL_OVERRIDE', + mockAssistantServer.endpoint + ); + }); + + after(async function () { + await mockAssistantServer.stop(); + await cleanup(compass); + }); + + afterEach(async function () { + await screenshotIfFailed(compass, this.currentTest); + try { + mockAssistantServer.clearRequests(); + } catch (err) { + await browser.screenshot(screenshotPathName('afterEach-GenAi-Query')); + throw err; + } + }); + + describe('when the ai model response is valid', function () { + beforeEach(async function () { + await setup(browser, dbName, collName); + mockAssistantServer.setResponse({ + status: 200, + body: '{i: {$gt: 50}}', + }); + }); + + it('makes request to the server and updates the query bar with the response', async function () { + await new Promise((resolve) => setTimeout(resolve, 10000)); + // Click the ai entry button. + await browser.clickVisible(Selectors.GenAIEntryButton); + + // Enter the ai prompt. + await browser.clickVisible(Selectors.GenAITextInput); + + const testUserInput = 'find all documents where i is greater than 50'; + await browser.setValueVisible(Selectors.GenAITextInput, testUserInput); + + // Click generate. + await browser.clickVisible(Selectors.GenAIGenerateQueryButton); + + // Wait for the ipc events to succeed. + await browser.waitUntil(async function () { + // Make sure the query bar was updated. + const queryBarFilterContent = await browser.getCodemirrorEditorText( + Selectors.queryBarOptionInputFilter('Documents') + ); + return queryBarFilterContent === '{"i":{"$gt":50}}'; + }); + + // Check that the request was made with the correct parameters. + const requests = mockAssistantServer.getRequests(); + expect(requests.length).to.equal(1); + + const queryRequest = requests[0]; + // 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.input).to.be.an('array').of.length(1); + + const message = queryRequest.content.input[0]; + expect(message.role).to.equal('user'); + expect(message.content).to.be.an('array').of.length(1); + expect(message.content[0]).to.have.property('type'); + expect(message.content[0]).to.have.property('text'); + + // Run it and check that the correct documents are shown. + await browser.runFind('Documents', true); + const modifiedResult = await browser.getFirstListDocument(); + expect(modifiedResult.i).to.be.equal('51'); + }); + }); + + describe.only('when the chatbot api request errors', function () { + beforeEach(async function () { + await setup(browser, dbName, collName); + mockAssistantServer.setResponse({ + status: 500, + body: '', + }); + }); + + it('the error is shown to the user', async function () { + // Click the ai entry button. + await browser.clickVisible(Selectors.GenAIEntryButton); + + // Enter the ai prompt. + await browser.clickVisible(Selectors.GenAITextInput); + + const testUserInput = 'find all documents where i is greater than 50'; + await browser.setValueVisible(Selectors.GenAITextInput, testUserInput); + + // Click generate. + await browser.clickVisible(Selectors.GenAIGenerateQueryButton); + + // Check that the error is shown. + const errorBanner = browser.$(Selectors.GenAIErrorMessageBanner); + await errorBanner.waitForDisplayed(); + expect(await errorBanner.getText()).to.equal( + 'Sorry, we were unable to generate the query, please try again. If the error persists, try changing your prompt.' + ); + }); + }); +}); From 058e950003387a49de5e9c7301d872590efe0176 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 13:57:07 +0300 Subject: [PATCH 09/23] fix error handling --- .../tests/collection-ai-query.test.ts | 3 +-- .../src/atlas-ai-errors.ts | 13 +----------- .../src/atlas-ai-service.ts | 4 ++-- .../src/chatbot-errors.ts | 20 +++++++++++++++++++ 4 files changed, 24 insertions(+), 16 deletions(-) create mode 100644 packages/compass-generative-ai/src/chatbot-errors.ts diff --git a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts index e84159481f1..67de96ed58d 100644 --- a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts +++ b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts @@ -239,7 +239,6 @@ describe.only('Collection ai query with chatbot (with mocked backend)', function }); it('makes request to the server and updates the query bar with the response', async function () { - await new Promise((resolve) => setTimeout(resolve, 10000)); // Click the ai entry button. await browser.clickVisible(Selectors.GenAIEntryButton); @@ -285,7 +284,7 @@ describe.only('Collection ai query with chatbot (with mocked backend)', function }); }); - describe.only('when the chatbot api request errors', function () { + describe('when the chatbot api request errors', function () { beforeEach(async function () { await setup(browser, dbName, collName); mockAssistantServer.setResponse({ diff --git a/packages/compass-generative-ai/src/atlas-ai-errors.ts b/packages/compass-generative-ai/src/atlas-ai-errors.ts index e8160b88324..992377cf18b 100644 --- a/packages/compass-generative-ai/src/atlas-ai-errors.ts +++ b/packages/compass-generative-ai/src/atlas-ai-errors.ts @@ -18,15 +18,4 @@ class AtlasAiServiceApiResponseParseError extends Error { } } -class AtlasAiServiceGenAiResponseError extends Error { - constructor(message: string) { - super(message); - this.name = 'AtlasAiServiceGenAiResponseError'; - } -} - -export { - AtlasAiServiceInvalidInputError, - AtlasAiServiceApiResponseParseError, - AtlasAiServiceGenAiResponseError, -}; +export { AtlasAiServiceInvalidInputError, AtlasAiServiceApiResponseParseError }; diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 521563118e2..54d5ffd7538 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -15,7 +15,6 @@ import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer'; import { AtlasAiServiceInvalidInputError, AtlasAiServiceApiResponseParseError, - AtlasAiServiceGenAiResponseError, } from './atlas-ai-errors'; import { createOpenAI } from '@ai-sdk/openai'; import { @@ -27,6 +26,7 @@ import { buildFindQueryPrompt, } from './utils/gen-ai-prompt'; import { parseXmlToMmsJsonResponse } from './utils/xml-to-mms-response'; +import { AiChatbotInvalidResponseError } from './chatbot-errors'; type GenerativeAiInput = { userInput: string; @@ -620,7 +620,7 @@ export class AtlasAiService { chunks.push(value.delta); } if (value.type === 'error') { - throw new AtlasAiServiceGenAiResponseError(value.errorText); + throw new AiChatbotInvalidResponseError(value.errorText); } } diff --git a/packages/compass-generative-ai/src/chatbot-errors.ts b/packages/compass-generative-ai/src/chatbot-errors.ts new file mode 100644 index 00000000000..2d778a46b8f --- /dev/null +++ b/packages/compass-generative-ai/src/chatbot-errors.ts @@ -0,0 +1,20 @@ +export class AiChatbotError extends Error { + statusCode: number; + errorCode: string; + detail: string; + + constructor(statusCode: number, detail: string, errorCode: string) { + super(`${errorCode}: ${detail}`); + this.name = 'ServerError'; + this.statusCode = statusCode; + this.errorCode = errorCode; + this.detail = detail; + } +} + +export class AiChatbotInvalidResponseError extends AiChatbotError { + constructor(message: string) { + super(500, message, 'INVALID_RESPONSE'); + this.name = 'AiChatbotInvalidResponseError'; + } +} From 7d54d4d0f19f5e128c03d57a7a54e18cff374dae Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 14:48:19 +0300 Subject: [PATCH 10/23] clean up transport --- .../src/atlas-ai-service.ts | 92 ++++--------------- .../src/utils/ai-chat-transport.ts | 81 ---------------- .../src/utils/gen-ai-response.ts | 42 +++++++++ .../src/utils/xml-to-mms-response.spec.ts | 77 ++++++++++++++-- .../src/utils/xml-to-mms-response.ts | 32 +++++-- 5 files changed, 156 insertions(+), 168 deletions(-) delete mode 100644 packages/compass-generative-ai/src/utils/ai-chat-transport.ts create mode 100644 packages/compass-generative-ai/src/utils/gen-ai-response.ts diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 54d5ffd7538..80ab92a6120 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -17,16 +17,14 @@ import { AtlasAiServiceApiResponseParseError, } from './atlas-ai-errors'; import { createOpenAI } from '@ai-sdk/openai'; -import { - AiChatTransport, - type SendMessagesPayload, -} from './utils/ai-chat-transport'; +import { type LanguageModel } from 'ai'; +import type { AiQueryPrompt } from './utils/gen-ai-prompt'; import { buildAggregateQueryPrompt, buildFindQueryPrompt, } from './utils/gen-ai-prompt'; import { parseXmlToMmsJsonResponse } from './utils/xml-to-mms-response'; -import { AiChatbotInvalidResponseError } from './chatbot-errors'; +import { getAiQueryResponse } from './utils/gen-ai-response'; type GenerativeAiInput = { userInput: string; @@ -103,40 +101,6 @@ function hasExtraneousKeys(obj: any, expectedKeys: string[]) { return Object.keys(obj).some((key) => !expectedKeys.includes(key)); } -/** - * For generating queries with the AI chatbot, currently we don't - * have a chat interface and only need to send a single message. - * This function builds the message so it can be sent to the AI model. - */ -function buildChatMessageForAiModel( - { signal, ...restOfInput }: Omit, - { type }: { type: 'find' | 'aggregate' } -): SendMessagesPayload { - const { prompt, metadata } = - type === 'aggregate' - ? buildAggregateQueryPrompt(restOfInput) - : buildFindQueryPrompt(restOfInput); - return { - trigger: 'submit-message', - chatId: crypto.randomUUID(), - messageId: undefined, - messages: [ - { - id: crypto.randomUUID(), - role: 'user', - parts: [ - { - type: 'text', - text: prompt, - }, - ], - metadata, - }, - ], - abortSignal: signal, - }; -} - export function validateAIQueryResponse( response: any ): asserts response is AIQuery { @@ -308,7 +272,7 @@ export class AtlasAiService { private preferences: PreferencesAccess; private logger: Logger; - private chatTransport: AiChatTransport; + private aiModel: LanguageModel; constructor({ apiURLPreset, @@ -328,7 +292,7 @@ export class AtlasAiService { this.initPromise = this.setupAIAccess(); const initialBaseUrl = 'http://PLACEHOLDER_BASE_URL_TO_BE_REPLACED.invalid'; - const model = createOpenAI({ + this.aiModel = createOpenAI({ apiKey: '', baseURL: initialBaseUrl, fetch: (url, init) => { @@ -344,7 +308,6 @@ export class AtlasAiService { // TODO(COMPASS-10125): Switch the model to `mongodb-slim-latest` when // enabling this feature (to use edu-chatbot for GenAI). }).responses('mongodb-chat-latest'); - this.chatTransport = new AiChatTransport({ model }); } /** @@ -481,10 +444,11 @@ export class AtlasAiService { connectionInfo: ConnectionInfo ) { if (this.preferences.getPreferences().enableChatbotEndpointForGenAI) { + const message = buildAggregateQueryPrompt(input); return this.generateQueryUsingChatbot( - input, + message, validateAIAggregationResponse, - { type: 'aggregate' } + { signal: input.signal } ); } return this.getQueryOrAggregationFromUserInput( @@ -502,8 +466,9 @@ export class AtlasAiService { connectionInfo: ConnectionInfo ) { if (this.preferences.getPreferences().enableChatbotEndpointForGenAI) { - return this.generateQueryUsingChatbot(input, validateAIQueryResponse, { - type: 'find', + const message = buildFindQueryPrompt(input); + return this.generateQueryUsingChatbot(message, validateAIQueryResponse, { + signal: input.signal, }); } return this.getQueryOrAggregationFromUserInput( @@ -597,39 +562,18 @@ export class AtlasAiService { } private async generateQueryUsingChatbot( - // TODO(COMPASS-10083): When storing data, we want have requestId as well. - input: Omit, + message: AiQueryPrompt, validateFn: (res: any) => asserts res is T, - options: { type: 'find' | 'aggregate' } + options: { signal: AbortSignal } ): Promise { this.throwIfAINotEnabled(); - const response = await this.chatTransport.sendMessages( - buildChatMessageForAiModel(input, options) - ); - - const chunks: string[] = []; - let done = false; - const reader = response.getReader(); - while (!done) { - const { done: _done, value } = await reader.read(); - if (_done) { - done = true; - break; - } - if (value.type === 'text-delta') { - chunks.push(value.delta); - } - if (value.type === 'error') { - throw new AiChatbotInvalidResponseError(value.errorText); - } - } - - const parsedResponse = parseXmlToMmsJsonResponse( - chunks.join(''), - this.logger + const response = await getAiQueryResponse( + this.aiModel, + message, + options.signal ); + const parsedResponse = parseXmlToMmsJsonResponse(response, this.logger); validateFn(parsedResponse); - return parsedResponse; } } diff --git a/packages/compass-generative-ai/src/utils/ai-chat-transport.ts b/packages/compass-generative-ai/src/utils/ai-chat-transport.ts deleted file mode 100644 index bc51640b627..00000000000 --- a/packages/compass-generative-ai/src/utils/ai-chat-transport.ts +++ /dev/null @@ -1,81 +0,0 @@ -import { - type ChatTransport, - type LanguageModel, - type UIMessageChunk, - type UIMessage, - convertToModelMessages, - streamText, -} from 'ai'; - -type ChatMessageMetadata = { - confirmation?: boolean; - instructions?: string; - sendWithoutHistory?: boolean; -}; -type ChatMessage = UIMessage; - -export type SendMessagesPayload = Parameters< - ChatTransport['sendMessages'] ->[0]; - -/** Returns true if the message should be excluded from being sent to the API. */ -export function shouldExcludeMessage({ metadata }: ChatMessage) { - if (metadata?.confirmation) { - return true; - } - return false; -} - -export class AiChatTransport implements ChatTransport { - private model: LanguageModel; - - constructor({ model }: { model: LanguageModel }) { - this.model = model; - } - - static emptyStream = new ReadableStream({ - start(controller) { - controller.close(); - }, - }); - - sendMessages({ messages, abortSignal }: SendMessagesPayload) { - // If the most recent message is a message that is meant to be excluded - // then we do not need to send this request to the assistant API as it's likely - // redundant otherwise. - if (shouldExcludeMessage(messages[messages.length - 1])) { - return Promise.resolve(AiChatTransport.emptyStream); - } - - const filteredMessages = messages.filter( - (message) => !shouldExcludeMessage(message) - ); - - // If no messages remain after filtering, return an empty stream - if (filteredMessages.length === 0) { - return Promise.resolve(AiChatTransport.emptyStream); - } - - const lastMessage = filteredMessages[filteredMessages.length - 1]; - const result = streamText({ - model: this.model, - messages: lastMessage.metadata?.sendWithoutHistory - ? convertToModelMessages([lastMessage]) - : convertToModelMessages(filteredMessages), - abortSignal: abortSignal, - providerOptions: { - openai: { - store: false, - instructions: lastMessage.metadata?.instructions ?? '', - }, - }, - }); - - return Promise.resolve(result.toUIMessageStream({ sendSources: true })); - } - - reconnectToStream(): Promise | null> { - // For this implementation, we don't support reconnecting to streams - return Promise.resolve(null); - } -} diff --git a/packages/compass-generative-ai/src/utils/gen-ai-response.ts b/packages/compass-generative-ai/src/utils/gen-ai-response.ts new file mode 100644 index 00000000000..13ca53c786a --- /dev/null +++ b/packages/compass-generative-ai/src/utils/gen-ai-response.ts @@ -0,0 +1,42 @@ +import { AiChatbotInvalidResponseError } from '../chatbot-errors'; +import { type AiQueryPrompt } from './gen-ai-prompt'; +import type { LanguageModel } from 'ai'; +import { streamText } from 'ai'; +/** + * Function that takes an LanguageModel and a message of type + * AiQueryPrompt and returns a Promise that resolves to a string response + */ +export async function getAiQueryResponse( + model: LanguageModel, + message: AiQueryPrompt, + abortSignal: AbortSignal +): Promise { + const response = streamText({ + model, + messages: [{ role: 'user', content: message.prompt }], + providerOptions: { + openai: { + store: false, + instructions: message.metadata.instructions || '', + }, + }, + abortSignal, + }).toUIMessageStream(); + const chunks: string[] = []; + let done = false; + const reader = response.getReader(); + while (!done) { + const { done: _done, value } = await reader.read(); + if (_done) { + done = true; + break; + } + if (value.type === 'text-delta') { + chunks.push(value.delta); + } + if (value.type === 'error') { + throw new AiChatbotInvalidResponseError(value.errorText); + } + } + return chunks.join(''); +} diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts index ca02709890e..c0b5571b208 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts @@ -11,13 +11,9 @@ const loggerMock = { mongoLogId: (id: number) => id, } as unknown as Logger; describe('parseXmlToMmsJsonResponse', function () { - it('should parse valid XML string to MMS JSON response', function () { + it('should prioritize aggregation over query fields', function () { const xmlString = ` { age: { $gt: 25 } } - { name: 1, age: 1 } - { age: -1 } - 5 - 10 [{ $match: { status: "A" } }] `; @@ -28,6 +24,75 @@ describe('parseXmlToMmsJsonResponse', function () { aggregation: { pipeline: '[{"$match":{"status":"A"}}]', }, + query: { + filter: null, + project: null, + sort: null, + skip: null, + limit: null, + }, + }, + }); + }); + + it('should return null for aggregation if not provided', function () { + const xmlString = ` + { age: { $gt: 25 } } + `; + + const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + expect(result).to.deep.equal({ + content: { + aggregation: null, + query: { + filter: '{"age":{"$gt":25}}', + project: null, + sort: null, + skip: null, + limit: null, + }, + }, + }); + }); + + it('should return null for query fields if not provided', function () { + const xmlString = ` + [{ $match: { status: "A" } }] + `; + + const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + + expect(result).to.deep.equal({ + content: { + aggregation: { + pipeline: '[{"$match":{"status":"A"}}]', + }, + query: { + filter: null, + project: null, + sort: null, + skip: null, + limit: null, + }, + }, + }); + }); + + it('should return all the query fields if provided', function () { + const xmlString = ` + { age: { $gt: 25 } } + { name: 1, age: 1 } + { age: -1 } + 5 + 10 + + `; + + const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + + expect(result).to.deep.equal({ + content: { + aggregation: null, query: { filter: '{"age":{"$gt":25}}', project: '{"name":1,"age":1}', @@ -59,7 +124,7 @@ describe('parseXmlToMmsJsonResponse', function () { `[]`, loggerMock ); - expect(result.content.aggregation.pipeline).to.equal(null); + expect(result.content.aggregation).to.equal(null); }); it('zero value', function () { const result = parseXmlToMmsJsonResponse(`0`, loggerMock); diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts index bd543aa1444..8dc83e9142c 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -40,16 +40,34 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { } } - // Keep the response same as we have from mms api + // Keep the response same as we have from mms api. If llm generated + // an aggregation, we want to return that instead of a query + if (result.aggregation) { + return { + content: { + aggregation: { + pipeline: result.aggregation, + }, + query: { + filter: null, + project: null, + sort: null, + skip: null, + limit: null, + }, + }, + }; + } + return { content: { - ...(result.aggregation ? { aggregation: result.aggregation } : {}), + aggregation: null, query: { - filter: result.filter, - project: result.project, - sort: result.sort, - skip: result.skip, - limit: result.limit, + filter: result.filter ?? null, + project: result.project ?? null, + sort: result.sort ?? null, + skip: result.skip ?? null, + limit: result.limit ?? null, }, }, }; From 309335acd1605f8ea8f1116ec8e64de59771deed Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 14:52:43 +0300 Subject: [PATCH 11/23] changes in field name --- .../src/utils/gen-ai-prompt.spec.ts | 6 +++--- .../src/utils/gen-ai-prompt.ts | 14 +++++++------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts index 526861c71af..3ae68818975 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts @@ -8,7 +8,7 @@ import { toJSString } from 'mongodb-query-parser'; import { ObjectId } from 'bson'; const OPTIONS: UserPromptForQueryOptions = { - userPrompt: 'Find all users older than 30', + userInput: 'Find all users older than 30', databaseName: 'airbnb', collectionName: 'listings', schema: { @@ -50,7 +50,7 @@ describe('GenAI Prompts', function () { expect(prompt).to.be.a('string'); expect(prompt).to.include( - `Write a query that does the following: "${OPTIONS.userPrompt}"`, + `Write a query that does the following: "${OPTIONS.userInput}"`, 'includes user prompt' ); expect(prompt).to.include( @@ -93,7 +93,7 @@ describe('GenAI Prompts', function () { expect(prompt).to.be.a('string'); expect(prompt).to.include( - `Generate an aggregation that does the following: "${OPTIONS.userPrompt}"`, + `Generate an aggregation that does the following: "${OPTIONS.userInput}"`, 'includes user prompt' ); expect(prompt).to.include( diff --git a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts index c235feaabf0..1d8e22d8d81 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts @@ -58,7 +58,7 @@ function buildInstructionsForAggregateQuery() { } export type UserPromptForQueryOptions = { - userPrompt: string; + userInput: string; databaseName?: string; collectionName?: string; schema?: unknown; @@ -76,7 +76,7 @@ function withCodeFence(code: string): string { function buildUserPromptForQuery({ type, - userPrompt, + userInput, databaseName, collectionName, schema, @@ -87,7 +87,7 @@ function buildUserPromptForQuery({ const queryPrompt = [ type === 'find' ? 'Write a query' : 'Generate an aggregation', 'that does the following:', - `"${userPrompt}"`, + `"${userInput}"`, ].join(' '); if (databaseName) { @@ -148,7 +148,7 @@ export type AiQueryPrompt = { }; export function buildFindQueryPrompt({ - userPrompt, + userInput, databaseName, collectionName, schema, @@ -156,7 +156,7 @@ export function buildFindQueryPrompt({ }: UserPromptForQueryOptions): AiQueryPrompt { const prompt = buildUserPromptForQuery({ type: 'find', - userPrompt, + userInput, databaseName, collectionName, schema, @@ -172,7 +172,7 @@ export function buildFindQueryPrompt({ } export function buildAggregateQueryPrompt({ - userPrompt, + userInput, databaseName, collectionName, schema, @@ -180,7 +180,7 @@ export function buildAggregateQueryPrompt({ }: UserPromptForQueryOptions): AiQueryPrompt { const prompt = buildUserPromptForQuery({ type: 'aggregate', - userPrompt, + userInput, databaseName, collectionName, schema, From feacddce085defcde5c85542e009fd974af6b23f Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 15:31:21 +0300 Subject: [PATCH 12/23] use query parser --- .../src/utils/xml-to-mms-response.spec.ts | 12 ++++++------ .../src/utils/xml-to-mms-response.ts | 7 ++++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts index c0b5571b208..d461c3ece44 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts @@ -22,7 +22,7 @@ describe('parseXmlToMmsJsonResponse', function () { expect(result).to.deep.equal({ content: { aggregation: { - pipeline: '[{"$match":{"status":"A"}}]', + pipeline: "[{$match:{status:'A'}}]", }, query: { filter: null, @@ -45,7 +45,7 @@ describe('parseXmlToMmsJsonResponse', function () { content: { aggregation: null, query: { - filter: '{"age":{"$gt":25}}', + filter: '{age:{$gt:25}}', project: null, sort: null, skip: null, @@ -65,7 +65,7 @@ describe('parseXmlToMmsJsonResponse', function () { expect(result).to.deep.equal({ content: { aggregation: { - pipeline: '[{"$match":{"status":"A"}}]', + pipeline: "[{$match:{status:'A'}}]", }, query: { filter: null, @@ -94,9 +94,9 @@ describe('parseXmlToMmsJsonResponse', function () { content: { aggregation: null, query: { - filter: '{"age":{"$gt":25}}', - project: '{"name":1,"age":1}', - sort: '{"age":-1}', + filter: '{age:{$gt:25}}', + project: '{name:1,age:1}', + sort: '{age:-1}', skip: '5', limit: '10', }, diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts index 8dc83e9142c..48b9facbffd 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -1,4 +1,5 @@ import type { Logger } from '@mongodb-js/compass-logging'; +import parse, { toJSString } from 'mongodb-query-parser'; export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { const expectedTags = [ @@ -18,15 +19,15 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { if (match && match[1]) { const value = match[1].trim(); try { - // Here the value is valid js string, but not valid json, so we use eval to parse it. - const tagValue = eval(`(${value})`); + const tagValue = parse(value); if ( !tagValue || (typeof tagValue === 'object' && Object.keys(tagValue).length === 0) ) { result[tag] = null; } else { - result[tag] = JSON.stringify(tagValue); + // No indentation + result[tag] = toJSString(tagValue, 0); } } catch (e) { logger.log.warn( From c32198335aa1c71fc0d3da09c1abfffccaed509a Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 15:32:28 +0300 Subject: [PATCH 13/23] fix check --- package-lock.json | 2 +- packages/compass-generative-ai/package.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/package-lock.json b/package-lock.json index d2dcbb22080..8a999c4c449 100644 --- a/package-lock.json +++ b/package-lock.json @@ -49692,6 +49692,7 @@ "bson": "^6.10.4", "compass-preferences-model": "^2.66.3", "mongodb": "^6.19.0", + "mongodb-query-parser": "^4.5.0", "mongodb-schema": "^12.6.3", "react": "^17.0.2", "react-redux": "^8.1.3", @@ -49714,7 +49715,6 @@ "depcheck": "^1.4.1", "electron-mocha": "^12.2.0", "mocha": "^10.2.0", - "mongodb-query-parser": "^4.5.0", "nyc": "^15.1.0", "p-queue": "^7.4.1", "sinon": "^9.2.3", diff --git a/packages/compass-generative-ai/package.json b/packages/compass-generative-ai/package.json index 751880946c6..1ba664073fb 100644 --- a/packages/compass-generative-ai/package.json +++ b/packages/compass-generative-ai/package.json @@ -63,6 +63,7 @@ "bson": "^6.10.4", "compass-preferences-model": "^2.66.3", "mongodb": "^6.19.0", + "mongodb-query-parser": "^4.5.0", "mongodb-schema": "^12.6.3", "react": "^17.0.2", "react-redux": "^8.1.3", @@ -87,7 +88,6 @@ "depcheck": "^1.4.1", "electron-mocha": "^12.2.0", "mocha": "^10.2.0", - "mongodb-query-parser": "^4.5.0", "nyc": "^15.1.0", "p-queue": "^7.4.1", "sinon": "^9.2.3", From f5a34dfe20b41d9e9ae5384c80446e0258efdbbb Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 15:38:40 +0300 Subject: [PATCH 14/23] fix test --- packages/compass-e2e-tests/tests/collection-ai-query.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts index 67de96ed58d..ac09eaae4c9 100644 --- a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts +++ b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts @@ -195,7 +195,7 @@ async function setup( await browser.setFeature('optInGenAIFeatures', true); } -describe.only('Collection ai query with chatbot (with mocked backend)', function () { +describe('Collection ai query with chatbot (with mocked backend)', function () { const dbName = 'test'; const collName = 'numbers'; let compass: Compass; @@ -257,7 +257,7 @@ describe.only('Collection ai query with chatbot (with mocked backend)', function const queryBarFilterContent = await browser.getCodemirrorEditorText( Selectors.queryBarOptionInputFilter('Documents') ); - return queryBarFilterContent === '{"i":{"$gt":50}}'; + return queryBarFilterContent === '{i:{$gt:50}}'; }); // Check that the request was made with the correct parameters. From 91b17d560eebcef297eeaca169de2edc2e6f61f6 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 15:43:28 +0300 Subject: [PATCH 15/23] clean up --- .../compass-generative-ai/src/utils/gen-ai-response.ts | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/packages/compass-generative-ai/src/utils/gen-ai-response.ts b/packages/compass-generative-ai/src/utils/gen-ai-response.ts index 13ca53c786a..50e5b956fef 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-response.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-response.ts @@ -2,10 +2,7 @@ import { AiChatbotInvalidResponseError } from '../chatbot-errors'; import { type AiQueryPrompt } from './gen-ai-prompt'; import type { LanguageModel } from 'ai'; import { streamText } from 'ai'; -/** - * Function that takes an LanguageModel and a message of type - * AiQueryPrompt and returns a Promise that resolves to a string response - */ + export async function getAiQueryResponse( model: LanguageModel, message: AiQueryPrompt, @@ -17,7 +14,7 @@ export async function getAiQueryResponse( providerOptions: { openai: { store: false, - instructions: message.metadata.instructions || '', + instructions: message.metadata.instructions, }, }, abortSignal, From ea4dfdf40ecabe42ace8d9ba502f201329c5bf98 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 15:53:18 +0300 Subject: [PATCH 16/23] copilot feedback --- packages/compass-generative-ai/src/atlas-ai-service.ts | 7 ++++--- .../compass-generative-ai/src/utils/xml-to-mms-response.ts | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 80ab92a6120..3001f6ec3ca 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -291,16 +291,17 @@ export class AtlasAiService { this.logger = logger; this.initPromise = this.setupAIAccess(); - const initialBaseUrl = 'http://PLACEHOLDER_BASE_URL_TO_BE_REPLACED.invalid'; + const PLACEHOLDER_BASE_URL = + 'http://PLACEHOLDER_BASE_URL_TO_BE_REPLACED.invalid'; this.aiModel = createOpenAI({ apiKey: '', - baseURL: initialBaseUrl, + baseURL: PLACEHOLDER_BASE_URL, fetch: (url, init) => { // The `baseUrl` can be dynamically changed, but `createOpenAI` // doesn't allow us to change it after initial call. Instead // we're going to update it every time the fetch call happens const uri = String(url).replace( - initialBaseUrl, + PLACEHOLDER_BASE_URL, this.atlasService.assistantApiEndpoint() ); return this.atlasService.authenticatedFetch(uri, init); diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts index 48b9facbffd..363ca616add 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -9,10 +9,11 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { 'skip', 'limit', 'aggregation', - ]; + ] as const; // Currently the prompt forces LLM to return xml-styled data - const result: any = {}; + const result: Partial> = + {}; for (const tag of expectedTags) { const regex = new RegExp(`<${tag}>([\\s\\S]*?)<\\/${tag}>`, 'i'); const match = xmlString.match(regex); @@ -27,7 +28,7 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { result[tag] = null; } else { // No indentation - result[tag] = toJSString(tagValue, 0); + result[tag] = toJSString(tagValue, 0) ?? null; } } catch (e) { logger.log.warn( From f8a4c43494868ac15da316220654efd4dfa12c98 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Wed, 3 Dec 2025 18:25:26 +0300 Subject: [PATCH 17/23] fix log id --- packages/compass-generative-ai/src/utils/xml-to-mms-response.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts index 363ca616add..ad5e5e74976 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -32,7 +32,7 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { } } catch (e) { logger.log.warn( - logger.mongoLogId(1_001_000_309), + logger.mongoLogId(1_001_000_384), 'AtlasAiService', `Failed to parse value for tag <${tag}>: ${value}`, { error: e } From 8190219d73c54561fc91b84cccbd40ea9e81dbef Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Thu, 4 Dec 2025 10:38:20 +0300 Subject: [PATCH 18/23] fix cors issue on e2e tests --- packages/compass-e2e-tests/helpers/assistant-service.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/compass-e2e-tests/helpers/assistant-service.ts b/packages/compass-e2e-tests/helpers/assistant-service.ts index 5abca2c7d5f..970f8edd6dd 100644 --- a/packages/compass-e2e-tests/helpers/assistant-service.ts +++ b/packages/compass-e2e-tests/helpers/assistant-service.ts @@ -170,12 +170,13 @@ export async function startMockAssistantServer( let response = _response; const server = http .createServer((req, res) => { - res.setHeader('Access-Control-Allow-Origin', '*'); + res.setHeader('Access-Control-Allow-Origin', req.headers.origin || '*'); res.setHeader('Access-Control-Allow-Methods', 'POST, OPTIONS'); res.setHeader( 'Access-Control-Allow-Headers', - 'Content-Type, Authorization, X-Request-Origin, User-Agent' + 'Content-Type, Authorization, X-Request-Origin, User-Agent, X-CSRF-Token, X-CSRF-Time' ); + res.setHeader('Access-Control-Allow-Credentials', 'true'); // Handle preflight requests if (req.method === 'OPTIONS') { From c444cb9fe1b0e7dbbfeb8a54b83d03f19d3e0f7b Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Thu, 4 Dec 2025 15:07:22 +0300 Subject: [PATCH 19/23] more tests --- .../src/atlas-ai-service.spec.ts | 353 ++++++++++++++++++ .../src/utils/gen-ai-prompt.spec.ts | 59 +++ .../src/utils/gen-ai-prompt.ts | 8 + .../src/utils/xml-to-mms-response.spec.ts | 25 +- .../src/utils/xml-to-mms-response.ts | 37 +- 5 files changed, 450 insertions(+), 32 deletions(-) diff --git a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts index a9dc4016f0f..e56ffbcef2f 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -737,4 +737,357 @@ describe('AtlasAiService', function () { }); }); } + + describe('with chatbot api', function () { + describe('getQueryFromUserInput and getAggregationFromUserInput', function () { + type Chunk = { type: 'text' | 'error'; content: string }; + let atlasAiService: AtlasAiService; + const mockConnectionInfo = getMockConnectionInfo(); + + function streamChunkResponse( + readableStreamController: ReadableStreamController, + chunks: Chunk[] + ) { + const responseId = `resp_${Date.now()}`; + const itemId = `item_${Date.now()}`; + let sequenceNumber = 0; + + const encoder = new TextEncoder(); + + // openai response format: + // https://github.com/vercel/ai/blob/811119c1808d7b62a4857bcad42353808cdba17c/packages/openai/src/responses/openai-responses-api.ts#L322 + + // Send response.created event + readableStreamController.enqueue( + encoder.encode( + `data: ${JSON.stringify({ + type: 'response.created', + response: { + id: responseId, + object: 'realtime.response', + status: 'in_progress', + output: [], + usage: { + input_tokens: 0, + output_tokens: 0, + total_tokens: 0, + }, + }, + sequence_number: sequenceNumber++, + })}\n\n` + ) + ); + + // Send output_item.added event + readableStreamController.enqueue( + encoder.encode( + `data: ${JSON.stringify({ + type: 'response.output_item.added', + response_id: responseId, + output_index: 0, + item: { + id: itemId, + object: 'realtime.item', + type: 'message', + role: 'assistant', + content: [], + }, + sequence_number: sequenceNumber++, + })}\n\n` + ) + ); + + for (const chunk of chunks) { + if (chunk.type === 'error') { + readableStreamController.enqueue( + encoder.encode( + `data: ${JSON.stringify({ + type: `error`, + response_id: responseId, + item_id: itemId, + output_index: 0, + error: { + type: 'model_error', + code: 'model_error', + message: chunk.content, + }, + sequence_number: sequenceNumber++, + })}\n\n` + ) + ); + } else { + readableStreamController.enqueue( + encoder.encode( + `data: ${JSON.stringify({ + type: 'response.output_text.delta', + response_id: responseId, + item_id: itemId, + output_index: 0, + delta: chunk.content, + sequence_number: sequenceNumber++, + })}\n\n` + ) + ); + } + } + + const content = chunks + .filter((c) => c.type === 'text') + .map((c) => c.content) + .join(''); + + // Send output_item.done event + readableStreamController.enqueue( + encoder.encode( + `data: ${JSON.stringify({ + type: 'response.output_item.done', + response_id: responseId, + output_index: 0, + item: { + id: itemId, + object: 'realtime.item', + type: 'message', + role: 'assistant', + content: [ + { + type: 'text', + text: content, + }, + ], + }, + sequence_number: sequenceNumber++, + })}\n\n` + ) + ); + + // Send response.completed event + const tokenCount = Math.ceil(content.length / 4); // assume 4 chars per token + readableStreamController.enqueue( + encoder.encode( + `data: ${JSON.stringify({ + type: 'response.completed', + response: { + id: responseId, + object: 'realtime.response', + status: 'completed', + output: [ + { + id: itemId, + object: 'realtime.item', + type: 'message', + role: 'assistant', + content: [ + { + type: 'text', + text: content, + }, + ], + }, + ], + usage: { + input_tokens: 10, + output_tokens: tokenCount, + total_tokens: 10 + tokenCount, + }, + }, + sequence_number: sequenceNumber++, + })}\n\n` + ) + ); + } + + function streamableFetchMock(chunks: Chunk[]) { + const readableStream = new ReadableStream({ + start(controller) { + streamChunkResponse(controller, chunks); + controller.close(); + }, + }); + return new Response(readableStream, { + headers: { 'Content-Type': 'text/event-stream' }, + }); + } + + beforeEach(async function () { + const mockAtlasService = new MockAtlasService(); + await preferences.savePreferences({ + enableChatbotEndpointForGenAI: true, + }); + atlasAiService = new AtlasAiService({ + apiURLPreset: 'cloud', + atlasService: mockAtlasService as any, + preferences, + logger: createNoopLogger(), + }); + // Enable the AI feature + const fetchStub = sandbox.stub().resolves( + makeResponse({ + features: { + GEN_AI_COMPASS: { + enabled: true, + }, + }, + }) + ); + global.fetch = fetchStub; + await atlasAiService['setupAIAccess'](); + }); + + after(function () { + global.fetch = initialFetch; + }); + + const testCases = [ + { + functionName: 'getQueryFromUserInput', + successResponse: { + request: [ + { type: 'text', content: 'Hello' }, + { type: 'text', content: ' world' }, + { + type: 'text', + content: '. This is some non relevant text in the output', + }, + { type: 'text', content: '{test: ' }, + { type: 'text', content: '"pineapple"' }, + { type: 'text', content: '}' }, + ] as Chunk[], + response: { + content: { + query: { + filter: "{test:'pineapple'}", + project: null, + sort: null, + skip: null, + limit: null, + }, + }, + }, + }, + invalidModelResponse: { + request: [ + { type: 'text', content: 'Hello' }, + { type: 'text', content: ' world.' }, + { type: 'text', content: '{test: ' }, + { type: 'text', content: '"pineapple"' }, + { type: 'text', content: '}' }, + { type: 'error', content: 'Model crashed!' }, + ] as Chunk[], + errorMessage: 'Model crashed!', + }, + }, + { + functionName: 'getAggregationFromUserInput', + successResponse: { + request: [ + { type: 'text', content: 'Hello' }, + { type: 'text', content: ' world' }, + { + type: 'text', + content: '. This is some non relevant text in the output', + }, + { type: 'text', content: '[{$count: ' }, + { type: 'text', content: '"pineapple"' }, + { type: 'text', content: '}]' }, + ] as Chunk[], + response: { + content: { + aggregation: { + pipeline: "[{$count:'pineapple'}]", + }, + }, + }, + }, + invalidModelResponse: { + request: [ + { type: 'text', content: 'Hello' }, + { type: 'text', content: ' world.' }, + { type: 'text', content: '[{test: ' }, + { type: 'text', content: '"pineapple"' }, + { type: 'text', content: '}]' }, + { type: 'error', content: 'Model crashed!' }, + ] as Chunk[], + errorMessage: 'Model crashed!', + }, + }, + ] as const; + + for (const { + functionName, + successResponse, + invalidModelResponse, + } of testCases) { + describe(functionName, function () { + it('makes a post request with the user input to the endpoint in the environment', async function () { + const fetchStub = sandbox + .stub() + .resolves(streamableFetchMock(successResponse.request)); + global.fetch = fetchStub; + + const input = { + userInput: 'test', + signal: new AbortController().signal, + collectionName: 'jam', + databaseName: 'peanut', + schema: { _id: { types: [{ bsonType: 'ObjectId' }] } }, + sampleDocuments: [ + { _id: new ObjectId('642d766b7300158b1f22e972') }, + ], + requestId: 'abc', + }; + + const res = await atlasAiService[functionName]( + input as any, + mockConnectionInfo + ); + + expect(fetchStub).to.have.been.calledOnce; + + const { args } = fetchStub.firstCall; + 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.instructions).to.be.a('string'); + expect(requestBody.input).to.be.an('array'); + + const { role, content } = requestBody.input[0]; + expect(role).to.equal('user'); + expect(content[0].text).to.include( + `Database name: "${input.databaseName}"` + ); + expect(content[0].text).to.include( + `Collection name: "${input.collectionName}"` + ); + expect(res).to.deep.eq(successResponse.response); + }); + + it('should throw an error when the stream contains an error chunk', async function () { + const fetchStub = sandbox + .stub() + .resolves(streamableFetchMock(invalidModelResponse.request)); + global.fetch = fetchStub; + + try { + await atlasAiService[functionName]( + { + userInput: 'test', + collectionName: 'test', + databaseName: 'peanut', + requestId: 'abc', + signal: new AbortController().signal, + }, + mockConnectionInfo + ); + expect.fail(`Expected ${functionName} to throw`); + } catch (err) { + expect((err as Error).message).to.match( + new RegExp(invalidModelResponse.errorMessage, 'i') + ); + } + }); + }); + } + }); + }); }); diff --git a/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts index 3ae68818975..ec55525fb69 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts @@ -121,4 +121,63 @@ describe('GenAI Prompts', function () { 'includes actual sample documents' ); }); + + it('throws if user prompt exceeds the max size', function () { + try { + buildFindQueryPrompt({ + ...OPTIONS, + userInput: 'a'.repeat(512001), + }); + expect.fail('Expected buildFindQueryPrompt to throw'); + } catch (err) { + expect(err).to.have.property( + 'message', + 'Sorry, your request is too large. Please use a smaller prompt or try using this feature on a collection with smaller documents.' + ); + } + }); + + context('handles large sample documents', function () { + it('sends all the sample docs if within limits', function () { + const sampleDocuments = [ + { a: '1' }, + { a: '2' }, + { a: '3' }, + { a: '4'.repeat(5120) }, + ]; + const prompt = buildFindQueryPrompt({ + ...OPTIONS, + sampleDocuments, + }).prompt; + + expect(prompt).to.include(toJSString(sampleDocuments)); + }); + it('sends only one sample doc if all exceed limits', function () { + const sampleDocuments = [ + { a: '1'.repeat(5120) }, + { a: '2'.repeat(5120001) }, + { a: '3'.repeat(5120001) }, + { a: '4'.repeat(5120001) }, + ]; + const prompt = buildFindQueryPrompt({ + ...OPTIONS, + sampleDocuments, + }).prompt; + expect(prompt).to.include(toJSString([sampleDocuments[0]])); + }); + it('should not send sample docs if even one exceeds limits', function () { + const sampleDocuments = [ + { a: '1'.repeat(5120001) }, + { a: '2'.repeat(5120001) }, + { a: '3'.repeat(5120001) }, + { a: '4'.repeat(5120001) }, + ]; + const prompt = buildFindQueryPrompt({ + ...OPTIONS, + sampleDocuments, + }).prompt; + expect(prompt).to.not.include('Sample document from the collection:'); + expect(prompt).to.not.include('Sample documents from the collection:'); + }); + }); }); diff --git a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts index 1d8e22d8d81..8ee6ec8bbe3 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts @@ -137,6 +137,14 @@ function buildUserPromptForQuery({ } } messages.push(queryPrompt); + + // If at this point we have exceeded the limit, throw an error. + const totalPromptLength = messages.join('\n').length; + if (totalPromptLength > MAX_TOTAL_PROMPT_LENGTH) { + throw new Error( + 'Sorry, your request is too large. Please use a smaller prompt or try using this feature on a collection with smaller documents.' + ); + } return messages.join('\n'); } diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts index d461c3ece44..863d42ad502 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts @@ -11,7 +11,7 @@ const loggerMock = { mongoLogId: (id: number) => id, } as unknown as Logger; describe('parseXmlToMmsJsonResponse', function () { - it('should prioritize aggregation over query fields', function () { + it('should return prioritize aggregation over query when available and valid', function () { const xmlString = ` { age: { $gt: 25 } } [{ $match: { status: "A" } }] @@ -35,7 +35,7 @@ describe('parseXmlToMmsJsonResponse', function () { }); }); - it('should return null for aggregation if not provided', function () { + it('should not return aggregation if its not available in the response', function () { const xmlString = ` { age: { $gt: 25 } } `; @@ -43,7 +43,6 @@ describe('parseXmlToMmsJsonResponse', function () { const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); expect(result).to.deep.equal({ content: { - aggregation: null, query: { filter: '{age:{$gt:25}}', project: null, @@ -55,7 +54,7 @@ describe('parseXmlToMmsJsonResponse', function () { }); }); - it('should return null for query fields if not provided', function () { + it('should not return query if its not available in the response', function () { const xmlString = ` [{ $match: { status: "A" } }] `; @@ -67,13 +66,6 @@ describe('parseXmlToMmsJsonResponse', function () { aggregation: { pipeline: "[{$match:{status:'A'}}]", }, - query: { - filter: null, - project: null, - sort: null, - skip: null, - limit: null, - }, }, }); }); @@ -92,7 +84,6 @@ describe('parseXmlToMmsJsonResponse', function () { expect(result).to.deep.equal({ content: { - aggregation: null, query: { filter: '{age:{$gt:25}}', project: '{name:1,age:1}', @@ -104,31 +95,31 @@ describe('parseXmlToMmsJsonResponse', function () { }); }); - context('it should return null values for invalid data', function () { + context('it should handle invalid data', function () { it('invalid json', function () { const result = parseXmlToMmsJsonResponse( `{ age: { $gt: 25 `, loggerMock ); - expect(result.content.query.filter).to.equal(null); + expect(result.content).to.not.have.property('query'); }); it('empty object', function () { const result = parseXmlToMmsJsonResponse( `{}`, loggerMock ); - expect(result.content.query.filter).to.equal(null); + expect(result.content).to.not.have.property('query'); }); it('empty array', function () { const result = parseXmlToMmsJsonResponse( `[]`, loggerMock ); - expect(result.content.aggregation).to.equal(null); + expect(result.content).to.not.have.property('aggregation'); }); it('zero value', function () { const result = parseXmlToMmsJsonResponse(`0`, loggerMock); - expect(result.content.query.limit).to.equal(null); + expect(result.content).to.not.have.property('query'); }); }); }); diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts index ad5e5e74976..0ce0058518b 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -12,8 +12,14 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { ] as const; // Currently the prompt forces LLM to return xml-styled data - const result: Partial> = - {}; + const result: Record<(typeof expectedTags)[number], string | null> = { + filter: null, + project: null, + sort: null, + skip: null, + limit: null, + aggregation: null, + }; for (const tag of expectedTags) { const regex = new RegExp(`<${tag}>([\\s\\S]*?)<\\/${tag}>`, 'i'); const match = xmlString.match(regex); @@ -42,13 +48,15 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { } } - // Keep the response same as we have from mms api. If llm generated - // an aggregation, we want to return that instead of a query - if (result.aggregation) { + const { aggregation, ...query } = result; + const isQueryEmpty = Object.values(query).every((v) => v === null); + + // It prioritizes aggregation over query if both are present + if (aggregation && !isQueryEmpty) { return { content: { aggregation: { - pipeline: result.aggregation, + pipeline: aggregation, }, query: { filter: null, @@ -60,17 +68,16 @@ export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { }, }; } - return { content: { - aggregation: null, - query: { - filter: result.filter ?? null, - project: result.project ?? null, - sort: result.sort ?? null, - skip: result.skip ?? null, - limit: result.limit ?? null, - }, + ...(aggregation + ? { + aggregation: { + pipeline: aggregation, + }, + } + : {}), + ...(isQueryEmpty ? {} : { query }), }, }; } From 6eec8db3e162b15d2812d1906105ba4e0c31d5a9 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Thu, 4 Dec 2025 15:20:34 +0300 Subject: [PATCH 20/23] add type --- .../src/utils/xml-to-mms-response.ts | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts index 0ce0058518b..9287d35f6e2 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts @@ -1,7 +1,25 @@ import type { Logger } from '@mongodb-js/compass-logging'; import parse, { toJSString } from 'mongodb-query-parser'; -export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { +type MmsJsonResponse = { + content: { + query?: { + filter: string | null; + project: string | null; + sort: string | null; + skip: string | null; + limit: string | null; + }; + aggregation?: { + pipeline: string; + }; + }; +}; + +export function parseXmlToMmsJsonResponse( + xmlString: string, + logger: Logger +): MmsJsonResponse { const expectedTags = [ 'filter', 'project', From 77a0eea9c0aad6fc977b9c1c1a7e3c398f21f165 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Mon, 8 Dec 2025 15:22:36 +0300 Subject: [PATCH 21/23] use noop logger and clean up api response handling --- .../src/utils/gen-ai-response.ts | 9 +----- .../src/utils/xml-to-mms-response.spec.ts | 29 ++++++++----------- 2 files changed, 13 insertions(+), 25 deletions(-) diff --git a/packages/compass-generative-ai/src/utils/gen-ai-response.ts b/packages/compass-generative-ai/src/utils/gen-ai-response.ts index 50e5b956fef..823921f0079 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-response.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-response.ts @@ -20,14 +20,7 @@ export async function getAiQueryResponse( abortSignal, }).toUIMessageStream(); const chunks: string[] = []; - let done = false; - const reader = response.getReader(); - while (!done) { - const { done: _done, value } = await reader.read(); - if (_done) { - done = true; - break; - } + for await (const value of response) { if (value.type === 'text-delta') { chunks.push(value.delta); } diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts index 863d42ad502..b50ba6de5dc 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts +++ b/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts @@ -1,15 +1,7 @@ import { expect } from 'chai'; import { parseXmlToMmsJsonResponse } from './xml-to-mms-response'; -import type { Logger } from '@mongodb-js/compass-logging'; +import { createNoopLogger } from '@mongodb-js/compass-logging/provider'; -const loggerMock = { - log: { - warn: () => { - /* noop */ - }, - }, - mongoLogId: (id: number) => id, -} as unknown as Logger; describe('parseXmlToMmsJsonResponse', function () { it('should return prioritize aggregation over query when available and valid', function () { const xmlString = ` @@ -17,7 +9,7 @@ describe('parseXmlToMmsJsonResponse', function () { [{ $match: { status: "A" } }] `; - const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { @@ -40,7 +32,7 @@ describe('parseXmlToMmsJsonResponse', function () { { age: { $gt: 25 } } `; - const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { query: { @@ -59,7 +51,7 @@ describe('parseXmlToMmsJsonResponse', function () { [{ $match: { status: "A" } }] `; - const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { @@ -80,7 +72,7 @@ describe('parseXmlToMmsJsonResponse', function () { `; - const result = parseXmlToMmsJsonResponse(xmlString, loggerMock); + const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { @@ -99,26 +91,29 @@ describe('parseXmlToMmsJsonResponse', function () { it('invalid json', function () { const result = parseXmlToMmsJsonResponse( `{ age: { $gt: 25 `, - loggerMock + createNoopLogger() ); expect(result.content).to.not.have.property('query'); }); it('empty object', function () { const result = parseXmlToMmsJsonResponse( `{}`, - loggerMock + createNoopLogger() ); expect(result.content).to.not.have.property('query'); }); it('empty array', function () { const result = parseXmlToMmsJsonResponse( `[]`, - loggerMock + createNoopLogger() ); expect(result.content).to.not.have.property('aggregation'); }); it('zero value', function () { - const result = parseXmlToMmsJsonResponse(`0`, loggerMock); + const result = parseXmlToMmsJsonResponse( + `0`, + createNoopLogger() + ); expect(result.content).to.not.have.property('query'); }); }); From aafe769d7c89af36271aeb63c4c3fead6341b524 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Mon, 8 Dec 2025 15:25:08 +0300 Subject: [PATCH 22/23] remove mms from name --- .../src/atlas-ai-service.ts | 4 ++-- ...nse.spec.ts => parse-xml-response.spec.ts} | 20 +++++++++---------- ...-mms-response.ts => parse-xml-response.ts} | 6 +++--- 3 files changed, 15 insertions(+), 15 deletions(-) rename packages/compass-generative-ai/src/utils/{xml-to-mms-response.spec.ts => parse-xml-response.spec.ts} (81%) rename packages/compass-generative-ai/src/utils/{xml-to-mms-response.ts => parse-xml-response.ts} (96%) diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 3001f6ec3ca..5baf8c62b60 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -23,7 +23,7 @@ import { buildAggregateQueryPrompt, buildFindQueryPrompt, } from './utils/gen-ai-prompt'; -import { parseXmlToMmsJsonResponse } from './utils/xml-to-mms-response'; +import { parseXmlToJsonResponse } from './utils/parse-xml-response'; import { getAiQueryResponse } from './utils/gen-ai-response'; type GenerativeAiInput = { @@ -573,7 +573,7 @@ export class AtlasAiService { message, options.signal ); - const parsedResponse = parseXmlToMmsJsonResponse(response, this.logger); + const parsedResponse = parseXmlToJsonResponse(response, this.logger); validateFn(parsedResponse); return parsedResponse; } diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts b/packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts similarity index 81% rename from packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts rename to packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts index b50ba6de5dc..87fa314ac13 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts +++ b/packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts @@ -1,15 +1,15 @@ import { expect } from 'chai'; -import { parseXmlToMmsJsonResponse } from './xml-to-mms-response'; +import { parseXmlToJsonResponse } from './parse-xml-response'; import { createNoopLogger } from '@mongodb-js/compass-logging/provider'; -describe('parseXmlToMmsJsonResponse', function () { +describe('parseXmlToJsonResponse', function () { it('should return prioritize aggregation over query when available and valid', function () { const xmlString = ` { age: { $gt: 25 } } [{ $match: { status: "A" } }] `; - const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); + const result = parseXmlToJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { @@ -32,7 +32,7 @@ describe('parseXmlToMmsJsonResponse', function () { { age: { $gt: 25 } } `; - const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); + const result = parseXmlToJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { query: { @@ -51,7 +51,7 @@ describe('parseXmlToMmsJsonResponse', function () { [{ $match: { status: "A" } }] `; - const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); + const result = parseXmlToJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { @@ -72,7 +72,7 @@ describe('parseXmlToMmsJsonResponse', function () { `; - const result = parseXmlToMmsJsonResponse(xmlString, createNoopLogger()); + const result = parseXmlToJsonResponse(xmlString, createNoopLogger()); expect(result).to.deep.equal({ content: { @@ -89,28 +89,28 @@ describe('parseXmlToMmsJsonResponse', function () { context('it should handle invalid data', function () { it('invalid json', function () { - const result = parseXmlToMmsJsonResponse( + const result = parseXmlToJsonResponse( `{ age: { $gt: 25 `, createNoopLogger() ); expect(result.content).to.not.have.property('query'); }); it('empty object', function () { - const result = parseXmlToMmsJsonResponse( + const result = parseXmlToJsonResponse( `{}`, createNoopLogger() ); expect(result.content).to.not.have.property('query'); }); it('empty array', function () { - const result = parseXmlToMmsJsonResponse( + const result = parseXmlToJsonResponse( `[]`, createNoopLogger() ); expect(result.content).to.not.have.property('aggregation'); }); it('zero value', function () { - const result = parseXmlToMmsJsonResponse( + const result = parseXmlToJsonResponse( `0`, createNoopLogger() ); diff --git a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts b/packages/compass-generative-ai/src/utils/parse-xml-response.ts similarity index 96% rename from packages/compass-generative-ai/src/utils/xml-to-mms-response.ts rename to packages/compass-generative-ai/src/utils/parse-xml-response.ts index 9287d35f6e2..c4886e74714 100644 --- a/packages/compass-generative-ai/src/utils/xml-to-mms-response.ts +++ b/packages/compass-generative-ai/src/utils/parse-xml-response.ts @@ -1,7 +1,7 @@ import type { Logger } from '@mongodb-js/compass-logging'; import parse, { toJSString } from 'mongodb-query-parser'; -type MmsJsonResponse = { +type JsonResponse = { content: { query?: { filter: string | null; @@ -16,10 +16,10 @@ type MmsJsonResponse = { }; }; -export function parseXmlToMmsJsonResponse( +export function parseXmlToJsonResponse( xmlString: string, logger: Logger -): MmsJsonResponse { +): JsonResponse { const expectedTags = [ 'filter', 'project', From 089b30ca558cf049290f8e847a0858ca67ec5112 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Mon, 8 Dec 2025 23:31:33 +0300 Subject: [PATCH 23/23] clean up error, names --- .../src/chatbot-errors.ts | 18 +++--------------- .../src/utils/gen-ai-prompt.spec.ts | 4 ++-- .../src/utils/gen-ai-prompt.ts | 15 ++++++++------- .../src/utils/parse-xml-response.ts | 16 ++++++++++------ 4 files changed, 23 insertions(+), 30 deletions(-) diff --git a/packages/compass-generative-ai/src/chatbot-errors.ts b/packages/compass-generative-ai/src/chatbot-errors.ts index 2d778a46b8f..f115482b1b0 100644 --- a/packages/compass-generative-ai/src/chatbot-errors.ts +++ b/packages/compass-generative-ai/src/chatbot-errors.ts @@ -1,20 +1,8 @@ -export class AiChatbotError extends Error { - statusCode: number; - errorCode: string; - detail: string; +import { AtlasServiceError } from '@mongodb-js/atlas-service/renderer'; - constructor(statusCode: number, detail: string, errorCode: string) { - super(`${errorCode}: ${detail}`); - this.name = 'ServerError'; - this.statusCode = statusCode; - this.errorCode = errorCode; - this.detail = detail; - } -} - -export class AiChatbotInvalidResponseError extends AiChatbotError { +export class AiChatbotInvalidResponseError extends AtlasServiceError { constructor(message: string) { - super(500, message, 'INVALID_RESPONSE'); + super('ServerError', 500, message, 'INVALID_RESPONSE'); this.name = 'AiChatbotInvalidResponseError'; } } diff --git a/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts index ec55525fb69..f1aa217c21d 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts @@ -2,12 +2,12 @@ import { expect } from 'chai'; import { buildFindQueryPrompt, buildAggregateQueryPrompt, - type UserPromptForQueryOptions, + type PromptContextOptions, } from './gen-ai-prompt'; import { toJSString } from 'mongodb-query-parser'; import { ObjectId } from 'bson'; -const OPTIONS: UserPromptForQueryOptions = { +const OPTIONS: PromptContextOptions = { userInput: 'Find all users older than 30', databaseName: 'airbnb', collectionName: 'listings', diff --git a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts index 8ee6ec8bbe3..ac45fdf4eaf 100644 --- a/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts +++ b/packages/compass-generative-ai/src/utils/gen-ai-prompt.ts @@ -57,7 +57,7 @@ function buildInstructionsForAggregateQuery() { ].join('\n'); } -export type UserPromptForQueryOptions = { +export type PromptContextOptions = { userInput: string; databaseName?: string; collectionName?: string; @@ -81,7 +81,7 @@ function buildUserPromptForQuery({ collectionName, schema, sampleDocuments, -}: UserPromptForQueryOptions & { type: 'find' | 'aggregate' }): string { +}: PromptContextOptions & { type: 'find' | 'aggregate' }): string { const messages = []; const queryPrompt = [ @@ -138,14 +138,15 @@ function buildUserPromptForQuery({ } messages.push(queryPrompt); + const prompt = messages.join('\n'); + // If at this point we have exceeded the limit, throw an error. - const totalPromptLength = messages.join('\n').length; - if (totalPromptLength > MAX_TOTAL_PROMPT_LENGTH) { + if (prompt.length > MAX_TOTAL_PROMPT_LENGTH) { throw new Error( 'Sorry, your request is too large. Please use a smaller prompt or try using this feature on a collection with smaller documents.' ); } - return messages.join('\n'); + return prompt; } export type AiQueryPrompt = { @@ -161,7 +162,7 @@ export function buildFindQueryPrompt({ collectionName, schema, sampleDocuments, -}: UserPromptForQueryOptions): AiQueryPrompt { +}: PromptContextOptions): AiQueryPrompt { const prompt = buildUserPromptForQuery({ type: 'find', userInput, @@ -185,7 +186,7 @@ export function buildAggregateQueryPrompt({ collectionName, schema, sampleDocuments, -}: UserPromptForQueryOptions): AiQueryPrompt { +}: PromptContextOptions): AiQueryPrompt { const prompt = buildUserPromptForQuery({ type: 'aggregate', userInput, diff --git a/packages/compass-generative-ai/src/utils/parse-xml-response.ts b/packages/compass-generative-ai/src/utils/parse-xml-response.ts index c4886e74714..44e34a6c971 100644 --- a/packages/compass-generative-ai/src/utils/parse-xml-response.ts +++ b/packages/compass-generative-ai/src/utils/parse-xml-response.ts @@ -1,7 +1,7 @@ import type { Logger } from '@mongodb-js/compass-logging'; import parse, { toJSString } from 'mongodb-query-parser'; -type JsonResponse = { +type ParsedXmlJsonResponse = { content: { query?: { filter: string | null; @@ -19,7 +19,7 @@ type JsonResponse = { export function parseXmlToJsonResponse( xmlString: string, logger: Logger -): JsonResponse { +): ParsedXmlJsonResponse { const expectedTags = [ 'filter', 'project', @@ -29,6 +29,12 @@ export function parseXmlToJsonResponse( 'aggregation', ] as const; + const parser = new DOMParser(); + const xmlDoc = parser.parseFromString( + `${xmlString}`, + 'text/xml' + ); + // Currently the prompt forces LLM to return xml-styled data const result: Record<(typeof expectedTags)[number], string | null> = { filter: null, @@ -39,10 +45,8 @@ export function parseXmlToJsonResponse( aggregation: null, }; for (const tag of expectedTags) { - const regex = new RegExp(`<${tag}>([\\s\\S]*?)<\\/${tag}>`, 'i'); - const match = xmlString.match(regex); - if (match && match[1]) { - const value = match[1].trim(); + const value = xmlDoc.querySelector(tag)?.textContent?.trim(); + if (value) { try { const tagValue = parse(value); if (