feat: custom prompt

android-com-pl · Apr 15, 2024 · 21fcfac · 21fcfac
1 parent 7597e32
commit 21fcfac
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 19 deletions.
diff --git a/includes/AltGenerator.php b/includes/AltGenerator.php
@@ -9,7 +9,7 @@ class AltGenerator {
 	const API_URL = 'https://api.openai.com/v1/chat/completions';
 	const MODEL   = 'gpt-4-vision-preview';
 
-	public static function generate_alt_text( int $attachment_id ): string|WP_Error {
+	public static function generate_alt_text( int $attachment_id, string $user_prompt = '' ): string|WP_Error {
 		if ( ! wp_attachment_is_image( $attachment_id ) ) {
 			return ErrorCodes::Not_image->to_wp_error( [ 'attachment_id' => $attachment_id ] );
 		}
@@ -24,13 +24,6 @@ public static function generate_alt_text( int $attachment_id ): string|WP_Error
 		$locale   = get_locale();
 		$language = locale_get_display_language( $locale );
 
-		$user_prompt = apply_filters(
-			'acpl/ai_alt_generator/user_prompt',
-			"Generate a high-quality and concise alt text in $language ($locale) for the provided image without adding any additional comments.",
-			$locale,
-			$language
-		);
-
 		$image_mime_type = get_post_mime_type( $attachment_id );
 		$image_base64    = self::get_image_as_base64( $attachment_id );
 
@@ -51,12 +44,28 @@ public static function generate_alt_text( int $attachment_id ): string|WP_Error
 					[
 						'model'      => self::MODEL,
 						'messages'   => [
+							[
+								'role'    => 'system',
+								'content' => apply_filters(
+									'acpl/ai_alt_generator/system_prompt',
+									"Generate a high-quality and concise alt text in $language ($locale) for the provided image without adding any additional comments.",
+									$attachment_id,
+									$locale,
+									$language
+								),
+							],
 							[
 								'role'    => 'user',
 								'content' => [
 									[
 										'type' => 'text',
-										'text' => $user_prompt,
+										'text' => apply_filters(
+											'acpl/ai_alt_generator/user_prompt',
+											$user_prompt,
+											$attachment_id,
+											$locale,
+											$language
+										),
 									],
 									[
 										'type'      => 'image_url',
@@ -90,8 +99,8 @@ public static function generate_alt_text( int $attachment_id ): string|WP_Error
 		return $completion['choices'][0]['message']['content'] ?? '';
 	}
 
-	public static function generate_and_set_alt_text( int $attachment_id ): string|WP_Error|null {
-		$alt_text = self::generate_alt_text( $attachment_id );
+	public static function generate_and_set_alt_text( int $attachment_id, string $user_prompt = '' ): string|WP_Error|null {
+		$alt_text = self::generate_alt_text( $attachment_id, $user_prompt );
 		if ( is_wp_error( $alt_text ) ) {
 			AltGeneratorPlugin::error_log( $alt_text );
 

diff --git a/includes/Api.php b/includes/Api.php
@@ -19,6 +19,11 @@ public function register_routes(): void {
 						'required' => true,
 						'type'     => 'integer',
 					],
+					'user_prompt'   => [
+						'required'          => false,
+						'type'              => 'string',
+						'sanitize_callback' => 'sanitize_text_field',
+					],
 					'save'          => [
 						'required'    => false,
 						'type'        => 'boolean',
@@ -35,11 +40,12 @@ public function register_routes(): void {
 	public function generate_alt_text( WP_REST_Request $request ): WP_REST_Response|WP_Error {
 		$attachment_id = $request->get_param( 'attachment_id' );
 		$save_alt      = $request->get_param( 'save' );
+		$user_prompt   = $request->get_param( 'user_prompt' ) ?? '';
 
 		if ( $save_alt ) {
-			$alt_text = AltGenerator::generate_and_set_alt_text( $attachment_id );
+			$alt_text = AltGenerator::generate_and_set_alt_text( $attachment_id, $user_prompt );
 		} else {
-			$alt_text = AltGenerator::generate_alt_text( $attachment_id );
+			$alt_text = AltGenerator::generate_alt_text( $attachment_id, $user_prompt );
 		}
 
 		if ( is_wp_error( $alt_text ) ) {

diff --git a/src/components/BulkGenerateModal.tsx b/src/components/BulkGenerateModal.tsx
@@ -84,7 +84,7 @@ export default function BulkGenerateModal({
           new Map(prevMap.set(id, { ...details, status: "generating" })),
       );
 
-      generateAltText(id, true)
+      generateAltText(id, true, customPrompt)
         .then((alt) => {
           setAltGenerationMap(
             (prevMap) =>

diff --git a/src/utils/generateAltText.ts b/src/utils/generateAltText.ts
@@ -1,18 +1,32 @@
 import apiFetch from "@wordpress/api-fetch";
 import { API_PATH } from "../constants";
 
-export default async (attachmentId: number, save: boolean = false) => {
+export default async (
+  attachmentId: number,
+  save: boolean = false,
+  userPrompt?: string,
+) => {
+  const requestData: {
+    attachment_id: number;
+    save: boolean;
+    user_prompt?: string;
+  } = {
+    attachment_id: attachmentId,
+    save,
+  };
+
+  if (userPrompt?.length) {
+    requestData.user_prompt = userPrompt;
+  }
+
   return apiFetch<{ alt: string; img_id: number }>({
     path: API_PATH,
     method: "POST",
     headers: {
       "Content-Type": "application/json",
     },
 
-    body: JSON.stringify({
-      attachment_id: attachmentId,
-      save,
-    }),
+    body: JSON.stringify(requestData),
   })
     .then((response) => {
       return response.alt;