Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
6.25%
11 / 176
11.11%
1 / 9
 CRAP
0.00%
0 / 1
WPCOM_REST_API_V2_Endpoint_AI
6.36%
11 / 173
11.11%
1 / 9
 254.30
0.00%
0 / 1
 __construct
0.00%
0 / 10
0.00%
0 / 1
 20
 register_routes
0.00%
0 / 49
0.00%
0 / 1
 2
 register_ai_chat_routes
0.00%
0 / 51
0.00%
0 / 1
 2
 register_basic_routes
100.00%
11 / 11
100.00%
1 / 1
 1
 request_chat_with_site
0.00%
0 / 22
0.00%
0 / 1
 12
 rank_response
0.00%
0 / 27
0.00%
0 / 1
 20
 request_gpt_completion
0.00%
0 / 1
0.00%
0 / 1
 2
 request_dalle_generation
0.00%
0 / 1
0.00%
0 / 1
 2
 request_get_ai_assistance_feature
0.00%
0 / 1
0.00%
0 / 1
 2
1<?php
2/**
3 * REST API endpoint for the Jetpack AI blocks.
4 *
5 * @package automattic/jetpack
6 * @since 11.8
7 */
8
9use Automattic\Jetpack\Connection\Client;
10
11if ( ! defined( 'ABSPATH' ) ) {
12    exit( 0 );
13}
14
15/**
16 * Class WPCOM_REST_API_V2_Endpoint_AI
17 */
18class WPCOM_REST_API_V2_Endpoint_AI extends WP_REST_Controller {
19    /**
20     * Namespace prefix.
21     *
22     * @var string
23     */
24    public $namespace = 'wpcom/v2';
25
26    /**
27     * Endpoint base route.
28     *
29     * @var string
30     */
31    public $rest_base = 'jetpack-ai';
32
33    /**
34     * WPCOM_REST_API_V2_Endpoint_AI constructor.
35     */
36    public function __construct() {
37        $this->is_wpcom                     = true;
38        $this->wpcom_is_wpcom_only_endpoint = true;
39
40        if ( ! class_exists( 'Jetpack_AI_Helper' ) ) {
41            require_once JETPACK__PLUGIN_DIR . '_inc/lib/class-jetpack-ai-helper.php';
42        }
43
44        // Register routes that don't require Jetpack AI to be enabled.
45        add_action( 'rest_api_init', array( $this, 'register_basic_routes' ) );
46
47        if ( Jetpack_AI_Helper::is_ai_chat_enabled() ) {
48            add_action( 'rest_api_init', array( $this, 'register_ai_chat_routes' ) );
49        }
50
51        if ( ! \Jetpack_AI_Helper::is_enabled() ) {
52            return;
53        }
54
55        // Register routes that require Jetpack AI to be enabled.
56        add_action( 'rest_api_init', array( $this, 'register_routes' ) );
57    }
58
59    /**
60     * Register routes.
61     */
62    public function register_routes() {
63        register_rest_route(
64            $this->namespace,
65            $this->rest_base . '/completions',
66            array(
67                array(
68                    'methods'             => WP_REST_Server::CREATABLE,
69                    'callback'            => array( $this, 'request_gpt_completion' ),
70                    'permission_callback' => array( 'Jetpack_AI_Helper', 'get_status_permission_check' ),
71                ),
72                'args' => array(
73                    'content'    => array(
74                        'type'              => 'string',
75                        'required'          => true,
76                        'sanitize_callback' => 'sanitize_textarea_field',
77                    ),
78                    'post_id'    => array(
79                        'required' => false,
80                        'type'     => 'integer',
81                    ),
82                    'skip_cache' => array(
83                        'required'    => false,
84                        'type'        => 'boolean',
85                        'description' => 'Whether to skip the cache and make a new request',
86                    ),
87                ),
88            )
89        );
90
91        register_rest_route(
92            $this->namespace,
93            $this->rest_base . '/images/generations',
94            array(
95                array(
96                    'methods'             => WP_REST_Server::CREATABLE,
97                    'callback'            => array( $this, 'request_dalle_generation' ),
98                    'permission_callback' => array( 'Jetpack_AI_Helper', 'get_status_permission_check' ),
99                ),
100                'args' => array(
101                    'prompt'  => array(
102                        'type'              => 'string',
103                        'required'          => true,
104                        'sanitize_callback' => 'sanitize_textarea_field',
105                    ),
106                    'post_id' => array(
107                        'required' => false,
108                        'type'     => 'integer',
109                    ),
110                ),
111            )
112        );
113    }
114
115    /**
116     * Register routes for the AI Chat block.
117     * Relies on a site connection and Jetpack Search.
118     */
119    public function register_ai_chat_routes() {
120        register_rest_route(
121            $this->namespace,
122            '/jetpack-search/ai/search',
123            array(
124                array(
125                    'methods'             => WP_REST_Server::READABLE,
126                    'callback'            => array( $this, 'request_chat_with_site' ),
127                    'permission_callback' => '__return_true',
128                ),
129                'args' => array(
130                    'query'         => array(
131                        'description'       => 'Your question to the site',
132                        'required'          => true,
133                        'sanitize_callback' => 'sanitize_text_field',
134                    ),
135                    'answer_prompt' => array(
136                        'description'       => 'Answer prompt override',
137                        'required'          => false,
138                        'sanitize_callback' => 'sanitize_text_field',
139                    ),
140                ),
141            )
142        );
143
144        register_rest_route(
145            $this->namespace,
146            '/jetpack-search/ai/rank',
147            array(
148                array(
149                    'methods'             => WP_REST_Server::CREATABLE,
150                    'callback'            => array( $this, 'rank_response' ),
151                    'permission_callback' => '__return_true',
152                ),
153                'args' => array(
154                    'cache_key' => array(
155                        'description'       => 'Cache key of your response',
156                        'required'          => true,
157                        'sanitize_callback' => 'sanitize_text_field',
158                    ),
159                    'comment'   => array(
160                        'description'       => 'Optional feedback',
161                        'required'          => false,
162                        'sanitize_callback' => 'sanitize_text_field',
163                    ),
164                    'rank'      => array(
165                        'description'       => 'How do you rank this response',
166                        'required'          => false,
167                        'sanitize_callback' => 'sanitize_text_field',
168                    ),
169                ),
170            )
171        );
172    }
173
174    /**
175     * Register routes that don't require Jetpack AI to be enabled.
176     */
177    public function register_basic_routes() {
178        register_rest_route(
179            $this->namespace,
180            $this->rest_base . '/ai-assistant-feature',
181            array(
182                array(
183                    'methods'             => WP_REST_Server::READABLE,
184                    'callback'            => array( $this, 'request_get_ai_assistance_feature' ),
185                    'permission_callback' => array( 'Jetpack_AI_Helper', 'get_status_permission_check' ),
186                ),
187            )
188        );
189    }
190
191    /**
192     * Get a response from chatting with the site.
193     * Uses Jetpack Search.
194     *
195     * @param  WP_REST_Request $request The request.
196     * @return mixed
197     */
198    public function request_chat_with_site( $request ) {
199        $question = $request->get_param( 'query' );
200        $blog_id  = \Jetpack_Options::get_option( 'id' );
201        $response = Client::wpcom_json_api_request_as_blog(
202            sprintf( '/sites/%d/jetpack-search/ai/search', $blog_id ) . '?force=wpcom',
203            2,
204            array(
205                'method'  => 'GET',
206                'headers' => array( 'content-type' => 'application/json' ),
207                'timeout' => MINUTE_IN_SECONDS,
208            ),
209            array(
210                'query'         => $question,
211                /**
212                 * Filter for an answer prompt override.
213                 * Example: "Talk like a cowboy."
214                 *
215                 * @param string $prompt_override The prompt override string.
216                 *
217                 * @since 12.6
218                 */
219                'answer_prompt' => apply_filters( 'jetpack_ai_chat_answer_prompt', false ),
220            ),
221            'wpcom'
222        );
223
224        if ( is_wp_error( $response ) ) {
225            return $response;
226        }
227
228        $data = json_decode( wp_remote_retrieve_body( $response ) );
229
230        if ( empty( $data->cache_key ) ) {
231            return new WP_Error( 'invalid_ask_response', __( 'Invalid response from the server.', 'jetpack' ), 400 );
232        }
233
234        return $data;
235    }
236
237    /**
238     * Rank a response from chatting with the site.
239     *
240     * @param  WP_REST_Request $request The request.
241     * @return mixed
242     */
243    public function rank_response( $request ) {
244        $rank      = $request->get_param( 'rank' );
245        $comment   = $request->get_param( 'comment' );
246        $cache_key = $request->get_param( 'cache_key' );
247
248        if ( strpos( $cache_key, 'jp-search-ai-' ) !== 0 ) {
249            return new WP_Error( 'invalid_cache_key', __( 'Invalid cached context for the answer feedback.', 'jetpack' ), 400 );
250        }
251
252        $blog_id  = \Jetpack_Options::get_option( 'id' );
253        $response = Client::wpcom_json_api_request_as_blog(
254            sprintf( '/sites/%d/jetpack-search/ai/rank', $blog_id ) . '?force=wpcom',
255            2,
256            array(
257                'method'  => 'GET',
258                'headers' => array( 'content-type' => 'application/json' ),
259                'timeout' => 30,
260            ),
261            array(
262                'rank'      => $rank,
263                'comment'   => $comment,
264                'cache_key' => $cache_key,
265            ),
266            'wpcom'
267        );
268
269        if ( is_wp_error( $response ) ) {
270            return $response;
271        }
272
273        $data = json_decode( wp_remote_retrieve_body( $response ) );
274
275        if ( 'ok' !== $data ) {
276            return new WP_Error( 'invalid_feedback_response', __( 'Invalid response from the server.', 'jetpack' ), 400 );
277        }
278
279        return $data;
280    }
281
282    /**
283     * Get completions for a given text.
284     *
285     * @param  WP_REST_Request $request The request.
286     */
287    public function request_gpt_completion( $request ) {
288        return Jetpack_AI_Helper::get_gpt_completion( $request['content'], $request['post_id'], $request['skip_cache'] );
289    }
290
291    /**
292     * Get image generations for a given prompt.
293     *
294     * @param  WP_REST_Request $request The request.
295     */
296    public function request_dalle_generation( $request ) {
297        return Jetpack_AI_Helper::get_dalle_generation( $request['prompt'], $request['post_id'] );
298    }
299
300    /**
301     * Collect and provide relevat data about the AI feature,
302     * such as the number of requests made.
303     */
304    public function request_get_ai_assistance_feature() {
305        return Jetpack_AI_Helper::get_ai_assistance_feature();
306    }
307}
308
309wpcom_rest_api_v2_load_plugin( 'WPCOM_REST_API_V2_Endpoint_AI' );