Skip to content

Commit b3f9f2a

Browse files
artpiartpi
committed
Add MariaDB vector RAG search
1 parent c71c9c9 commit b3f9f2a

21 files changed

Lines changed: 3912 additions & 2 deletions

‎includes/CLI/RAG_Command.php‎

Lines changed: 291 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,291 @@
1+
<?php
2+
/**
3+
* WP-CLI utilities for RAG indexing.
4+
*
5+
* @package WordPress\AI
6+
*
7+
* @since 1.1.0
8+
*/
9+
10+
declare( strict_types=1 );
11+
12+
namespace WordPress\AI\CLI;
13+
14+
use WP_CLI;
15+
use WP_CLI\Utils;
16+
use WordPress\AI\RAG\Availability;
17+
use WordPress\AI\RAG\Index_Manager;
18+
19+
defined( 'ABSPATH' ) || exit;
20+
21+
/**
22+
* Manages MariaDB vector RAG indexing utilities.
23+
*
24+
* @since 1.1.0
25+
*/
26+
class RAG_Command {
27+
/**
28+
* Availability service.
29+
*
30+
* @var \WordPress\AI\RAG\Availability
31+
*/
32+
private Availability $availability;
33+
34+
/**
35+
* Index manager.
36+
*
37+
* @var \WordPress\AI\RAG\Index_Manager
38+
*/
39+
private Index_Manager $index_manager;
40+
41+
/**
42+
* Constructor.
43+
*
44+
* @since 1.1.0
45+
*
46+
* @param \WordPress\AI\RAG\Availability|null $availability Availability service.
47+
* @param \WordPress\AI\RAG\Index_Manager|null $index_manager Index manager.
48+
*/
49+
public function __construct( ?Availability $availability = null, ?Index_Manager $index_manager = null ) {
50+
$this->availability = $availability ?? new Availability();
51+
$this->index_manager = $index_manager ?? new Index_Manager( $this->availability );
52+
}
53+
54+
/**
55+
* Shows RAG index status.
56+
*
57+
* ## EXAMPLES
58+
*
59+
* $ wp ai rag status
60+
*
61+
* @when after_wp_load
62+
*
63+
* @param array<int, string> $args Positional arguments.
64+
* @param array<string, mixed> $assoc_args Associative arguments.
65+
*/
66+
public function status( $args, $assoc_args ): void {
67+
unset( $args, $assoc_args );
68+
69+
$available = $this->availability->is_available();
70+
$table_ok = $available && $this->index_manager->ensure_index_table();
71+
$counts = $this->index_manager->get_status_counts();
72+
$scheduled = $this->index_manager->get_next_scheduled_indexing();
73+
74+
WP_CLI::log( 'RAG search status:' );
75+
WP_CLI::log( sprintf( ' Available: %s', $available ? 'yes' : 'no' ) );
76+
77+
if ( ! $available ) {
78+
WP_CLI::log( sprintf( ' Reason: %s', $this->availability->get_unavailable_reason() ) );
79+
}
80+
81+
WP_CLI::log( sprintf( ' Index table: %s', $table_ok ? 'present' : 'missing' ) );
82+
WP_CLI::log( sprintf( ' Dirty: %d', $counts[ Index_Manager::STATUS_DIRTY ] ?? 0 ) );
83+
WP_CLI::log( sprintf( ' Processing: %d', $counts[ Index_Manager::STATUS_PROCESSING ] ?? 0 ) );
84+
WP_CLI::log( sprintf( ' Clean: %d', $counts[ Index_Manager::STATUS_CLEAN ] ?? 0 ) );
85+
WP_CLI::log( sprintf( ' Error: %d', $counts[ Index_Manager::STATUS_ERROR ] ?? 0 ) );
86+
WP_CLI::log(
87+
sprintf(
88+
' Next scheduled run: %s',
89+
null === $scheduled ? 'none' : gmdate( 'Y-m-d H:i:s', $scheduled ) . ' UTC'
90+
)
91+
);
92+
}
93+
94+
/**
95+
* Marks eligible posts as dirty.
96+
*
97+
* ## OPTIONS
98+
*
99+
* [--ids=<ids>]
100+
* : Comma-separated list of post IDs to mark dirty. Defaults to all eligible posts.
101+
*
102+
* [--limit=<number>]
103+
* : Maximum number of eligible posts to mark when --ids is not supplied.
104+
*
105+
* ## EXAMPLES
106+
*
107+
* # Mark all eligible posts/pages dirty
108+
* $ wp ai rag mark-dirty
109+
*
110+
* # Mark selected posts dirty
111+
* $ wp ai rag mark-dirty --ids=42,99
112+
*
113+
* @when after_wp_load
114+
*
115+
* @param array<int, string> $args Positional arguments.
116+
* @param array<string, mixed> $assoc_args Associative arguments.
117+
*/
118+
public function mark_dirty( $args, $assoc_args ): void {
119+
unset( $args );
120+
121+
$this->ensure_available();
122+
123+
$ids = $this->parse_ids_flag( (string) Utils\get_flag_value( $assoc_args, 'ids', '' ) );
124+
$limit = max( 0, (int) Utils\get_flag_value( $assoc_args, 'limit', 0 ) );
125+
$stats = $this->index_manager->mark_posts_dirty_for_indexing( $ids, $limit );
126+
127+
WP_CLI::success(
128+
sprintf(
129+
'Marked %d post(s) dirty. Skipped %d, removed %d from the index.',
130+
$stats['marked'],
131+
$stats['skipped'],
132+
$stats['removed']
133+
)
134+
);
135+
}
136+
137+
/**
138+
* Runs RAG indexing immediately.
139+
*
140+
* Processes dirty posts synchronously. Use `mark-dirty` first to bootstrap
141+
* an initial index for existing content.
142+
*
143+
* ## OPTIONS
144+
*
145+
* [--batch-size=<number>]
146+
* : Number of dirty posts to process per batch.
147+
* ---
148+
* default: 50
149+
* ---
150+
*
151+
* [--all]
152+
* : Continue running batches until no dirty posts remain.
153+
*
154+
* [--schedule]
155+
* : Schedule a follow-up cron event if dirty posts remain.
156+
* ---
157+
* default: true
158+
* ---
159+
*
160+
* ## EXAMPLES
161+
*
162+
* # Process one dirty-post batch now
163+
* $ wp ai rag index
164+
*
165+
* # Process all dirty posts now
166+
* $ wp ai rag index --all
167+
*
168+
* @when after_wp_load
169+
*
170+
* @param array<int, string> $args Positional arguments.
171+
* @param array<string, mixed> $assoc_args Associative arguments.
172+
*/
173+
public function index( $args, $assoc_args ): void {
174+
unset( $args );
175+
176+
$this->ensure_available();
177+
178+
$batch_size = max( 1, min( 200, (int) Utils\get_flag_value( $assoc_args, 'batch-size', 50 ) ) );
179+
$process_all = (bool) Utils\get_flag_value( $assoc_args, 'all', false );
180+
$schedule_tail = (bool) Utils\get_flag_value( $assoc_args, 'schedule', true );
181+
182+
$totals = array(
183+
'processed' => 0,
184+
'clean' => 0,
185+
'error' => 0,
186+
'removed' => 0,
187+
'remaining' => 0,
188+
);
189+
190+
do {
191+
$stats = $this->index_manager->run_indexing_batch( $batch_size, $schedule_tail && ! $process_all );
192+
193+
if ( is_wp_error( $stats ) ) {
194+
WP_CLI::error( $stats->get_error_message() );
195+
return;
196+
}
197+
198+
foreach ( array( 'processed', 'clean', 'error', 'removed' ) as $key ) {
199+
$totals[ $key ] += $stats[ $key ];
200+
}
201+
$totals['remaining'] = $stats['remaining'];
202+
203+
WP_CLI::log(
204+
sprintf(
205+
'Processed %d post(s): %d clean, %d error, %d removed. %d dirty post(s) remain.',
206+
$stats['processed'],
207+
$stats['clean'],
208+
$stats['error'],
209+
$stats['removed'],
210+
$stats['remaining']
211+
)
212+
);
213+
} while ( $process_all && $totals['remaining'] > 0 );
214+
215+
WP_CLI::success(
216+
sprintf(
217+
'RAG indexing complete for %d post(s): %d clean, %d error, %d removed. %d dirty post(s) remain.',
218+
$totals['processed'],
219+
$totals['clean'],
220+
$totals['error'],
221+
$totals['removed'],
222+
$totals['remaining']
223+
)
224+
);
225+
}
226+
227+
/**
228+
* Schedules the RAG indexing cron event.
229+
*
230+
* ## OPTIONS
231+
*
232+
* [--delay=<seconds>]
233+
* : Delay before the event runs.
234+
* ---
235+
* default: 1
236+
* ---
237+
*
238+
* ## EXAMPLES
239+
*
240+
* $ wp ai rag schedule
241+
*
242+
* @when after_wp_load
243+
*
244+
* @param array<int, string> $args Positional arguments.
245+
* @param array<string, mixed> $assoc_args Associative arguments.
246+
*/
247+
public function schedule( $args, $assoc_args ): void {
248+
unset( $args );
249+
250+
$this->ensure_available();
251+
252+
$delay = max( 1, (int) Utils\get_flag_value( $assoc_args, 'delay', 1 ) );
253+
$this->index_manager->schedule_indexing( $delay );
254+
255+
WP_CLI::success( sprintf( 'Scheduled RAG indexing in %d second(s).', $delay ) );
256+
}
257+
258+
/**
259+
* Ensures RAG indexing can run.
260+
*
261+
* @since 1.1.0
262+
*/
263+
private function ensure_available(): void {
264+
if ( ! $this->availability->is_available() ) {
265+
WP_CLI::error( $this->availability->get_unavailable_reason() );
266+
return;
267+
}
268+
269+
if ( $this->index_manager->ensure_index_table() ) {
270+
return;
271+
}
272+
273+
WP_CLI::error( 'The RAG index table could not be created.' );
274+
}
275+
276+
/**
277+
* Parses a comma-separated ID flag.
278+
*
279+
* @since 1.1.0
280+
*
281+
* @param string $ids_flag Comma-separated IDs.
282+
* @return list<int> IDs.
283+
*/
284+
private function parse_ids_flag( string $ids_flag ): array {
285+
if ( '' === trim( $ids_flag ) ) {
286+
return array();
287+
}
288+
289+
return array_values( array_filter( array_map( 'absint', explode( ',', $ids_flag ) ) ) );
290+
}
291+
}

‎includes/Experiments/Experiments.php‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ final class Experiments {
4040
\WordPress\AI\Experiments\Summarization\Summarization::class,
4141
\WordPress\AI\Experiments\Title_Generation\Title_Generation::class,
4242
\WordPress\AI\Experiments\Comment_Moderation\Comment_Moderation::class,
43+
\WordPress\AI\Experiments\RAG_Search\RAG_Search::class,
4344
);
4445

4546
/**

0 commit comments

Comments
 (0)