@automattic/jetpack-ai-client 0.1.16 → 0.2.0

package/build/jwt/index.js

@@ -0,0 +1,74 @@
+/**
+ * External dependencies
+ */
+import { isSimpleSite } from '@automattic/jetpack-shared-extension-utils';
+import apiFetch from '@wordpress/api-fetch';
+import debugFactory from 'debug';
+const debug = debugFactory('jetpack-ai-client:jwt');
+const JWT_TOKEN_ID = 'jetpack-ai-jwt';
+const JWT_TOKEN_EXPIRATION_TIME = 2 * 60 * 1000; // 2 minutes
+/**
+ * Request a token from the Jetpack site.
+ *
+ * @param {RequestTokenOptions} options - Options
+ * @returns {Promise<TokenDataProps>}     The token and the blogId
+ */
+export default async function requestJwt({ apiNonce, siteId, expirationTime, } = {}) {
+    // Default values
+    apiNonce = apiNonce || window.JP_CONNECTION_INITIAL_STATE.apiNonce;
+    siteId = siteId || window.JP_CONNECTION_INITIAL_STATE.siteSuffix;
+    expirationTime = expirationTime || JWT_TOKEN_EXPIRATION_TIME;
+    const isSimple = isSimpleSite();
+    // Trying to pick the token from localStorage
+    const token = localStorage.getItem(JWT_TOKEN_ID);
+    let tokenData = null;
+    if (token) {
+        try {
+            tokenData = JSON.parse(token);
+        }
+        catch (err) {
+            debug('Error parsing token', err);
+        }
+    }
+    if (tokenData && tokenData?.expire > Date.now()) {
+        debug('Using cached token');
+        return tokenData;
+    }
+    let data;
+    if (!isSimple) {
+        data = await apiFetch({
+            /*
+             * This endpoint is registered in the Jetpack plugin.
+             * Provably we should move it to another package, but for now it's here.
+             * issue: https://github.com/Automattic/jetpack/issues/31938
+             */
+            path: '/jetpack/v4/jetpack-ai-jwt?_cacheBuster=' + Date.now(),
+            credentials: 'same-origin',
+            headers: {
+                'X-WP-Nonce': apiNonce,
+            },
+            method: 'POST',
+        });
+    }
+    else {
+        data = await apiFetch({
+            path: '/wpcom/v2/sites/' + siteId + '/jetpack-openai-query/jwt',
+            method: 'POST',
+        });
+    }
+    const newTokenData = {
+        token: data.token,
+        /**
+         * TODO: make sure we return id from the .com token acquisition endpoint too
+         */
+        blogId: !isSimple ? data.blog_id : siteId,
+        /**
+         * Let's expire the token in 2 minutes
+         */
+        expire: Date.now() + expirationTime,
+    };
+    // Store the token in localStorage
+    debug('Storing new token');
+    localStorage.setItem(JWT_TOKEN_ID, JSON.stringify(newTokenData));
+    return newTokenData;
+}

package/build/suggestions-event-source/index.d.ts

