From 498c893d12684bfbbc3be199041e6e364dc6cdff Mon Sep 17 00:00:00 2001 From: pallpointben Date: Mon, 12 May 2025 15:20:17 -0400 Subject: [PATCH 1/2] setup basic doctor command --- src/participant/constants.ts | 6 ++ src/participant/participant.ts | 120 ++++++++++++++++++++++++++++ src/participant/participantTypes.ts | 3 +- src/participant/prompts/doctor.ts | 53 ++++++++++++ src/participant/prompts/index.ts | 2 + 5 files changed, 183 insertions(+), 1 deletion(-) create mode 100644 src/participant/prompts/doctor.ts diff --git a/src/participant/constants.ts b/src/participant/constants.ts index 5c4921976..93fa0e400 100644 --- a/src/participant/constants.ts +++ b/src/participant/constants.ts @@ -89,6 +89,12 @@ export function genericRequestChatResult( return createChatResult('generic', history); } +export function doctorRequestChatResult( + history: ReadonlyArray, +): ChatResult { + return createChatResult('doctor', history); +} + export function queryRequestChatResult( history: ReadonlyArray, ): ChatResult { diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 42f95c6b7..e637fd5e8 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -19,6 +19,7 @@ import { genericRequestChatResult, namespaceRequestChatResult, queryRequestChatResult, + doctorRequestChatResult, docsRequestChatResult, schemaRequestChatResult, createCancelledRequestChatResult, @@ -509,6 +510,8 @@ export default class ParticipantController { return this.handleSchemaRequest(request, context, stream, token); case 'Code': return this.handleQueryRequest(request, context, stream, token); + case 'Doctor': + return this.handleDoctorRequest(request, context, stream, token); default: return this._handleRoutedGenericRequest( request, @@ -1456,6 +1459,123 @@ export default class ParticipantController { return schemaRequestChatResult(context.history); } + // @MongoDB /doctor what's wrong with my query to find offices with at least 100 employees? + // eslint-disable-next-line complexity + async handleDoctorRequest( + request: vscode.ChatRequest, + context: vscode.ChatContext, + stream: vscode.ChatResponseStream, + token: vscode.CancellationToken, + ): Promise { + if (!this._connectionController.getActiveDataService()) { + return this._askToConnect({ + command: '/doctor', + context, + stream, + }); + } + + if (Prompts.isPromptEmpty(request)) { + if (this._doesLastMessageAskForNamespace(context.history)) { + return this.handleEmptyNamespaceMessage({ + command: '/doctor', + context, + stream, + }); + } + + stream.markdown(Prompts.doctor.emptyRequestResponse); + return emptyRequestChatResult(context.history); + } + + // 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 namespace = await this._getNamespaceFromChat({ + request, + context, + token, + }); + const { databaseName, collectionName } = + await this._getOrAskForMissingNamespace({ + ...namespace, + context, + stream, + command: '/doctor', + }); + + // If either the database or collection name could not be automatically picked + // then the user has been prompted to select one manually. + if (databaseName === undefined || collectionName === undefined) { + return namespaceRequestChatResult({ + databaseName, + collectionName, + history: context.history, + }); + } + + if (token.isCancellationRequested) { + return this._handleCancelledRequest({ + context, + stream, + }); + } + + // TODO try to access the user's Atlas credentials including an identifier for the connected + // Atlas project. If not available, ask the user to provide. If they say no, we can try + // to help without it. + + let schema: string | undefined; + let sampleDocuments: Document[] | undefined; + try { + ({ schema, sampleDocuments } = + await this._fetchCollectionSchemaAndSampleDocuments({ + databaseName, + collectionName, + token, + stream, + })); + } catch (e) { + // When an error fetching the collection schema or sample docs occurs, + // we still want to continue as it isn't critical, however, + // we do want to notify the user. + stream.markdown( + vscode.l10n.t( + 'An error occurred while fetching the collection schema and sample documents.', + ), + ); + } + + const modelInput = await Prompts.doctor.buildMessages({ + request, + context, + databaseName, + collectionName, + schema, + connectionNames: this._getConnectionNames(), + ...(sampleDocuments ? { sampleDocuments } : {}), + }); + + await this.streamChatResponseContentWithCodeActions({ + modelInput, + stream, + token, + }); + + // We don't want telemetry for this very unofficial command. + // this._telemetryService.track( + // new ParticipantResponseGeneratedTelemetryEvent({ + // command: 'doctor', + // hasCta: false, + // foundNamespace: true, + // hasRunnableContent: hasCodeBlock, + // outputLength: outputLength, + // }), + // ); + + return doctorRequestChatResult(context.history); + } + // @MongoDB /query find all documents where the "address" has the word Broadway in it. async handleQueryRequest( request: vscode.ChatRequest, diff --git a/src/participant/participantTypes.ts b/src/participant/participantTypes.ts index 3b65ceef0..969f21343 100644 --- a/src/participant/participantTypes.ts +++ b/src/participant/participantTypes.ts @@ -1,7 +1,7 @@ import type * as vscode from 'vscode'; import type { ParticipantTelemetryMetadata } from '../telemetry'; -export type ParticipantCommandType = 'query' | 'schema' | 'docs'; +export type ParticipantCommandType = 'query' | 'schema' | 'docs' | 'doctor'; export type ParticipantCommand = `/${ParticipantCommandType}`; export type ParticipantRequestType = ParticipantCommandType | 'generic'; @@ -12,6 +12,7 @@ export type ParticipantResponseType = | 'docs' | 'docs/chatbot' | 'docs/copilot' + | 'doctor' | 'exportToPlayground' | 'generic' | 'emptyRequest' diff --git a/src/participant/prompts/doctor.ts b/src/participant/prompts/doctor.ts new file mode 100644 index 000000000..fd43dd785 --- /dev/null +++ b/src/participant/prompts/doctor.ts @@ -0,0 +1,53 @@ +import * as vscode from 'vscode'; +import type { Document } from 'bson'; + +import { getStringifiedSampleDocuments } from '../sampleDocuments'; +import type { PromptArgsBase, UserPromptResponse } from './promptBase'; +import { codeBlockIdentifier } from '../constants'; +import { PromptBase } from './promptBase'; + +interface DoctorPromptArgs extends PromptArgsBase { + databaseName: string; + collectionName: string; + schema?: string; + sampleDocuments?: Document[]; + connectionNames?: string[]; +} + +export class DoctorPrompt extends PromptBase { + protected getAssistantPrompt(): string { + return `You are a MongoDB expert. +Your task is to respond that this prompt has not been implemented yet.`; + } + + async getUserPrompt({ + databaseName = 'mongodbVSCodeCopilotDB', + collectionName = 'test', + request, + schema, + sampleDocuments, + }: DoctorPromptArgs): Promise { + let prompt = request.prompt; + prompt += `\nDatabase name: ${databaseName}\n`; + prompt += `Collection name: ${collectionName}\n`; + if (schema) { + prompt += `Collection schema: ${schema}\n`; + } + + const sampleDocumentsPrompt = await getStringifiedSampleDocuments({ + sampleDocuments, + prompt, + }); + + return { + prompt: `${prompt}${sampleDocumentsPrompt}`, + hasSampleDocs: !!sampleDocumentsPrompt, + }; + } + + get emptyRequestResponse(): string { + return vscode.l10n.t( + 'Please specify a question when using this command. Usage: @MongoDB /query find documents where "name" contains "database".', + ); + } +} diff --git a/src/participant/prompts/index.ts b/src/participant/prompts/index.ts index 59d0f8718..50c9dda82 100644 --- a/src/participant/prompts/index.ts +++ b/src/participant/prompts/index.ts @@ -9,6 +9,7 @@ import { ExportToPlaygroundPrompt } from './exportToPlayground'; import { ExportToLanguagePrompt } from './exportToLanguage'; import { isContentEmpty } from './promptBase'; import { DocsPrompt } from './docs'; +import { DoctorPrompt } from './doctor'; export { getContentLength } from './promptBase'; @@ -18,6 +19,7 @@ export class Prompts { public static intent = new IntentPrompt(); public static namespace = new NamespacePrompt(); public static query = new QueryPrompt(); + public static doctor = new DoctorPrompt(); public static schema = new SchemaPrompt(); public static exportToPlayground = new ExportToPlaygroundPrompt(); public static exportToLanguage = new ExportToLanguagePrompt(); From 39afbcc46e16e71c19e8f665dc1fc4788dfed0f8 Mon Sep 17 00:00:00 2001 From: pallpointben Date: Mon, 12 May 2025 15:23:14 -0400 Subject: [PATCH 2/2] lint --- src/participant/prompts/doctor.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/participant/prompts/doctor.ts b/src/participant/prompts/doctor.ts index fd43dd785..fc6077e8b 100644 --- a/src/participant/prompts/doctor.ts +++ b/src/participant/prompts/doctor.ts @@ -3,7 +3,6 @@ import type { Document } from 'bson'; import { getStringifiedSampleDocuments } from '../sampleDocuments'; import type { PromptArgsBase, UserPromptResponse } from './promptBase'; -import { codeBlockIdentifier } from '../constants'; import { PromptBase } from './promptBase'; interface DoctorPromptArgs extends PromptArgsBase {