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