@@ -0,0 +1,55 @@
+/**
+ * External dependencies
+ */
+import { EventSourceMessage } from '@microsoft/fetch-event-source';
+import type { AiModelTypeProp, PromptProp } from '../types';
+type SuggestionsEventSourceConstructorArgs = {
+    url?: string;
+    question: PromptProp;
+    token?: string;
+    options?: {
+        postId?: number;
+        feature?: 'ai-assistant-experimental' | string | undefined;
+        fromCache?: boolean;
+        functions?: Array<object>;
+        model?: AiModelTypeProp;
+    };
+};
+type FunctionCallProps = {
+    name?: string;
+    arguments?: string;
+};
+/**
+ * SuggestionsEventSource is a wrapper around EvenTarget that emits
+ * a 'chunk' event for each chunk of data received, and a 'done' event
+ * when the stream is closed.
+ * It also emits a 'suggestion' event with the full suggestion received so far
+ *
+ * @returns {EventSource} The event source
+ * @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;
+    fullFunctionCall: FunctionCallProps;
+    isPromptClear: boolean;
+    controller: AbortController;
+    errorUnclearPromptTriggered: boolean;
+    constructor(data: SuggestionsEventSourceConstructorArgs);
+    initEventSource({ url, question, token, options, }: SuggestionsEventSourceConstructorArgs): Promise<void>;
+    checkForUnclearPrompt(): void;
+    close(): void;
+    processEvent(e: EventSourceMessage): void;
+    processConnectionError(response: any): void;
+    processErrorEvent(e: any): void;
+}
+export {};

package/build/suggestions-event-source/index.js

@@ -0,0 +1,282 @@
+/**
+ * External dependencies
+ */
+import { fetchEventSource } from '@microsoft/fetch-event-source';
+import debugFactory from 'debug';
+/**
+ * Internal dependencies
+ */
+import { getErrorData } from '../hooks/use-ai-suggestions';
+import requestJwt from '../jwt';
+/*
+ * Types & constants
+ */
+import { ERROR_CONTEXT_TOO_LARGE, ERROR_MODERATION, ERROR_NETWORK, ERROR_QUOTA_EXCEEDED, ERROR_RESPONSE, ERROR_SERVICE_UNAVAILABLE, ERROR_UNCLEAR_PROMPT, } from '../types';
+const debug = debugFactory('jetpack-ai-client:suggestions-event-source');
+/**
+ * SuggestionsEventSource is a wrapper around EvenTarget that emits
+ * a 'chunk' event for each chunk of data received, and a 'done' event
+ * when the stream is closed.
+ * It also emits a 'suggestion' event with the full suggestion received so far
+ *
+ * @returns {EventSource} The event source
+ * @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;
+    fullFunctionCall;
+    isPromptClear;
+    controller;
+    // Flag to detect if the unclear prompt event was already dispatched
+    errorUnclearPromptTriggered;
+    constructor(data) {
+        super();
+        this.fullMessage = '';
+        this.fullFunctionCall = {
+            name: '',
+            arguments: '',
+        };
+        this.isPromptClear = false;
+        // The AbortController is used to close the fetchEventSource connection
+        this.controller = new AbortController();
+        this.initEventSource(data);
+    }
+    async initEventSource({ url, question, token, options = {}, }) {
+        // If the token is not provided, try to get one
+        if (!token) {
+            try {
+                debug('Token was not provided, requesting one...');
+                token = (await requestJwt()).token;
+            }
+            catch (err) {
+                this.processErrorEvent(err);
+                return;
+            }
+        }
+        const bodyData = {};
+        // 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) {
+            const urlHandler = new URL('https://public-api.wordpress.com/wpcom/v2/jetpack-ai-query');
+            // Support response from cache option
+            if (options?.fromCache) {
+                urlHandler.searchParams.append('stream_cache', 'true');
+            }
+            url = urlHandler.toString();
+            debug('URL not provided, using default: %o', url);
+        }
+        // question can be a string or an array of PromptMessagesProp
+        if (Array.isArray(question)) {
+            bodyData.messages = question;
+        }
+        else {
+            bodyData.question = question;
+        }
+        // Propagate the feature option
+        if (options?.feature?.length) {
+            debug('Feature: %o', options.feature);
+            bodyData.feature = options.feature;
+        }
+        // Propagate the functions option
+        if (options?.functions?.length) {
+            debug('Functions: %o', options.functions);
+            bodyData.functions = options.functions;
+        }
+        // Model
+        if (options?.model?.length) {
+            debug('Model: %o', options.model);
+            bodyData.model = options.model;
+        }
+        // Clean the unclear prompt trigger flag
+        this.errorUnclearPromptTriggered = false;
+        await fetchEventSource(url, {
+            openWhenHidden: true,
+            method: 'POST',
+            headers: {
+                'Content-type': 'application/json',
+                Authorization: 'Bearer ' + token,
+            },
+            body: JSON.stringify(bodyData),
+            onclose: () => {
+                debug('Stream closed');
+            },
+            onerror: err => {
+                this.processErrorEvent(err);
+                throw err; // rethrow to stop the operation otherwise it will retry forever
+            },
+            onmessage: ev => {
+                this.processEvent(ev);
+            },
+            onopen: async (response) => {
+                if (response.ok) {
+                    return;
+                }
+                let errorCode;
+                if (response.status >= 400 &&
+                    response.status <= 500 &&
+                    ![413, 422, 429].includes(response.status)) {
+                    this.processConnectionError(response);
+                }
+                /*
+                 * error code 503
+                 * service unavailable
+                 */
+                if (response.status === 503) {
+                    errorCode = ERROR_SERVICE_UNAVAILABLE;
+                    this.dispatchEvent(new CustomEvent(ERROR_SERVICE_UNAVAILABLE));
+                }
+                /*
+                 * error code 413
+                 * request context too large
+                 */
+                if (response.status === 413) {
+                    errorCode = ERROR_CONTEXT_TOO_LARGE;
+                    this.dispatchEvent(new CustomEvent(ERROR_CONTEXT_TOO_LARGE));
+                }
+                /*
+                 * error code 422
+                 * request flagged by moderation system
+                 */
+                if (response.status === 422) {
+                    errorCode = ERROR_MODERATION;
+                    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, {
+                    detail: getErrorData(errorCode),
+                }));
+                throw new Error();
+            },
+            signal: this.controller.signal,
+        });
+    }
+    checkForUnclearPrompt() {
+        if (this.isPromptClear) {
+            return;
+        }
+        /*
+         * Sometimes the first token of the message is not received,
+         * so we check only for JETPACK_AI_ERROR, ignoring:
+         * - the double underscores (italic markdown)
+         * - the double asterisks (bold markdown)
+         */
+        const replacedMessage = this.fullMessage.replace(/__|(\*\*)/g, '');
+        if (replacedMessage.startsWith('JETPACK_AI_ERROR')) {
+            /*
+             * Check if the unclear prompt event was already dispatched,
+             * to ensure that it is dispatched only once per request.
+             */
+            if (this.errorUnclearPromptTriggered) {
+                return;
+            }
+            this.errorUnclearPromptTriggered = true;
+            // The unclear prompt marker was found, so we dispatch an error event
+            this.dispatchEvent(new CustomEvent(ERROR_UNCLEAR_PROMPT));
+            debug('Unclear error prompt dispatched');
+            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);
+        }
+        else {
+            // Mark the prompt as clear
+            this.isPromptClear = true;
+        }
+    }
+    close() {
+        this.controller.abort();
+    }
+    processEvent(e) {
+        if (e.data === '[DONE]') {
+            /*
+             * Check if the unclear prompt event was already dispatched,
+             * to ensure that it is dispatched only once per request.
+             */
+            if (this.errorUnclearPromptTriggered) {
+                return;
+            }
+            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;
+        if (chunk) {
+            this.fullMessage += chunk;
+            this.checkForUnclearPrompt();
+            if (this.isPromptClear) {
+                // 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_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_RESPONSE, {
+            detail: getErrorData(ERROR_NETWORK),
+        }));
+    }
+}

