From 66d5627952ca35e80558117f2dd470a4fc9c26f1 Mon Sep 17 00:00:00 2001 From: Rhys Howell Date: Fri, 27 Sep 2024 15:29:30 -0400 Subject: [PATCH 1/3] fix: add message content when blank for namespace request --- src/participant/participant.ts | 85 ++++++++- src/participant/prompts/index.ts | 20 ++- .../suite/participant/participant.test.ts | 169 +++++++++++------- 3 files changed, 202 insertions(+), 72 deletions(-) diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 8d57f0726..4088a1281 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -2,6 +2,7 @@ import * as vscode from 'vscode'; import { getSimplifiedSchema, parseSchema } from 'mongodb-schema'; import type { Document } from 'bson'; import type { Reference } from 'mongodb-rag-core'; +import util from 'util'; import { createLogger } from '../logging'; import type ConnectionController from '../connectionController'; @@ -137,6 +138,13 @@ export default class ParticipantController { errorName = ParticipantErrorTypes.OTHER; } + log.error('Participant encountered an error', { + command, + err, + error_code: errorCode, + error_name: errorName, + }); + this._telemetryService.track( TelemetryEventTypes.PARTICIPANT_RESPONSE_FAILED, { @@ -176,7 +184,21 @@ export default class ParticipantController { throw new Error('Copilot model not found'); } - return await model.sendRequest(messages, {}, token); + log.info('Sending request to model', { + truncatedMessages: messages.map( + (message: vscode.LanguageModelChatMessage) => + util.inspect({ + role: message.role, + content: message.content.slice(0, 50), + }) + ), + }); + + const modelResponse = await model.sendRequest(messages, {}, token); + + log.info('Model response received'); + + return modelResponse; } async streamChatResponse({ @@ -239,16 +261,25 @@ export default class ParticipantController { token, }); + let responseText = ''; + let codeBlocksInResponse = 0; await processStreamWithIdentifiers({ processStreamFragment: (fragment: string) => { + responseText += fragment; stream.markdown(fragment); }, onStreamIdentifier: (content: string) => { + codeBlocksInResponse++; this._streamCodeBlockActions({ runnableContent: content, stream }); }, inputIterable: chatResponse.text, identifier: codeBlockIdentifier, }); + + log.info('Streamed response to chat', { + responseText, + codeBlocksInResponse, + }); } // This will stream all of the response content and create a string from it. @@ -345,6 +376,10 @@ export default class ParticipantController { token, }); + log.info('Received intent response from model', { + responseContent, + }); + return Prompts.intent.getIntentFromModelResponse(responseContent); } @@ -707,14 +742,39 @@ export default class ParticipantController { request, connectionNames: this._getConnectionNames(), }); - const responseContentWithNamespace = await this.getChatResponseContent({ - messages: messagesWithNamespace, - token, - }); - const { databaseName, collectionName } = - Prompts.namespace.extractDatabaseAndCollectionNameFromResponse( - responseContentWithNamespace - ); + + let { + databaseName, + collectionName, + }: { + databaseName: string | undefined; + collectionName: string | undefined; + } = { + databaseName: undefined, + collectionName: undefined, + }; + + // When there's no user message content we can + // skip the request to the model. This would happen with /schema. + if (Prompts.doMessagesContainUserInput(messagesWithNamespace)) { + // VSCODE-626: When there's an empty message sent to the ai model, + // it currently errors (not on insiders, only main VSCode). + // Here we're defaulting to have some content as a workaround. + // TODO: Remove this when the issue is fixed. + messagesWithNamespace[messagesWithNamespace.length - 1].content = + messagesWithNamespace[ + messagesWithNamespace.length - 1 + ].content.trim() || 'see previous messages'; + + const responseContentWithNamespace = await this.getChatResponseContent({ + messages: messagesWithNamespace, + token, + }); + ({ databaseName, collectionName } = + Prompts.namespace.extractDatabaseAndCollectionNameFromResponse( + responseContentWithNamespace + )); + } // See if there's a namespace set in the // chat metadata we can fallback to if the model didn't find it. @@ -726,6 +786,11 @@ export default class ParticipantController { collectionName: collectionNameFromMetadata, } = this._chatMetadataStore.getChatMetadata(chatId) ?? {}; + log.info('Namespaces found in chat', { + databaseName: databaseName || databaseNameFromMetadata, + collectionName: collectionName || collectionNameFromMetadata, + }); + return { databaseName: databaseName || databaseNameFromMetadata, collectionName: collectionName || collectionNameFromMetadata, @@ -800,6 +865,8 @@ export default class ParticipantController { context: vscode.ChatContext; stream: vscode.ChatResponseStream; }): ChatResult { + log.info('Participant asked user to connect'); + stream.markdown( "Looks like you aren't currently connected, first let's get you connected to the cluster we'd like to create this query to run against.\n\n" ); diff --git a/src/participant/prompts/index.ts b/src/participant/prompts/index.ts index 867b057dd..18e4150af 100644 --- a/src/participant/prompts/index.ts +++ b/src/participant/prompts/index.ts @@ -1,4 +1,4 @@ -import type * as vscode from 'vscode'; +import * as vscode from 'vscode'; import { GenericPrompt } from './generic'; import { IntentPrompt } from './intent'; @@ -16,4 +16,22 @@ export class Prompts { public static isPromptEmpty(request: vscode.ChatRequest): boolean { return !request.prompt || request.prompt.trim().length === 0; } + + // Check if any of the messages contain user input. + // This is useful since when there's no user input in any + // messages, we can skip some additional processing. + public static doMessagesContainUserInput( + messages: vscode.LanguageModelChatMessage[] + ): boolean { + for (const message of messages) { + if ( + message.role === vscode.LanguageModelChatMessageRole.User && + message.content.trim().length > 0 + ) { + return true; + } + } + + return false; + } } diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts index 2610ccddf..09c7cc93d 100644 --- a/src/test/suite/participant/participant.test.ts +++ b/src/test/suite/participant/participant.test.ts @@ -1147,26 +1147,70 @@ suite('Participant Controller Test Suite', function () { }); suite('schema command', function () { - suite('known namespace from running namespace LLM', function () { + suite('no namespace provided', function () { beforeEach(function () { sendRequestStub.onCall(0).resolves({ - text: ['DATABASE_NAME: dbOne\n', 'COLLECTION_NAME: collOne\n`'], + text: ['none'], }); }); - test('shows a button to view the json output', async function () { + test('without a prompt it asks for the database name without pinging ai', async function () { const chatRequestMock = { prompt: '', command: 'schema', references: [], }; - sampleStub.resolves([ - { - _id: new ObjectId('63ed1d522d8573fa5c203660'), - }, - ]); await invokeChatHandler(chatRequestMock); - const expectedSchema = `{ + + expect(sendRequestStub.called).to.be.false; + const askForDBMessage = chatStreamStub.markdown.getCall(0).args[0]; + expect(askForDBMessage).to.include( + 'What is the name of the database you would like to run against?' + ); + }); + + test('with a prompt it asks the ai for the namespace', async function () { + const chatRequestMock = { + prompt: 'pineapple', + command: 'schema', + references: [], + }; + await invokeChatHandler(chatRequestMock); + + expect(sendRequestStub.calledOnce).to.be.true; + expect(sendRequestStub.firstCall.args[0][0].content).to.include( + 'Parse all user messages to find a database name and a collection name.' + ); + + const askForDBMessage = chatStreamStub.markdown.getCall(0).args[0]; + expect(askForDBMessage).to.include( + 'What is the name of the database you would like to run against?' + ); + }); + }); + + suite( + 'with a prompt and a known namespace from running namespace LLM', + function () { + beforeEach(function () { + sendRequestStub.onCall(0).resolves({ + text: ['DATABASE_NAME: dbOne\n', 'COLLECTION_NAME: collOne\n`'], + }); + }); + + test('shows a button to view the json output', async function () { + const chatRequestMock = { + prompt: 'what is my schema', + command: 'schema', + references: [], + }; + sampleStub.resolves([ + { + _id: new ObjectId('63ed1d522d8573fa5c203660'), + }, + ]); + await invokeChatHandler(chatRequestMock); + const expectedSchema = `{ "count": 1, "fields": [ { @@ -1192,73 +1236,74 @@ suite('Participant Controller Test Suite', function () { } ] }`; - expect(chatStreamStub?.button.getCall(0).args[0]).to.deep.equal({ - command: 'mdb.participantViewRawSchemaOutput', - title: 'Open JSON Output', - arguments: [ - { - schema: expectedSchema, - }, - ], + expect(chatStreamStub?.button.getCall(0).args[0]).to.deep.equal({ + command: 'mdb.participantViewRawSchemaOutput', + title: 'Open JSON Output', + arguments: [ + { + schema: expectedSchema, + }, + ], + }); }); - }); - test("includes the collection's schema in the request", async function () { - sampleStub.resolves([ - { - _id: new ObjectId('63ed1d522d8573fa5c203660'), - field: { - stringField: - 'There was a house cat who finally got the chance to do what it had always wanted to do.', - arrayField: [new Int32('1')], + test("includes the collection's schema in the request", async function () { + sampleStub.resolves([ + { + _id: new ObjectId('63ed1d522d8573fa5c203660'), + field: { + stringField: + 'There was a house cat who finally got the chance to do what it had always wanted to do.', + arrayField: [new Int32('1')], + }, }, - }, - { - _id: new ObjectId('63ed1d522d8573fa5c203660'), - field: { - stringField: 'Pineapple.', - arrayField: [new Int32('166')], + { + _id: new ObjectId('63ed1d522d8573fa5c203660'), + field: { + stringField: 'Pineapple.', + arrayField: [new Int32('166')], + }, }, - }, - ]); - const chatRequestMock = { - prompt: '', - command: 'schema', - references: [], - }; - await invokeChatHandler(chatRequestMock); - const messages = sendRequestStub.secondCall.args[0]; - expect(messages[0].content).to.include( - 'Amount of documents sampled: 2' - ); - expect(messages[1].content).to.include( - `Database name: dbOne + ]); + const chatRequestMock = { + prompt: 'what is my schema', + command: 'schema', + references: [], + }; + await invokeChatHandler(chatRequestMock); + const messages = sendRequestStub.secondCall.args[0]; + expect(messages[0].content).to.include( + 'Amount of documents sampled: 2' + ); + expect(messages[1].content).to.include( + `Database name: dbOne Collection name: collOne Schema: { "count": 2, "fields": [` - ); - expect(messages[1].content).to.include(`"name": "arrayField", + ); + expect(messages[1].content).to.include(`"name": "arrayField", "path": [ "field", "arrayField" ],`); - }); + }); - test('prints a message when no documents are found', async function () { - sampleStub.resolves([]); - const chatRequestMock = { - prompt: '', - command: 'schema', - references: [], - }; - await invokeChatHandler(chatRequestMock); - expect(chatStreamStub?.markdown.getCall(0).args[0]).to.include( - 'Unable to generate a schema from the collection, no documents found.' - ); - }); - }); + test('prints a message when no documents are found', async function () { + sampleStub.resolves([]); + const chatRequestMock = { + prompt: 'what is my schema', + command: 'schema', + references: [], + }; + await invokeChatHandler(chatRequestMock); + expect(chatStreamStub?.markdown.getCall(0).args[0]).to.include( + 'Unable to generate a schema from the collection, no documents found.' + ); + }); + } + ); }); suite('docs command', function () { From 155d35e8b5d66572d8a5abeb572d3bc1a0d87675 Mon Sep 17 00:00:00 2001 From: Rhys Howell Date: Fri, 27 Sep 2024 16:28:02 -0400 Subject: [PATCH 2/3] fixup: remove possibly sensitive info from logs --- src/participant/participant.ts | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 4088a1281..d58335f7a 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -140,7 +140,6 @@ export default class ParticipantController { log.error('Participant encountered an error', { command, - err, error_code: errorCode, error_name: errorName, }); @@ -185,12 +184,11 @@ export default class ParticipantController { } log.info('Sending request to model', { - truncatedMessages: messages.map( - (message: vscode.LanguageModelChatMessage) => - util.inspect({ - role: message.role, - content: message.content.slice(0, 50), - }) + messages: messages.map((message: vscode.LanguageModelChatMessage) => + util.inspect({ + role: message.role, + contentLength: message.content.length, + }) ), }); @@ -261,11 +259,11 @@ export default class ParticipantController { token, }); - let responseText = ''; + let responseLength = ''; let codeBlocksInResponse = 0; await processStreamWithIdentifiers({ processStreamFragment: (fragment: string) => { - responseText += fragment; + responseLength += fragment.length; stream.markdown(fragment); }, onStreamIdentifier: (content: string) => { @@ -277,7 +275,7 @@ export default class ParticipantController { }); log.info('Streamed response to chat', { - responseText, + responseLength: responseLength, codeBlocksInResponse, }); } From 191831c1f68c05af292bd13989742f9ad6b811c2 Mon Sep 17 00:00:00 2001 From: Rhys Howell Date: Fri, 27 Sep 2024 18:08:15 -0400 Subject: [PATCH 3/3] fixup: remove content log --- src/participant/participant.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 05522a0c2..519464fab 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -406,7 +406,7 @@ export default class ParticipantController { }); log.info('Received intent response from model', { - responseContent, + responseContentLength: responseContent.length, }); return Prompts.intent.getIntentFromModelResponse(responseContent);