feat(gen-ai): handle model limits COMPASS-10129 (#7655)

* handle model limits and throw error with code

* copilot feedback

Author: mabaasit

packages/compass-generative-ai/src/chatbot-errors.ts

@@ -6,3 +6,10 @@ export class AiChatbotInvalidResponseError extends AtlasServiceError {
     this.name = 'AiChatbotInvalidResponseError';
   }
 }
+
+export class AiChatbotPromptTooLargeError extends AtlasServiceError {
+  constructor(message: string) {
+    super('NetworkError', 400, message, 'PROMPT_TOO_LARGE');
+    this.name = 'AiChatbotPromptTooLargeError';
+  }
+}

packages/compass-generative-ai/src/utils/gen-ai-prompt.spec.ts

@@ -7,6 +7,7 @@ import {
 } from './gen-ai-prompt';
 import { toJSString } from 'mongodb-query-parser';
 import { ObjectId } from 'bson';
+import { AiChatbotPromptTooLargeError } from '../chatbot-errors';
 
 const OPTIONS: PromptContextOptions = {
   userInput: 'Find all users older than 30',
@@ -149,13 +150,13 @@ describe('GenAI Prompts', function () {
     try {
       buildFindQueryPrompt({
         ...OPTIONS,
-        userInput: 'a'.repeat(512001),
+        userInput: 'a'.repeat(250_001),
       });
       expect.fail('Expected buildFindQueryPrompt to throw');
     } catch (err) {
-      expect(err).to.have.property(
-        'message',
-        'Sorry, your request is too large. Please use a smaller prompt or try using this feature on a collection with smaller documents.'
+      expect(err).to.be.instanceOf(AiChatbotPromptTooLargeError);
+      expect((err as AiChatbotPromptTooLargeError).message).to.equal(
+        'PROMPT_TOO_LARGE: Sorry, your request is too large. Please use a smaller prompt or try using this feature on a collection with smaller documents.'
       );
     }
   });
@@ -178,9 +179,9 @@ describe('GenAI Prompts', function () {
     it('sends only one sample doc if all exceed limits', function () {
       const sampleDocuments = [
         { a: '1'.repeat(5120) },
-        { a: '2'.repeat(5120001) },
-        { a: '3'.repeat(5120001) },
-        { a: '4'.repeat(5120001) },
+        { a: '2'.repeat(250_001) },
+        { a: '3'.repeat(250_001) },
+        { a: '4'.repeat(250_001) },
       ];
       const prompt = buildFindQueryPrompt({
         ...OPTIONS,
@@ -190,10 +191,10 @@ describe('GenAI Prompts', function () {
     });
     it('should not send sample docs if even one exceeds limits', function () {
       const sampleDocuments = [
-        { a: '1'.repeat(5120001) },
-        { a: '2'.repeat(5120001) },
-        { a: '3'.repeat(5120001) },
-        { a: '4'.repeat(5120001) },
+        { a: '1'.repeat(250_001) },
+        { a: '2'.repeat(250_001) },
+        { a: '3'.repeat(250_001) },
+        { a: '4'.repeat(250_001) },
       ];
       const prompt = buildFindQueryPrompt({
         ...OPTIONS,

packages/compass-generative-ai/src/utils/gen-ai-prompt.ts

@@ -1,12 +1,9 @@
 import { toJSString } from 'mongodb-query-parser';
 import { flattenSchemaToObject } from './util';
+import { AiChatbotPromptTooLargeError } from '../chatbot-errors';
 
-// When including sample documents, we want to ensure that we do not
-// attach large documents and exceed the limit. OpenAI roughly estimates
-// 4 characters = 1 token and we should not exceed context window limits.
-// This roughly translates to 128k tokens.
-// TODO(COMPASS-10129): Adjust this limit based on the model's context window.
-const MAX_TOTAL_PROMPT_LENGTH = 512000;
+// chatbot slim model support ~250k characters
+const MAX_TOTAL_PROMPT_LENGTH = 250_000;
 const MIN_SAMPLE_DOCUMENTS = 1;
 
 function getCurrentTimeString() {
@@ -165,7 +162,7 @@ function buildUserPromptForQuery({
 
   // If at this point we have exceeded the limit, throw an error.
   if (prompt.length > MAX_TOTAL_PROMPT_LENGTH) {
-    throw new Error(
+    throw new AiChatbotPromptTooLargeError(
       'Sorry, your request is too large. Please use a smaller prompt or try using this feature on a collection with smaller documents.'
     );
   }