package/build/types.d.ts

@@ -0,0 +1,34 @@
+export declare const ERROR_SERVICE_UNAVAILABLE: "error_service_unavailable";
+export declare const ERROR_QUOTA_EXCEEDED: "error_quota_exceeded";
+export declare const ERROR_MODERATION: "error_moderation";
+export declare const ERROR_CONTEXT_TOO_LARGE: "error_context_too_large";
+export declare const ERROR_NETWORK: "error_network";
+export declare const ERROR_UNCLEAR_PROMPT: "error_unclear_prompt";
+export declare const ERROR_RESPONSE: "error_response";
+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;
+export type PromptItemProps = {
+    role: 'system' | 'user' | 'assistant' | 'jetpack-ai';
+    content?: string;
+    context?: object;
+};
+export type PromptMessagesProp = Array<PromptItemProps>;
+export type PromptProp = PromptMessagesProp | string;
+export type { UseAiContextOptions } from './data-flow/use-ai-context';
+export type { RequestingErrorProps } from './hooks/use-ai-suggestions';
+export declare const REQUESTING_STATES: readonly ["init", "requesting", "suggesting", "done", "error"];
+export type RequestingStateProp = (typeof REQUESTING_STATES)[number];
+export declare const AI_MODEL_GPT_3_5_Turbo_16K: "gpt-3.5-turbo-16k";
+export declare const AI_MODEL_GPT_4: "gpt-4";
+export type AiModelTypeProp = typeof AI_MODEL_GPT_3_5_Turbo_16K | typeof AI_MODEL_GPT_4;
+interface JPConnectionInitialState {
+    apiNonce: string;
+    siteSuffix: string;
+    connectionStatus: {
+        isActive: boolean;
+    };
+}
+declare global {
+    interface Window {
+        JP_CONNECTION_INITIAL_STATE: JPConnectionInitialState;
+    }
+}

