| 1 |
<?php |
| 2 |
|
| 3 |
namespace WPDataAccess\API; |
| 4 |
|
| 5 |
use WPDataAccess\Connection\WPDADB; |
| 6 |
use WPDataAccess\Data_Dictionary\WPDA_Dictionary_Lists; |
| 7 |
use WPDataAccess\WPDA; |
| 8 |
class WPDA_AI extends WPDA_API_Core { |
| 9 |
const CIPHER = 'AES-256-CBC'; |
| 10 |
|
| 11 |
const AI_API_KEY = 'wpda_ai_key'; |
| 12 |
|
| 13 |
const SUPPORTED_MODELS = ['gpt-3.5-turbo', 'gpt-4-turbo']; |
| 14 |
|
| 15 |
public function register_rest_routes() { |
| 16 |
register_rest_route( WPDA_API::WPDA_NAMESPACE, 'ai/sql', array( |
| 17 |
'methods' => array('POST'), |
| 18 |
'callback' => array($this, 'ai_sql'), |
| 19 |
'permission_callback' => function () { |
| 20 |
return $this->current_user_can_access(); |
| 21 |
}, |
| 22 |
'args' => array( |
| 23 |
'prompt' => array( |
| 24 |
'required' => true, |
| 25 |
'type' => 'string', |
| 26 |
'description' => __( 'Prompt', 'wp-data-access' ), |
| 27 |
'sanitize_callback' => 'sanitize_text_field', |
| 28 |
'validate_callback' => 'rest_validate_request_arg', |
| 29 |
), |
| 30 |
'model' => array( |
| 31 |
'required' => true, |
| 32 |
'type' => 'string', |
| 33 |
'description' => __( 'Model', 'wp-data-access' ), |
| 34 |
'sanitize_callback' => 'sanitize_text_field', |
| 35 |
'validate_callback' => function ( $param ) { |
| 36 |
return in_array( $param, self::SUPPORTED_MODELS ); |
| 37 |
}, |
| 38 |
), |
| 39 |
'explain' => array( |
| 40 |
'required' => true, |
| 41 |
'type' => 'boolean', |
| 42 |
'description' => __( 'Add explanations', 'wp-data-access' ), |
| 43 |
'sanitize_callback' => 'sanitize_text_field', |
| 44 |
'validate_callback' => 'rest_validate_request_arg', |
| 45 |
), |
| 46 |
), |
| 47 |
) ); |
| 48 |
register_rest_route( WPDA_API::WPDA_NAMESPACE, 'ai/hints', array( |
| 49 |
'methods' => array('POST'), |
| 50 |
'callback' => array($this, 'hints'), |
| 51 |
'permission_callback' => '__return_true', |
| 52 |
'args' => array( |
| 53 |
'dbs' => $this->get_param( 'dbs' ), |
| 54 |
), |
| 55 |
) ); |
| 56 |
register_rest_route( WPDA_API::WPDA_NAMESPACE, 'ai/enabled', array( |
| 57 |
'methods' => array('POST'), |
| 58 |
'callback' => array($this, 'enabled'), |
| 59 |
'permission_callback' => '__return_true', |
| 60 |
'args' => array(), |
| 61 |
) ); |
| 62 |
register_rest_route( WPDA_API::WPDA_NAMESPACE, 'ai/enable', array( |
| 63 |
'methods' => array('POST'), |
| 64 |
'callback' => array($this, 'enable'), |
| 65 |
'permission_callback' => '__return_true', |
| 66 |
'args' => array( |
| 67 |
'key' => array( |
| 68 |
'required' => true, |
| 69 |
'type' => 'string', |
| 70 |
'description' => __( 'API Key', 'wp-data-access' ), |
| 71 |
'sanitize_callback' => 'sanitize_text_field', |
| 72 |
'validate_callback' => 'rest_validate_request_arg', |
| 73 |
), |
| 74 |
'encrypt' => array( |
| 75 |
'required' => true, |
| 76 |
'type' => 'boolean', |
| 77 |
'description' => __( 'Encrypt API key', 'wp-data-access' ), |
| 78 |
'sanitize_callback' => 'sanitize_text_field', |
| 79 |
'validate_callback' => 'rest_validate_request_arg', |
| 80 |
), |
| 81 |
), |
| 82 |
) ); |
| 83 |
} |
| 84 |
|
| 85 |
public function ai_sql( $request ) { |
| 86 |
$timeout = 30; |
| 87 |
$prompt = $request['prompt']; |
| 88 |
$model = $request['model']; |
| 89 |
$explain = ( $request['explain'] ? 'Write the SQL query first in a code block. After the code block, provide a clear, concise explanation of what the query does.' : 'Provide only the query without further explanation.' ); |
| 90 |
$prompt = "\nYou are a professional MySQL consultant helping developers write SQL queries.\nAlways respond with clean, optimized MySQL code.\nPlace the query inside a single Markdown code block using triple backticks (```sql).\nDo not include a semicolon at the end of the query.\nImportant: Do not include a LIMIT clause unless the user specifically requests limiting the number of results.\nAssume the system will handle limits automatically if needed.\n{$explain}\nExample input for a user:\nWrite a join between tables dept and emp and show the average and total salaries per department.\nExpected Output:\n```sql\nSELECT d.dname AS department_name, \n AVG(e.sal) AS average_salary, \n SUM(e.sal) AS total_salary\nFROM dept d\nJOIN emp e ON d.deptno = e.deptno\nGROUP BY d.dname\n```\n{$prompt}\n"; |
| 91 |
WPDA::wpda_log_wp_error( $prompt ); |
| 92 |
$api_key_saved = get_user_meta( get_current_user_id(), self::AI_API_KEY, true ); |
| 93 |
if ( false === $api_key_saved || '' === $api_key_saved ) { |
| 94 |
$api_key = ''; |
| 95 |
} else { |
| 96 |
$api_key = substr( $api_key_saved, 0, -2 ); |
| 97 |
$is_encrypted = substr( $api_key_saved, -1 ); |
| 98 |
if ( '1' === $is_encrypted ) { |
| 99 |
// Decrypt API key |
| 100 |
$api_key = $this->decrypt( $api_key ); |
| 101 |
} |
| 102 |
} |
| 103 |
if ( '' === trim( $api_key ) ) { |
| 104 |
return new \WP_Error('error', 'Invalid or missing API Key', array( |
| 105 |
'status' => 403, |
| 106 |
)); |
| 107 |
} |
| 108 |
return $this->ask_ai( |
| 109 |
$api_key, |
| 110 |
$model, |
| 111 |
$prompt, |
| 112 |
$timeout |
| 113 |
); |
| 114 |
} |
| 115 |
|
| 116 |
private function ask_ai( |
| 117 |
$api_key, |
| 118 |
$model, |
| 119 |
$prompt, |
| 120 |
$timeout, |
| 121 |
$msg = '' |
| 122 |
) { |
| 123 |
$response = wp_remote_post( 'https://api.openai.com/v1/chat/completions', array( |
| 124 |
'headers' => array( |
| 125 |
'Authorization' => 'Bearer ' . $api_key, |
| 126 |
'Content-Type' => 'application/json', |
| 127 |
), |
| 128 |
'body' => json_encode( array( |
| 129 |
'model' => $model, |
| 130 |
'messages' => array(array( |
| 131 |
'role' => 'user', |
| 132 |
'content' => $prompt, |
| 133 |
)), |
| 134 |
) ), |
| 135 |
'timeout' => $timeout, |
| 136 |
) ); |
| 137 |
if ( !is_wp_error( $response ) ) { |
| 138 |
$body = json_decode( wp_remote_retrieve_body( $response ), true ); |
| 139 |
if ( '' !== $msg ) { |
| 140 |
$body['msg'] = $msg; |
| 141 |
} |
| 142 |
return rest_ensure_response( $body ); |
| 143 |
} |
| 144 |
if ( self::SUPPORTED_MODELS[1] === $model ) { |
| 145 |
// Try 'gpt-3.5-turbo' if 'gpt-4-turbo' failed |
| 146 |
return $this->ask_ai( |
| 147 |
$api_key, |
| 148 |
self::SUPPORTED_MODELS[0], |
| 149 |
$prompt, |
| 150 |
$timeout, |
| 151 |
'Note: This result was generated using gpt-3.5-turbo due to a timeout using gpt-4-turbo.' |
| 152 |
); |
| 153 |
} |
| 154 |
return new \WP_Error('error', $response->get_error_message(), array( |
| 155 |
'status' => 403, |
| 156 |
)); |
| 157 |
} |
| 158 |
|
| 159 |
public function hints( $request ) { |
| 160 |
if ( !$this->current_user_can_access() ) { |
| 161 |
return $this->unauthorized(); |
| 162 |
} |
| 163 |
if ( !$this->current_user_token_valid( $request ) ) { |
| 164 |
return $this->invalid_nonce(); |
| 165 |
} |
| 166 |
$dbs = $request->get_param( 'dbs' ); |
| 167 |
$tables = WPDA_Dictionary_Lists::get_tables( true, $dbs ); |
| 168 |
$wpdadb = WPDADB::get_db_connection( $dbs ); |
| 169 |
if ( null === $wpdadb ) { |
| 170 |
// Error connecting. |
| 171 |
return new \WP_Error('error', "Error connecting to database {$dbs}", array( |
| 172 |
'status' => 420, |
| 173 |
)); |
| 174 |
} |
| 175 |
$hints = array(); |
| 176 |
foreach ( $tables as $table ) { |
| 177 |
$table_name = WPDA::remove_backticks( $table['table_name'] ); |
| 178 |
$sql_cmd = $wpdadb->get_results( "SHOW CREATE TABLE `{$table_name}`", 'ARRAY_N' ); |
| 179 |
if ( '' === $wpdadb->last_error && isset( $sql_cmd[0][1] ) ) { |
| 180 |
$hints[$table_name] = $sql_cmd[0][1]; |
| 181 |
} |
| 182 |
} |
| 183 |
return $this->WPDA_Rest_Response( '', $hints ); |
| 184 |
} |
| 185 |
|
| 186 |
public function enabled( $request ) { |
| 187 |
if ( !$this->current_user_can_access() ) { |
| 188 |
return $this->unauthorized(); |
| 189 |
} |
| 190 |
if ( !$this->current_user_token_valid( $request ) ) { |
| 191 |
return $this->invalid_nonce(); |
| 192 |
} |
| 193 |
$api_key = get_user_meta( WPDA::get_current_user_id(), self::AI_API_KEY, true ); |
| 194 |
$is_enabled = false !== $api_key && '' !== $api_key; |
| 195 |
return $this->WPDA_Rest_Response( '', array( |
| 196 |
'enabled' => $is_enabled, |
| 197 |
'encryption' => $this->get_encryption_key() !== null, |
| 198 |
) ); |
| 199 |
} |
| 200 |
|
| 201 |
public function enable( $request ) { |
| 202 |
if ( !$this->current_user_can_access() ) { |
| 203 |
return $this->unauthorized(); |
| 204 |
} |
| 205 |
if ( !$this->current_user_token_valid( $request ) ) { |
| 206 |
return $this->invalid_nonce(); |
| 207 |
} |
| 208 |
$key = $request['key']; |
| 209 |
$encrypt = ( $request['encrypt'] ? 1 : 0 ); |
| 210 |
if ( $encrypt ) { |
| 211 |
$key = $this->encrypt( $key ); |
| 212 |
} |
| 213 |
update_user_meta( WPDA::get_current_user_id(), self::AI_API_KEY, "{$key}|{$encrypt}" ); |
| 214 |
return $this->WPDA_Rest_Response( '' ); |
| 215 |
} |
| 216 |
|
| 217 |
private function encrypt( $string ) { |
| 218 |
$key = $this->get_encryption_key(); |
| 219 |
if ( null === $key ) { |
| 220 |
return $string; |
| 221 |
} |
| 222 |
$ivlen = openssl_cipher_iv_length( self::CIPHER ); |
| 223 |
$iv = openssl_random_pseudo_bytes( $ivlen ); |
| 224 |
$ciphertext_raw = openssl_encrypt( |
| 225 |
$string, |
| 226 |
self::CIPHER, |
| 227 |
$key, |
| 228 |
OPENSSL_RAW_DATA, |
| 229 |
$iv |
| 230 |
); |
| 231 |
$hmac = hash_hmac( |
| 232 |
'sha256', |
| 233 |
$ciphertext_raw, |
| 234 |
$key, |
| 235 |
true |
| 236 |
); |
| 237 |
return base64_encode( $iv . $hmac . $ciphertext_raw ); |
| 238 |
} |
| 239 |
|
| 240 |
private function decrypt( $string ) { |
| 241 |
$key = $this->get_encryption_key(); |
| 242 |
if ( null === $key ) { |
| 243 |
return $string; |
| 244 |
} |
| 245 |
$c = base64_decode( $string ); |
| 246 |
$ivlen = openssl_cipher_iv_length( self::CIPHER ); |
| 247 |
$iv = substr( $c, 0, $ivlen ); |
| 248 |
$hmac = substr( $c, $ivlen, 32 ); |
| 249 |
$ciphertext_raw = substr( $c, $ivlen + 32 ); |
| 250 |
$calculated_hmac = hash_hmac( |
| 251 |
'sha256', |
| 252 |
$ciphertext_raw, |
| 253 |
$key, |
| 254 |
true |
| 255 |
); |
| 256 |
if ( !hash_equals( $hmac, $calculated_hmac ) ) { |
| 257 |
return false; |
| 258 |
} |
| 259 |
return openssl_decrypt( |
| 260 |
$ciphertext_raw, |
| 261 |
self::CIPHER, |
| 262 |
$key, |
| 263 |
OPENSSL_RAW_DATA, |
| 264 |
$iv |
| 265 |
); |
| 266 |
} |
| 267 |
|
| 268 |
private function get_encryption_key() { |
| 269 |
if ( defined( 'WPDA_ENCRYPT_AI_KEY' ) && '' !== trim( constant( 'WPDA_ENCRYPT_AI_KEY' ) ) ) { |
| 270 |
return constant( 'WPDA_ENCRYPT_AI_KEY' ); |
| 271 |
} else { |
| 272 |
return null; |
| 273 |
} |
| 274 |
} |
| 275 |
|
| 276 |
} |
| 277 |
|
