Skip to content
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

Jump to bottom

feat(chat): automatically pick the collection if there exists only one in the database VSCODE-610 #863

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

feat(chat): automatically pick the collection if there exists only one in the database VSCODE-610 #863

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ff0c1cc
WIP
gagik Oct 31, 2024
f9dd536
Move around dependencies
gagik Nov 1, 2024
09c7414
Remove grep
gagik Nov 1, 2024
aada8e1
Use firstCall
gagik Nov 1, 2024
8582193
Add schema tests
gagik Nov 1, 2024
a8bc30d
align tests and use a stub
gagik Nov 1, 2024
5ddc7fb
Add saving to metadata
gagik Nov 1, 2024
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
248 changes: 183 additions & 65 deletions src/participant/participant.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ import formatError from '../utils/formatError';
import type { ModelInput } from './prompts/promptBase';
import { processStreamWithIdentifiers } from './streamParsing';
import type { PromptIntent } from './prompts/intent';
import type { DataService } from 'mongodb-data-service';

const log = createLogger('participant');

Expand Down Expand Up @@ -671,64 +672,49 @@ export default class ParticipantController {
}
}

async renderCollectionsTree({
renderCollectionsTree({
collections,
command,
context,
databaseName,
stream,
}: {
collections: Awaited<ReturnType<DataService['listCollections']>>;
command: ParticipantCommand;
databaseName: string;
context: vscode.ChatContext;
stream: vscode.ChatResponseStream;
}): Promise<void> {
const dataService = this._connectionController.getActiveDataService();
if (!dataService) {
return;
}

stream.push(
new vscode.ChatResponseProgressPart('Fetching collection names...')
}): void {
collections.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((coll) =>
stream.markdown(
createMarkdownLink({
commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
data: {
command,
chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
context.history
),
databaseName,
collectionName: coll.name,
},
name: coll.name,
})
)
);

try {
const collections = await dataService.listCollections(databaseName);
collections.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((coll) =>
stream.markdown(
createMarkdownLink({
commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
data: {
command,
chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
context.history
),
databaseName,
collectionName: coll.name,
},
name: coll.name,
})
)
if (collections.length > MAX_MARKDOWN_LIST_LENGTH) {
stream.markdown(
createMarkdownLink({
commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
data: {
command,
chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
context.history
),
databaseName,
},
name: 'Show more',
})
);
if (collections.length > MAX_MARKDOWN_LIST_LENGTH) {
stream.markdown(
createMarkdownLink({
commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
data: {
command,
chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
context.history
),
databaseName,
},
name: 'Show more',
})
);
}
} catch (error) {
log.error('Unable to fetch collections:', error);

// Users can always do this manually when asked to provide a collection name.
return;
Copy link
Author

Choose a reason for hiding this comment

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

I am not sure what "this" is referring to here so I am not sure if there's something more to this but the behavior should be largely same with _getCollections, I just didn't transfer the comment over as I didn't understand it.

Copy link
Member

Choose a reason for hiding this comment

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

The comment here was alluding to if we can't fetch the collection names then a user can still input it by typing. It could definitely have been a bit more descriptive. I think that's why we're silently returning here instead of throwing an error.

}
}

Expand Down Expand Up @@ -811,7 +797,67 @@ export default class ParticipantController {
};
}

