a3s_memory/repository/
query.rs

1use super::validation::{
2    validate_count, validate_required_text, MAX_CONTENT_BYTES, MAX_QUERY_LIMIT,
3};
4use super::{DurableMemoryKind, MemoryNamespace, MemoryNode, MemoryRepositoryError, MemoryStatus};
5use serde::{Deserialize, Serialize};
6use std::collections::{BTreeMap, BTreeSet};
7
8/// Stable identity of the deterministic V2 lexical query algorithm.
9///
10/// Hosts that persist serving policy can bind this value and reject silent
11/// query-semantics drift after a binary upgrade.
12pub const MEMORY_LEXICAL_QUERY_PROFILE_V1: &str = "a3s.memory.lexical.word-cjk-bigram.v1";
13
14/// Exact-namespace repository query. Active status is the default view.
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16#[serde(rename_all = "camelCase")]
17pub struct MemoryQuery {
18    pub namespace: MemoryNamespace,
19    pub text: Option<String>,
20    pub kinds: BTreeSet<DurableMemoryKind>,
21    pub statuses: BTreeSet<MemoryStatus>,
22    pub limit: usize,
23}
24
25impl MemoryQuery {
26    pub fn new(namespace: MemoryNamespace) -> Self {
27        Self {
28            namespace,
29            text: None,
30            kinds: BTreeSet::new(),
31            statuses: BTreeSet::from([MemoryStatus::Active]),
32            limit: 20,
33        }
34    }
35
36    pub fn with_text(mut self, text: impl Into<String>) -> Self {
37        self.text = Some(text.into());
38        self
39    }
40
41    pub fn with_kinds(mut self, kinds: impl IntoIterator<Item = DurableMemoryKind>) -> Self {
42        self.kinds = kinds.into_iter().collect();
43        self
44    }
45
46    pub fn with_statuses(mut self, statuses: impl IntoIterator<Item = MemoryStatus>) -> Self {
47        self.statuses = statuses.into_iter().collect();
48        self
49    }
50
51    pub fn with_limit(mut self, limit: usize) -> Self {
52        self.limit = limit;
53        self
54    }
55
56    pub(crate) fn validate(&self) -> Result<(), MemoryRepositoryError> {
57        self.namespace.validate()?;
58        if self.limit == 0 {
59            return Err(MemoryRepositoryError::invalid(
60                "query.limit",
61                "must be greater than zero",
62            ));
63        }
64        validate_count("query limit", self.limit, MAX_QUERY_LIMIT)?;
65        if self.statuses.is_empty() {
66            return Err(MemoryRepositoryError::invalid(
67                "query.statuses",
68                "must contain at least one status",
69            ));
70        }
71        if let Some(text) = &self.text {
72            validate_required_text("query.text", text, MAX_CONTENT_BYTES)?;
73            if !text.chars().any(char::is_alphanumeric) {
74                return Err(MemoryRepositoryError::invalid(
75                    "query.text",
76                    "must contain at least one alphanumeric term",
77                ));
78            }
79        }
80        Ok(())
81    }
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
85#[serde(rename_all = "camelCase")]
86pub struct MemoryScore {
87    pub lexical: f32,
88    pub total: f32,
89}
90
91#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
92#[serde(rename_all = "camelCase")]
93pub struct MemoryQueryHit {
94    pub node: MemoryNode,
95    pub score: MemoryScore,
96}
97
98#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
99#[serde(rename_all = "camelCase")]
100pub struct MemoryQueryResult {
101    pub hits: Vec<MemoryQueryHit>,
102}
103
104pub(crate) fn query_nodes(
105    nodes: Option<&BTreeMap<String, MemoryNode>>,
106    query: &MemoryQuery,
107) -> MemoryQueryResult {
108    let query_tokens = query.text.as_deref().map(tokens);
109    let mut hits = nodes
110        .into_iter()
111        .flat_map(BTreeMap::values)
112        .filter(|node| query.statuses.contains(&node.status))
113        .filter(|node| query.kinds.is_empty() || query.kinds.contains(&node.kind))
114        .filter_map(|node| {
115            let lexical = match &query_tokens {
116                Some(terms) => lexical_score(terms, &node.content),
117                None => 0.0,
118            };
119            if query_tokens.is_some() && lexical == 0.0 {
120                return None;
121            }
122            Some(MemoryQueryHit {
123                node: node.clone(),
124                score: MemoryScore {
125                    lexical,
126                    total: lexical,
127                },
128            })
129        })
130        .collect::<Vec<_>>();
131
132    hits.sort_by(|left, right| {
133        right
134            .score
135            .total
136            .total_cmp(&left.score.total)
137            .then_with(|| right.node.updated_at.cmp(&left.node.updated_at))
138            .then_with(|| left.node.id.cmp(&right.node.id))
139    });
140    hits.truncate(query.limit);
141    MemoryQueryResult { hits }
142}
143
144fn lexical_score(query_tokens: &BTreeSet<String>, content: &str) -> f32 {
145    let content_tokens = tokens(content);
146    let matches = query_tokens.intersection(&content_tokens).count();
147    matches as f32 / query_tokens.len() as f32
148}
149
150fn tokens(value: &str) -> BTreeSet<String> {
151    let mut tokens = BTreeSet::new();
152    let mut word = String::new();
153    let mut cjk_run = Vec::new();
154
155    for character in value.to_lowercase().chars() {
156        if is_cjk_character(character) {
157            flush_word(&mut tokens, &mut word);
158            cjk_run.push(character);
159        } else {
160            flush_cjk_run(&mut tokens, &mut cjk_run);
161            if character.is_alphanumeric() {
162                word.push(character);
163            } else {
164                flush_word(&mut tokens, &mut word);
165            }
166        }
167    }
168    flush_word(&mut tokens, &mut word);
169    flush_cjk_run(&mut tokens, &mut cjk_run);
170    tokens
171}
172
173fn flush_word(tokens: &mut BTreeSet<String>, word: &mut String) {
174    if !word.is_empty() {
175        tokens.insert(std::mem::take(word));
176    }
177}
178
179fn flush_cjk_run(tokens: &mut BTreeSet<String>, run: &mut Vec<char>) {
180    if run.is_empty() {
181        return;
182    }
183    tokens.insert(run.iter().collect());
184    for pair in run.windows(2) {
185        tokens.insert(pair.iter().collect());
186    }
187    run.clear();
188}
189
190fn is_cjk_character(character: char) -> bool {
191    matches!(
192        u32::from(character),
193        0x1100..=0x11FF
194            | 0x2E80..=0x2FFF
195            | 0x3005..=0x3007
196            | 0x3031..=0x3035
197            | 0x3040..=0x30FF
198            | 0x3100..=0x312F
199            | 0x3130..=0x318F
200            | 0x3190..=0x319F
201            | 0x31A0..=0x31BF
202            | 0x31F0..=0x31FF
203            | 0x3400..=0x4DBF
204            | 0x4E00..=0x9FFF
205            | 0xA960..=0xA97F
206            | 0xAC00..=0xD7AF
207            | 0xD7B0..=0xD7FF
208            | 0xF900..=0xFAFF
209            | 0xFF66..=0xFF9F
210            | 0x20000..=0x2FA1F
211            | 0x30000..=0x323AF
212    )
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    #[test]
220    fn tokenizer_preserves_words_and_adds_cjk_bigrams() {
221        assert_eq!(
222            tokens("Rust 数据库迁移 cargo_fmt"),
223            BTreeSet::from([
224                "rust".to_string(),
225                "数据库迁移".to_string(),
226                "数据".to_string(),
227                "据库".to_string(),
228                "库迁".to_string(),
229                "迁移".to_string(),
230                "cargo".to_string(),
231                "fmt".to_string(),
232            ])
233        );
234    }
235
236    #[test]
237    fn cjk_bigrams_cover_chinese_japanese_and_korean_without_unigrams() {
238        for (input, expected) in [
239            ("数据库", "数据"),
240            ("確認手順", "確認"),
241            ("배포절차", "배포"),
242        ] {
243            let actual = tokens(input);
244            assert!(actual.contains(expected), "{input}: {actual:?}");
245            assert!(!actual.contains(&input.chars().next().unwrap().to_string()));
246        }
247    }
248}
a3s_memory/repository/query.rs

a3s_memory/repository/
query.rs
