context_harness/
document_ingestor.rs

1//! Shared per-item document ingestion used by connector-driven full sync.
2//!
3//! Discovery, filtering, progress and checkpoints remain owned by `ingest`.
4
5use crate::app_store::AppStore;
6use crate::chunk::chunk_text;
7use crate::config::Config;
8use crate::embed_cmd;
9use crate::extract;
10use crate::models::SourceItem;
11use anyhow::Result;
12
13const DEFAULT_MAX_EXTRACT_BYTES: u64 = 50_000_000;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub(crate) struct IngestedDocument {
17    pub(crate) chunks_written: u64,
18    pub(crate) embeddings_written: u64,
19    pub(crate) embeddings_pending: u64,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub(crate) enum DocumentIngestOutcome {
24    Ingested(IngestedDocument),
25    ExtractionSkipped,
26}
27
28pub(crate) enum DocumentPreparationError {
29    TooLarge { size: usize, limit: u64 },
30    Extraction(extract::ExtractError),
31}
32
33pub(crate) fn prepare_source_item(
34    item: &mut SourceItem,
35    max_extract_bytes: u64,
36) -> std::result::Result<(), DocumentPreparationError> {
37    let Some(bytes) = item.raw_bytes.as_ref() else {
38        return Ok(());
39    };
40    if bytes.len() as u64 > max_extract_bytes {
41        return Err(DocumentPreparationError::TooLarge {
42            size: bytes.len(),
43            limit: max_extract_bytes,
44        });
45    }
46    let body = extract::extract_text(bytes, &item.content_type)
47        .map_err(DocumentPreparationError::Extraction)?;
48    item.body = body;
49    item.raw_bytes = None;
50    Ok(())
51}
52
53/// Executes the existing canonical write path for one already-discovered item.
54pub(crate) struct DocumentIngestor<'a, S> {
55    config: &'a Config,
56    store: &'a S,
57    max_extract_bytes: u64,
58}
59
60impl<'a, S: AppStore> DocumentIngestor<'a, S> {
61    pub(crate) fn new(config: &'a Config, store: &'a S, source_label: &str) -> Self {
62        Self {
63            config,
64            store,
65            max_extract_bytes: max_extract_bytes_for_source(config, source_label),
66        }
67    }
68
69    pub(crate) async fn ingest(&self, item: &mut SourceItem) -> Result<DocumentIngestOutcome> {
70        if let Err(error) = prepare_source_item(item, self.max_extract_bytes) {
71            match error {
72                DocumentPreparationError::TooLarge { size, limit } => eprintln!(
73                    "Warning: skipping {} (size {} > max_extract_bytes {})",
74                    item.source_id, size, limit
75                ),
76                DocumentPreparationError::Extraction(error) => {
77                    eprintln!(
78                        "Warning: extraction failed for {}: {}",
79                        item.source_id, error
80                    );
81                }
82            }
83            return Ok(DocumentIngestOutcome::ExtractionSkipped);
84        }
85
86        let document_id = self.store.upsert_source_item(item).await?;
87        let chunks = chunk_text(&document_id, &item.body, self.config.chunking.max_tokens);
88        let chunks_written = chunks.len() as u64;
89        self.store
90            .replace_chunks(&document_id, &chunks, None)
91            .await?;
92        let (embeddings_written, embeddings_pending) =
93            embed_cmd::embed_chunks_inline(self.config, self.store, &chunks).await;
94
95        Ok(DocumentIngestOutcome::Ingested(IngestedDocument {
96            chunks_written,
97            embeddings_written,
98            embeddings_pending,
99        }))
100    }
101}
102
103fn max_extract_bytes_for_source(config: &Config, source_label: &str) -> u64 {
104    if let Some(name) = source_label.strip_prefix("filesystem:") {
105        config
106            .connectors
107            .filesystem
108            .get(name)
109            .map(|connector| connector.max_extract_bytes)
110            .unwrap_or(DEFAULT_MAX_EXTRACT_BYTES)
111    } else {
112        DEFAULT_MAX_EXTRACT_BYTES
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use crate::app_store::SqliteAppStore;
120    use chrono::{TimeZone, Utc};
121    use sqlx::Row;
122    use tempfile::TempDir;
123
124    fn config(temp: &TempDir) -> Config {
125        toml::from_str(&format!(
126            r#"
127[db]
128path = "{}"
129
130[chunking]
131max_tokens = 3
132
133[retrieval]
134final_limit = 12
135
136[server]
137bind = "127.0.0.1:0"
138"#,
139            temp.path().join("ctx.sqlite").display()
140        ))
141        .unwrap()
142    }
143
144    fn item(body: &str) -> SourceItem {
145        SourceItem {
146            source: "custom:fixture".into(),
147            source_id: "doc-1".into(),
148            source_url: None,
149            title: Some("Fixture".into()),
150            author: None,
151            created_at: Utc.timestamp_opt(1_700_000_000, 0).unwrap(),
152            updated_at: Utc.timestamp_opt(1_700_000_001, 0).unwrap(),
153            content_type: "text/plain".into(),
154            body: body.into(),
155            metadata_json: "{}".into(),
156            raw_json: None,
157            raw_bytes: None,
158        }
159    }
160
161    #[tokio::test]
162    async fn replaces_canonical_chunks_without_touching_discovery_checkpoint() {
163        let temp = TempDir::new().unwrap();
164        let config = config(&temp);
165        SqliteAppStore::initialize_config(&config).await.unwrap();
166        let store = SqliteAppStore::connect(&config).await.unwrap();
167        let ingestor = DocumentIngestor::new(&config, &store, "custom:fixture");
168
169        let mut original = item("obsolete alpha paragraph\n\nsecond paragraph");
170        let first = ingestor.ingest(&mut original).await.unwrap();
171        assert!(matches!(first, DocumentIngestOutcome::Ingested(_)));
172
173        let mut replacement = item("replacement beta text");
174        let second = ingestor.ingest(&mut replacement).await.unwrap();
175        let DocumentIngestOutcome::Ingested(second) = second else {
176            panic!("text input should be ingested");
177        };
178        assert_eq!(second.embeddings_written, 0);
179        assert_eq!(second.embeddings_pending, 0);
180        assert_eq!(store.get_checkpoint("custom:fixture").await.unwrap(), None);
181
182        let rows = sqlx::query(
183            "SELECT c.text FROM chunks c JOIN documents d ON d.id = c.document_id WHERE d.source = ? AND d.source_id = ? ORDER BY c.chunk_index",
184        )
185        .bind("custom:fixture")
186        .bind("doc-1")
187        .fetch_all(store.pool())
188        .await
189        .unwrap();
190        assert_eq!(rows.len() as u64, second.chunks_written);
191        let text = rows
192            .iter()
193            .map(|row| row.get::<String, _>("text"))
194            .collect::<Vec<_>>()
195            .join(" ");
196        assert!(text.contains("replacement beta text"));
197        assert!(!text.contains("obsolete alpha"));
198    }
199}
context_harness/document_ingestor.rs

context_harness/
document_ingestor.rs
