Skip to content

Commit df56a36

Browse files
felipeeliafelipeelia
committed
Formatting and basic prompt validation
1 parent 8c1c940 commit df56a36

File tree

3 files changed

+160
-14
lines changed

3 files changed

+160
-14
lines changed

‎assets/js/blocks/rag/frontend.js

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ if (modelUrl) {
2525
* @returns {WPElement} App component.
2626
*/
2727
const App = () => {
28+
const [className, setClassName] = useState('');
2829
const [message, setMessage] = useState('');
2930
const [isLoading, setIsLoading] = useState(true);
3031

@@ -41,6 +42,7 @@ const App = () => {
4142
},
4243
})
4344
.then((response) => {
45+
setClassName(response.class);
4446
setMessage(response.html);
4547
})
4648
.finally(() => {
@@ -53,6 +55,7 @@ const App = () => {
5355
path: `${restApiEndpoint}?search_query=${searchQuery}`,
5456
})
5557
.then((response) => {
58+
setClassName(response.class);
5659
setMessage(response.html);
5760
})
5861
.finally(() => {
@@ -66,8 +69,11 @@ const App = () => {
6669
<Skeleton count={5} />
6770
</Placeholder>
6871
) : (
69-
// eslint-disable-next-line react/no-danger
70-
<div className="ep-rag-generated" dangerouslySetInnerHTML={{ __html: message }} />
72+
<div
73+
className={`ep-rag-generated ${className}`}
74+
// eslint-disable-next-line react/no-danger
75+
dangerouslySetInnerHTML={{ __html: message }}
76+
/>
7177
);
7278
};
7379

‎includes/classes/Feature/RAG.php

Lines changed: 97 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
/**
33
* RAG Feature
44
*
5-
* @since 2.4.0
5+
* @since 2.5.0
66
* @package ElasticPressLabs
77
*/
88

@@ -18,7 +18,7 @@
1818
/**
1919
* RAG feature
2020
*
21-
* @since 2.4.0
21+
* @since 2.5.0
2222
*/
2323
class RAG extends Feature {
2424
/**
@@ -213,13 +213,41 @@ public function setup_endpoint() {
213213
*
214214
* @param string $search_term Search term
215215
* @param null|array $search_vectors Search term vectors
216-
* @return string
216+
* @return string|\WP_Error
217217
*/
218218
public function get_ai_response( $search_term, $search_vectors = null ) {
219219
if ( ! $search_term ) {
220220
return '';
221221
}
222222

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+
223251
/**
224252
* Filters the AI response before it is returned.
225253
*
@@ -258,10 +286,11 @@ public function get_ai_response( $search_term, $search_vectors = null ) {
258286
* Fires after receiving the response for a RAG (Retrieval-Augmented Generation) post request.
259287
*
260288
* @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.
265294
*/
266295
do_action( 'ep_rag_post_response', $response, $search_term, $search_vectors, $prompt );
267296

@@ -357,7 +386,7 @@ public function get_prompt( $posts_representations ) {
357386
*
358387
* @param string $prompt Prompt for the AI model
359388
* @param string $search_term Search query
360-
* @return string
389+
* @return string|\WP_Error
361390
*/
362391
public function ai_api_request( $prompt, $search_term ) {
363392
$headers = [
@@ -372,6 +401,10 @@ public function ai_api_request( $prompt, $search_term ) {
372401
'role' => 'system',
373402
'content' => $prompt,
374403
],
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+
],
375408
[
376409
'role' => 'user',
377410
'content' => $search_term,
@@ -403,17 +436,20 @@ public function ai_api_request( $prompt, $search_term ) {
403436
);
404437

405438
$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+
}
406442

407443
$code = wp_remote_retrieve_response_code( $response );
408444
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' ) );
410446
}
411447

412448
$body = wp_remote_retrieve_body( $response );
413449
$body = json_decode( $body, true );
414450
return isset( $body['choices'], $body['choices'][0], $body['choices'][0]['message'], $body['choices'][0]['message']['content'] )
415451
? $body['choices'][0]['message']['content']
416-
: false;
452+
: new \WP_Error( 'ep_rag_unformatted_response', __( 'An error occurred. Try again later.', 'elasticpress-labs' ) );
417453
}
418454

419455
/**
@@ -503,4 +539,55 @@ public function requirements_status() {
503539

504540
return $status;
505541
}
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+
}
506593
}

‎includes/classes/REST/RAG.php

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88

99
namespace ElasticPressLabs\REST;
1010

11-
use ElasticPress\Utils;
1211
use ElasticPressLabs\Feature\RAG as RAGFeature;
1312

1413
/**
@@ -83,7 +82,20 @@ public function register_routes() {
8382
* @return object|\WP_Error
8483
*/
8584
public function get_rag_response( \WP_REST_Request $request ) {
86-
return [ 'html' => $this->feature->get_ai_response( $request['search_query'], $request['search_vectors'] ?? null ) ];
85+
$ai_response = $this->feature->get_ai_response( $request['search_query'], $request['search_vectors'] ?? null );
86+
87+
if ( ! is_wp_error( $ai_response ) ) {
88+
$class = 'ep-rag-success';
89+
$html = $this->format_response( $ai_response );
90+
} else {
91+
$class = str_replace( '_', '-', $ai_response->get_error_code() );
92+
$html = $ai_response->get_error_message();
93+
}
94+
95+
return [
96+
'class' => $class,
97+
'html' => $html,
98+
];
8799
}
88100

89101
/**
@@ -104,4 +116,45 @@ public function sanitize_vectors_array( array $array ): array {
104116
protected function get_embed_method(): string {
105117
return (string) $this->feature->get_setting( 'ep_rag_search_term_embed_method' );
106118
}
119+
120+
/**
121+
* Formats the AI response into a string.
122+
*
123+
* @param mixed $ai_response The AI response data to be formatted.
124+
* @return string The formatted response as a string.
125+
*/
126+
protected function format_response( $ai_response ): string {
127+
$ai_response = trim( $ai_response, "'" );
128+
$ai_response_array = json_decode( $ai_response, true );
129+
130+
$html = '';
131+
if ( json_last_error() === JSON_ERROR_NONE && isset( $ai_response_array['response'] ) ) {
132+
$html = wp_kses_post( str_replace( '\\\\', '\\', $ai_response_array['response'] ) );
133+
if ( ! empty( $ai_response_array['references'] ) ) {
134+
$html .= '<div class="ep-rag--references">';
135+
$html .= '<p>' . esc_html__( 'Sources:', 'elasticpress-labs' ) . '</p>';
136+
$html .= '<ul>';
137+
foreach ( $ai_response_array['references'] as $reference ) {
138+
$html .= '<li>';
139+
$html .= '<a href="' . esc_url( $reference['url'] ) . '" target="_blank" rel="noopener noreferrer">';
140+
$html .= esc_html( $reference['title'] );
141+
$html .= '</a>';
142+
$html .= '</li>';
143+
}
144+
$html .= '</ul>';
145+
$html .= '</div>';
146+
}
147+
}
148+
149+
/**
150+
* Filters the formatted HTML response generated from the AI response.
151+
*
152+
* @since 2.5.0
153+
* @hook ep_rag_formatted_response
154+
* @param {string} $html The formatted HTML response.
155+
* @param {mixed} $ai_response The raw AI response data.
156+
* @return {string} Filtered HTML response.
157+
*/
158+
return apply_filters( 'ep_rag_formatted_response', $html, $ai_response );
159+
}
107160
}

0 commit comments

Comments
 (0)