package/build/types.js

@@ -0,0 +1,27 @@
+export const ERROR_SERVICE_UNAVAILABLE = 'error_service_unavailable';
+export const ERROR_QUOTA_EXCEEDED = 'error_quota_exceeded';
+export const ERROR_MODERATION = 'error_moderation';
+export const ERROR_CONTEXT_TOO_LARGE = 'error_context_too_large';
+export const ERROR_NETWORK = 'error_network';
+export const ERROR_UNCLEAR_PROMPT = 'error_unclear_prompt';
+export const ERROR_RESPONSE = 'error_response';
+/*
+ * Requests types
+ */
+const REQUESTING_STATE_INIT = 'init';
+const REQUESTING_STATE_REQUESTING = 'requesting';
+const REQUESTING_STATE_SUGGESTING = 'suggesting';
+const REQUESTING_STATE_DONE = 'done';
+const REQUESTING_STATE_ERROR = 'error';
+export const REQUESTING_STATES = [
+    REQUESTING_STATE_INIT,
+    REQUESTING_STATE_REQUESTING,
+    REQUESTING_STATE_SUGGESTING,
+    REQUESTING_STATE_DONE,
+    REQUESTING_STATE_ERROR,
+];
+/*
+ * Model types and constants
+ */
+export const AI_MODEL_GPT_3_5_Turbo_16K = 'gpt-3.5-turbo-16k';
+export const AI_MODEL_GPT_4 = 'gpt-4';

package/package.json

