From 6ad479181ca5760a4b73572631cc6d311efa8396 Mon Sep 17 00:00:00 2001 From: Basit Chonka Date: Fri, 12 Dec 2025 17:30:06 +0300 Subject: [PATCH] parse responses for find/aggregate --- .../src/atlas-ai-service.spec.ts | 3 + .../src/atlas-ai-service.ts | 10 +- .../src/utils/parse-xml-response.spec.ts | 202 +++++++++--------- .../src/utils/parse-xml-response.ts | 40 ++-- 4 files changed, 127 insertions(+), 128 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 e56ffbcef2f..1866a616f9d 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -954,6 +954,9 @@ describe('AtlasAiService', function () { ] as Chunk[], response: { content: { + aggregation: { + pipeline: '', + }, query: { filter: "{test:'pineapple'}", project: null, diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 5baf8c62b60..950940a6daa 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -449,7 +449,7 @@ export class AtlasAiService { return this.generateQueryUsingChatbot( message, validateAIAggregationResponse, - { signal: input.signal } + { signal: input.signal, type: 'aggregate' } ); } return this.getQueryOrAggregationFromUserInput( @@ -470,6 +470,7 @@ export class AtlasAiService { const message = buildFindQueryPrompt(input); return this.generateQueryUsingChatbot(message, validateAIQueryResponse, { signal: input.signal, + type: 'find', }); } return this.getQueryOrAggregationFromUserInput( @@ -565,7 +566,7 @@ export class AtlasAiService { private async generateQueryUsingChatbot( message: AiQueryPrompt, validateFn: (res: any) => asserts res is T, - options: { signal: AbortSignal } + options: { signal: AbortSignal; type: 'find' | 'aggregate' } ): Promise { this.throwIfAINotEnabled(); const response = await getAiQueryResponse( @@ -573,7 +574,10 @@ export class AtlasAiService { message, options.signal ); - const parsedResponse = parseXmlToJsonResponse(response, this.logger); + const parsedResponse = parseXmlToJsonResponse(response, { + logger: this.logger, + type: options.type, + }); validateFn(parsedResponse); return parsedResponse; } diff --git a/packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts b/packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts index 87fa314ac13..6b59fc39e2e 100644 --- a/packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts +++ b/packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts @@ -3,118 +3,116 @@ import { parseXmlToJsonResponse } from './parse-xml-response'; import { createNoopLogger } from '@mongodb-js/compass-logging/provider'; 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 = parseXmlToJsonResponse(xmlString, createNoopLogger()); - - expect(result).to.deep.equal({ - content: { - aggregation: { - pipeline: "[{$match:{status:'A'}}]", - }, - query: { - filter: null, - project: null, - sort: null, - skip: null, - limit: null, + context('handles find', function () { + const options = { logger: createNoopLogger(), type: 'find' as const }; + const NULL_QUERY = { + filter: null, + project: null, + sort: null, + skip: null, + limit: null, + }; + it('should return aggregation and query if provided', function () { + const xmlString = ` + { age: { $gt: 25 } } + [{ $match: { status: "A" } }] + `; + const result = parseXmlToJsonResponse(xmlString, options); + expect(result).to.deep.equal({ + content: { + aggregation: { + pipeline: "[{$match:{status:'A'}}]", + }, + query: { + filter: '{age:{$gt:25}}', + project: null, + sort: null, + skip: null, + limit: null, + }, }, - }, + }); }); - }); - - it('should not return aggregation if its not available in the response', function () { - const xmlString = ` - { age: { $gt: 25 } } - `; - - const result = parseXmlToJsonResponse(xmlString, createNoopLogger()); - expect(result).to.deep.equal({ - content: { - query: { - filter: '{age:{$gt:25}}', - 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 = parseXmlToJsonResponse(xmlString, options); + expect(result).to.deep.equal({ + content: { + aggregation: { + pipeline: '', + }, + query: { + filter: '{age:{$gt:25}}', + project: '{name:1,age:1}', + sort: '{age:-1}', + skip: '5', + limit: '10', + }, }, - }, + }); }); - }); - - it('should not return query if its not available in the response', function () { - const xmlString = ` - [{ $match: { status: "A" } }] - `; - - const result = parseXmlToJsonResponse(xmlString, createNoopLogger()); - - expect(result).to.deep.equal({ - content: { - aggregation: { - pipeline: "[{$match:{status:'A'}}]", - }, - }, + context('it should handle invalid data', function () { + it('invalid json', function () { + const result = parseXmlToJsonResponse( + `{ age: { $gt: 25 `, + options + ); + expect(result.content.query).to.deep.equal(NULL_QUERY); + }); + it('empty object', function () { + const result = parseXmlToJsonResponse(`{}`, options); + expect(result.content.query).to.deep.equal(NULL_QUERY); + }); + it('zero value', function () { + const result = parseXmlToJsonResponse(`0`, options); + expect(result.content.query).to.deep.equal(NULL_QUERY); + }); }); }); - 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 = parseXmlToJsonResponse(xmlString, createNoopLogger()); - - expect(result).to.deep.equal({ - content: { - query: { - filter: '{age:{$gt:25}}', - project: '{name:1,age:1}', - sort: '{age:-1}', - skip: '5', - limit: '10', + context('handles aggregate', function () { + const options = { logger: createNoopLogger(), type: 'aggregate' as const }; + it('returns empty pipeline its not available in the response', function () { + const xmlString = ``; + const result = parseXmlToJsonResponse(xmlString, options); + expect(result).to.deep.equal({ + content: { + aggregation: { + pipeline: '', + }, }, - }, + }); }); - }); - - context('it should handle invalid data', function () { - it('invalid json', function () { - const result = parseXmlToJsonResponse( - `{ age: { $gt: 25 `, - createNoopLogger() - ); - expect(result.content).to.not.have.property('query'); - }); - it('empty object', function () { - const result = parseXmlToJsonResponse( - `{}`, - createNoopLogger() - ); - expect(result.content).to.not.have.property('query'); - }); - it('empty array', function () { - const result = parseXmlToJsonResponse( - `[]`, - createNoopLogger() - ); - expect(result.content).to.not.have.property('aggregation'); + it('handles empty array', function () { + const xmlString = `[]`; + const result = parseXmlToJsonResponse(xmlString, options); + expect(result).to.deep.equal({ + content: { + aggregation: { + pipeline: '', + }, + }, + }); }); - it('zero value', function () { - const result = parseXmlToJsonResponse( - `0`, - createNoopLogger() - ); - expect(result.content).to.not.have.property('query'); + it('returns aggregation pipeline if available', function () { + const xmlString = ` + [{ $match: { status: "A" } }] + `; + const result = parseXmlToJsonResponse(xmlString, options); + expect(result).to.deep.equal({ + content: { + aggregation: { + pipeline: "[{$match:{status:'A'}}]", + }, + }, + }); }); }); }); 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 44e34a6c971..dc3531aa0a1 100644 --- a/packages/compass-generative-ai/src/utils/parse-xml-response.ts +++ b/packages/compass-generative-ai/src/utils/parse-xml-response.ts @@ -18,7 +18,13 @@ type ParsedXmlJsonResponse = { export function parseXmlToJsonResponse( xmlString: string, - logger: Logger + { + logger, + type, + }: { + logger: Logger; + type: 'find' | 'aggregate'; + } ): ParsedXmlJsonResponse { const expectedTags = [ 'filter', @@ -70,36 +76,24 @@ export function parseXmlToJsonResponse( } } - const { aggregation, ...query } = result; - const isQueryEmpty = Object.values(query).every((v) => v === null); + const { aggregation: pipeline, ...query } = result; - // It prioritizes aggregation over query if both are present - if (aggregation && !isQueryEmpty) { + const aggregation = { + pipeline: pipeline ?? '', + }; + // For aggregation, we only return aggregation field + if (type === 'aggregate') { return { content: { - aggregation: { - pipeline: aggregation, - }, - query: { - filter: null, - project: null, - sort: null, - skip: null, - limit: null, - }, + aggregation, }, }; } + return { content: { - ...(aggregation - ? { - aggregation: { - pipeline: aggregation, - }, - } - : {}), - ...(isQueryEmpty ? {} : { query }), + query, + aggregation, }, }; }