1use anyhow::{bail, Result};
2use rusqlite::{params, Connection};
3
4use crate::retrieval::embedding::{
5 embedding_provider_status, EmbeddingBackfillTarget, EmbeddingProviderStatus,
6};
7
8#[derive(Debug, Clone, PartialEq)]
9pub struct ActiveEmbeddingCoverage {
10 pub embedded: i64,
11 pub total: i64,
12 pub percent: f64,
13 pub mixed_profile_count: i64,
14}
15
16#[derive(Debug, Clone, PartialEq)]
17pub struct InactiveEmbeddingPruneReport {
18 pub pruned: i64,
19 pub active_model: String,
20 pub active_dimensions: usize,
21 pub coverage: ActiveEmbeddingCoverage,
22}
23
24pub fn active_embedding_coverage(conn: &Connection) -> Result<ActiveEmbeddingCoverage> {
25 let status = embedding_provider_status()?;
26 active_embedding_coverage_for_status(conn, &status)
27}
28
29pub fn active_embedding_coverage_for_status(
30 conn: &Connection,
31 status: &EmbeddingProviderStatus,
32) -> Result<ActiveEmbeddingCoverage> {
33 if !super::table_exists(conn, "memories")? {
34 return Ok(ActiveEmbeddingCoverage {
35 embedded: 0,
36 total: 0,
37 percent: 0.0,
38 mixed_profile_count: 0,
39 });
40 }
41 let total = searchable_memory_count(conn)?;
42 if status.disabled || !super::table_exists(conn, "memory_embeddings")? {
43 return Ok(ActiveEmbeddingCoverage {
44 embedded: 0,
45 total,
46 percent: percent(0, total),
47 mixed_profile_count: 0,
48 });
49 }
50 let Some(model) = status.active_model_id.as_deref() else {
51 return Ok(ActiveEmbeddingCoverage {
52 embedded: 0,
53 total,
54 percent: percent(0, total),
55 mixed_profile_count: embedding_profile_count(conn)?,
56 });
57 };
58 let embedded = match status.active_dimensions {
59 Some(dimensions) => conn.query_row(
60 "SELECT COUNT(DISTINCT m.id)
61 FROM memories m
62 JOIN memory_embeddings e ON e.memory_id = m.id
63 WHERE m.status IN ('active', 'stale', 'archived')
64 AND e.model = ?1
65 AND e.dimensions = ?2",
66 params![model, dimensions as i64],
67 |row| row.get(0),
68 )?,
69 None => conn.query_row(
70 "SELECT COUNT(DISTINCT m.id)
71 FROM memories m
72 JOIN memory_embeddings e ON e.memory_id = m.id
73 WHERE m.status IN ('active', 'stale', 'archived')
74 AND e.model = ?1",
75 [model],
76 |row| row.get(0),
77 )?,
78 };
79 Ok(ActiveEmbeddingCoverage {
80 embedded,
81 total,
82 percent: percent(embedded, total),
83 mixed_profile_count: embedding_profile_count(conn)?,
84 })
85}
86
87fn searchable_memory_count(conn: &Connection) -> Result<i64> {
88 Ok(conn.query_row(
89 "SELECT COUNT(*) FROM memories WHERE status IN ('active', 'stale', 'archived')",
90 [],
91 |row| row.get(0),
92 )?)
93}
94
95fn embedding_profile_count(conn: &Connection) -> Result<i64> {
96 if !super::table_exists(conn, "memory_embeddings")? {
97 return Ok(0);
98 }
99 Ok(conn.query_row(
100 "SELECT COUNT(*)
101 FROM (
102 SELECT model, dimensions
103 FROM memory_embeddings
104 GROUP BY model, dimensions
105 )",
106 [],
107 |row| row.get(0),
108 )?)
109}
110
111pub fn prune_inactive_memory_embeddings(
112 conn: &Connection,
113 target: &EmbeddingBackfillTarget,
114) -> Result<InactiveEmbeddingPruneReport> {
115 if !super::table_exists(conn, "memories")? || !super::table_exists(conn, "memory_embeddings")? {
116 return Ok(InactiveEmbeddingPruneReport {
117 pruned: 0,
118 active_model: target.model.clone(),
119 active_dimensions: target.dimensions,
120 coverage: ActiveEmbeddingCoverage {
121 embedded: 0,
122 total: 0,
123 percent: 0.0,
124 mixed_profile_count: 0,
125 },
126 });
127 }
128 let coverage = embedding_coverage_for_target(conn, target)?;
129 if coverage.embedded < coverage.total {
130 bail!(
131 "refusing to prune inactive embedding profiles before active coverage reaches 100%: {}/{} ({:.1}%)",
132 coverage.embedded,
133 coverage.total,
134 coverage.percent
135 );
136 }
137 let stale_or_missing = pending_reindex_count_for_target(conn, target)?;
138 if stale_or_missing > 0 {
139 bail!(
140 "refusing to prune inactive embedding profiles while active profile has {stale_or_missing} missing or stale rows; run embedding backfill without --limit before pruning"
141 );
142 }
143 let pruned = conn.execute(
144 "DELETE FROM memory_embeddings
145 WHERE rowid IN (
146 SELECT e.rowid
147 FROM memory_embeddings e
148 JOIN memories m ON m.id = e.memory_id
149 WHERE m.status IN ('active', 'stale', 'archived')
150 AND NOT (e.model = ?1 AND e.dimensions = ?2)
151 )",
152 params![target.model.as_str(), target.dimensions as i64],
153 )? as i64;
154 Ok(InactiveEmbeddingPruneReport {
155 pruned,
156 active_model: target.model.clone(),
157 active_dimensions: target.dimensions,
158 coverage,
159 })
160}
161
162fn embedding_coverage_for_target(
163 conn: &Connection,
164 target: &EmbeddingBackfillTarget,
165) -> Result<ActiveEmbeddingCoverage> {
166 let total = searchable_memory_count(conn)?;
167 let embedded = conn.query_row(
168 "SELECT COUNT(DISTINCT m.id)
169 FROM memories m
170 JOIN memory_embeddings e ON e.memory_id = m.id
171 WHERE m.status IN ('active', 'stale', 'archived')
172 AND e.model = ?1
173 AND e.dimensions = ?2",
174 params![target.model.as_str(), target.dimensions as i64],
175 |row| row.get(0),
176 )?;
177 Ok(ActiveEmbeddingCoverage {
178 embedded,
179 total,
180 percent: percent(embedded, total),
181 mixed_profile_count: embedding_profile_count(conn)?,
182 })
183}
184
185fn pending_reindex_count_for_target(
186 conn: &Connection,
187 target: &EmbeddingBackfillTarget,
188) -> Result<i64> {
189 Ok(conn.query_row(
190 "SELECT COUNT(*)
191 FROM memories m
192 LEFT JOIN memory_embeddings e
193 ON e.memory_id = m.id
194 AND e.model = ?1
195 AND e.dimensions = ?2
196 WHERE (e.memory_id IS NULL
197 OR e.updated_at_epoch < m.updated_at_epoch)
198 AND m.status IN ('active', 'stale', 'archived')",
199 params![target.model.as_str(), target.dimensions as i64],
200 |row| row.get(0),
201 )?)
202}
203
204fn percent(numerator: i64, denominator: i64) -> f64 {
205 if denominator <= 0 {
206 0.0
207 } else {
208 (numerator as f64 * 100.0) / denominator as f64
209 }
210}