mempal_runtime/ingest/
reindex.rs1use std::{
2 collections::{HashMap, HashSet},
3 path::PathBuf,
4};
5
6use thiserror::Error;
7
8use crate::core::{db::Database, types::ReindexSource};
9use crate::embed::Embedder;
10
11use super::{
12 IngestError, IngestOptions, ingest_file_with_options, normalize::CURRENT_NORMALIZE_VERSION,
13};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum ReindexMode {
17 Stale,
18 Force,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct ReindexOptions {
23 pub mode: ReindexMode,
24 pub dry_run: bool,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq, Default)]
28pub struct ReindexReport {
29 pub candidate_drawers: u64,
30 pub candidate_sources: u64,
31 pub processed_sources: u64,
32 pub reingested_files: usize,
33 pub reingested_chunks: usize,
34 pub skipped_existing_chunks: usize,
35 pub skipped_missing_sources: u64,
36 pub skipped_missing_drawers: u64,
37 pub skipped_protected_sources: u64,
38 pub skipped_protected_drawers: u64,
39 pub protecting_references: u64,
40}
41
42#[derive(Debug, Error)]
43pub enum ReindexError {
44 #[error(transparent)]
45 Db(#[from] crate::core::db::DbError),
46 #[error("failed to reindex source {source_file}")]
47 Ingest {
48 source_file: String,
49 #[source]
50 source: IngestError,
51 },
52}
53
54pub async fn reindex_sources<E: Embedder + ?Sized>(
55 db: &Database,
56 embedder: &E,
57 options: ReindexOptions,
58) -> Result<ReindexReport, ReindexError> {
59 let sources = match options.mode {
60 ReindexMode::Stale => db.reindex_sources_stale(CURRENT_NORMALIZE_VERSION)?,
61 ReindexMode::Force => db.reindex_sources_force()?,
62 };
63
64 let mut report = ReindexReport {
65 candidate_drawers: sources.iter().map(|source| source.drawer_count).sum(),
66 candidate_sources: sources.len() as u64,
67 ..ReindexReport::default()
68 };
69
70 let mut reindexable = Vec::new();
74 let mut reference_summaries = HashMap::new();
75 let mut counted_protected_sources = HashSet::new();
76 for source in sources {
77 let Some(source_file) = source.source_file.as_deref() else {
78 report.skipped_missing_sources += 1;
79 report.skipped_missing_drawers += source.drawer_count;
80 continue;
81 };
82 if !PathBuf::from(source_file).is_file() {
83 report.skipped_missing_sources += 1;
84 report.skipped_missing_drawers += source.drawer_count;
85 continue;
86 }
87
88 let source_key = (source_file.to_string(), source.wing.clone());
89 let reference_summary = if let Some(summary) = reference_summaries.get(&source_key) {
90 *summary
91 } else {
92 let summary = db.source_knowledge_reference_summary(source_file, &source.wing)?;
93 reference_summaries.insert(source_key.clone(), summary);
94 summary
95 };
96 if reference_summary.referenced_drawers > 0 {
97 report.skipped_protected_sources += 1;
98 report.skipped_protected_drawers += source.drawer_count;
99 if counted_protected_sources.insert(source_key) {
100 report.protecting_references += reference_summary.references;
101 }
102 continue;
103 }
104
105 reindexable.push(source);
106 }
107
108 if options.dry_run {
109 return Ok(report);
110 }
111
112 for source in reindexable {
113 let source_file = source
114 .source_file
115 .as_deref()
116 .expect("checked Some above")
117 .to_string();
118 let source_path = PathBuf::from(&source_file);
119 let stats = reindex_one_source(db, embedder, &source, &source_file, source_path).await?;
120 report.processed_sources += 1;
121 report.reingested_files += stats.files;
122 report.reingested_chunks += stats.chunks;
123 report.skipped_existing_chunks += stats.skipped;
124 }
125
126 Ok(report)
127}
128
129async fn reindex_one_source<E: Embedder + ?Sized>(
130 db: &Database,
131 embedder: &E,
132 source: &ReindexSource,
133 source_file: &str,
134 source_path: PathBuf,
135) -> Result<super::IngestStats, ReindexError> {
136 ingest_file_with_options(
137 db,
138 embedder,
139 &source_path,
140 &source.wing,
141 IngestOptions {
142 room: source.room.as_deref(),
143 source_root: source_path.parent(),
144 dry_run: false,
145 source_file_override: Some(source_file),
146 replace_existing_source: true,
147 replace_across_rooms: true,
148 no_strip_noise: false,
149 ..IngestOptions::default()
150 },
151 )
152 .await
153 .map_err(|source| ReindexError::Ingest {
154 source_file: source_file.to_string(),
155 source,
156 })
157}