Skip to main content

mempal_runtime/ingest/
reindex.rs

1use 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    // P117: the source-file existence scan runs in dry-run too, so the
71    // report is a real feasibility statement (sources ingested without an
72    // on-disk file — e.g. MCP-ingested content — are not reindexable).
73    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}