AI Assistant: Add error handling when request is too long (#33157)

* emit error_context_too_large error for 413 status on SuggestionsEventSource * changelog * show different erro message when the tokens limit is exceeded * changelog Committed via a GitHub action: https://github.com/Automattic/jetpack/actions/runs/6236930696
Automattic · Sep 19, 2023 · 87b3c18 · 87b3c18
1 parent 94755ac
commit 87b3c18
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 16 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,7 @@ This is an alpha version! The changes listed here are not final.
 ### Added
 - AI Client: Add support for the jetpack-ai role on the prompt messages.
 - AI Client: add `model` param to request helpers
+- AI Client: Emit specific error for large context error on SuggestionsEventSource
 - AI Client: Introduce blockListBlockWithAiDataProvider() function
 
 ### Changed

diff --git a/src/suggestions-event-source/index.ts b/src/suggestions-event-source/index.ts
@@ -12,6 +12,7 @@ import requestJwt from '../jwt';
  * Types & constants
  */
 import {
+	ERROR_CONTEXT_TOO_LARGE,
 	ERROR_MODERATION,
 	ERROR_NETWORK,
 	ERROR_QUOTA_EXCEEDED,
@@ -53,16 +54,17 @@ const debug = debugFactory( 'jetpack-ai-client:suggestions-event-source' );
  * It also emits a 'suggestion' event with the full suggestion received so far
  *
  * @returns {EventSource} The event source
- * @fires suggestion                - The full suggestion has been received so far
- * @fires message                   - A message has been received
- * @fires chunk                     - A chunk of data has been received
- * @fires done                      - The stream has been closed. No more data will be received
- * @fires error                     - An error has occurred
- * @fires error_network             - The EventSource connection to the server returned some error
- * @fires error_service_unavailable - The server returned a 503 error
- * @fires error_quota_exceeded      - The server returned a 429 error
- * @fires error_moderation          - The server returned a 422 error
- * @fires error_unclear_prompt      - The server returned a message starting with JETPACK_AI_ERROR
+ * @fires SuggestionsEventSource#suggestion                - The full suggestion has been received so far
+ * @fires SuggestionsEventSource#message                   - A message has been received
+ * @fires SuggestionsEventSource#chunk                     - A chunk of data has been received
+ * @fires SuggestionsEventSource#done                      - The stream has been closed. No more data will be received
+ * @fires SuggestionsEventSource#error                     - An error has occurred
+ * @fires SuggestionsEventSource#error_network             - The EventSource connection to the server returned some error
+ * @fires SuggestionsEventSource#error_context_too_large   - The server returned a 413 error
+ * @fires SuggestionsEventSource#error_moderation          - The server returned a 422 error
+ * @fires SuggestionsEventSource#error_quota_exceeded      - The server returned a 429 error
+ * @fires SuggestionsEventSource#error_service_unavailable - The server returned a 503 error
+ * @fires SuggestionsEventSource#error_unclear_prompt      - The server returned a message starting with JETPACK_AI_ERROR
  */
 export default class SuggestionsEventSource extends EventTarget {
 	fullMessage: string;
@@ -187,7 +189,7 @@ export default class SuggestionsEventSource extends EventTarget {
 				if (
 					response.status >= 400 &&
 					response.status <= 500 &&
-					! [ 422, 429 ].includes( response.status )
+					! [ 413, 422, 429 ].includes( response.status )
 				) {
 					this.processConnectionError( response );
 				}
@@ -202,12 +204,12 @@ export default class SuggestionsEventSource extends EventTarget {
 				}
 
 				/*
-				 * error code 429
-				 * you exceeded your current quota please check your plan and billing details
+				 * error code 413
+				 * request context too large
 				 */
-				if ( response.status === 429 ) {
-					errorCode = ERROR_QUOTA_EXCEEDED;
-					this.dispatchEvent( new CustomEvent( ERROR_QUOTA_EXCEEDED ) );
+				if ( response.status === 413 ) {
+					errorCode = ERROR_CONTEXT_TOO_LARGE;
+					this.dispatchEvent( new CustomEvent( ERROR_CONTEXT_TOO_LARGE ) );
 				}
 
 				/*
@@ -219,6 +221,15 @@ export default class SuggestionsEventSource extends EventTarget {
 					this.dispatchEvent( new CustomEvent( ERROR_MODERATION ) );
 				}
 
+				/*
+				 * error code 429
+				 * you exceeded your current quota please check your plan and billing details
+				 */
+				if ( response.status === 429 ) {
+					errorCode = ERROR_QUOTA_EXCEEDED;
+					this.dispatchEvent( new CustomEvent( ERROR_QUOTA_EXCEEDED ) );
+				}
+
 				// Always dispatch a global ERROR_RESPONSE event
 				this.dispatchEvent(
 					new CustomEvent( ERROR_RESPONSE, {

diff --git a/src/types.ts b/src/types.ts
@@ -1,6 +1,7 @@
 export const ERROR_SERVICE_UNAVAILABLE = 'error_service_unavailable' as const;
 export const ERROR_QUOTA_EXCEEDED = 'error_quota_exceeded' as const;
 export const ERROR_MODERATION = 'error_moderation' as const;
+export const ERROR_CONTEXT_TOO_LARGE = 'error_context_too_large' as const;
 export const ERROR_NETWORK = 'error_network' as const;
 export const ERROR_UNCLEAR_PROMPT = 'error_unclear_prompt' as const;
 export const ERROR_RESPONSE = 'error_response' as const;
@@ -9,6 +10,7 @@ export type SuggestionErrorCode =
 	| typeof ERROR_SERVICE_UNAVAILABLE
 	| typeof ERROR_QUOTA_EXCEEDED
 	| typeof ERROR_MODERATION
+	| typeof ERROR_CONTEXT_TOO_LARGE
 	| typeof ERROR_NETWORK
 	| typeof ERROR_UNCLEAR_PROMPT
 	| typeof ERROR_RESPONSE;