AI Client: add model param to request helpers (#33083)

* register model consts and types

* add `model` option to askQuestion fn

* add `model` to SuggestionErrorCode fn

* changelog

Committed via a GitHub action: https://github.com/Automattic/jetpack/actions/runs/6195933723

CHANGELOG.md

@@ -11,6 +11,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: Introduce blockListBlockWithAiDataProvider() function
 
 ### Changed

src/ask-question/Readme.md

@@ -22,6 +22,7 @@ function askQuestion(
     - **postId** (**number**, optional): ID of the post where the question is asked.
     - **fromCache** (**boolean**, optional): If set to true, the answer will be fetched from the cache. Default value is false.
     - **feature** (**string**, optional): Allows to use a specific AI assistant feature.
+    - **model**( **AiModelTypeProp** optional): Allows to use a specific AI model.
 
 ## Returns
 

src/ask-question/index.ts

@@ -6,7 +6,7 @@ import SuggestionsEventSource from '../suggestions-event-source';
 /*
  * Types & constants
  */
-import type { PromptProp } from '../types';
+import type { AiModelTypeProp, PromptProp } from '../types';
 
 export type AskQuestionOptionsArgProps = {
 	/*
@@ -24,6 +24,11 @@ export type AskQuestionOptionsArgProps = {
 	 */
 	feature?: 'ai-assistant-experimental' | string | undefined;
 
+	/*
+	 * Allows to use a specific AI model.
+	 */
+	model?: AiModelTypeProp;
+
 	/*
 	 * Allows the use of function calling. Default value is undefined.
 	 */
@@ -57,12 +62,18 @@ const debug = debugFactory( 'jetpack-ai-client:ask-question' );
  */
 export default async function askQuestion(
 	question: PromptProp,
-	{ postId = null, fromCache = false, feature, functions }: AskQuestionOptionsArgProps = {}
+	{ postId = null, fromCache = false, feature, functions, model }: AskQuestionOptionsArgProps = {}
 ): Promise< SuggestionsEventSource > {
-	debug( 'Asking question: %o. options: %o', question, { postId, fromCache, feature, functions } );
+	debug( 'Asking question: %o. options: %o', question, {
+		postId,
+		fromCache,
+		feature,
+		functions,
+		model,
+	} );
 
 	return new SuggestionsEventSource( {
 		question,
-		options: { postId, feature, fromCache, functions },
+		options: { postId, feature, fromCache, functions, model },
 	} );
 }

src/suggestions-event-source/Readme.md

@@ -39,6 +39,7 @@ The constructor takes an object with the following properties:
     - `postId` (optional): The post ID.
     - `feature` (optional): A string that specifies the AI model to use (default or 'ai-assistant-experimental').
     - `fromCache` (optional): A boolean to indicate whether to fetch the response from cache.
+    - `model` (optional): Allows to use a specific AI model.
 
 ### Usage
 

src/suggestions-event-source/index.ts

@@ -19,7 +19,12 @@ import {
 	ERROR_SERVICE_UNAVAILABLE,
 	ERROR_UNCLEAR_PROMPT,
 } from '../types';
-import type { PromptMessagesProp, PromptProp, SuggestionErrorCode } from '../types';
+import type {
+	AiModelTypeProp,
+	PromptMessagesProp,
+	PromptProp,
+	SuggestionErrorCode,
+} from '../types';
 
 type SuggestionsEventSourceConstructorArgs = {
 	url?: string;
@@ -30,6 +35,7 @@ type SuggestionsEventSourceConstructorArgs = {
 		feature?: 'ai-assistant-experimental' | string | undefined;
 		fromCache?: boolean;
 		functions?: Array< object >;
+		model?: AiModelTypeProp;
 	};
 };
 
@@ -103,6 +109,7 @@ export default class SuggestionsEventSource extends EventTarget {
 			question?: PromptProp;
 			feature?: string;
 			functions?: Array< object >;
+			model?: AiModelTypeProp;
 		} = {};
 
 		// Populate body data with post id
@@ -142,6 +149,12 @@ export default class SuggestionsEventSource extends EventTarget {
 			bodyData.functions = options.functions;
 		}
 
+		// Model
+		if ( options?.model?.length ) {
+			debug( 'Model: %o', options.model );
+			bodyData.model = options.model;
+		}
+
 		await fetchEventSource( url, {
 			openWhenHidden: true,
 			method: 'POST',

src/types.ts

@@ -35,7 +35,6 @@ 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;
@@ -51,3 +50,11 @@ export const REQUESTING_STATES = [
 ] as const;
 
 export type RequestingStateProp = ( typeof REQUESTING_STATES )[ number ];
+
+/*
+ * Model types and constants
+ */
+export const AI_MODEL_GPT_3_5_Turbo_16K = 'gpt-3.5-turbo-16k' as const;
+export const AI_MODEL_GPT_4 = 'gpt-4' as const;
+
+export type AiModelTypeProp = typeof AI_MODEL_GPT_3_5_Turbo_16K | typeof AI_MODEL_GPT_4;