llm_embedding/
preprocessor.rs

1//! Query text preprocessing before embedding.
2//!
3//! Provides generic prefix and template-based preprocessing. Model-specific
4//! prefixes are supplied as configuration data by the caller.
5
6use serde::{Deserialize, Serialize};
7
8/// Preprocessor configuration, deserializable from service config files.
9#[derive(Debug, Clone, Serialize, Deserialize, Default)]
10#[serde(tag = "type", rename_all = "snake_case")]
11pub enum PreprocessorConfig {
12    /// No preprocessing (default).
13    #[default]
14    None,
15    /// Prepend a fixed prefix to every text.
16    Prefix { prefix: String },
17    /// Substitute `{{text}}` into a template.
18    Template { template: String },
19}
20
21/// Concrete preprocessor built from [`PreprocessorConfig`].
22#[derive(Debug, Clone)]
23pub enum PreprocessorImpl {
24    /// Pass text through unchanged.
25    None,
26    /// Prepend a fixed prefix.
27    Prefix { prefix: String },
28    /// Substitute `{{text}}` into a template.
29    Template { template: String },
30}
31
32impl PreprocessorImpl {
33    /// Builds the concrete preprocessor for a config value.
34    pub fn from_config(config: &PreprocessorConfig) -> Self {
35        match config {
36            PreprocessorConfig::None => Self::None,
37            PreprocessorConfig::Prefix { prefix } => Self::Prefix {
38                prefix: prefix.clone(),
39            },
40            PreprocessorConfig::Template { template } => Self::Template {
41                template: template.clone(),
42            },
43        }
44    }
45
46    /// Applies the preprocessing rule to a single text.
47    pub fn preprocess(&self, text: &str) -> String {
48        match self {
49            Self::None => text.to_string(),
50            Self::Prefix { prefix } => format!("{prefix}{text}"),
51            Self::Template { template } => template.replace("{{text}}", text),
52        }
53    }
54
55    /// Applies the preprocessing rule to a batch of texts, preserving order.
56    pub fn process_batch(&self, texts: &[&str]) -> Vec<String> {
57        texts.iter().map(|text| self.preprocess(text)).collect()
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn none_passes_text_through() {
67        assert_eq!(PreprocessorImpl::None.preprocess("hello"), "hello");
68    }
69
70    #[test]
71    fn prefix_prepends_literal() {
72        let preprocessor = PreprocessorImpl::Prefix {
73            prefix: "query: ".into(),
74        };
75        assert_eq!(preprocessor.preprocess("rust"), "query: rust");
76    }
77
78    #[test]
79    fn template_substitutes_placeholder() {
80        let preprocessor = PreprocessorImpl::Template {
81            template: "classify: {{text}}".into(),
82        };
83        assert_eq!(preprocessor.preprocess("hello"), "classify: hello");
84    }
85
86    #[test]
87    fn from_config_round_trips_each_variant() {
88        let none = PreprocessorImpl::from_config(&PreprocessorConfig::None);
89        assert_eq!(none.preprocess("hello"), "hello");
90
91        let prefixed = PreprocessorImpl::from_config(&PreprocessorConfig::Prefix {
92            prefix: ">> ".into(),
93        });
94        assert_eq!(prefixed.preprocess("x"), ">> x");
95
96        let templated = PreprocessorImpl::from_config(&PreprocessorConfig::Template {
97            template: "[{{text}}]".into(),
98        });
99        assert_eq!(templated.preprocess("x"), "[x]");
100    }
101
102    #[test]
103    fn batch_preserves_order() {
104        let preprocessor = PreprocessorImpl::Prefix {
105            prefix: "> ".into(),
106        };
107        assert_eq!(
108            preprocessor.process_batch(&["a", "b", "c"]),
109            vec!["> a", "> b", "> c"]
110        );
111    }
112}
llm_embedding/preprocessor.rs

llm_embedding/
preprocessor.rs
