triplets_core/source/backends/
file_source.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use crate::config::{NegativeStrategy, SamplerConfig, Selector, TripletRecipe};
6use crate::data::{DataRecord, QualityScore, RecordSection, SectionRole};
7use crate::errors::SamplerError;
8use crate::source::indexing::file_corpus::FileCorpusIndex;
9use crate::source::{DataSource, SourceCursor, SourceSnapshot};
10use crate::types::{CategoryId, SourceId, TaxonomyValue};
11use crate::utils::{file_times, is_text_file};
12use crate::utils::{make_section, normalize_inline_whitespace};
13
14const FILE_RECIPE_TITLE_CONTEXT_WRONG_DATE: &str = "title_context_wrong_date";
15const FILE_RECIPE_TITLE_ANCHOR_WRONG_DATE: &str = "title_anchor_wrong_date";
16const FILE_RECIPE_TITLE_CONTEXT_WRONG_ARTICLE: &str = "title_context_wrong_article";
17const FILE_RECIPE_TITLE_ANCHOR_WRONG_ARTICLE: &str = "title_anchor_wrong_article";
18
19/// Builds taxonomy values from a root path and file path.
20pub type TaxonomyBuilder =
21    Arc<dyn Fn(&Path, &Path, &SourceId) -> Vec<TaxonomyValue> + Send + Sync + 'static>;
22
23/// Builds record sections from a normalized title and body.
24pub type SectionBuilder = Arc<dyn Fn(&str, &str) -> Vec<RecordSection> + Send + Sync + 'static>;
25
26/// Configuration for a generic filesystem-backed data source.
27#[derive(Clone)]
28pub struct FileSourceConfig {
29    /// Stable source identifier used in records and persistence keys.
30    pub source_id: SourceId,
31    /// Root directory containing source files.
32    pub root: PathBuf,
33    /// Default quality trust score applied to generated records.
34    pub trust: f32,
35    /// Optional trust overrides keyed by taxonomy segment.
36    pub category_trust: HashMap<CategoryId, f32>,
37    /// Whether to follow symlinks during index walking.
38    pub follow_links: bool,
39    /// Whether indexing should include only text files.
40    pub text_files_only: bool,
41    /// Whether deterministic directory grouping is enabled.
42    pub group_by_directory: bool,
43    /// Whether title extraction should replace underscores with spaces.
44    pub title_replace_underscores: bool,
45    /// Whether default recipe set includes the date-aware negative lane.
46    pub include_date_aware_default_recipe: bool,
47    /// Optional directory used for persisted file-corpus index stores.
48    ///
49    /// When `None`, file-corpus indexing uses the managed cache discovery root.
50    /// Set this in tests to keep index writes inside temporary directories.
51    pub index_dir: Option<PathBuf>,
52    /// Optional default recipes returned by this source.
53    pub default_triplet_recipes: Vec<TripletRecipe>,
54    /// Taxonomy builder invoked per file.
55    pub taxonomy_builder: TaxonomyBuilder,
56    /// Section builder invoked per file.
57    pub section_builder: SectionBuilder,
58}
59
60impl FileSourceConfig {
61    /// Create a config for a filesystem source with explicit id and root.
62    pub fn new(source_id: impl Into<SourceId>, root: impl Into<PathBuf>) -> Self {
63        Self {
64            source_id: source_id.into(),
65            root: root.into(),
66            trust: 0.85,
67            category_trust: HashMap::new(),
68            follow_links: true,
69            text_files_only: false,
70            group_by_directory: true,
71            title_replace_underscores: true,
72            include_date_aware_default_recipe: false,
73            index_dir: None,
74            default_triplet_recipes: default_title_context_triplet_recipes(false),
75            taxonomy_builder: Arc::new(taxonomy_from_path),
76            section_builder: Arc::new(anchor_context_sections),
77        }
78    }
79
80    /// Override default trust score.
81    pub fn with_trust(mut self, trust: f32) -> Self {
82        self.trust = trust;
83        self
84    }
85
86    /// Add a taxonomy-segment trust override.
87    pub fn with_category_trust(mut self, category: impl Into<String>, trust: f32) -> Self {
88        self.category_trust
89            .insert(category.into().to_lowercase(), trust);
90        self
91    }
92
93    /// Override whether symlinks are followed during index walk.
94    pub fn with_follow_links(mut self, follow_links: bool) -> Self {
95        self.follow_links = follow_links;
96        self
97    }
98
99    /// Override whether index walk includes only text files.
100    pub fn with_text_files_only(mut self, text_files_only: bool) -> Self {
101        self.text_files_only = text_files_only;
102        self
103    }
104
105    /// Enable or disable deterministic directory grouping.
106    pub fn with_directory_grouping(mut self, group_by_directory: bool) -> Self {
107        self.group_by_directory = group_by_directory;
108        self
109    }
110
111    /// Set whether title extraction replaces underscores with spaces.
112    pub fn with_title_replace_underscores(mut self, replace_underscores: bool) -> Self {
113        self.title_replace_underscores = replace_underscores;
114        self
115    }
116
117    /// Enable/disable the date-aware default recipe lane (`WrongPublicationDate`).
118    ///
119    /// "Date-aware" here uses publication-date metadata on records (for example
120    /// taxonomy/meta date fields), not filesystem timestamps from source files.
121    pub fn with_date_aware_default_recipe(mut self, include: bool) -> Self {
122        self.include_date_aware_default_recipe = include;
123        self.default_triplet_recipes = default_title_context_triplet_recipes(include);
124        self
125    }
126
127    /// Override the directory used to persist file-corpus index stores.
128    pub fn with_index_dir(mut self, index_dir: impl Into<PathBuf>) -> Self {
129        self.index_dir = Some(index_dir.into());
130        self
131    }
132
133    /// Set source-provided default triplet recipes.
134    pub fn with_default_triplet_recipes(mut self, recipes: Vec<TripletRecipe>) -> Self {
135        self.default_triplet_recipes = recipes;
136        self
137    }
138
139    /// Set a custom taxonomy builder.
140    pub fn with_taxonomy_builder(mut self, taxonomy_builder: TaxonomyBuilder) -> Self {
141        self.taxonomy_builder = taxonomy_builder;
142        self
143    }
144
145    /// Set a custom section builder.
146    pub fn with_section_builder(mut self, section_builder: SectionBuilder) -> Self {
147        self.section_builder = section_builder;
148        self
149    }
150}
151
152/// Default mixed-negative recipes used by `FileSource` title/body corpora.
153///
154/// When `include_date_aware` is enabled, the date-aware lane compares metadata
155/// publication dates, not filesystem mtime/ctime/atime values.
156pub fn default_title_context_triplet_recipes(include_date_aware: bool) -> Vec<TripletRecipe> {
157    let mut recipes = Vec::new();
158    if include_date_aware {
159        // Make date-aware summary negatives nearly as common as summary wrong-article
160        // negatives so temporal contrast is meaningfully represented.
161        recipes.push(TripletRecipe {
162            name: FILE_RECIPE_TITLE_CONTEXT_WRONG_DATE.into(),
163            anchor: Selector::Role(SectionRole::Anchor),
164            positive_selector: Selector::Role(SectionRole::Context),
165            negative_selector: Selector::Role(SectionRole::Context),
166            negative_strategy: NegativeStrategy::WrongPublicationDate,
167            weight: 0.30,
168            instruction: None,
169            allow_same_anchor_positive: false,
170        });
171        // Keep a smaller anchor-negative date-aware lane for harder examples
172        // without overwhelming the primary summary-driven objectives.
173        // Date-aware means publication-date metadata comparison, not file mtime.
174        recipes.push(TripletRecipe {
175            name: FILE_RECIPE_TITLE_ANCHOR_WRONG_DATE.into(),
176            anchor: Selector::Role(SectionRole::Anchor),
177            positive_selector: Selector::Role(SectionRole::Context),
178            negative_selector: Selector::Role(SectionRole::Anchor),
179            negative_strategy: NegativeStrategy::WrongPublicationDate,
180            weight: 0.10,
181            instruction: None,
182            allow_same_anchor_positive: false,
183        });
184    }
185    // Rebalance summary wrong-article depending on whether date-aware lanes are
186    // enabled so the full pool stays intentionally weighted in both modes.
187    recipes.push(TripletRecipe {
188        name: FILE_RECIPE_TITLE_CONTEXT_WRONG_ARTICLE.into(),
189        anchor: Selector::Role(SectionRole::Anchor),
190        positive_selector: Selector::Role(SectionRole::Context),
191        negative_selector: Selector::Role(SectionRole::Context),
192        negative_strategy: NegativeStrategy::WrongArticle,
193        weight: if include_date_aware { 0.35 } else { 0.75 },
194        instruction: None,
195        allow_same_anchor_positive: false,
196    });
197    // Medium-hard lane adds anchor-as-negative pressure to improve
198    // discrimination among title-like anchor fields.
199    recipes.push(TripletRecipe {
200        name: FILE_RECIPE_TITLE_ANCHOR_WRONG_ARTICLE.into(),
201        anchor: Selector::Role(SectionRole::Anchor),
202        positive_selector: Selector::Role(SectionRole::Context),
203        negative_selector: Selector::Role(SectionRole::Anchor),
204        negative_strategy: NegativeStrategy::WrongArticle,
205        weight: 0.25,
206        instruction: None,
207        allow_same_anchor_positive: false,
208    });
209    recipes
210}
211
212/// Generic filesystem-backed source with configurable taxonomy and section mapping.
213pub struct FileSource {
214    config: FileSourceConfig,
215}
216
217impl FileSource {
218    /// Create a generic file source from configuration.
219    pub fn new(config: FileSourceConfig) -> Self {
220        Self { config }
221    }
222
223    fn file_corpus_index(&self, sampler_seed: u64) -> FileCorpusIndex {
224        let mut index = FileCorpusIndex::new(&self.config.root, &self.config.source_id)
225            .with_sampler_seed(sampler_seed)
226            .with_follow_links(self.config.follow_links)
227            .with_text_files_only(self.config.text_files_only)
228            .with_directory_grouping(self.config.group_by_directory);
229
230        if let Some(index_dir) = &self.config.index_dir {
231            index = index.with_index_dir(index_dir.clone());
232        }
233
234        index
235    }
236
237    fn trust_for_taxonomy(&self, taxonomy: &[String]) -> f32 {
238        for segment in taxonomy.iter().skip(1) {
239            if let Some(weight) = self.config.category_trust.get(&segment.to_lowercase()) {
240                return *weight;
241            }
242        }
243        self.config.trust
244    }
245
246    fn build_record(&self, path: &Path) -> Result<Option<DataRecord>, SamplerError> {
247        if !is_text_file(path) {
248            return Ok(None);
249        }
250        let title = FileCorpusIndex::normalized_title_from_stem(
251            path,
252            &self.config.source_id,
253            self.config.title_replace_underscores,
254        )?;
255        if title.is_empty() {
256            return Ok(None);
257        }
258
259        let body_raw = std::fs::read_to_string(path)?;
260        let body = normalize_inline_whitespace(body_raw);
261        if body.is_empty() {
262            return Ok(None);
263        }
264
265        let taxonomy =
266            (self.config.taxonomy_builder)(&self.config.root, path, &self.config.source_id);
267        let sections = (self.config.section_builder)(&title, &body);
268        let trust = self.trust_for_taxonomy(&taxonomy);
269        let (created_at, updated_at) = file_times(path);
270
271        Ok(Some(DataRecord {
272            id: FileCorpusIndex::source_scoped_record_id(
273                &self.config.source_id,
274                &self.config.root,
275                path,
276            ),
277            source: self.config.source_id.clone(),
278            created_at,
279            updated_at,
280            quality: QualityScore { trust },
281            taxonomy,
282            sections,
283            meta_prefix: None,
284            label: None,
285        }))
286    }
287}
288
289impl DataSource for FileSource {
290    fn id(&self) -> &str {
291        &self.config.source_id
292    }
293
294    fn refresh(
295        &self,
296        config: &SamplerConfig,
297        cursor: Option<&SourceCursor>,
298        limit: Option<usize>,
299    ) -> Result<SourceSnapshot, SamplerError> {
300        self.file_corpus_index(config.seed)
301            .refresh_indexable(cursor, limit, |path| self.build_record(path))
302    }
303
304    fn reported_record_count(&self, config: &SamplerConfig) -> Result<u128, SamplerError> {
305        self.file_corpus_index(config.seed)
306            .indexed_record_count()
307            .map(|count| count as u128)
308    }
309
310    fn default_triplet_recipes(&self) -> Vec<TripletRecipe> {
311        self.config.default_triplet_recipes.clone()
312    }
313}
314
315/// Build default taxonomy from the file path relative to `root`.
316///
317/// Output shape is `[source_id, <parent segments...>]`.
318pub fn taxonomy_from_path(root: &Path, path: &Path, source_id: &SourceId) -> Vec<TaxonomyValue> {
319    let mut taxonomy = vec![source_id.to_string()];
320    if let Ok(rel) = path.strip_prefix(root)
321        && let Some(parent) = rel.parent()
322    {
323        for segment in parent.iter() {
324            taxonomy.push(segment.to_string_lossy().to_string());
325        }
326    }
327    taxonomy
328}
329
330/// Build a default two-section payload of title anchor and body context.
331pub fn anchor_context_sections(title: &str, body: &str) -> Vec<RecordSection> {
332    vec![
333        make_section(SectionRole::Anchor, None, title),
334        make_section(SectionRole::Context, None, body),
335    ]
336}
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341    use crate::config::{NegativeStrategy, Selector};
342    use tempfile::tempdir;
343
344    fn sampler_config(seed: u64) -> SamplerConfig {
345        SamplerConfig {
346            seed,
347            ..SamplerConfig::default()
348        }
349    }
350
351    #[test]
352    fn reads_records_without_default_source_id() {
353        let temp = tempdir().unwrap();
354        let category = temp.path().join("factual");
355        std::fs::create_dir_all(&category).unwrap();
356        std::fs::write(
357            category.join("What_is_alpha.txt"),
358            "Alpha measures risk-adjusted outperformance.",
359        )
360        .unwrap();
361
362        let source = FileSource::new(FileSourceConfig::new("qa_custom", temp.path()));
363        let snapshot = source.refresh(&sampler_config(101), None, None).unwrap();
364
365        assert_eq!(snapshot.records.len(), 1);
366        assert_eq!(snapshot.records[0].source, "qa_custom");
367    }
368
369    #[test]
370    fn applies_category_trust_overrides() {
371        let temp = tempdir().unwrap();
372        let factual = temp.path().join("factual");
373        let opinion = temp.path().join("opinionated");
374        std::fs::create_dir_all(&factual).unwrap();
375        std::fs::create_dir_all(&opinion).unwrap();
376        std::fs::write(
377            factual.join("What_is_beta.txt"),
378            "Beta compares volatility.",
379        )
380        .unwrap();
381        std::fs::write(
382            opinion.join("Will_rates_fall.txt"),
383            "Probably not this year.",
384        )
385        .unwrap();
386
387        let source = FileSource::new(
388            FileSourceConfig::new("qa_weighted", temp.path())
389                .with_category_trust("factual", 0.95)
390                .with_category_trust("opinionated", 0.6),
391        );
392        let snapshot = source.refresh(&sampler_config(101), None, None).unwrap();
393
394        let factual_record = snapshot
395            .records
396            .iter()
397            .find(|record| record.taxonomy.iter().any(|value| value == "factual"))
398            .unwrap();
399        let opinion_record = snapshot
400            .records
401            .iter()
402            .find(|record| record.taxonomy.iter().any(|value| value == "opinionated"))
403            .unwrap();
404        assert_eq!(factual_record.quality.trust, 0.95);
405        assert_eq!(opinion_record.quality.trust, 0.6);
406    }
407
408    #[test]
409    fn supports_custom_sections_and_default_recipes() {
410        let temp = tempdir().unwrap();
411        std::fs::write(
412            temp.path().join("What_is_gamma.txt"),
413            "Gamma measures convexity.",
414        )
415        .unwrap();
416
417        let sections: SectionBuilder = Arc::new(|question, answer| {
418            vec![
419                make_section(SectionRole::Anchor, Some("Question"), question),
420                make_section(SectionRole::Context, Some("Answer"), answer),
421            ]
422        });
423
424        let recipes = vec![TripletRecipe {
425            name: "question_answer".into(),
426            anchor: Selector::Role(SectionRole::Anchor),
427            positive_selector: Selector::Role(SectionRole::Context),
428            negative_selector: Selector::Role(SectionRole::Context),
429            negative_strategy: NegativeStrategy::QuestionAnswerMismatch,
430            weight: 1.0,
431            instruction: None,
432            allow_same_anchor_positive: false,
433        }];
434
435        let source = FileSource::new(
436            FileSourceConfig::new("qa_sections", temp.path())
437                .with_section_builder(sections)
438                .with_default_triplet_recipes(recipes.clone()),
439        );
440
441        let snapshot = source.refresh(&sampler_config(101), None, None).unwrap();
442        assert_eq!(snapshot.records.len(), 1);
443        assert_eq!(snapshot.records[0].sections.len(), 2);
444        assert_eq!(source.default_triplet_recipes().len(), recipes.len());
445    }
446
447    #[test]
448    fn file_source_config_new_has_explicit_default_triplet_recipes() {
449        let temp = tempdir().unwrap();
450        let source = FileSource::new(FileSourceConfig::new("qa_defaults", temp.path()));
451        let defaults = source.default_triplet_recipes();
452        assert!(!defaults.is_empty());
453        let names: Vec<&str> = defaults.iter().map(|recipe| recipe.name.as_ref()).collect();
454        assert!(!names.contains(&FILE_RECIPE_TITLE_CONTEXT_WRONG_DATE));
455        assert!(!names.contains(&FILE_RECIPE_TITLE_ANCHOR_WRONG_DATE));
456        assert!(names.contains(&FILE_RECIPE_TITLE_CONTEXT_WRONG_ARTICLE));
457        assert!(names.contains(&FILE_RECIPE_TITLE_ANCHOR_WRONG_ARTICLE));
458        let summary_wrong_article = defaults
459            .iter()
460            .find(|recipe| recipe.name == FILE_RECIPE_TITLE_CONTEXT_WRONG_ARTICLE)
461            .unwrap();
462        let anchor_wrong_article = defaults
463            .iter()
464            .find(|recipe| recipe.name == FILE_RECIPE_TITLE_ANCHOR_WRONG_ARTICLE)
465            .unwrap();
466        assert_eq!(summary_wrong_article.weight, 0.75);
467        assert_eq!(anchor_wrong_article.weight, 0.25);
468    }
469
470    #[test]
471    fn file_source_config_can_enable_date_aware_default_recipe() {
472        let temp = tempdir().unwrap();
473        let source = FileSource::new(
474            FileSourceConfig::new("qa_defaults_with_date", temp.path())
475                .with_date_aware_default_recipe(true),
476        );
477        let defaults = source.default_triplet_recipes();
478        let names: Vec<&str> = defaults.iter().map(|recipe| recipe.name.as_ref()).collect();
479        assert!(names.contains(&FILE_RECIPE_TITLE_CONTEXT_WRONG_DATE));
480        assert!(names.contains(&FILE_RECIPE_TITLE_ANCHOR_WRONG_DATE));
481        assert!(names.contains(&FILE_RECIPE_TITLE_CONTEXT_WRONG_ARTICLE));
482        assert!(names.contains(&FILE_RECIPE_TITLE_ANCHOR_WRONG_ARTICLE));
483        let summary_wrong_date = defaults
484            .iter()
485            .find(|recipe| recipe.name == FILE_RECIPE_TITLE_CONTEXT_WRONG_DATE)
486            .unwrap();
487        let anchor_wrong_date = defaults
488            .iter()
489            .find(|recipe| recipe.name == FILE_RECIPE_TITLE_ANCHOR_WRONG_DATE)
490            .unwrap();
491        let summary_wrong_article = defaults
492            .iter()
493            .find(|recipe| recipe.name == FILE_RECIPE_TITLE_CONTEXT_WRONG_ARTICLE)
494            .unwrap();
495        let anchor_wrong_article = defaults
496            .iter()
497            .find(|recipe| recipe.name == FILE_RECIPE_TITLE_ANCHOR_WRONG_ARTICLE)
498            .unwrap();
499        assert_eq!(summary_wrong_date.weight, 0.30);
500        assert_eq!(anchor_wrong_date.weight, 0.10);
501        assert_eq!(summary_wrong_article.weight, 0.35);
502        assert_eq!(anchor_wrong_article.weight, 0.25);
503    }
504
505    #[test]
506    fn file_source_config_override_replaces_default_triplet_recipes() {
507        let temp = tempdir().unwrap();
508        let custom = vec![TripletRecipe {
509            name: "custom_only".into(),
510            anchor: Selector::Role(SectionRole::Context),
511            positive_selector: Selector::Role(SectionRole::Context),
512            negative_selector: Selector::Role(SectionRole::Context),
513            negative_strategy: NegativeStrategy::WrongArticle,
514            weight: 1.0,
515            instruction: None,
516            allow_same_anchor_positive: false,
517        }];
518        let source = FileSource::new(
519            FileSourceConfig::new("qa_defaults_override", temp.path())
520                .with_default_triplet_recipes(custom.clone()),
521        );
522        let recipes = source.default_triplet_recipes();
523        assert_eq!(recipes.len(), 1);
524        assert_eq!(recipes[0].name.as_ref(), "custom_only");
525    }
526
527    #[test]
528    fn taxonomy_from_path_handles_nested_and_non_descendant_paths() {
529        let temp = tempdir().unwrap();
530        let root = temp.path().join("root");
531        std::fs::create_dir_all(root.join("topic/subtopic")).unwrap();
532
533        let nested = root.join("topic/subtopic/doc.txt");
534        let taxonomy = taxonomy_from_path(&root, &nested, &"qa_tax".to_string());
535        assert_eq!(taxonomy, vec!["qa_tax", "topic", "subtopic"]);
536
537        let outside = temp.path().join("outside.txt");
538        let outside_taxonomy = taxonomy_from_path(&root, &outside, &"qa_tax".to_string());
539        assert_eq!(outside_taxonomy, vec!["qa_tax"]);
540    }
541
542    #[test]
543    fn anchor_context_sections_build_expected_roles_and_text() {
544        let sections = anchor_context_sections("What is delta", "Delta is change over time.");
545        assert_eq!(sections.len(), 2);
546        assert_eq!(sections[0].role, SectionRole::Anchor);
547        assert_eq!(sections[0].text, "What is delta");
548        assert_eq!(sections[1].role, SectionRole::Context);
549        assert_eq!(sections[1].text, "Delta is change over time.");
550    }
551
552    #[test]
553    fn title_replace_underscores_toggle_changes_anchor_title_text() {
554        let temp = tempdir().unwrap();
555        std::fs::write(
556            temp.path().join("What_is_delta.txt"),
557            "Delta captures directional change.",
558        )
559        .unwrap();
560
561        let source_default =
562            FileSource::new(FileSourceConfig::new("qa_title_default", temp.path()));
563        let default_snapshot = source_default
564            .refresh(&sampler_config(101), None, Some(1))
565            .unwrap();
566        assert_eq!(default_snapshot.records.len(), 1);
567        assert_eq!(
568            default_snapshot.records[0].sections[0].text,
569            "What is delta"
570        );
571
572        let source_preserve = FileSource::new(
573            FileSourceConfig::new("qa_title_preserve", temp.path())
574                .with_title_replace_underscores(false),
575        );
576        let preserve_snapshot = source_preserve
577            .refresh(&sampler_config(101), None, Some(1))
578            .unwrap();
579        assert_eq!(preserve_snapshot.records.len(), 1);
580        assert_eq!(
581            preserve_snapshot.records[0].sections[0].text,
582            "What_is_delta"
583        );
584    }
585
586    #[test]
587    fn refresh_skips_non_txt_files_even_when_text_only_disabled() {
588        let temp = tempdir().unwrap();
589        std::fs::write(temp.path().join("notes.md"), "markdown should be skipped").unwrap();
590        std::fs::write(temp.path().join("doc.txt"), "plain text should be indexed").unwrap();
591
592        let source = FileSource::new(
593            FileSourceConfig::new("qa_filtering", temp.path()).with_text_files_only(false),
594        );
595        let snapshot = source.refresh(&sampler_config(101), None, None).unwrap();
596        assert_eq!(snapshot.records.len(), 1);
597        assert!(snapshot.records[0].id.contains("doc.txt"));
598    }
599
600    #[test]
601    fn trust_falls_back_to_default_and_count_and_id_are_exposed() {
602        let temp = tempdir().unwrap();
603        let docs = temp.path().join("docs");
604        std::fs::create_dir_all(&docs).unwrap();
605        std::fs::write(docs.join("alpha.txt"), "Alpha body.").unwrap();
606
607        let source = FileSource::new(
608            FileSourceConfig::new("qa_count", temp.path())
609                .with_trust(0.42)
610                .with_category_trust("factual", 0.95)
611                .with_taxonomy_builder(Arc::new(|_, _, source_id| {
612                    vec![source_id.clone(), "UNMATCHED".to_string()]
613                })),
614        );
615
616        let seed_101 = sampler_config(101);
617        let snapshot = source.refresh(&seed_101, None, None).unwrap();
618        assert_eq!(snapshot.records.len(), 1);
619        assert_eq!(snapshot.records[0].quality.trust, 0.42);
620        assert_eq!(source.id(), "qa_count");
621        assert_eq!(source.reported_record_count(&seed_101).unwrap(), 1);
622    }
623
624    #[test]
625    fn sampler_seed_controls_file_source_refresh_order() {
626        let temp = tempdir().unwrap();
627        for idx in 0..12 {
628            std::fs::write(
629                temp.path().join(format!("doc_{idx:02}.txt")),
630                format!("Body text for {idx}"),
631            )
632            .unwrap();
633        }
634
635        let source_a = FileSource::new(FileSourceConfig::new("seeded_a", temp.path()));
636        let source_b = FileSource::new(FileSourceConfig::new("seeded_a", temp.path()));
637        let source_c = FileSource::new(FileSourceConfig::new("seeded_a", temp.path()));
638
639        let ids_a: Vec<String> = source_a
640            .refresh(&sampler_config(11), None, Some(8))
641            .unwrap()
642            .records
643            .into_iter()
644            .map(|record| record.id)
645            .collect();
646        let ids_b: Vec<String> = source_b
647            .refresh(&sampler_config(11), None, Some(8))
648            .unwrap()
649            .records
650            .into_iter()
651            .map(|record| record.id)
652            .collect();
653        let ids_c: Vec<String> = source_c
654            .refresh(&sampler_config(29), None, Some(8))
655            .unwrap()
656            .records
657            .into_iter()
658            .map(|record| record.id)
659            .collect();
660
661        assert_eq!(ids_a, ids_b);
662        assert_ne!(ids_a, ids_c);
663    }
664
665    #[test]
666    fn with_index_dir_persists_refresh_index_store_under_custom_directory() {
667        let temp = tempdir().unwrap();
668        let index_temp = tempdir().unwrap();
669        std::fs::write(temp.path().join("alpha.txt"), "Alpha body").unwrap();
670
671        let custom_index_dir = index_temp.path().join("custom_index_store");
672        std::fs::create_dir_all(&custom_index_dir).unwrap();
673
674        let source_id = "qa_custom_index_refresh".to_string();
675        let source = FileSource::new(
676            FileSourceConfig::new(source_id.clone(), temp.path())
677                .with_index_dir(custom_index_dir.clone()),
678        );
679
680        let snapshot = source.refresh(&sampler_config(101), None, None).unwrap();
681        assert_eq!(snapshot.records.len(), 1);
682
683        let expected_store_path = FileCorpusIndex::index_store_path_for(
684            Some(custom_index_dir.as_path()),
685            temp.path(),
686            &source_id,
687        );
688        assert!(
689            expected_store_path.is_file(),
690            "expected index store to exist at {}",
691            expected_store_path.display()
692        );
693    }
694
695    #[test]
696    fn with_index_dir_persists_count_index_store_under_custom_directory() {
697        let temp = tempdir().unwrap();
698        let index_temp = tempdir().unwrap();
699        std::fs::write(temp.path().join("alpha.txt"), "Alpha body").unwrap();
700
701        let custom_index_dir = index_temp.path().join("custom_index_store");
702        std::fs::create_dir_all(&custom_index_dir).unwrap();
703
704        let source_id = "qa_custom_index_count".to_string();
705        let source = FileSource::new(
706            FileSourceConfig::new(source_id.clone(), temp.path())
707                .with_index_dir(custom_index_dir.clone()),
708        );
709
710        assert_eq!(
711            source.reported_record_count(&sampler_config(101)).unwrap(),
712            1
713        );
714
715        let expected_store_path = FileCorpusIndex::index_store_path_for(
716            Some(custom_index_dir.as_path()),
717            temp.path(),
718            &source_id,
719        );
720        assert!(
721            expected_store_path.is_file(),
722            "expected index store to exist at {}",
723            expected_store_path.display()
724        );
725    }
726}
triplets_core/source/backends/file_source.rs

triplets_core/source/backends/
file_source.rs
