adjust the tests

mongodb-js · Dec 3, 2024 · 895e306 · 895e306
1 parent 545207a
commit 895e306
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/src/participant/prompts/promptBase.ts b/src/participant/prompts/promptBase.ts
@@ -101,9 +101,9 @@ export abstract class PromptBase<PromptArgs extends PromptArgsBase> {
     return undefined;
   }
 
-  protected getUserPrompt(
-    request: PromptArgsBase['request']
-  ): Promise<UserPromptResponse> {
+  protected getUserPrompt({
+    request,
+  }: PromptArgs): Promise<UserPromptResponse> {
     return Promise.resolve({
       prompt: request.prompt,
       hasSampleDocs: false,
@@ -121,7 +121,7 @@ export abstract class PromptBase<PromptArgs extends PromptArgsBase> {
       this.getAssistantPrompt(args)
     );
 
-    const { prompt, hasSampleDocs } = await this.getUserPrompt(request);
+    const { prompt, hasSampleDocs } = await this.getUserPrompt(args);
     // eslint-disable-next-line new-cap
     const userPrompt = vscode.LanguageModelChatMessage.User(prompt);
 

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
@@ -2038,7 +2038,7 @@ Schema:
   suite('prompt builders', function () {
     suite('prompt history', function () {
       test('gets filtered once history goes over maxInputTokens', async function () {
-        const expectedMaxMessages = 8;
+        const expectedMaxMessages = 10;
 
         const mockedMessages = Array.from(
           { length: 20 },
@@ -2065,11 +2065,18 @@ Schema:
           connectionNames: [],
         });
 
-        // Should include the limit and the initial generic prompt and the newly sent request
-        expect(messages.length + 2).equals(expectedMaxMessages);
+        expect(messages.length).equals(expectedMaxMessages);
+
+        // Should consist of the assistant prompt (1 token), 8 history messages (8 tokens),
+        // and the new request (1 token)
         expect(
           messages.slice(1).map((message) => getMessageContent(message))
-        ).deep.equal([...mockedMessages, chatRequestMock.prompt]);
+        ).deep.equal([
+          ...mockedMessages.slice(
+            mockedMessages.length - (expectedMaxMessages - 2)
+          ),
+          chatRequestMock.prompt,
+        ]);
       });
     });