use std::path::PathBuf;
use crate::Engine;
use crate::auth::ScopedSpace;
use crate::chunker::chunk_text;
use crate::error::{Result, SconeError};
pub(crate) const CHUNK_TARGET_BYTES: usize = 1000;
const KINDS: [&str; 4] = ["note", "file", "conversation", "observation"];
#[derive(Debug)]
pub enum IngestInput {
Note { text: String },
File { path: PathBuf },
}
#[derive(Debug)]
pub enum IngestOutcome {
Ingested { episode_id: i64, chunks: usize },
Deduplicated { episode_id: i64 },
}
impl Engine {
pub fn ingest(&mut self, space: &ScopedSpace, input: IngestInput) -> Result<IngestOutcome> {
let (kind, content, source) = match input {
IngestInput::Note { text } => ("note", text, None),
IngestInput::File { path } => {
let bytes = std::fs::read(&path)?;
let text = String::from_utf8(bytes).map_err(|_| {
SconeError::InvalidInput(format!("{} is not valid UTF-8", path.display()))
})?;
("file", text, Some(path.display().to_string()))
}
};
self.ingest_raw(space, kind, &content, source.as_deref(), None)
}
pub(crate) fn import_episode(
&mut self,
space: &ScopedSpace,
kind: &str,
content: &str,
source: Option<&str>,
created_at: Option<&str>,
) -> Result<(i64, bool)> {
match self.ingest_raw(space, kind, content, source, created_at)? {
IngestOutcome::Ingested { episode_id, .. } => Ok((episode_id, true)),
IngestOutcome::Deduplicated { episode_id } => Ok((episode_id, false)),
}
}
fn ingest_raw(
&mut self,
space: &ScopedSpace,
kind: &str,
content: &str,
source: Option<&str>,
created_at: Option<&str>,
) -> Result<IngestOutcome> {
if !KINDS.contains(&kind) {
return Err(SconeError::InvalidInput(format!(
"kind must be one of {KINDS:?}, got {kind:?}"
)));
}
if content.trim().is_empty() {
return Err(SconeError::InvalidInput("content is empty".into()));
}
self.require_writable()?;
let hash = blake3::hash(content.as_bytes()).to_hex().to_string();
let spans = chunk_text(content, self.chunk_target);
let texts: Vec<&str> = spans.iter().map(|s| &content[s.start..s.end]).collect();
let embeddings = self.embedder.embed(&texts)?;
let tx = self.conn.transaction()?;
let inserted = tx.execute(
"INSERT INTO episodes (space_id, kind, content, hash, source, created_at)
VALUES (?1, ?2, ?3, ?4, ?5,
COALESCE(?6, strftime('%Y-%m-%dT%H:%M:%fZ','now')))
ON CONFLICT (space_id, hash) DO NOTHING",
rusqlite::params![space.id(), kind, content, hash, source, created_at],
)?;
if inserted == 0 {
let episode_id = tx.query_row(
"SELECT id FROM episodes WHERE space_id = ?1 AND hash = ?2",
rusqlite::params![space.id(), hash],
|r| r.get(0),
)?;
return Ok(IngestOutcome::Deduplicated { episode_id });
}
let episode_id = tx.last_insert_rowid();
let mut chunk_ids = Vec::with_capacity(spans.len());
{
let mut stmt = tx.prepare(
"INSERT INTO chunks (episode_id, pos, start_byte, end_byte, embedding)
VALUES (?1, ?2, ?3, ?4, ?5)",
)?;
for (pos, (span, emb)) in spans.iter().zip(&embeddings).enumerate() {
let blob: Vec<u8> = emb.iter().flat_map(|f| f.to_le_bytes()).collect();
stmt.execute(rusqlite::params![
episode_id,
pos as i64,
span.start as i64,
span.end as i64,
blob
])?;
chunk_ids.push(tx.last_insert_rowid());
}
}
tx.execute(
"UPDATE spaces SET revision = revision + 1 WHERE id = ?1",
[space.id()],
)?;
tx.execute(
"INSERT OR IGNORE INTO distill_queue (episode_id) VALUES (?1)",
[episode_id],
)?;
tx.execute(
"INSERT OR IGNORE INTO meta (key, value) VALUES ('embedder_id', ?1)",
[self.embedder.id()],
)?;
tx.execute(
"INSERT OR IGNORE INTO meta (key, value) VALUES ('embedder_dim', ?1)",
[self.embedder.dim().to_string()],
)?;
tx.commit()?;
let fts_rows: Vec<(u64, u64, &str)> = chunk_ids
.iter()
.zip(&texts)
.map(|(id, text)| (*id as u64, space.id() as u64, *text))
.collect();
let vec_rows: Vec<(u64, &[f32])> = chunk_ids
.iter()
.zip(&embeddings)
.map(|(id, emb)| (*id as u64, emb.as_slice()))
.collect();
let index_result = self
.fts
.add(&fts_rows)
.and_then(|()| self.vectors.add(&vec_rows));
match index_result {
Ok(()) => self.indexes_dirty = true,
Err(_) => self.set_meta("index_dirty", "1")?,
}
Ok(IngestOutcome::Ingested {
episode_id,
chunks: spans.len(),
})
}
}
#[derive(Debug, Default)]
pub struct ScanReport {
pub ingested: usize,
pub deduplicated: usize,
pub skipped: usize,
}
impl Engine {
pub fn ingest_directory(
&mut self,
space: &ScopedSpace,
dir: &std::path::Path,
max_file_bytes: u64,
) -> Result<ScanReport> {
const SKIP_DIRS: [&str; 4] = ["node_modules", "target", ".git", "__pycache__"];
let mut report = ScanReport::default();
let mut stack = vec![dir.to_path_buf()];
while let Some(current) = stack.pop() {
for entry in std::fs::read_dir(¤t)? {
let entry = entry?;
let path = entry.path();
let name = entry.file_name();
let name = name.to_string_lossy();
if name.starts_with('.') {
continue;
}
let file_type = entry.file_type()?;
if file_type.is_dir() {
if !SKIP_DIRS.contains(&name.as_ref()) {
stack.push(path);
}
continue;
}
if !file_type.is_file() {
continue;
}
if entry.metadata()?.len() > max_file_bytes {
report.skipped += 1;
continue;
}
match self.ingest(space, IngestInput::File { path }) {
Ok(IngestOutcome::Ingested { .. }) => report.ingested += 1,
Ok(IngestOutcome::Deduplicated { .. }) => report.deduplicated += 1,
Err(SconeError::InvalidInput(_)) => report.skipped += 1,
Err(other) => return Err(other),
}
}
}
Ok(report)
}
}