@automattic/jetpack-ai-client 0.1.0 → 0.1.2

package/src/suggestions-event-source/index.ts

@@ -6,19 +6,34 @@ import debugFactory from 'debug';
 /*
  * Types & constants
  */
-import type { PromptItemProps } from '../types';
+import { getErrorData } from '../hooks/use-ai-suggestions';
+import {
+	ERROR_MODERATION,
+	ERROR_NETWORK,
+	ERROR_QUOTA_EXCEEDED,
+	ERROR_RESPONSE,
+	ERROR_SERVICE_UNAVAILABLE,
+	ERROR_UNCLEAR_PROMPT,
+} from '../types';
+import type { PromptMessagesProp, PromptProp, SuggestionErrorCode } from '../types';
 
 type SuggestionsEventSourceConstructorArgs = {
 	url?: string;
-	question: string | PromptItemProps[];
+	question: PromptProp;
 	token: string;
 	options?: {
 		postId?: number;
 		feature?: 'ai-assistant-experimental' | string | undefined;
 		fromCache?: boolean;
+		functions?: Array< object >;
 	};
 };
 
+type FunctionCallProps = {
+	name?: string;
+	arguments?: string;
+};
+
 const debug = debugFactory( 'jetpack-ai-client:suggestions-event-source' );
 
 /**
@@ -41,12 +56,17 @@ const debug = debugFactory( 'jetpack-ai-client:suggestions-event-source' );
  */
 export default class SuggestionsEventSource extends EventTarget {
 	fullMessage: string;
+	fullFunctionCall: FunctionCallProps;
 	isPromptClear: boolean;
 	controller: AbortController;
 
 	constructor( data: SuggestionsEventSourceConstructorArgs ) {
 		super();
 		this.fullMessage = '';
+		this.fullFunctionCall = {
+			name: '',
+			arguments: '',
+		};
 		this.isPromptClear = false;
 
 		// The AbortController is used to close the fetchEventSource connection
@@ -61,7 +81,18 @@ export default class SuggestionsEventSource extends EventTarget {
 		token,
 		options = {},
 	}: SuggestionsEventSourceConstructorArgs ) {
-		const bodyData = { post_id: options?.postId, messages: [], question: '', feature: '' };
+		const bodyData: {
+			post_id?: number;
+			messages?: PromptMessagesProp;
+			question?: PromptProp;
+			feature?: string;
+			functions?: Array< object >;
+		} = {};
+
+		// Populate body data with post id
+		if ( options?.postId ) {
+			bodyData.post_id = options.postId;
+		}
 
 		// If the url is not provided, we use the default one
 		if ( ! url ) {
@@ -76,7 +107,7 @@ export default class SuggestionsEventSource extends EventTarget {
 			debug( 'URL not provided, using default: %o', url );
 		}
 
-		// question can be a string or an array of PromptItemProps
+		// question can be a string or an array of PromptMessagesProp
 		if ( Array.isArray( question ) ) {
 			bodyData.messages = question;
 		} else {
@@ -89,6 +120,12 @@ export default class SuggestionsEventSource extends EventTarget {
 			bodyData.feature = options.feature;
 		}
 
+		// Propagate the functions option
+		if ( options?.functions?.length ) {
+			debug( 'Functions: %o', options.functions );
+			bodyData.functions = options.functions;
+		}
+
 		await fetchEventSource( url, {
 			openWhenHidden: true,
 			method: 'POST',
@@ -115,6 +152,9 @@ export default class SuggestionsEventSource extends EventTarget {
 				if ( response.ok ) {
 					return;
 				}
+
+				let errorCode: SuggestionErrorCode;
+
 				if (
 					response.status >= 400 &&
 					response.status <= 500 &&
@@ -128,7 +168,8 @@ export default class SuggestionsEventSource extends EventTarget {
 				 * service unavailable
 				 */
 				if ( response.status === 503 ) {
-					this.dispatchEvent( new CustomEvent( 'error_service_unavailable' ) );
+					errorCode = ERROR_SERVICE_UNAVAILABLE;
+					this.dispatchEvent( new CustomEvent( ERROR_SERVICE_UNAVAILABLE ) );
 				}
 
 				/*
@@ -136,7 +177,8 @@ export default class SuggestionsEventSource extends EventTarget {
 				 * you exceeded your current quota please check your plan and billing details
 				 */
 				if ( response.status === 429 ) {
-					this.dispatchEvent( new CustomEvent( 'error_quota_exceeded' ) );
+					errorCode = ERROR_QUOTA_EXCEEDED;
+					this.dispatchEvent( new CustomEvent( ERROR_QUOTA_EXCEEDED ) );
 				}
 
 				/*
@@ -144,9 +186,17 @@ export default class SuggestionsEventSource extends EventTarget {
 				 * request flagged by moderation system
 				 */
 				if ( response.status === 422 ) {
-					this.dispatchEvent( new CustomEvent( 'error_moderation' ) );
+					errorCode = ERROR_MODERATION;
+					this.dispatchEvent( new CustomEvent( ERROR_MODERATION ) );
 				}
 
+				// Always dispatch a global ERROR_RESPONSE event
+				this.dispatchEvent(
+					new CustomEvent( ERROR_RESPONSE, {
+						detail: getErrorData( errorCode ),
+					} )
+				);
+
 				throw new Error();
 			},
 
@@ -168,7 +218,12 @@ export default class SuggestionsEventSource extends EventTarget {
 		const replacedMessage = this.fullMessage.replace( /__|(\*\*)/g, '' );
 		if ( replacedMessage.startsWith( 'JETPACK_AI_ERROR' ) ) {
 			// The unclear prompt marker was found, so we dispatch an error event
-			this.dispatchEvent( new CustomEvent( 'error_unclear_prompt' ) );
+			this.dispatchEvent( new CustomEvent( ERROR_UNCLEAR_PROMPT ) );
+			this.dispatchEvent(
+				new CustomEvent( ERROR_RESPONSE, {
+					detail: getErrorData( ERROR_UNCLEAR_PROMPT ),
+				} )
+			);
 		} else if ( 'JETPACK_AI_ERROR'.startsWith( replacedMessage ) ) {
 			// Partial unclear prompt marker was found, so we wait for more data and print a debug message without dispatching an event
 			debug( this.fullMessage );
@@ -184,14 +239,31 @@ export default class SuggestionsEventSource extends EventTarget {
 
 	processEvent( e: EventSourceMessage ) {
 		if ( e.data === '[DONE]' ) {
-			// Dispatch an event with the full content
-			this.dispatchEvent( new CustomEvent( 'done', { detail: this.fullMessage } ) );
-			debug( 'Done: %o', this.fullMessage );
+			if ( this.fullMessage.length ) {
+				// Dispatch an event with the full content
+				this.dispatchEvent( new CustomEvent( 'done', { detail: this.fullMessage } ) );
+				debug( 'Done: %o', this.fullMessage );
+				return;
+			}
+
+			if ( this.fullFunctionCall.name.length ) {
+				this.dispatchEvent( new CustomEvent( 'function_done', { detail: this.fullFunctionCall } ) );
+				debug( 'Done: %o', this.fullFunctionCall );
+				return;
+			}
+		}
+
+		let data;
+		try {
+			data = JSON.parse( e.data );
+		} catch ( err ) {
+			debug( 'Error parsing JSON', e, err );
 			return;
 		}
+		const { delta } = data?.choices?.[ 0 ] ?? { delta: { content: null, function_call: null } };
+		const chunk = delta.content;
+		const functionCallChunk = delta.function_call;
 
-		const data = JSON.parse( e.data );
-		const chunk = data.choices[ 0 ].delta.content;
 		if ( chunk ) {
 			this.fullMessage += chunk;
 			this.checkForUnclearPrompt();
@@ -200,20 +272,46 @@ export default class SuggestionsEventSource extends EventTarget {
 				// Dispatch an event with the chunk
 				this.dispatchEvent( new CustomEvent( 'chunk', { detail: chunk } ) );
 				// Dispatch an event with the full message
+				debug( 'suggestion: %o', this.fullMessage );
 				this.dispatchEvent( new CustomEvent( 'suggestion', { detail: this.fullMessage } ) );
 			}
 		}
+
+		if ( functionCallChunk ) {
+			if ( functionCallChunk.name != null ) {
+				this.fullFunctionCall.name += functionCallChunk.name;
+			}
+
+			if ( functionCallChunk.arguments != null ) {
+				this.fullFunctionCall.arguments += functionCallChunk.arguments;
+			}
+
+			// Dispatch an event with the function call
+			this.dispatchEvent(
+				new CustomEvent( 'functionCallChunk', { detail: this.fullFunctionCall } )
+			);
+		}
 	}
 
 	processConnectionError( response ) {
 		debug( 'Connection error: %o', response );
-		this.dispatchEvent( new CustomEvent( 'error_network', { detail: response } ) );
+		this.dispatchEvent( new CustomEvent( ERROR_NETWORK, { detail: response } ) );
+		this.dispatchEvent(
+			new CustomEvent( ERROR_RESPONSE, {
+				detail: getErrorData( ERROR_NETWORK ),
+			} )
+		);
 	}
 
 	processErrorEvent( e ) {
 		debug( 'onerror: %o', e );
 
 		// Dispatch a generic network error event
-		this.dispatchEvent( new CustomEvent( 'error_network', { detail: e } ) );
+		this.dispatchEvent( new CustomEvent( ERROR_NETWORK, { detail: e } ) );
+		this.dispatchEvent(
+			new CustomEvent( ERROR_RESPONSE, {
+				detail: getErrorData( ERROR_NETWORK ),
+			} )
+		);
 	}
 }

package/src/types.ts

@@ -1,4 +1,52 @@
+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_NETWORK = 'error_network' as const;
+export const ERROR_UNCLEAR_PROMPT = 'error_unclear_prompt' as const;
+export const ERROR_RESPONSE = 'error_response' as const;
+
+export type SuggestionErrorCode =
+	| typeof ERROR_SERVICE_UNAVAILABLE
+	| typeof ERROR_QUOTA_EXCEEDED
+	| typeof ERROR_MODERATION
+	| typeof ERROR_NETWORK
+	| typeof ERROR_UNCLEAR_PROMPT
+	| typeof ERROR_RESPONSE;
+
+/*
+ * Prompt types
+ */
 export type PromptItemProps = {
 	role: 'system' | 'user' | 'assistant';
 	content: string;
 };
+
+export type PromptMessagesProp = Array< PromptItemProps >;
+
+export type PromptProp = PromptMessagesProp | string;
+
+/*
+ * Data Flow types
+ */
+export type { UseAiContextOptions } from './data-flow/use-ai-context';
+export type { RequestingErrorProps } from './hooks/use-ai-suggestions';
+
+/*
+ * Requests types
+ */
+
+const REQUESTING_STATE_INIT = 'init' as const;
+const REQUESTING_STATE_REQUESTING = 'requesting' as const;
+const REQUESTING_STATE_SUGGESTING = 'suggesting' as const;
+const REQUESTING_STATE_DONE = 'done' as const;
+const REQUESTING_STATE_ERROR = 'error' as const;
+
+export const REQUESTING_STATES = [
+	REQUESTING_STATE_INIT,
+	REQUESTING_STATE_REQUESTING,
+	REQUESTING_STATE_SUGGESTING,
+	REQUESTING_STATE_DONE,
+	REQUESTING_STATE_ERROR,
+] as const;
+
+export type RequestingStateProp = ( typeof REQUESTING_STATES )[ number ];

package/src/index.js

@@ -1,26 +0,0 @@
-import apiFetch from '@wordpress/api-fetch';
-
-/**
- * Fetches images from Jetpack AI
- *
- * It's up to the consumer to catch errors
- *
- * @param {*} prompt  - The prompt to send to Jetpack AI
- * @param {*} postId - The post ID where the completion is being requested
- * @returns {Promise<Array<string>>} The images
- */
-export function requestImages( prompt, postId ) {
-	return apiFetch( {
-		path: '/wpcom/v2/jetpack-ai/images/generations',
-		method: 'POST',
-		data: {
-			prompt,
-			post_id: postId,
-		},
-	} ).then( res => {
-		const images = res.data.map( image => {
-			return 'data:image/png;base64,' + image.b64_json;
-		} );
-		return images;
-	} );
-}