@@ -1,7 +1,7 @@
 {
 	"private": false,
 	"name": "@automattic/jetpack-ai-client",
-	"version": "0.1.16",
+	"version": "0.2.0",
 	"description": "A JS client for consuming Jetpack AI services",
 	"homepage": "https://github.com/Automattic/jetpack/tree/HEAD/projects/js-packages/ai-client/#readme",
 	"bugs": {
@@ -15,7 +15,8 @@
 	"license": "GPL-2.0-or-later",
 	"author": "Automattic",
 	"scripts": {
-		"build": "pnpm run compile-ts",
+		"build": "pnpm run clean && pnpm run compile-ts",
+		"clean": "rm -rf build/",
 		"compile-ts": "tsc --pretty",
 		"test": "NODE_OPTIONS=--experimental-vm-modules jest"
 	},
@@ -29,13 +30,19 @@
 		"typescript": "5.0.4"
 	},
 	"exports": {
-		".": "./index.ts"
+		".": {
+			"types": "./build/index.d.ts",
+			"default": "./build/index.js"
+		}
 	},
+	"main": "./build/index.js",
+	"types": "./build/index.d.ts",
 	"dependencies": {
-		"@automattic/jetpack-base-styles": "^0.6.12",
-		"@automattic/jetpack-connection": "^0.30.7",
-		"@automattic/jetpack-shared-extension-utils": "^0.13.1",
+		"@automattic/jetpack-base-styles": "^0.6.13",
+		"@automattic/jetpack-connection": "^0.30.8",
+		"@automattic/jetpack-shared-extension-utils": "^0.13.2",
 		"@microsoft/fetch-event-source": "2.0.1",
+		"@types/react": "18.2.33",
 		"@wordpress/api-fetch": "6.42.0",
 		"@wordpress/block-editor": "12.13.0",
 		"@wordpress/components": "25.11.0",

package/src/components/ai-control/index.tsx

@@ -19,6 +19,22 @@ import { GuidelineMessage } from './message';
  * Types
  */
 import type { RequestingStateProp } from '../../types';
+type AIControlProps = {
+	disabled?: boolean;
+	value: string;
+	placeholder?: string;
+	showAccept?: boolean;
+	acceptLabel?: string;
+	showButtonLabels?: boolean;
+	isTransparent?: boolean;
+	state?: RequestingStateProp;
+	showClearButton?: boolean;
+	showGuideLine?: boolean;
+	onChange?: ( newValue: string ) => void;
+	onSend?: ( currentValue: string ) => void;
+	onStop?: () => void;
+	onAccept?: () => void;
+};
 
 // eslint-disable-next-line @typescript-eslint/no-empty-function
 const noop = () => {};
@@ -26,23 +42,9 @@ const noop = () => {};
 /**
  * AI Control component.
  *
- * @param {object} props                   - Component props
- * @param {boolean} props.disabled         - Input disabled state
- * @param {string} props.value             - The input value
- * @param {string} props.placeholder       - The input placeholder
- * @param {boolean} props.showAccept       - Whether to show the accept button
- * @param {string} props.acceptLabel       - The accept button label
- * @param {boolean} props.showButtonLabels - Whether to show the button labels
- * @param {boolean} props.isTransparent    - Whether the component has low opacity
- * @param {string} props.state             - The request state
- * @param {boolean} props.showClearButton  - Whether to show the clear button when the input has a value
- * @param {boolean} props.showGuideLine    - WHether to show the guideline message
- * @param {Function} props.onChange        - Input change handler
- * @param {Function} props.onSend          - Request send handler
- * @param {Function} props.onStop          - Request stop handler
- * @param {Function} props.onAccept        - Response accept handler
- * @param {object} ref                     - Auto injected ref from react
- * @returns {object} - AI Control component
+ * @param {AIControlProps} props       - Component props.
+ * @param {React.MutableRefObject} ref - Ref to the component.
+ * @returns {React.ReactElement}         Rendered component.
  */
 export function AIControl(
 	{
@@ -76,8 +78,8 @@ export function AIControl(
 		onStop?: () => void;
 		onAccept?: () => void;
 	},
-	ref
-) {
+	ref: React.MutableRefObject< null > // eslint-disable-line @typescript-eslint/ban-types
+): React.ReactElement {
 	const promptUserInputRef = useRef( null );
 	const loading = state === 'requesting' || state === 'suggesting';
 

package/src/data-flow/context.tsx

@@ -53,7 +53,7 @@ type AiDataContextProviderProps = {
 	/*
 	 * Children
 	 */
-	children: React.ReactNode;
+	children: React.ReactElement;
 };
 
 /**
@@ -67,7 +67,7 @@ export const AiDataContext = createContext( {} as AiDataContextProps );
  * AI Data Context Provider
  *
  * @param {AiDataContextProviderProps} props - Component props.
- * @returns {React.ReactNode}                           Context provider.
+ * @returns {React.ReactElement}                           Context provider.
  * @example
  * <AiDataContextProvider value={ value }>
  * 	{ children }
@@ -76,6 +76,6 @@ export const AiDataContext = createContext( {} as AiDataContextProps );
 export const AiDataContextProvider = ( {
 	value,
 	children,
-}: AiDataContextProviderProps ): React.ReactNode => (
+}: AiDataContextProviderProps ): React.ReactElement => (
 	<AiDataContext.Provider value={ value } children={ children } />
 );