1pub mod auth;
7pub mod chunker;
8mod db;
9pub mod distill;
10pub mod embed;
11mod error;
12pub mod index;
13mod ingest;
14pub mod llm;
15mod portability;
16pub mod profile;
17mod recall;
18pub mod rerank;
19mod tags;
20pub mod timeparse;
21
22use std::path::{Path, PathBuf};
23
24use rusqlite::Connection;
25
26pub use distill::{ApplyReport, DistillReport, ProvenanceItem};
27pub use error::{Result, SconeError};
28pub use ingest::{IngestInput, IngestOutcome, ScanReport};
29pub use portability::ImportReport;
30pub use profile::Profile;
31pub use recall::{ContextPack, FactItem, RecallItem, RecallOpts};
32
33#[derive(Debug)]
34pub struct DoctorReport {
35 pub episodes: usize,
36 pub chunks: usize,
37 pub reembedded: usize,
38}
39
40impl std::fmt::Debug for Engine {
41 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42 f.debug_struct("Engine")
43 .field("data_dir", &self.data_dir)
44 .field("embedder", &self.embedder.id())
45 .finish_non_exhaustive()
46 }
47}
48
49#[derive(Debug)]
50pub struct SpaceStatus {
51 pub name: String,
52 pub episodes: i64,
53 pub chunks: i64,
54 pub revision: i64,
55}
56
57#[derive(Debug)]
58pub struct StatusReport {
59 pub read_only: bool,
60 pub spaces: Vec<SpaceStatus>,
61 pub embedder_id: String,
62 pub embedder_dim: usize,
63 pub index_dirty: bool,
64 pub pending_distill: i64,
65 pub failed_distill: i64,
66 pub llm_id: Option<String>,
67}
68
69pub struct Engine {
70 conn: Connection,
71 data_dir: PathBuf,
72 embedder: Box<dyn embed::EmbeddingProvider>,
73 llm: Option<Box<dyn llm::LlmProvider>>,
74 reranker: Option<Box<dyn rerank::Reranker>>,
75 fts: index::fts::FtsIndex,
76 vectors: index::vectors::VectorIndex,
77 indexes_dirty: bool,
79 chunk_target: usize,
82}
83
84impl Drop for Engine {
85 fn drop(&mut self) {
86 let _ = self.flush_indexes();
89 }
90}
91
92impl Engine {
93 pub fn open(data_dir: &Path, embedder: Box<dyn embed::EmbeddingProvider>) -> Result<Engine> {
99 Self::open_inner(data_dir, embedder, false)
100 }
101
102 pub fn open_for_repair(
104 data_dir: &Path,
105 embedder: Box<dyn embed::EmbeddingProvider>,
106 ) -> Result<Engine> {
107 Self::open_inner(data_dir, embedder, true)
108 }
109
110 fn open_inner(
111 data_dir: &Path,
112 embedder: Box<dyn embed::EmbeddingProvider>,
113 repair: bool,
114 ) -> Result<Engine> {
115 std::fs::create_dir_all(data_dir)?;
116 let conn = db::open(&data_dir.join("scone.db"))?;
117 if !repair {
118 let pinned: Option<String> = match conn.query_row(
119 "SELECT value FROM meta WHERE key = 'embedder_id'",
120 [],
121 |r| r.get(0),
122 ) {
123 Ok(v) => Some(v),
124 Err(rusqlite::Error::QueryReturnedNoRows) => None,
125 Err(e) => return Err(SconeError::Db(e)),
126 };
127 if let Some(pinned) = pinned
128 && !(pinned == embedder.id() && Self::pinned_dim(&conn)? == Some(embedder.dim()))
129 {
130 return Err(SconeError::Index(format!(
131 "store is pinned to embedder {pinned}; opening with {} ({} dims) requires `scone doctor --rebuild`",
132 embedder.id(),
133 embedder.dim()
134 )));
135 }
136 }
137 let vectors_dir = data_dir.join("vectors");
138 let vectors = if repair {
139 index::vectors::VectorIndex::open_or_reset(&vectors_dir, embedder.dim())?
140 } else {
141 index::vectors::VectorIndex::open(&vectors_dir, embedder.dim())?
142 };
143 let fts = index::fts::FtsIndex::open(&data_dir.join("fts"))?;
144 let mut engine = Engine {
145 conn,
146 data_dir: data_dir.to_path_buf(),
147 embedder,
148 llm: None,
149 reranker: None,
150 fts,
151 vectors,
152 indexes_dirty: false,
153 chunk_target: ingest::CHUNK_TARGET_BYTES,
154 };
155 if engine.fts.writable() {
156 engine.catch_up_indexes()?;
157 }
158 Ok(engine)
159 }
160
161 pub fn set_chunk_target(&mut self, bytes: usize) {
163 self.chunk_target = bytes.max(64);
164 }
165
166 pub fn is_read_only(&self) -> bool {
170 !self.fts.writable()
171 }
172
173 fn require_writable(&self) -> Result<()> {
174 if self.is_read_only() {
175 return Err(SconeError::InvalidInput(
176 "this store is read-only: another scone process holds the write lock".into(),
177 ));
178 }
179 Ok(())
180 }
181
182 fn catch_up_indexes(&mut self) -> Result<()> {
186 let indexed_max: i64 = self
187 .get_meta("indexed_max")?
188 .and_then(|v| v.parse().ok())
189 .unwrap_or(0);
190 let chunk_max: i64 =
191 self.conn
192 .query_row("SELECT coalesce(max(id), 0) FROM chunks", [], |r| r.get(0))?;
193 if chunk_max <= indexed_max {
194 return Ok(());
195 }
196 let rows: Vec<(i64, i64, String, Option<Vec<u8>>)> = {
197 let mut stmt = self.conn.prepare(
198 "SELECT c.id, e.space_id, e.content, c.start_byte, c.end_byte, c.embedding
199 FROM chunks c JOIN episodes e ON e.id = c.episode_id
200 WHERE c.id > ?1 ORDER BY c.id",
201 )?;
202 let mapped = stmt.query_map([indexed_max], |r| {
203 let content: String = r.get(2)?;
204 let start: i64 = r.get(3)?;
205 let end: i64 = r.get(4)?;
206 let text = content
207 .get(start as usize..end as usize)
208 .unwrap_or_default()
209 .to_owned();
210 Ok((r.get(0)?, r.get(1)?, text, r.get(5)?))
211 })?;
212 mapped.collect::<std::result::Result<Vec<_>, _>>()?
213 };
214 let dim = self.embedder.dim();
215 let fts_rows: Vec<(u64, u64, &str)> = rows
216 .iter()
217 .map(|(id, space, text, _)| (*id as u64, *space as u64, text.as_str()))
218 .collect();
219 self.fts.add(&fts_rows)?;
220 let mut vec_rows: Vec<(u64, Vec<f32>)> = Vec::new();
221 for (id, _, _, blob) in &rows {
222 if let Some(b) = blob
223 && b.len() == dim * 4
224 {
225 let v = b
226 .as_chunks::<4>()
227 .0
228 .iter()
229 .map(|c| f32::from_le_bytes(*c))
230 .collect();
231 vec_rows.push((*id as u64, v));
232 }
233 }
234 let borrowed: Vec<(u64, &[f32])> =
235 vec_rows.iter().map(|(id, v)| (*id, v.as_slice())).collect();
236 self.vectors.add(&borrowed)?;
237 self.indexes_dirty = true;
238 self.flush_indexes()
239 }
240
241 pub(crate) fn flush_indexes(&mut self) -> Result<()> {
244 if !self.indexes_dirty || self.is_read_only() {
245 return Ok(());
246 }
247 self.fts.commit()?;
248 self.vectors.flush()?;
249 let chunk_max: i64 =
250 self.conn
251 .query_row("SELECT coalesce(max(id), 0) FROM chunks", [], |r| r.get(0))?;
252 self.set_meta("indexed_max", &chunk_max.to_string())?;
253 self.indexes_dirty = false;
254 Ok(())
255 }
256
257 pub fn set_llm(&mut self, llm: Option<Box<dyn llm::LlmProvider>>) {
260 self.llm = llm;
261 }
262
263 pub fn has_llm(&self) -> bool {
264 self.llm.is_some()
265 }
266
267 pub fn set_reranker(&mut self, reranker: Option<Box<dyn rerank::Reranker>>) {
269 self.reranker = reranker;
270 }
271
272 pub fn llm_answer(&self, question: &str, context: &str) -> Result<String> {
274 match &self.llm {
275 Some(llm) => llm.answer(question, context),
276 None => Err(SconeError::Llm("no LLM configured".into())),
277 }
278 }
279
280 pub fn llm_answer_with_system(
282 &self,
283 system: &str,
284 question: &str,
285 context: &str,
286 ) -> Result<String> {
287 match &self.llm {
288 Some(llm) => llm.answer_with_system(system, question, context),
289 None => Err(SconeError::Llm("no LLM configured".into())),
290 }
291 }
292
293 fn pinned_dim(conn: &Connection) -> Result<Option<usize>> {
294 match conn.query_row(
295 "SELECT value FROM meta WHERE key = 'embedder_dim'",
296 [],
297 |r| r.get::<_, String>(0),
298 ) {
299 Ok(v) => Ok(v.parse().ok()),
300 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
301 Err(e) => Err(SconeError::Db(e)),
302 }
303 }
304
305 pub fn doctor_rebuild(&mut self) -> Result<DoctorReport> {
308 self.require_writable()?;
309 self.fts.wipe()?;
310 self.vectors.wipe()?;
311 let rows: Vec<(i64, i64, String, Option<Vec<u8>>)> = {
314 let mut stmt = self.conn.prepare(
315 "SELECT c.id, e.space_id, e.content, c.start_byte, c.end_byte, c.embedding
316 FROM chunks c JOIN episodes e ON e.id = c.episode_id
317 ORDER BY c.id",
318 )?;
319 let mapped = stmt.query_map([], |r| {
320 let content: String = r.get(2)?;
321 let start: i64 = r.get(3)?;
322 let end: i64 = r.get(4)?;
323 let text = content
324 .get(start as usize..end as usize)
325 .unwrap_or_default()
326 .to_owned();
327 Ok((r.get(0)?, r.get(1)?, text, r.get(5)?))
328 })?;
329 mapped.collect::<std::result::Result<Vec<_>, _>>()?
330 };
331 let dim = self.embedder.dim();
332 let mut reembedded = 0usize;
333 let mut fts_rows = Vec::with_capacity(rows.len());
334 let mut vec_data = Vec::with_capacity(rows.len());
335 for (chunk_id, space_id, text, blob) in &rows {
336 let vector = match blob {
337 Some(b) if b.len() == dim * 4 => b
338 .as_chunks::<4>()
339 .0
340 .iter()
341 .map(|c| f32::from_le_bytes(*c))
342 .collect::<Vec<f32>>(),
343 _ => {
344 let embedded = self.embedder.embed(&[text.as_str()])?;
345 let v = embedded
346 .into_iter()
347 .next()
348 .ok_or_else(|| SconeError::Embed("embedder returned no vector".into()))?;
349 let new_blob: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
350 self.conn.execute(
351 "UPDATE chunks SET embedding = ?1 WHERE id = ?2",
352 rusqlite::params![new_blob, chunk_id],
353 )?;
354 reembedded += 1;
355 v
356 }
357 };
358 fts_rows.push((*chunk_id as u64, *space_id as u64, text.as_str()));
359 vec_data.push((*chunk_id as u64, vector));
360 }
361 self.fts.add(&fts_rows)?;
362 let vec_rows: Vec<(u64, &[f32])> =
363 vec_data.iter().map(|(id, v)| (*id, v.as_slice())).collect();
364 self.vectors.add(&vec_rows)?;
365 self.indexes_dirty = true;
366 self.flush_indexes()?;
367 self.set_meta("embedder_id", self.embedder.id())?;
368 self.set_meta("embedder_dim", &dim.to_string())?;
369 self.set_meta("index_dirty", "0")?;
370 let episodes: i64 = self
371 .conn
372 .query_row("SELECT count(*) FROM episodes", [], |r| r.get(0))?;
373 Ok(DoctorReport {
374 episodes: episodes as usize,
375 chunks: rows.len(),
376 reembedded,
377 })
378 }
379
380 pub(crate) fn set_meta(&self, key: &str, value: &str) -> Result<()> {
381 self.conn.execute(
382 "INSERT INTO meta (key, value) VALUES (?1, ?2)
383 ON CONFLICT (key) DO UPDATE SET value = excluded.value",
384 [key, value],
385 )?;
386 Ok(())
387 }
388
389 pub(crate) fn get_meta(&self, key: &str) -> Result<Option<String>> {
390 match self
391 .conn
392 .query_row("SELECT value FROM meta WHERE key = ?1", [key], |r| r.get(0))
393 {
394 Ok(v) => Ok(Some(v)),
395 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
396 Err(e) => Err(SconeError::Db(e)),
397 }
398 }
399
400 pub fn space_revision(&self, space: &auth::ScopedSpace) -> Result<i64> {
401 Ok(self.conn.query_row(
402 "SELECT revision FROM spaces WHERE id = ?1",
403 [space.id()],
404 |r| r.get(0),
405 )?)
406 }
407
408 pub fn status(&self) -> Result<StatusReport> {
411 let mut stmt = self.conn.prepare(
412 "SELECT s.name, s.revision,
413 (SELECT count(*) FROM episodes e WHERE e.space_id = s.id),
414 (SELECT count(*) FROM chunks c JOIN episodes e ON e.id = c.episode_id
415 WHERE e.space_id = s.id)
416 FROM spaces s ORDER BY s.name",
417 )?;
418 let spaces = stmt
419 .query_map([], |r| {
420 Ok(SpaceStatus {
421 name: r.get(0)?,
422 revision: r.get(1)?,
423 episodes: r.get(2)?,
424 chunks: r.get(3)?,
425 })
426 })?
427 .collect::<std::result::Result<Vec<_>, _>>()?;
428 let (pending_distill, failed_distill) = self.conn.query_row(
429 "SELECT sum(state = 'pending'), sum(state = 'failed') FROM distill_queue",
430 [],
431 |r| {
432 Ok((
433 r.get::<_, Option<i64>>(0)?.unwrap_or(0),
434 r.get::<_, Option<i64>>(1)?.unwrap_or(0),
435 ))
436 },
437 )?;
438 Ok(StatusReport {
439 read_only: self.is_read_only(),
440 spaces,
441 embedder_id: self.embedder.id().to_owned(),
442 embedder_dim: self.embedder.dim(),
443 index_dirty: self.get_meta("index_dirty")?.as_deref() == Some("1"),
444 pending_distill,
445 failed_distill,
446 llm_id: self.llm.as_ref().map(|l| l.id().to_owned()),
447 })
448 }
449
450 pub fn episode_content(&self, episode_id: i64) -> Result<(String, String)> {
452 self.conn
453 .query_row(
454 "SELECT content, kind FROM episodes WHERE id = ?1",
455 [episode_id],
456 |r| Ok((r.get(0)?, r.get(1)?)),
457 )
458 .map_err(|e| match e {
459 rusqlite::Error::QueryReturnedNoRows => {
460 SconeError::NotFound(format!("episode {episode_id}"))
461 }
462 other => SconeError::Db(other),
463 })
464 }
465
466 pub fn data_dir(&self) -> &Path {
467 &self.data_dir
468 }
469
470 pub(crate) fn conn_mut(&mut self) -> &mut Connection {
471 &mut self.conn
472 }
473
474 pub fn schema_version(&self) -> Result<i64> {
475 let v: String = self.conn.query_row(
476 "SELECT value FROM meta WHERE key = 'schema_version'",
477 [],
478 |r| r.get(0),
479 )?;
480 v.parse()
481 .map_err(|_| SconeError::Index("schema_version is not a number".into()))
482 }
483}