async _askForNamespace({
async _getCollections({
stream,
databaseName,
}: {
stream: vscode.ChatResponseStream;
databaseName: string;
}): Promise<ReturnType<DataService['listCollections']> | undefined> {
stream.push(
new vscode.ChatResponseProgressPart('Fetching collection names...')
);

const dataService = this._connectionController.getActiveDataService();

if (!dataService) {
return;
}
try {
return await dataService.listCollections(databaseName);
} catch (error) {
log.error('Unable to fetch collections:', error);
return;
}
}

/** Gets the collection name if there is only one collection.
* Otherwise returns undefined and asks the user to select the collection. */
async _getOrAskForCollectionName({
context,
databaseName,
stream,
command,
}: {
command: ParticipantCommand;
context: vscode.ChatContext;
databaseName: string;
stream: vscode.ChatResponseStream;
}): Promise<string | undefined> {
const collections = await this._getCollections({ stream, databaseName });

if (collections !== undefined) {
if (collections.length === 1) {
return collections[0].name;
}

stream.markdown(
`Which collection would you like to use within ${databaseName}?\n\n`
);

this.renderCollectionsTree({
collections,
command,
databaseName,
context,
stream,
});
}

return;
}

async _askForDatabaseName({
command,
context,
databaseName,
Expand All @@ -838,16 +884,6 @@ export default class ParticipantController {
context,
stream,
});
} else if (!collectionName) {
stream.markdown(
`Which collection would you like to use within ${databaseName}?\n\n`
);
await this.renderCollectionsTree({
command,
databaseName,
context,
stream,
});
}

return namespaceRequestChatResult({
Expand Down Expand Up @@ -1011,12 +1047,33 @@ export default class ParticipantController {
// we re-ask the question.
const databaseName = lastMessage.metadata.databaseName;
if (databaseName) {
const collections = await this._getCollections({
stream,
databaseName,
});

if (!collections) {
return namespaceRequestChatResult({
databaseName,
collectionName: undefined,
history: context.history,
});
}

// If there is only 1 collection in the database, we can pick it automatically
if (collections.length === 1) {
stream.markdown(Prompts.generic.getEmptyRequestResponse());
return emptyRequestChatResult(context.history);
}

stream.markdown(
vscode.l10n.t(
'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.'
)
);
await this.renderCollectionsTree({

this.renderCollectionsTree({
collections,
command,
databaseName,
context,
Expand All @@ -1043,6 +1100,7 @@ export default class ParticipantController {
}

// @MongoDB /schema
// eslint-disable-next-line complexity
async handleSchemaRequest(
request: vscode.ChatRequest,
context: vscode.ChatContext,
Expand All @@ -1068,14 +1126,16 @@ export default class ParticipantController {
});
}

const { databaseName, collectionName } = await this._getNamespaceFromChat({
const namespace = await this._getNamespaceFromChat({
request,
context,
token,
});
const { databaseName } = namespace;
let { collectionName } = namespace;

if (!databaseName || !collectionName) {
return await this._askForNamespace({
if (!databaseName) {
return await this._askForDatabaseName({
command: '/schema',
context,
databaseName,
Expand All @@ -1084,6 +1144,34 @@ export default class ParticipantController {
});
}

if (!collectionName) {
collectionName = await this._getOrAskForCollectionName({
command: '/schema',
context,
databaseName,
stream,
});

// If the collection name could not get automatically selected,
// then the user has been prompted for it.
if (!collectionName) {
return namespaceRequestChatResult({
databaseName,
collectionName: undefined,
history: context.history,
});
}

// Save the collection name in the metadata.
const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
context.history
);
this._chatMetadataStore.setChatMetadata(chatId, {
Copy link
Author

@gagik gagik Nov 1, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

Does this make sense to do? If I understand right this means this will effectively cache it but I haven't tested it yet, will soon.

...this._chatMetadataStore.getChatMetadata(chatId),
collectionName,
});
}

if (token.isCancellationRequested) {
return this._handleCancelledRequest({
context,
Expand Down Expand Up @@ -1194,20 +1282,50 @@ export default class ParticipantController {
// First we ask the model to parse for the database and collection name.
// If they exist, we can then use them in our final completion.
// When they don't exist we ask the user for them.
const { databaseName, collectionName } = await this._getNamespaceFromChat({
const namespace = await this._getNamespaceFromChat({
request,
context,
token,
});
if (!databaseName || !collectionName) {
return await this._askForNamespace({
const { databaseName } = namespace;
let { collectionName } = namespace;

if (!databaseName) {
return await this._askForDatabaseName({
command: '/query',
context,
databaseName,
collectionName,
stream,
});
}
if (!collectionName) {
collectionName = await this._getOrAskForCollectionName({
command: '/query',
context,
databaseName,
stream,
});

// If the collection name could not get automatically selected,
// then the user has been prompted for it.
if (!collectionName) {
return namespaceRequestChatResult({
databaseName,
collectionName: undefined,
history: context.history,
});
}

// Save the collection name in the metadata.
const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
context.history
);
this._chatMetadataStore.setChatMetadata(chatId, {
...this._chatMetadataStore.getChatMetadata(chatId),
collectionName,
});
}

if (token.isCancellationRequested) {
return this._handleCancelledRequest({
Expand Down
Loading
Loading