Skip to content

Commit cd5bf6e

Browse files
mabaasitmabaasit
authored
fix(gen-ai): parse responses individually for find and aggregate COMPASS-10176 (#7645)
parse responses for find/aggregate
1 parent cbc379b commit cd5bf6e

File tree

4 files changed

+127
-128
lines changed

4 files changed

+127
-128
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -954,6 +954,9 @@ describe('AtlasAiService', function () {
954954
] as Chunk[],
955955
response: {
956956
content: {
957+
aggregation: {
958+
pipeline: '',
959+
},
957960
query: {
958961
filter: "{test:'pineapple'}",
959962
project: null,

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -449,7 +449,7 @@ export class AtlasAiService {
449449
return this.generateQueryUsingChatbot(
450450
message,
451451
validateAIAggregationResponse,
452-
{ signal: input.signal }
452+
{ signal: input.signal, type: 'aggregate' }
453453
);
454454
}
455455
return this.getQueryOrAggregationFromUserInput(
@@ -470,6 +470,7 @@ export class AtlasAiService {
470470
const message = buildFindQueryPrompt(input);
471471
return this.generateQueryUsingChatbot(message, validateAIQueryResponse, {
472472
signal: input.signal,
473+
type: 'find',
473474
});
474475
}
475476
return this.getQueryOrAggregationFromUserInput(
@@ -565,15 +566,18 @@ export class AtlasAiService {
565566
private async generateQueryUsingChatbot<T>(
566567
message: AiQueryPrompt,
567568
validateFn: (res: any) => asserts res is T,
568-
options: { signal: AbortSignal }
569+
options: { signal: AbortSignal; type: 'find' | 'aggregate' }
569570
): Promise<T> {
570571
this.throwIfAINotEnabled();
571572
const response = await getAiQueryResponse(
572573
this.aiModel,
573574
message,
574575
options.signal
575576
);
576-
const parsedResponse = parseXmlToJsonResponse(response, this.logger);
577+
const parsedResponse = parseXmlToJsonResponse(response, {
578+
logger: this.logger,
579+
type: options.type,
580+
});
577581
validateFn(parsedResponse);
578582
return parsedResponse;
579583
}

‎packages/compass-generative-ai/src/utils/parse-xml-response.spec.ts‎

Lines changed: 100 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -3,118 +3,116 @@ import { parseXmlToJsonResponse } from './parse-xml-response';
33
import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
44

55
describe('parseXmlToJsonResponse', function () {
6-
it('should return prioritize aggregation over query when available and valid', function () {
7-
const xmlString = `
8-
<filter>{ age: { $gt: 25 } }</filter>
9-
<aggregation>[{ $match: { status: "A" } }]</aggregation>
10-
`;
11-
12-
const result = parseXmlToJsonResponse(xmlString, createNoopLogger());
13-
14-
expect(result).to.deep.equal({
15-
content: {
16-
aggregation: {
17-
pipeline: "[{$match:{status:'A'}}]",
18-
},
19-
query: {
20-
filter: null,
21-
project: null,
22-
sort: null,
23-
skip: null,
24-
limit: null,
6+
context('handles find', function () {
7+
const options = { logger: createNoopLogger(), type: 'find' as const };
8+
const NULL_QUERY = {
9+
filter: null,
10+
project: null,
11+
sort: null,
12+
skip: null,
13+
limit: null,
14+
};
15+
it('should return aggregation and query if provided', function () {
16+
const xmlString = `
17+
<filter>{ age: { $gt: 25 } }</filter>
18+
<aggregation>[{ $match: { status: "A" } }]</aggregation>
19+
`;
20+
const result = parseXmlToJsonResponse(xmlString, options);
21+
expect(result).to.deep.equal({
22+
content: {
23+
aggregation: {
24+
pipeline: "[{$match:{status:'A'}}]",
25+
},
26+
query: {
27+
filter: '{age:{$gt:25}}',
28+
project: null,
29+
sort: null,
30+
skip: null,
31+
limit: null,
32+
},
2533
},
26-
},
34+
});
2735
});
28-
});
29-
30-
it('should not return aggregation if its not available in the response', function () {
31-
const xmlString = `
32-
<filter>{ age: { $gt: 25 } }</filter>
33-
`;
34-
35-
const result = parseXmlToJsonResponse(xmlString, createNoopLogger());
36-
expect(result).to.deep.equal({
37-
content: {
38-
query: {
39-
filter: '{age:{$gt:25}}',
40-
project: null,
41-
sort: null,
42-
skip: null,
43-
limit: null,
36+
it('should return all the query fields if provided', function () {
37+
const xmlString = `
38+
<filter>{ age: { $gt: 25 } }</filter>
39+
<project>{ name: 1, age: 1 }</project>
40+
<sort>{ age: -1 }</sort>
41+
<skip>5</skip>
42+
<limit>10</limit>
43+
<aggregation></aggregation>
44+
`;
45+
const result = parseXmlToJsonResponse(xmlString, options);
46+
expect(result).to.deep.equal({
47+
content: {
48+
aggregation: {
49+
pipeline: '',
50+
},
51+
query: {
52+
filter: '{age:{$gt:25}}',
53+
project: '{name:1,age:1}',
54+
sort: '{age:-1}',
55+
skip: '5',
56+
limit: '10',
57+
},
4458
},
45-
},
59+
});
4660
});
47-
});
48-
49-
it('should not return query if its not available in the response', function () {
50-
const xmlString = `
51-
<aggregation>[{ $match: { status: "A" } }]</aggregation>
52-
`;
53-
54-
const result = parseXmlToJsonResponse(xmlString, createNoopLogger());
55-
56-
expect(result).to.deep.equal({
57-
content: {
58-
aggregation: {
59-
pipeline: "[{$match:{status:'A'}}]",
60-
},
61-
},
61+
context('it should handle invalid data', function () {
62+
it('invalid json', function () {
63+
const result = parseXmlToJsonResponse(
64+
`<filter>{ age: { $gt: 25 </filter>`,
65+
options
66+
);
67+
expect(result.content.query).to.deep.equal(NULL_QUERY);
68+
});
69+
it('empty object', function () {
70+
const result = parseXmlToJsonResponse(`<filter>{}</filter>`, options);
71+
expect(result.content.query).to.deep.equal(NULL_QUERY);
72+
});
73+
it('zero value', function () {
74+
const result = parseXmlToJsonResponse(`<limit>0</limit>`, options);
75+
expect(result.content.query).to.deep.equal(NULL_QUERY);
76+
});
6277
});
6378
});
6479

65-
it('should return all the query fields if provided', function () {
66-
const xmlString = `
67-
<filter>{ age: { $gt: 25 } }</filter>
68-
<project>{ name: 1, age: 1 }</project>
69-
<sort>{ age: -1 }</sort>
70-
<skip>5</skip>
71-
<limit>10</limit>
72-
<aggregation></aggregation>
73-
`;
74-
75-
const result = parseXmlToJsonResponse(xmlString, createNoopLogger());
76-
77-
expect(result).to.deep.equal({
78-
content: {
79-
query: {
80-
filter: '{age:{$gt:25}}',
81-
project: '{name:1,age:1}',
82-
sort: '{age:-1}',
83-
skip: '5',
84-
limit: '10',
80+
context('handles aggregate', function () {
81+
const options = { logger: createNoopLogger(), type: 'aggregate' as const };
82+
it('returns empty pipeline its not available in the response', function () {
83+
const xmlString = ``;
84+
const result = parseXmlToJsonResponse(xmlString, options);
85+
expect(result).to.deep.equal({
86+
content: {
87+
aggregation: {
88+
pipeline: '',
89+
},
8590
},
86-
},
91+
});
8792
});
88-
});
89-
90-
context('it should handle invalid data', function () {
91-
it('invalid json', function () {
92-
const result = parseXmlToJsonResponse(
93-
`<filter>{ age: { $gt: 25 </filter>`,
94-
createNoopLogger()
95-
);
96-
expect(result.content).to.not.have.property('query');
97-
});
98-
it('empty object', function () {
99-
const result = parseXmlToJsonResponse(
100-
`<filter>{}</filter>`,
101-
createNoopLogger()
102-
);
103-
expect(result.content).to.not.have.property('query');
104-
});
105-
it('empty array', function () {
106-
const result = parseXmlToJsonResponse(
107-
`<aggregation>[]</aggregation>`,
108-
createNoopLogger()
109-
);
110-
expect(result.content).to.not.have.property('aggregation');
93+
it('handles empty array', function () {
94+
const xmlString = `<aggregation>[]</aggregation>`;
95+
const result = parseXmlToJsonResponse(xmlString, options);
96+
expect(result).to.deep.equal({
97+
content: {
98+
aggregation: {
99+
pipeline: '',
100+
},
101+
},
102+
});
111103
});
112-
it('zero value', function () {
113-
const result = parseXmlToJsonResponse(
114-
`<limit>0</limit>`,
115-
createNoopLogger()
116-
);
117-
expect(result.content).to.not.have.property('query');
104+
it('returns aggregation pipeline if available', function () {
105+
const xmlString = `
106+
<aggregation>[{ $match: { status: "A" } }]</aggregation>
107+
`;
108+
const result = parseXmlToJsonResponse(xmlString, options);
109+
expect(result).to.deep.equal({
110+
content: {
111+
aggregation: {
112+
pipeline: "[{$match:{status:'A'}}]",
113+
},
114+
},
115+
});
118116
});
119117
});
120118
});

‎packages/compass-generative-ai/src/utils/parse-xml-response.ts‎

Lines changed: 17 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,13 @@ type ParsedXmlJsonResponse = {
1818

1919
export function parseXmlToJsonResponse(
2020
xmlString: string,
21-
logger: Logger
21+
{
22+
logger,
23+
type,
24+
}: {
25+
logger: Logger;
26+
type: 'find' | 'aggregate';
27+
}
2228
): ParsedXmlJsonResponse {
2329
const expectedTags = [
2430
'filter',
@@ -70,36 +76,24 @@ export function parseXmlToJsonResponse(
7076
}
7177
}
7278

73-
const { aggregation, ...query } = result;
74-
const isQueryEmpty = Object.values(query).every((v) => v === null);
79+
const { aggregation: pipeline, ...query } = result;
7580

76-
// It prioritizes aggregation over query if both are present
77-
if (aggregation && !isQueryEmpty) {
81+
const aggregation = {
82+
pipeline: pipeline ?? '',
83+
};
84+
// For aggregation, we only return aggregation field
85+
if (type === 'aggregate') {
7886
return {
7987
content: {
80-
aggregation: {
81-
pipeline: aggregation,
82-
},
83-
query: {
84-
filter: null,
85-
project: null,
86-
sort: null,
87-
skip: null,
88-
limit: null,
89-
},
88+
aggregation,
9089
},
9190
};
9291
}
92+
9393
return {
9494
content: {
95-
...(aggregation
96-
? {
97-
aggregation: {
98-
pipeline: aggregation,
99-
},
100-
}
101-
: {}),
102-
...(isQueryEmpty ? {} : { query }),
95+
query,
96+
aggregation,
10397
},
10498
};
10599
}

0 commit comments

Comments
 (0)