From e8ba8d18f7bf14f8d876d474e55ab5129f76a9ea Mon Sep 17 00:00:00 2001 From: Ben Perlmutter Date: Thu, 28 Aug 2025 12:25:00 -0400 Subject: [PATCH] draft of changes --- package-lock.json | 111 +++--------------- packages/compass-assistant/package.json | 2 +- .../src/compass-assistant-provider.tsx | 11 +- .../src/docs-provider-transport.ts | 60 ++++++++-- packages/compass-assistant/src/index.tsx | 4 + packages/compass-assistant/src/prompts.ts | 19 ++- .../compass-assistant/test/assistant.eval.ts | 77 ++++++------ ...aggregation-pipeline.ts => aggregation.ts} | 24 ++-- .../test/eval-cases/atlas-search.ts | 16 +++ .../test/eval-cases/connection-error.ts | 21 ++++ .../test/eval-cases/explain-plan.ts | 21 ++++ .../eval-cases/filter-docs-before-search.ts | 13 -- .../test/eval-cases/index.ts | 10 +- .../test/eval-cases/model-data.ts | 22 ++-- 14 files changed, 235 insertions(+), 176 deletions(-) rename packages/compass-assistant/test/eval-cases/{aggregation-pipeline.ts => aggregation.ts} (52%) create mode 100644 packages/compass-assistant/test/eval-cases/atlas-search.ts create mode 100644 packages/compass-assistant/test/eval-cases/connection-error.ts create mode 100644 packages/compass-assistant/test/eval-cases/explain-plan.ts delete mode 100644 packages/compass-assistant/test/eval-cases/filter-docs-before-search.ts diff --git a/package-lock.json b/package-lock.json index 80ae040abf8..c62d15b2fcf 100644 --- a/package-lock.json +++ b/package-lock.json @@ -47112,7 +47112,7 @@ "version": "1.2.0", "license": "SSPL", "dependencies": { - "@ai-sdk/openai": "^2.0.4", + "@ai-sdk/openai": "^2.0.22", "@mongodb-js/atlas-service": "^0.56.0", "@mongodb-js/compass-app-registry": "^9.4.20", "@mongodb-js/compass-components": "^1.49.0", @@ -47165,31 +47165,14 @@ "zod": "^3.25.76 || ^4" } }, - "packages/compass-assistant/node_modules/@ai-sdk/gateway/node_modules/@ai-sdk/provider-utils": { - "version": "3.0.7", - "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.7.tgz", - "integrity": "sha512-o3BS5/t8KnBL3ubP8k3w77AByOypLm+pkIL/DCw0qKkhDbvhCy+L3hRTGPikpdb8WHcylAeKsjgwOxhj4cqTUA==", - "license": "Apache-2.0", - "dependencies": { - "@ai-sdk/provider": "2.0.0", - "@standard-schema/spec": "^1.0.0", - "eventsource-parser": "^3.0.5" - }, - "engines": { - "node": ">=18" - }, - "peerDependencies": { - "zod": "^3.25.76 || ^4" - } - }, "packages/compass-assistant/node_modules/@ai-sdk/openai": { - "version": "2.0.10", - "resolved": "https://registry.npmjs.org/@ai-sdk/openai/-/openai-2.0.10.tgz", - "integrity": "sha512-vnB6Jk2Qb245fajaWjG3q6N0QQy/uej7kZ0QR9xxq09x++3Tx/UPOcgAKMhFFA2fnuGpkFSRKoiDCyp/E3RARQ==", + "version": "2.0.22", + "resolved": "https://registry.npmjs.org/@ai-sdk/openai/-/openai-2.0.22.tgz", + "integrity": "sha512-qjSIPL5+LNM9flcBPeR64ZWeAZdYg4XWkAK34H3FaY61dSbuIaeqFPSzmQUrxotVcphAzgfL5tuYRqRYP2ZYyg==", "license": "Apache-2.0", "dependencies": { "@ai-sdk/provider": "2.0.0", - "@ai-sdk/provider-utils": "3.0.1" + "@ai-sdk/provider-utils": "3.0.7" }, "engines": { "node": ">=18" @@ -47199,15 +47182,14 @@ } }, "packages/compass-assistant/node_modules/@ai-sdk/provider-utils": { - "version": "3.0.1", - "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.1.tgz", - "integrity": "sha512-/iP1sKc6UdJgGH98OCly7sWJKv+J9G47PnTjIj40IJMUQKwDrUMyf7zOOfRtPwSuNifYhSoJQ4s1WltI65gJ/g==", + "version": "3.0.7", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.7.tgz", + "integrity": "sha512-o3BS5/t8KnBL3ubP8k3w77AByOypLm+pkIL/DCw0qKkhDbvhCy+L3hRTGPikpdb8WHcylAeKsjgwOxhj4cqTUA==", "license": "Apache-2.0", "dependencies": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", - "eventsource-parser": "^3.0.3", - "zod-to-json-schema": "^3.24.1" + "eventsource-parser": "^3.0.5" }, "engines": { "node": ">=18" @@ -47275,23 +47257,6 @@ "zod": "^3.25.76 || ^4" } }, - "packages/compass-assistant/node_modules/ai/node_modules/@ai-sdk/provider-utils": { - "version": "3.0.7", - "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.7.tgz", - "integrity": "sha512-o3BS5/t8KnBL3ubP8k3w77AByOypLm+pkIL/DCw0qKkhDbvhCy+L3hRTGPikpdb8WHcylAeKsjgwOxhj4cqTUA==", - "license": "Apache-2.0", - "dependencies": { - "@ai-sdk/provider": "2.0.0", - "@standard-schema/spec": "^1.0.0", - "eventsource-parser": "^3.0.5" - }, - "engines": { - "node": ">=18" - }, - "peerDependencies": { - "zod": "^3.25.76 || ^4" - } - }, "packages/compass-assistant/node_modules/eventsource-parser": { "version": "3.0.5", "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.0.5.tgz", @@ -47358,15 +47323,6 @@ "url": "https://github.com/sponsors/colinhacks" } }, - "packages/compass-assistant/node_modules/zod-to-json-schema": { - "version": "3.24.6", - "resolved": "https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.24.6.tgz", - "integrity": "sha512-h/z3PKvcTcTetyjl1fkj79MHNEjm+HpD6NXheWjzOekY7kV+lwDYnHw+ivHkijnCSMz1yJaWBD9vu/Fcmk+vEg==", - "license": "ISC", - "peerDependencies": { - "zod": "^3.24.1" - } - }, "packages/compass-collection": { "name": "@mongodb-js/compass-collection", "version": "4.70.0", @@ -60545,7 +60501,7 @@ "@mongodb-js/compass-assistant": { "version": "file:packages/compass-assistant", "requires": { - "@ai-sdk/openai": "^2.0.4", + "@ai-sdk/openai": "^2.0.22", "@mongodb-js/atlas-service": "^0.56.0", "@mongodb-js/compass-app-registry": "^9.4.20", "@mongodb-js/compass-components": "^1.49.0", @@ -60587,38 +60543,25 @@ "requires": { "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.7" - }, - "dependencies": { - "@ai-sdk/provider-utils": { - "version": "3.0.7", - "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.7.tgz", - "integrity": "sha512-o3BS5/t8KnBL3ubP8k3w77AByOypLm+pkIL/DCw0qKkhDbvhCy+L3hRTGPikpdb8WHcylAeKsjgwOxhj4cqTUA==", - "requires": { - "@ai-sdk/provider": "2.0.0", - "@standard-schema/spec": "^1.0.0", - "eventsource-parser": "^3.0.5" - } - } } }, "@ai-sdk/openai": { - "version": "2.0.10", - "resolved": "https://registry.npmjs.org/@ai-sdk/openai/-/openai-2.0.10.tgz", - "integrity": "sha512-vnB6Jk2Qb245fajaWjG3q6N0QQy/uej7kZ0QR9xxq09x++3Tx/UPOcgAKMhFFA2fnuGpkFSRKoiDCyp/E3RARQ==", + "version": "2.0.22", + "resolved": "https://registry.npmjs.org/@ai-sdk/openai/-/openai-2.0.22.tgz", + "integrity": "sha512-qjSIPL5+LNM9flcBPeR64ZWeAZdYg4XWkAK34H3FaY61dSbuIaeqFPSzmQUrxotVcphAzgfL5tuYRqRYP2ZYyg==", "requires": { "@ai-sdk/provider": "2.0.0", - "@ai-sdk/provider-utils": "3.0.1" + "@ai-sdk/provider-utils": "3.0.7" } }, "@ai-sdk/provider-utils": { - "version": "3.0.1", - "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.1.tgz", - "integrity": "sha512-/iP1sKc6UdJgGH98OCly7sWJKv+J9G47PnTjIj40IJMUQKwDrUMyf7zOOfRtPwSuNifYhSoJQ4s1WltI65gJ/g==", + "version": "3.0.7", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.7.tgz", + "integrity": "sha512-o3BS5/t8KnBL3ubP8k3w77AByOypLm+pkIL/DCw0qKkhDbvhCy+L3hRTGPikpdb8WHcylAeKsjgwOxhj4cqTUA==", "requires": { "@ai-sdk/provider": "2.0.0", "@standard-schema/spec": "^1.0.0", - "eventsource-parser": "^3.0.3", - "zod-to-json-schema": "^3.24.1" + "eventsource-parser": "^3.0.5" } }, "@sinonjs/commons": { @@ -60666,18 +60609,6 @@ "@ai-sdk/provider": "2.0.0", "@ai-sdk/provider-utils": "3.0.7", "@opentelemetry/api": "1.9.0" - }, - "dependencies": { - "@ai-sdk/provider-utils": { - "version": "3.0.7", - "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-3.0.7.tgz", - "integrity": "sha512-o3BS5/t8KnBL3ubP8k3w77AByOypLm+pkIL/DCw0qKkhDbvhCy+L3hRTGPikpdb8WHcylAeKsjgwOxhj4cqTUA==", - "requires": { - "@ai-sdk/provider": "2.0.0", - "@standard-schema/spec": "^1.0.0", - "eventsource-parser": "^3.0.5" - } - } } }, "eventsource-parser": { @@ -60729,12 +60660,6 @@ "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", "peer": true - }, - "zod-to-json-schema": { - "version": "3.24.6", - "resolved": "https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.24.6.tgz", - "integrity": "sha512-h/z3PKvcTcTetyjl1fkj79MHNEjm+HpD6NXheWjzOekY7kV+lwDYnHw+ivHkijnCSMz1yJaWBD9vu/Fcmk+vEg==", - "requires": {} } } }, diff --git a/packages/compass-assistant/package.json b/packages/compass-assistant/package.json index f7b4613601d..381d2cf17ea 100644 --- a/packages/compass-assistant/package.json +++ b/packages/compass-assistant/package.json @@ -49,7 +49,7 @@ "reformat": "npm run eslint . -- --fix && npm run prettier -- --write ." }, "dependencies": { - "@ai-sdk/openai": "^2.0.4", + "@ai-sdk/openai": "^2.0.22", "@mongodb-js/atlas-service": "^0.56.0", "@mongodb-js/compass-app-registry": "^9.4.20", "@mongodb-js/compass-components": "^1.49.0", diff --git a/packages/compass-assistant/src/compass-assistant-provider.tsx b/packages/compass-assistant/src/compass-assistant-provider.tsx index 92531f74e45..8faf8653ca0 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.tsx @@ -9,7 +9,7 @@ import { import { atlasServiceLocator } from '@mongodb-js/atlas-service/provider'; import { DocsProviderTransport } from './docs-provider-transport'; import { useDrawerActions } from '@mongodb-js/compass-components'; -import { buildConnectionErrorPrompt, buildExplainPlanPrompt } from './prompts'; +import { buildPrompts } from './prompts'; import { usePreference } from 'compass-preferences-model/provider'; import { createLoggerLocator } from '@mongodb-js/compass-logging/provider'; import type { ConnectionInfo } from '@mongodb-js/connection-info'; @@ -21,6 +21,7 @@ export type AssistantMessage = UIMessage & { metadata?: { /** The text to display instead of the message text. */ displayText?: string; + type?: keyof typeof buildPrompts; }; }; @@ -91,7 +92,7 @@ export const AssistantProvider: React.FunctionComponent< const assistantActionsContext = useRef({ interpretExplainPlan: ({ explainPlan }) => { openDrawer(ASSISTANT_DRAWER_ID); - const { prompt, displayText } = buildExplainPlanPrompt({ + const { prompt, displayText } = buildPrompts['explain-plan'].user({ explainPlan, }); void chat.sendMessage( @@ -99,6 +100,7 @@ export const AssistantProvider: React.FunctionComponent< text: prompt, metadata: { displayText, + type: 'explain-plan', }, }, {} @@ -112,13 +114,16 @@ export const AssistantProvider: React.FunctionComponent< ); const connectionError = error.toString(); - const { prompt } = buildConnectionErrorPrompt({ + const { prompt } = buildPrompts['connection-error'].user({ connectionString, connectionError, }); void chat.sendMessage( { text: prompt, + metadata: { + type: 'connection-error', + }, }, {} ); diff --git a/packages/compass-assistant/src/docs-provider-transport.ts b/packages/compass-assistant/src/docs-provider-transport.ts index 44261e1124e..cf2cdcbd976 100644 --- a/packages/compass-assistant/src/docs-provider-transport.ts +++ b/packages/compass-assistant/src/docs-provider-transport.ts @@ -1,30 +1,39 @@ +import { createOpenAI } from '@ai-sdk/openai'; +import { buildPrompts } from './prompts'; import { type ChatTransport, + convertToModelMessages, type UIMessage, type UIMessageChunk, - convertToModelMessages, + type LanguageModel, + type ModelMessage, streamText, } from 'ai'; -import { createOpenAI } from '@ai-sdk/openai'; export class DocsProviderTransport implements ChatTransport { - private openai: ReturnType; + private createResponse: ReturnType; constructor({ baseUrl }: { baseUrl: string }) { - this.openai = createOpenAI({ + const openai = createOpenAI({ baseURL: baseUrl, apiKey: '', }); + this.createResponse = makeCreateResponse({ + languageModel: openai.responses('mongodb-chat-latest'), + }); } sendMessages({ messages, + // TODO: pass the metadata correctly in a strongly typed manner. + // I'm not sure how to do this. + metadata, abortSignal, }: Parameters['sendMessages']>[0]) { - const result = streamText({ - model: this.openai.responses('mongodb-chat-latest'), + const result = this.createResponse({ messages: convertToModelMessages(messages), - abortSignal: abortSignal, + abortSignal, + systemPromptType: metadata?.type as keyof typeof buildPrompts, }); return Promise.resolve(result.toUIMessageStream()); @@ -35,3 +44,40 @@ export class DocsProviderTransport implements ChatTransport { return Promise.resolve(null); } } + +export function makeCreateResponse({ + languageModel, +}: { + languageModel: LanguageModel; +}) { + return function createResponse({ + messages, + abortSignal, + systemPromptType, + headers, + }: { + messages: ModelMessage[]; + abortSignal?: AbortSignal; + systemPromptType?: keyof typeof buildPrompts; + headers?: Record; + }) { + const instructions = + systemPromptType && buildPrompts[systemPromptType]?.system + ? buildPrompts[systemPromptType].system + : undefined; + + const result = streamText({ + model: languageModel, + messages: messages, + abortSignal: abortSignal, + headers, + providerOptions: { + openai: { + // Have to pass like this because it only accepts JSONValue (not instructions: undefined). + ...(instructions ? { instructions } : {}), + }, + }, + }); + return result; + }; +} diff --git a/packages/compass-assistant/src/index.tsx b/packages/compass-assistant/src/index.tsx index 2a7bcb6e5e5..d8ea6b76d64 100644 --- a/packages/compass-assistant/src/index.tsx +++ b/packages/compass-assistant/src/index.tsx @@ -5,3 +5,7 @@ export { compassAssistantServiceLocator, } from './compass-assistant-provider'; export type { CompassAssistantService } from './compass-assistant-provider'; +export type { + DocsProviderTransport, + makeCreateResponse, +} from './docs-provider-transport'; diff --git a/packages/compass-assistant/src/prompts.ts b/packages/compass-assistant/src/prompts.ts index a421228796b..1a9b914501d 100644 --- a/packages/compass-assistant/src/prompts.ts +++ b/packages/compass-assistant/src/prompts.ts @@ -1,4 +1,6 @@ -export const buildExplainPlanPrompt = ({ +const explainPlanSystemPrompt = `TODO:....add custom prompt stuff here`; + +const buildExplainPlanUserPrompt = ({ explainPlan, }: { explainPlan: string; @@ -12,7 +14,9 @@ ${explainPlan}`, }; }; -export const buildConnectionErrorPrompt = ({ +const connectionErrorSystemPrompt = `TODO:....add custom prompt stuff here`; + +const buildConnectionUserErrorPrompt = ({ connectionString, connectionError, }: { @@ -29,3 +33,14 @@ Error message: ${connectionError}`, }; }; + +export const buildPrompts = { + 'explain-plan': { + user: buildExplainPlanUserPrompt, + system: explainPlanSystemPrompt, + }, + 'connection-error': { + user: buildConnectionUserErrorPrompt, + system: connectionErrorSystemPrompt, + }, +} as const; diff --git a/packages/compass-assistant/test/assistant.eval.ts b/packages/compass-assistant/test/assistant.eval.ts index 8ef4a340426..b1092d97325 100644 --- a/packages/compass-assistant/test/assistant.eval.ts +++ b/packages/compass-assistant/test/assistant.eval.ts @@ -1,14 +1,16 @@ /* eslint-disable no-console */ import { createOpenAI } from '@ai-sdk/openai'; -import { streamText } from 'ai'; +import { defaultSettingsMiddleware, wrapLanguageModel } from 'ai'; import { init, Factuality as _Factuality } from 'autoevals'; -import { Eval } from 'braintrust'; +import { Eval, BraintrustMiddleware } from 'braintrust'; import type { EvalCase, EvalScorer } from 'braintrust'; import { OpenAI } from 'openai'; import { evalCases } from './eval-cases'; import { fuzzyLinkMatch } from './fuzzylinkmatch'; import { binaryNdcgAtK } from './binaryndcgatk'; import { makeEntrypointCases } from './entrypoints'; +import { makeCreateResponse } from '../src/docs-provider-transport'; +import { type buildPrompts } from '../src/prompts'; const client = new OpenAI({ baseURL: 'https://api.braintrust.dev/v1/proxy', @@ -20,8 +22,13 @@ init({ client }); export type SimpleEvalCase = { name?: string; input: string; + /** + * Include this in eval cases where you want to test the custom system prompt behavior. + */ + systemPromptType?: keyof typeof buildPrompts; expected: string; expectedSources?: string[]; + tags?: string[]; }; type Message = { @@ -33,6 +40,7 @@ type ExpectedMessage = OutputMessage; type ConversationEvalCaseInput = { messages: InputMessage[]; + systemPromptType?: keyof typeof buildPrompts; }; type ConversationEvalCaseExpected = { @@ -78,6 +86,29 @@ function getScorerTemperature(): number | undefined { return undefined; } +const createResponse = makeCreateResponse({ + languageModel: wrapLanguageModel({ + model: createOpenAI({ + baseURL: + process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE ?? + 'https://knowledge.staging.corp.mongodb.com/api/v1', + apiKey: '', + headers: { 'User-Agent': 'mongodb-compass/x.x.x' }, + }).responses('mongodb-chat-latest'), + middleware: [ + // Traces calls to Assistant API + BraintrustMiddleware({ + debug: true, + }), + defaultSettingsMiddleware({ + settings: { + temperature: getChatTemperature(), + }, + }), + ], + }), +}); + function makeEvalCases(): ConversationEvalCase[] { const entrypointCases: ConversationEvalCase[] = makeEntrypointCases(); @@ -90,6 +121,7 @@ function makeEvalCases(): ConversationEvalCase[] { expected: { messages: [{ text: c.expected, sources: c.expectedSources || [] }], }, + tags: c.tags ?? [], metadata: {}, }; }); @@ -100,42 +132,17 @@ function makeEvalCases(): ConversationEvalCase[] { async function makeAssistantCall( input: ConversationEvalCaseInput ): Promise { - const openai = createOpenAI({ - baseURL: - process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE ?? - 'https://knowledge.staging.corp.mongodb.com/api/v1', - apiKey: '', - headers: { - 'User-Agent': 'mongodb-compass/x.x.x', - }, - }); const prompt = allText(input.messages); - const result = streamText({ - model: openai.responses('mongodb-chat-latest'), - temperature: getChatTemperature(), - prompt, + const result = createResponse({ + messages: [{ content: prompt, role: 'user' }], + systemPromptType: input.systemPromptType, + }); + const text = await result.text; + // TODO: validate that this works as expected. If not, we must pull the sources out of the stream + const sources = (await result.sources).map((source) => { + return source.id; }); - - const chunks: string[] = []; - - for await (const chunk of result.toUIMessageStream()) { - const t = ((chunk as any).delta as string) || ''; - if (t) { - chunks.push(t); - } - } - const text = chunks.join(''); - - // TODO: something's up with this type. url does exist on it. - const resolvedSources = (await result.sources) as { url: string }[]; - - const sources = resolvedSources - .map((source) => { - console.log(source); - return source.url; - }) - .filter((url) => !!url); return { messages: [{ text, sources }], diff --git a/packages/compass-assistant/test/eval-cases/aggregation-pipeline.ts b/packages/compass-assistant/test/eval-cases/aggregation.ts similarity index 52% rename from packages/compass-assistant/test/eval-cases/aggregation-pipeline.ts rename to packages/compass-assistant/test/eval-cases/aggregation.ts index 9f0486bcf5a..9cf819931d0 100644 --- a/packages/compass-assistant/test/eval-cases/aggregation-pipeline.ts +++ b/packages/compass-assistant/test/eval-cases/aggregation.ts @@ -1,8 +1,11 @@ import type { SimpleEvalCase } from '../assistant.eval'; -const evalCase: SimpleEvalCase = { - input: 'What is an aggregation pipeline?', - expected: `The aggregation pipeline in MongoDB is a framework for data processing and transformation. It consists of a sequence of stages, where each stage performs an operation on the input documents and passes the results to the next stage. Common operations include filtering, grouping, projecting, joining, and calculating values. Aggregation pipelines are powerful for data analysis, reporting, and transformation tasks in MongoDB. +export const aggregationTag = 'aggregation'; + +const evalCases: SimpleEvalCase[] = [ + { + input: 'What is an aggregation pipeline?', + expected: `The aggregation pipeline in MongoDB is a framework for data processing and transformation. It consists of a sequence of stages, where each stage performs an operation on the input documents and passes the results to the next stage. Common operations include filtering, grouping, projecting, joining, and calculating values. Aggregation pipelines are powerful for data analysis, reporting, and transformation tasks in MongoDB. Compass makes it easy to create and run aggregation pipelines under the Aggregations tab. You may generate an aggregation pipeline with natural language, utilize the visual stage editor, or edit aggregations in the text view. @@ -30,10 +33,11 @@ db.orders.aggregate([ { $unset: ["_id"] } ]) `, - expectedSources: [ - 'https://www.mongodb.com/docs/manual/core/aggregation-pipeline/', - 'https://www.mongodb.com/docs/compass/create-agg-pipeline/', - ], -}; - -export default evalCase; + expectedSources: [ + 'https://www.mongodb.com/docs/manual/core/aggregation-pipeline/', + 'https://www.mongodb.com/docs/compass/create-agg-pipeline/', + ], + }, +].map((c) => ({ ...c, tags: [aggregationTag] })); + +export default evalCases; diff --git a/packages/compass-assistant/test/eval-cases/atlas-search.ts b/packages/compass-assistant/test/eval-cases/atlas-search.ts new file mode 100644 index 00000000000..ba1d6dc3783 --- /dev/null +++ b/packages/compass-assistant/test/eval-cases/atlas-search.ts @@ -0,0 +1,16 @@ +import type { SimpleEvalCase } from '../assistant.eval'; + +export const atlasSearchTag = 'atlas-search'; +const evalCases: SimpleEvalCase[] = [ + { + input: 'How can I filter docs before running a $search query?', + expected: + 'Because the $search stage must be the first stage in an aggregation pipeline, you cannot pre-filter documents with a preceding $match stage. Instead, filtering should be performed within the $search stage using the filter clause of the compound operator. This allows you to apply predicate queries (e.g., on ranges, dates, or specific terms) to narrow down the dataset before the main query clauses (must or should) are executed. Alternatively, you can filter documents by creating a View—a partial index of your collection that pre-queries and filters out unwanted documents. Note that users need createCollection privileges to build views.', + expectedSources: [ + 'https://www.mongodb.com/docs/atlas/atlas-search/compound/#options', + 'https://www.mongodb.com/docs/atlas/atlas-search/transform-documents-collections/#example--filter-documents', + ], + }, +].map((c) => ({ ...c, tags: [atlasSearchTag] })); + +export default evalCases; diff --git a/packages/compass-assistant/test/eval-cases/connection-error.ts b/packages/compass-assistant/test/eval-cases/connection-error.ts new file mode 100644 index 00000000000..cae797d3f01 --- /dev/null +++ b/packages/compass-assistant/test/eval-cases/connection-error.ts @@ -0,0 +1,21 @@ +import type { SimpleEvalCase } from '../assistant.eval'; + +export const connectionErrorTag = 'connection-error'; + +// TODO: Julian/whoever fills this in +const evalCases: SimpleEvalCase[] = [ + { + input: 'TODO:...', + + expected: 'TODO:...', + expectedSources: [ + // TODO: add sources + ], + }, +].map((c) => ({ + ...c, + tags: [connectionErrorTag], + systemPromptType: 'connection-error', +})); + +export default evalCases; diff --git a/packages/compass-assistant/test/eval-cases/explain-plan.ts b/packages/compass-assistant/test/eval-cases/explain-plan.ts new file mode 100644 index 00000000000..77371df9771 --- /dev/null +++ b/packages/compass-assistant/test/eval-cases/explain-plan.ts @@ -0,0 +1,21 @@ +import type { SimpleEvalCase } from '../assistant.eval'; + +export const explainPlanTag = 'explain-plan'; + +// TODO: Julian fills this in +const evalCases: SimpleEvalCase[] = [ + { + input: 'TODO:...', + + expected: 'TODO:...', + expectedSources: [ + // TODO: add sources + ], + }, +].map((c) => ({ + ...c, + tags: [explainPlanTag], + systemPromptType: 'explain-plan', +})); + +export default evalCases; diff --git a/packages/compass-assistant/test/eval-cases/filter-docs-before-search.ts b/packages/compass-assistant/test/eval-cases/filter-docs-before-search.ts deleted file mode 100644 index 8b5495fb13c..00000000000 --- a/packages/compass-assistant/test/eval-cases/filter-docs-before-search.ts +++ /dev/null @@ -1,13 +0,0 @@ -import type { SimpleEvalCase } from '../assistant.eval'; - -const evalCase: SimpleEvalCase = { - input: 'How can I filter docs before running a $search query?', - expected: - 'Because the $search stage must be the first stage in an aggregation pipeline, you cannot pre-filter documents with a preceding $match stage. Instead, filtering should be performed within the $search stage using the filter clause of the compound operator. This allows you to apply predicate queries (e.g., on ranges, dates, or specific terms) to narrow down the dataset before the main query clauses (must or should) are executed. Alternatively, you can filter documents by creating a View—a partial index of your collection that pre-queries and filters out unwanted documents. Note that users need createCollection privileges to build views.', - expectedSources: [ - 'https://www.mongodb.com/docs/atlas/atlas-search/compound/#options', - 'https://www.mongodb.com/docs/atlas/atlas-search/transform-documents-collections/#example--filter-documents', - ], -}; - -export default evalCase; diff --git a/packages/compass-assistant/test/eval-cases/index.ts b/packages/compass-assistant/test/eval-cases/index.ts index 816f0976391..72a21e77b1d 100644 --- a/packages/compass-assistant/test/eval-cases/index.ts +++ b/packages/compass-assistant/test/eval-cases/index.ts @@ -1,10 +1,14 @@ import type { SimpleEvalCase } from '../assistant.eval'; -import filterDocsBeforeSearch from './filter-docs-before-search'; -import aggregationPipeline from './aggregation-pipeline'; +import filterDocsBeforeSearch from './atlas-search'; +import aggregationPipeline from './aggregation'; import modelData from './model-data'; +import explainPlan from './explain-plan'; +import connectionError from './connection-error'; export const evalCases: SimpleEvalCase[] = [ filterDocsBeforeSearch, aggregationPipeline, modelData, -]; + explainPlan, + connectionError, +].flat(); diff --git a/packages/compass-assistant/test/eval-cases/model-data.ts b/packages/compass-assistant/test/eval-cases/model-data.ts index 5c5a7b9ad24..2d183f55d10 100644 --- a/packages/compass-assistant/test/eval-cases/model-data.ts +++ b/packages/compass-assistant/test/eval-cases/model-data.ts @@ -1,18 +1,22 @@ import type { SimpleEvalCase } from '../assistant.eval'; -const evalCase: SimpleEvalCase = { - input: 'How do I model data with MongoDB?', - expected: `Data modeling in MongoDB is highly dependent on how you access your data. To ensure that your data model has a logical structure and achieves optimal performance, plan your schema prior to using your database at a production scale. To determine your data model, use the following schema design process: +export const modelDataTag = 'model-data'; + +const evalCases: SimpleEvalCase[] = [ + { + input: 'How do I model data with MongoDB?', + expected: `Data modeling in MongoDB is highly dependent on how you access your data. To ensure that your data model has a logical structure and achieves optimal performance, plan your schema prior to using your database at a production scale. To determine your data model, use the following schema design process: Identify your workload: Identify the operations that your application runs most frequently Map relationships: Identify the relationships in your application's data and decide whether to link or embed related data. Apply design patterns: Apply schema design patterns to optimize reads and writes. Create indexes: Create indexes to support common query patterns. `, - expectedSources: [ - 'https://www.mongodb.com/docs/manual/data-modeling/#plan-your-schema', - 'https://www.mongodb.com/docs/manual/data-modeling/schema-design-process/#designing-your-schema', - ], -}; + expectedSources: [ + 'https://www.mongodb.com/docs/manual/data-modeling/#plan-your-schema', + 'https://www.mongodb.com/docs/manual/data-modeling/schema-design-process/#designing-your-schema', + ], + }, +].map((c) => ({ ...c, tags: [modelDataTag] })); -export default evalCase; +export default evalCases;