Skip to content

fix(gen-ai): parse responses individually for find and aggregate COMPASS-10176 #7645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mabaasit merged 3 commits into from
Dec 15, 2025
Merged

fix(gen-ai): parse responses individually for find and aggregate COMPASS-10176 #7645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6ad4791
parse responses for find/aggregate
mabaasit Dec 12, 2025
e7041fc
Merge branch 'main' into fix-gen-ai-response
mabaasit Dec 15, 2025
2ec1f8c
Merge branch 'main' into fix-gen-ai-response
mabaasit Dec 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions packages/compass-generative-ai/src/atlas-ai-service.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -954,6 +954,9 @@ describe('AtlasAiService', function () {
] as Chunk[],
response: {
content: {
aggregation: {
pipeline: '',
},
query: {
filter: "{test:'pineapple'}",
project: null,
Expand Down
10 changes: 7 additions & 3 deletions packages/compass-generative-ai/src/atlas-ai-service.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -449,7 +449,7 @@ export class AtlasAiService {
return this.generateQueryUsingChatbot(
message,
validateAIAggregationResponse,
{ signal: input.signal }
{ signal: input.signal, type: 'aggregate' }
);
}
return this.getQueryOrAggregationFromUserInput(
Expand All @@ -470,6 +470,7 @@ export class AtlasAiService {
const message = buildFindQueryPrompt(input);
return this.generateQueryUsingChatbot(message, validateAIQueryResponse, {
signal: input.signal,
type: 'find',
});
}
return this.getQueryOrAggregationFromUserInput(
Expand Down Expand Up @@ -565,15 +566,18 @@ export class AtlasAiService {
private async generateQueryUsingChatbot<T>(
message: AiQueryPrompt,
validateFn: (res: any) => asserts res is T,
options: { signal: AbortSignal }
options: { signal: AbortSignal; type: 'find' | 'aggregate' }
): Promise<T> {
this.throwIfAINotEnabled();
const response = await getAiQueryResponse(
this.aiModel,
message,
options.signal
);
const parsedResponse = parseXmlToJsonResponse(response, this.logger);
const parsedResponse = parseXmlToJsonResponse(response, {
logger: this.logger,
type: options.type,
});
validateFn(parsedResponse);
return parsedResponse;
}
Expand Down
202 changes: 100 additions & 102 deletions packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = `
<filter>{ age: { $gt: 25 } }</filter>
<aggregation>[{ $match: { status: "A" } }]</aggregation>
`;

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 = `
<filter>{ age: { $gt: 25 } }</filter>
<aggregation>[{ $match: { status: "A" } }]</aggregation>
`;
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 = `
<filter>{ age: { $gt: 25 } }</filter>
`;

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 = `
<filter>{ age: { $gt: 25 } }</filter>
<project>{ name: 1, age: 1 }</project>
<sort>{ age: -1 }</sort>
<skip>5</skip>
<limit>10</limit>
<aggregation></aggregation>
`;
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 = `
<aggregation>[{ $match: { status: "A" } }]</aggregation>
`;

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(
`<filter>{ age: { $gt: 25 </filter>`,
options
);
expect(result.content.query).to.deep.equal(NULL_QUERY);
});
it('empty object', function () {
const result = parseXmlToJsonResponse(`<filter>{}</filter>`, options);
expect(result.content.query).to.deep.equal(NULL_QUERY);
});
it('zero value', function () {
const result = parseXmlToJsonResponse(`<limit>0</limit>`, options);
expect(result.content.query).to.deep.equal(NULL_QUERY);
});
});
});

it('should return all the query fields if provided', function () {
const xmlString = `
<filter>{ age: { $gt: 25 } }</filter>
<project>{ name: 1, age: 1 }</project>
<sort>{ age: -1 }</sort>
<skip>5</skip>
<limit>10</limit>
<aggregation></aggregation>
`;

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 () {
Copy link

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Corrected spelling: 'its' should be 'if it's' or 'when it's' for grammatical correctness.

Suggested change
it('returns empty pipeline its not available in the response', function () {
it("returns empty pipeline if it's not available in the response", function () {

Copilot uses AI. Check for mistakes.
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(
`<filter>{ age: { $gt: 25 </filter>`,
createNoopLogger()
);
expect(result.content).to.not.have.property('query');
});
it('empty object', function () {
const result = parseXmlToJsonResponse(
`<filter>{}</filter>`,
createNoopLogger()
);
expect(result.content).to.not.have.property('query');
});
it('empty array', function () {
const result = parseXmlToJsonResponse(
`<aggregation>[]</aggregation>`,
createNoopLogger()
);
expect(result.content).to.not.have.property('aggregation');
it('handles empty array', function () {
const xmlString = `<aggregation>[]</aggregation>`;
const result = parseXmlToJsonResponse(xmlString, options);
expect(result).to.deep.equal({
content: {
aggregation: {
pipeline: '',
},
},
});
});
it('zero value', function () {
const result = parseXmlToJsonResponse(
`<limit>0</limit>`,
createNoopLogger()
);
expect(result.content).to.not.have.property('query');
it('returns aggregation pipeline if available', function () {
const xmlString = `
<aggregation>[{ $match: { status: "A" } }]</aggregation>
`;
const result = parseXmlToJsonResponse(xmlString, options);
expect(result).to.deep.equal({
content: {
aggregation: {
pipeline: "[{$match:{status:'A'}}]",
},
},
});
});
});
});
40 changes: 17 additions & 23 deletions packages/compass-generative-ai/src/utils/parse-xml-response.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,13 @@ type ParsedXmlJsonResponse = {

export function parseXmlToJsonResponse(
xmlString: string,
logger: Logger
{
logger,
type,
}: {
logger: Logger;
type: 'find' | 'aggregate';
}
): ParsedXmlJsonResponse {
const expectedTags = [
'filter',
Expand Down Expand Up @@ -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,
Copy link

Copilot AI Dec 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the 'find' type, returning both query and aggregation fields means aggregation will always be included even when not relevant to find operations. Consider if the find path should only return query-related fields to maintain a clearer separation between operation types, similar to how aggregate only returns aggregation.

Suggested change
aggregation,

Copilot uses AI. Check for mistakes.
},
};
}
Loading