2
2
/**
3
3
* RAG Feature
4
4
*
5
- * @since 2.4 .0
5
+ * @since 2.5 .0
6
6
* @package ElasticPressLabs
7
7
*/
8
8
18
18
/**
19
19
* RAG feature
20
20
*
21
- * @since 2.4 .0
21
+ * @since 2.5 .0
22
22
*/
23
23
class RAG extends Feature {
24
24
/**
@@ -213,13 +213,41 @@ public function setup_endpoint() {
213
213
*
214
214
* @param string $search_term Search term
215
215
* @param null|array $search_vectors Search term vectors
216
- * @return string
216
+ * @return string|\WP_Error
217
217
*/
218
218
public function get_ai_response ( $ search_term , $ search_vectors = null ) {
219
219
if ( ! $ search_term ) {
220
220
return '' ;
221
221
}
222
222
223
+ $ is_valid_search_term = $ this ->validate_search_term ( $ search_term );
224
+
225
+ /**
226
+ * Filter to determine if a search term is valid for RAG feature.
227
+ *
228
+ * This filter allows customization of the validation logic for search terms
229
+ * used in the RAG feature. Developers can use this filter to override the
230
+ * default validation behavior.
231
+ *
232
+ * @since 2.5.0
233
+ * @hook ep_rag_is_valid_search_term
234
+ * @param {bool} $is_valid_search_term Whether the search term is valid. Default is determined by internal logic.
235
+ * @param {string} $search_term The search term being validated.
236
+ * @return {bool} Whether the search term is valid.
237
+ */
238
+ if ( ! apply_filters ( 'ep_rag_is_valid_search_term ' , $ is_valid_search_term , $ search_term ) ) {
239
+ /**
240
+ * Filter the response for an invalid search term in the RAG feature.
241
+ *
242
+ * @since 2.5.0
243
+ * @hook ep_rag_invalid_search_term_response
244
+ * @param {string} $response The response to return for an invalid search term. Default is an empty string.
245
+ * @param {string} $search_term The invalid search term that triggered the response.
246
+ * @return {\WP_Error} Response.
247
+ */
248
+ return apply_filters ( 'ep_rag_invalid_search_term_response ' , new \WP_Error ( 'ep-rag-invalid-search-term ' , '' ), $ search_term );
249
+ }
250
+
223
251
/**
224
252
* Filters the AI response before it is returned.
225
253
*
@@ -258,10 +286,11 @@ public function get_ai_response( $search_term, $search_vectors = null ) {
258
286
* Fires after receiving the response for a RAG (Retrieval-Augmented Generation) post request.
259
287
*
260
288
* @since 2.5.0
261
- * @param array $response The response from the RAG post request.
262
- * @param string $search_term The search term.
263
- * @param array $search_vectors The search vectors used for the RAG request.
264
- * @param string $prompt The prompt.
289
+ * @hook ep_rag_post_response
290
+ * @param array|\WP_error $response The response from the RAG post request.
291
+ * @param string $search_term The search term.
292
+ * @param array $search_vectors The search vectors used for the RAG request.
293
+ * @param string $prompt The prompt.
265
294
*/
266
295
do_action ( 'ep_rag_post_response ' , $ response , $ search_term , $ search_vectors , $ prompt );
267
296
@@ -357,7 +386,7 @@ public function get_prompt( $posts_representations ) {
357
386
*
358
387
* @param string $prompt Prompt for the AI model
359
388
* @param string $search_term Search query
360
- * @return string
389
+ * @return string|\WP_Error
361
390
*/
362
391
public function ai_api_request ( $ prompt , $ search_term ) {
363
392
$ headers = [
@@ -372,6 +401,10 @@ public function ai_api_request( $prompt, $search_term ) {
372
401
'role ' => 'system ' ,
373
402
'content ' => $ prompt ,
374
403
],
404
+ [
405
+ 'role ' => 'system ' ,
406
+ 'content ' => 'Send your response as a JSON object with the following keys: "response" (the asnwer, in HTML format) and "references" (an array of objects with the URLs you used to build the response, having "url" and "title" as attributes). Do not wrap the response in any other tags or limiters like "```json". Make sure the JSON object returned is properly escaped. ' ,
407
+ ],
375
408
[
376
409
'role ' => 'user ' ,
377
410
'content ' => $ search_term ,
@@ -403,17 +436,20 @@ public function ai_api_request( $prompt, $search_term ) {
403
436
);
404
437
405
438
$ response = wp_remote_post ( $ url , $ options );
439
+ if ( is_wp_error ( $ response ) ) {
440
+ return new \WP_Error ( 'ep_rag_request_failed ' , __ ( 'An error occurred. Try again later. ' , 'elasticpress-labs ' ) );
441
+ }
406
442
407
443
$ code = wp_remote_retrieve_response_code ( $ response );
408
444
if ( 200 !== $ code ) {
409
- return false ;
445
+ return new \ WP_Error ( ' ep_rag_non_200_code_ ' . $ code , __ ( ' An error occurred. Try again later. ' , ' elasticpress-labs ' ) ) ;
410
446
}
411
447
412
448
$ body = wp_remote_retrieve_body ( $ response );
413
449
$ body = json_decode ( $ body , true );
414
450
return isset ( $ body ['choices ' ], $ body ['choices ' ][0 ], $ body ['choices ' ][0 ]['message ' ], $ body ['choices ' ][0 ]['message ' ]['content ' ] )
415
451
? $ body ['choices ' ][0 ]['message ' ]['content ' ]
416
- : false ;
452
+ : new \ WP_Error ( ' ep_rag_unformatted_response ' , __ ( ' An error occurred. Try again later. ' , ' elasticpress-labs ' ) ) ;
417
453
}
418
454
419
455
/**
@@ -503,4 +539,55 @@ public function requirements_status() {
503
539
504
540
return $ status ;
505
541
}
542
+
543
+ /**
544
+ * Validates the provided search term.
545
+ *
546
+ * This function checks the validity of the given search term
547
+ * to ensure it meets the required criteria for processing.
548
+ *
549
+ * @param string $search_term The search term to validate.
550
+ * @return bool True if the search term is valid, false otherwise.
551
+ */
552
+ protected function validate_search_term ( string $ search_term ): bool {
553
+ $ attack_patterns = [
554
+ // Instruction override patterns
555
+ '/ignore previous (instructions|rule|prompt)/i ' ,
556
+ '/disregard (your|all|previous) (instructions|prompt)/i ' ,
557
+ '/forget (your|all) (instructions|prompt)/i ' ,
558
+
559
+ // Delimiter exploitation patterns
560
+ '/\<\/?system\>/i ' ,
561
+ '/\<\/?admin\>/i ' ,
562
+ '/\<\/?prompt\>/i ' ,
563
+ '/\<\/?instructions?\>/i ' ,
564
+
565
+ // Jailbreak attempts
566
+ '/DAN|Do Anything Now/i ' ,
567
+ '/you are a helpful assistant that only responds/i ' ,
568
+ '/you are in developer mode/i ' ,
569
+
570
+ // Role play exploitation
571
+ '/pretend to be/i ' ,
572
+ '/act as if/i ' ,
573
+ '/you are now/i ' ,
574
+
575
+ // Payload embedding attempts
576
+ '/\{\{[^}]+\}\}/i ' ,
577
+ '/\[\[[^]]+\]\]/i ' ,
578
+ '/```(system|exec|prompt)/i ' ,
579
+
580
+ // Coding context breaking
581
+ '/`\/\/ignore previous code`/i ' ,
582
+ '/\/\*\s*ignore previous\s*\*\//i ' ,
583
+ ];
584
+
585
+ foreach ( $ attack_patterns as $ pattern ) {
586
+ if ( preg_match ( $ pattern , $ search_term ) ) {
587
+ return false ;
588
+ }
589
+ }
590
+
591
+ return true ;
592
+ }
506
593
}
0 commit comments