-
Notifications
You must be signed in to change notification settings - Fork 246
feat(gen-ai): use chatbot api for queries COMPASS-10081 #7623
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.
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
Conversation
…s://github.com/mongodb-js/compass into COMPASS-10082-add-mms-prompts-and-feature-flag
…PASS-10081-switch-to-edu-api
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
This PR integrates the chatbot API for generative AI query functionality in MongoDB Compass, replacing the previous implementation with a new endpoint. The changes introduce XML response parsing, update authentication headers for service requests, and add comprehensive E2E testing for the new chatbot integration.
Key Changes:
- Introduces new chatbot API client using OpenAI SDK to generate MongoDB queries
- Adds XML-to-JSON response parser to handle chatbot responses
- Updates service providers with request origin headers for tracking
Reviewed changes
Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| packages/compass/src/app/components/entrypoint.tsx | Adds 'X-Request-Origin' header for Compass desktop |
| packages/compass-web/src/entrypoint.tsx | Adds 'X-Request-Origin' header for Atlas Data Explorer |
| packages/compass-generative-ai/src/utils/xml-to-mms-response.ts | Implements XML parsing logic to convert chatbot responses to expected format |
| packages/compass-generative-ai/src/utils/xml-to-mms-response.spec.ts | Adds comprehensive test coverage for XML parser |
| packages/compass-generative-ai/src/utils/gen-ai-response.ts | Implements chatbot API response streaming handler |
| packages/compass-generative-ai/src/utils/gen-ai-prompt.ts | Renames parameter from 'userPrompt' to 'userInput' for consistency |
| packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts | Updates tests to reflect parameter rename |
| packages/compass-generative-ai/src/chatbot-errors.ts | Defines error classes for chatbot API failures |
| packages/compass-generative-ai/src/atlas-ai-service.ts | Integrates chatbot client and adds feature flag for new endpoint |
| packages/compass-generative-ai/src/atlas-ai-service.spec.ts | Updates mock service with assistant endpoint method |
| packages/compass-generative-ai/package.json | Moves mongodb-query-parser to dependencies and adds AI SDK packages |
| packages/compass-e2e-tests/tests/collection-ai-query.test.ts | Adds E2E tests for chatbot-based query generation |
| packages/compass-e2e-tests/helpers/assistant-service.ts | Fixes header order in mock server response |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
packages/compass-generative-ai/src/utils/xml-to-mms-response.ts
Outdated
Show resolved
Hide resolved
|
Assigned |
| import type { LanguageModel } from 'ai'; | ||
| import { streamText } from 'ai'; | ||
|
|
||
| export async function getAiQueryResponse( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I know packages/compass-assistant/src/docs-provider-transport.ts is in entirely different path of code but wondering if one or the other should live closer together since there is some tie-in there.
I'm fine as is too though
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, good point. Initially I adopted that here as well and then remove it in clean up transport. I want to do a follow-up when this project ends and use the class removed in clean up transport commit. That will be exported by this package and will be used by the assistant.
| import type { Logger } from '@mongodb-js/compass-logging'; | ||
| import parse, { toJSString } from 'mongodb-query-parser'; | ||
|
|
||
| export function parseXmlToMmsJsonResponse(xmlString: string, logger: Logger) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
MmsJsonResponse feels like a technical detail to me; without context of the old assistant API I wouldn't understand what this meant.
Can we also add add a return type to this?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
added in add type
gribnoysup
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A few nits, but otherwise LGTM. We're planning to re-export the transport later down the road so that compass-assistant can depend on it isntead of having a copy, is that right?
| export type UserPromptForQueryOptions = { | ||
| userPrompt: string; | ||
| userInput: string; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nit: naming is slightly inconsistent, if you're changing userPrompt to userInput, then the type name should also change maybe?
| const reader = response.getReader(); | ||
| while (!done) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nit: the for await loop would be a more idiomatic way to deal with an async iterator I think
| const reader = response.getReader(); | |
| while (!done) { | |
| for await (const chunk of response) { |
| import { parseXmlToMmsJsonResponse } from './xml-to-mms-response'; | ||
| import type { Logger } from '@mongodb-js/compass-logging'; | ||
|
|
||
| const loggerMock = { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nit: this is what we have createNoopLogger for
| for (const tag of expectedTags) { | ||
| const regex = new RegExp(`<${tag}>([\\s\\S]*?)<\\/${tag}>`, 'i'); | ||
| const match = xmlString.match(regex); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Probably fine to roll our own "parser" here, but just FYI DOMParser browser API exists
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added this in clean up error, names. This however requires me to add top-level tag to the string so that its parsable.
|
|
||
| export class AiChatbotInvalidResponseError extends AiChatbotError { | ||
| constructor(message: string) { | ||
| super(500, message, 'INVALID_RESPONSE'); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Wouldn't we want to preserve the original server response code instead of hardcoding 500?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is a bit tricky one. This is thrown when the server returns an error while streaming. I was unable to recreate this use-case locally and I am not how the server responds with error in this case (what's contained in chunk.errorText). I think its fine if we throw 500 because we are not streaming response to the query bar.
| @@ -0,0 +1,20 @@ | |||
| export class AiChatbotError extends Error { | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is this a copy of AtlasServiceError that we export from the atlas-service package? This package depends on service already, can we just import it?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Naming nit: we don't usually use mms for naming in the codebase, it's an outdated term that is not really relevant here, especially as strictly speaking I guess this is not even an "mms response", it's just xml with selected fields parsed to json. There's probably a better name for that
Yes that's correct. I addressed the comments. Regarding using |
|
@gribnoysup I'll merge it for now and if there's something more that needs to be addressed, I'll handle that in upcoming PR. |
Description
Checklist
Motivation and Context
Open Questions
Dependents
Types of changes