1use anyhow::{bail, Result};
2use rusqlite::{params, Connection};
3
4use crate::retrieval::embedding::{
5 configured_backfill_target_with_fallback_cache, embedding_provider_status,
6 embedding_provider_status_without_probe, resolve_embedding_config, EmbeddingBackfillTarget,
7 EmbeddingConfig, EmbeddingFallbackCache, EmbeddingProviderStatus,
8};
9#[cfg(feature = "local-onnx")]
10use crate::retrieval::embedding::{with_configured_model_read_lock, EmbeddingProvider};
11
12#[derive(Debug, Clone, PartialEq)]
13pub struct ActiveEmbeddingCoverage {
14 pub embedded: i64,
15 pub total: i64,
16 pub percent: f64,
17 pub mixed_profile_count: i64,
18}
19
20#[derive(Debug, Clone, PartialEq)]
21pub struct InactiveEmbeddingPruneReport {
22 pub pruned: i64,
23 pub active_model: String,
24 pub active_dimensions: usize,
25 pub coverage: ActiveEmbeddingCoverage,
26}
27
28pub fn active_embedding_coverage(conn: &Connection) -> Result<ActiveEmbeddingCoverage> {
29 let status = embedding_provider_status()?;
30 active_embedding_coverage_for_status(conn, &status)
31}
32
33pub fn active_embedding_coverage_for_status(
34 conn: &Connection,
35 status: &EmbeddingProviderStatus,
36) -> Result<ActiveEmbeddingCoverage> {
37 if !super::table_exists(conn, "memories")? {
38 return Ok(ActiveEmbeddingCoverage {
39 embedded: 0,
40 total: 0,
41 percent: 0.0,
42 mixed_profile_count: 0,
43 });
44 }
45 let total = searchable_memory_count(conn)?;
46 if status.disabled || !super::table_exists(conn, "memory_embeddings")? {
47 return Ok(ActiveEmbeddingCoverage {
48 embedded: 0,
49 total,
50 percent: percent(0, total),
51 mixed_profile_count: 0,
52 });
53 }
54 let Some(model) = status.active_model_id.as_deref() else {
55 return Ok(ActiveEmbeddingCoverage {
56 embedded: 0,
57 total,
58 percent: percent(0, total),
59 mixed_profile_count: embedding_profile_count(conn)?,
60 });
61 };
62 let embedded = match status.active_dimensions {
63 Some(dimensions) => conn.query_row(
64 "SELECT COUNT(DISTINCT m.id)
65 FROM memories m
66 JOIN memory_embeddings e ON e.memory_id = m.id
67 WHERE m.status IN ('active', 'stale', 'archived')
68 AND e.model = ?1
69 AND e.dimensions = ?2",
70 params![model, dimensions as i64],
71 |row| row.get(0),
72 )?,
73 None => conn.query_row(
74 "SELECT COUNT(DISTINCT m.id)
75 FROM memories m
76 JOIN memory_embeddings e ON e.memory_id = m.id
77 WHERE m.status IN ('active', 'stale', 'archived')
78 AND e.model = ?1",
79 [model],
80 |row| row.get(0),
81 )?,
82 };
83 Ok(ActiveEmbeddingCoverage {
84 embedded,
85 total,
86 percent: percent(embedded, total),
87 mixed_profile_count: embedding_profile_count(conn)?,
88 })
89}
90
91fn searchable_memory_count(conn: &Connection) -> Result<i64> {
92 Ok(conn.query_row(
93 "SELECT COUNT(*) FROM memories WHERE status IN ('active', 'stale', 'archived')",
94 [],
95 |row| row.get(0),
96 )?)
97}
98
99fn embedding_profile_count(conn: &Connection) -> Result<i64> {
100 if !super::table_exists(conn, "memory_embeddings")? {
101 return Ok(0);
102 }
103 Ok(conn.query_row(
104 "SELECT COUNT(*)
105 FROM (
106 SELECT model, dimensions
107 FROM memory_embeddings
108 GROUP BY model, dimensions
109 )",
110 [],
111 |row| row.get(0),
112 )?)
113}
114
115pub fn prune_inactive_memory_embeddings(
116 conn: &Connection,
117 target: &EmbeddingBackfillTarget,
118) -> Result<InactiveEmbeddingPruneReport> {
119 let config = resolve_embedding_config()?;
120 #[cfg(feature = "local-onnx")]
121 let pin_local_model_state = match config.provider {
122 EmbeddingProvider::Local => true,
123 EmbeddingProvider::Auto => {
124 embedding_provider_status_without_probe()?.active_provider
125 != EmbeddingProvider::OpenAi.label()
126 }
127 EmbeddingProvider::FeatureHash | EmbeddingProvider::OpenAi | EmbeddingProvider::Off => {
128 false
129 }
130 };
131 #[cfg(feature = "local-onnx")]
132 if pin_local_model_state {
133 return with_configured_model_read_lock(&config, || {
134 prune_inactive_memory_embeddings_pinned(conn, target, &config)
135 });
136 }
137 prune_inactive_memory_embeddings_pinned(conn, target, &config)
138}
139
140fn prune_inactive_memory_embeddings_pinned(
141 conn: &Connection,
142 target: &EmbeddingBackfillTarget,
143 pinned_config: &EmbeddingConfig,
144) -> Result<InactiveEmbeddingPruneReport> {
145 ensure_current_prune_target(target, pinned_config)?;
146 if !super::table_exists(conn, "memories")? || !super::table_exists(conn, "memory_embeddings")? {
147 return Ok(InactiveEmbeddingPruneReport {
148 pruned: 0,
149 active_model: target.model.clone(),
150 active_dimensions: target.dimensions,
151 coverage: ActiveEmbeddingCoverage {
152 embedded: 0,
153 total: 0,
154 percent: 0.0,
155 mixed_profile_count: 0,
156 },
157 });
158 }
159 let coverage = active_embedding_coverage_for_target(conn, target)?;
160 if coverage.embedded < coverage.total {
161 bail!(
162 "refusing to prune inactive embedding profiles before active coverage reaches 100%: {}/{} ({:.1}%)",
163 coverage.embedded,
164 coverage.total,
165 coverage.percent
166 );
167 }
168 let stale_or_missing = super::pending_memory_embedding_reindex_count_for_target(conn, target)?;
169 if stale_or_missing > 0 {
170 bail!(
171 "refusing to prune inactive embedding profiles while active profile has {stale_or_missing} missing or stale rows; run embedding backfill without --limit before pruning"
172 );
173 }
174 let pruned = conn.execute(
175 "DELETE FROM memory_embeddings
176 WHERE rowid IN (
177 SELECT e.rowid
178 FROM memory_embeddings e
179 JOIN memories m ON m.id = e.memory_id
180 WHERE m.status IN ('active', 'stale', 'archived')
181 AND NOT (e.model = ?1 AND e.dimensions = ?2)
182 )",
183 params![target.model.as_str(), target.dimensions as i64],
184 )? as i64;
185 super::vec_index::sync_vec_keep_only_profile(conn, target.model.as_str(), target.dimensions)?;
186 Ok(InactiveEmbeddingPruneReport {
187 pruned,
188 active_model: target.model.clone(),
189 active_dimensions: target.dimensions,
190 coverage,
191 })
192}
193
194pub fn active_embedding_coverage_for_target(
195 conn: &Connection,
196 target: &EmbeddingBackfillTarget,
197) -> Result<ActiveEmbeddingCoverage> {
198 if !super::table_exists(conn, "memories")? {
199 return Ok(ActiveEmbeddingCoverage {
200 embedded: 0,
201 total: 0,
202 percent: 0.0,
203 mixed_profile_count: 0,
204 });
205 }
206 let total = searchable_memory_count(conn)?;
207 if target.dimensions == 0 || !super::table_exists(conn, "memory_embeddings")? {
208 return Ok(ActiveEmbeddingCoverage {
209 embedded: 0,
210 total,
211 percent: percent(0, total),
212 mixed_profile_count: embedding_profile_count(conn)?,
213 });
214 }
215 let embedded = conn.query_row(
216 "SELECT COUNT(DISTINCT m.id)
217 FROM memories m
218 JOIN memory_embeddings e ON e.memory_id = m.id
219 WHERE m.status IN ('active', 'stale', 'archived')
220 AND e.model = ?1
221 AND e.dimensions = ?2",
222 params![target.model.as_str(), target.dimensions as i64],
223 |row| row.get(0),
224 )?;
225 Ok(ActiveEmbeddingCoverage {
226 embedded,
227 total,
228 percent: percent(embedded, total),
229 mixed_profile_count: embedding_profile_count(conn)?,
230 })
231}
232
233fn ensure_current_prune_target(
234 target: &EmbeddingBackfillTarget,
235 pinned_config: &EmbeddingConfig,
236) -> Result<()> {
237 let config_before = resolve_embedding_config()?;
238 if &config_before != pinned_config {
239 bail!(
240 "refusing to prune embedding profiles because the embedding configuration changed before the model-state pin was acquired"
241 );
242 }
243 let status_before = embedding_provider_status_without_probe()?;
244 if status_before.disabled {
245 bail!("cannot prune embedding profiles while embedding provider is off");
246 }
247 if status_before.degraded {
248 bail!(
249 "refusing to prune embedding profiles while the current provider is degraded: {}",
250 status_before
251 .degradation_reason
252 .as_deref()
253 .or(status_before.unavailable_reason.as_deref())
254 .unwrap_or("unknown provider degradation")
255 );
256 }
257
258 let mut fallback_cache = EmbeddingFallbackCache::default();
259 let current = configured_backfill_target_with_fallback_cache(&mut fallback_cache)?;
260 let config_after = resolve_embedding_config()?;
261 let status_after = embedding_provider_status_without_probe()?;
262 if config_after != config_before || status_after != status_before {
263 bail!(
264 "refusing to prune embedding profiles because the embedding configuration or active profile changed while resolving the current target"
265 );
266 }
267 if let Some(fallback_target) = fallback_cache.call_failure_fallback_target() {
268 bail!(
269 "refusing to prune embedding profiles after typed provider fallback selected model={} dimensions={}",
270 fallback_target.model,
271 fallback_target.dimensions
272 );
273 }
274 if ¤t != target {
275 bail!(
276 "refusing to prune stale target model={} dimensions={}; current embedding profile is model={} dimensions={}",
277 target.model,
278 target.dimensions,
279 current.model,
280 current.dimensions
281 );
282 }
283 Ok(())
284}
285
286fn percent(numerator: i64, denominator: i64) -> f64 {
287 if denominator <= 0 {
288 0.0
289 } else {
290 (numerator as f64 * 100.0) / denominator as f64
291 }
292}