telos_agent/knowledge/memory/index/
query.rs

1use std::cmp::Ordering;
2
3use crate::knowledge::memory::format::{MemoryEntry, MemoryStatus};
4use crate::knowledge::memory::index::MemoryStore;
5
6/// Sort order for memory queries.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum MemorySort {
9    Relevance,
10    RecentlyUpdated,
11    MostUsed,
12}
13
14/// Query options for selecting memories.
15#[derive(Debug, Clone)]
16pub struct MemoryQuery {
17    pub status: Option<MemoryStatus>,
18    pub tags: Vec<String>,
19    pub limit: Option<usize>,
20    pub include_body: bool,
21    pub sort: MemorySort,
22}
23
24impl Default for MemoryQuery {
25    fn default() -> Self {
26        Self {
27            status: None,
28            tags: Vec::new(),
29            limit: None,
30            include_body: true,
31            sort: MemorySort::Relevance,
32        }
33    }
34}
35
36impl MemoryStore {
37    /// Query memories by metadata. Uses the in-memory cache — no disk I/O.
38    pub fn query(&self, query: MemoryQuery) -> Vec<MemoryEntry> {
39        let tags: Vec<String> = query.tags.iter().map(|tag| tag.to_lowercase()).collect();
40        let mut entries: Vec<MemoryEntry> = self
41            .cache
42            .values()
43            .filter(|entry| {
44                query.status.as_ref().is_none_or(|status| &entry.status == status)
45                    && tags.iter().all(|tag| entry.tags.iter().any(|t| t.to_lowercase() == *tag))
46            })
47            .cloned()
48            .collect();
49
50        match query.sort {
51            MemorySort::Relevance => entries.sort_by(|a, b| {
52                status_rank(&a.status)
53                    .cmp(&status_rank(&b.status))
54                    .then_with(|| {
55                        std::cmp::Reverse(a.times_used).cmp(&std::cmp::Reverse(b.times_used))
56                    })
57                    .then_with(|| b.updated.cmp(&a.updated))
58                    .then_with(|| a.name.cmp(&b.name))
59            }),
60            MemorySort::RecentlyUpdated => {
61                entries.sort_by(|a, b| b.updated.cmp(&a.updated).then_with(|| a.name.cmp(&b.name)));
62            }
63            MemorySort::MostUsed => entries.sort_by_key(|b| std::cmp::Reverse(b.times_used)),
64        }
65
66        if let Some(limit) = query.limit {
67            entries.truncate(limit);
68        }
69        if !query.include_body {
70            for entry in &mut entries {
71                entry.body.clear();
72            }
73        }
74        entries
75    }
76
77    /// Search memories by keyword — matches name, description, tags, and body.
78    /// Uses the in-memory cache — no disk I/O.
79    pub fn search(&self, query: &str) -> Vec<MemoryEntry> {
80        let query_lower = query.to_lowercase();
81        self.cache
82            .values()
83            .filter(|entry| {
84                entry.name.to_lowercase().contains(&query_lower)
85                    || entry.description.to_lowercase().contains(&query_lower)
86                    || entry.tags.iter().any(|t| t.to_lowercase().contains(&query_lower))
87                    || entry.body.to_lowercase().contains(&query_lower)
88            })
89            .cloned()
90            .collect()
91    }
92
93    /// Get the top N most-used memories for prompt injection.
94    pub fn top_by_usage(&self, n: usize) -> Vec<MemoryEntry> {
95        self.query(MemoryQuery {
96            status: Some(MemoryStatus::Working),
97            limit: Some(n),
98            sort: MemorySort::MostUsed,
99            ..MemoryQuery::default()
100        })
101    }
102
103    /// Score cached memories against a user query and return the top-K by
104    /// keyword-overlap relevance. Non-deprecated entries only.
105    ///
106    /// `min_relevance` (0.0-1.0) discards entries whose composite score falls
107    /// below the threshold.  A value around 0.10 filters out false-positive
108    /// matches from very common tokens.
109    ///
110    /// This is the primary method for dynamic memory injection — unlike
111    /// `query()` + `MemorySort::Relevance` (which sorts by metadata), this
112    /// uses the actual content of the user's prompt to find semantically
113    /// relevant memories via simple token overlap.
114    pub fn search_relevant(
115        &self,
116        query: &str,
117        limit: usize,
118        min_relevance: f64,
119    ) -> Vec<MemoryEntry> {
120        let query_tokens = tokenize(query);
121        if query_tokens.is_empty() || self.cache.is_empty() {
122            return Vec::new();
123        }
124
125        let mut scored: Vec<(f64, &MemoryEntry)> = self
126            .cache
127            .values()
128            .filter(|e| e.status != MemoryStatus::Deprecated)
129            .map(|entry| {
130                let raw_score = compute_relevance(&query_tokens, entry);
131                (raw_score, entry)
132            })
133            .collect();
134
135        // Drop entries below the relevance floor before ranking.
136        scored.retain(|(score, _)| *score >= min_relevance);
137
138        // Sort descending by score, stable for deterministic output.
139        scored.sort_by(|a, b| {
140            b.0.partial_cmp(&a.0).unwrap_or(Ordering::Equal).then_with(|| a.1.name.cmp(&b.1.name))
141        });
142        scored.truncate(limit);
143        scored.into_iter().map(|(_, e)| e.clone()).collect()
144    }
145}
146
147pub(super) fn status_rank(status: &MemoryStatus) -> u8 {
148    match status {
149        MemoryStatus::NeedsFix => 0,
150        MemoryStatus::Working => 1,
151        MemoryStatus::Deprecated => 2,
152    }
153}
154
155// ── Relevance scoring ───────────────────────────────────────
156
157/// Tokenize text into lowercase words, skipping single-char and empty tokens.
158fn tokenize(text: &str) -> Vec<String> {
159    text.to_lowercase()
160        .split(|c: char| !c.is_alphanumeric() && c != '-' && c != '_')
161        .filter(|s| s.len() > 1)
162        .map(String::from)
163        .collect()
164}
165
166/// Tokenize an entry's searchable fields into a single token list.
167#[allow(dead_code)]
168fn tokenize_entry(entry: &MemoryEntry) -> Vec<String> {
169    let mut tokens = tokenize(&entry.name);
170    tokens.extend(tokenize(&entry.description));
171    for tag in &entry.tags {
172        tokens.extend(tokenize(tag));
173    }
174    tokens
175}
176
177/// Fraction of `query_tokens` that exactly match a token in `entry_tokens`.
178///
179/// Exact match only, with one exception: a query token that is the prefix of a
180/// longer entry token and the boundary character is `-` or `_` counts as a match
181/// (e.g. `error` matches `error-handling`).  This prevents noisy substring
182/// matches like `hi` matching `think`.
183fn token_overlap(query_tokens: &[String], entry_tokens: &[String]) -> f64 {
184    if query_tokens.is_empty() || entry_tokens.is_empty() {
185        return 0.0;
186    }
187    let matches =
188        query_tokens.iter().filter(|qt| entry_tokens.iter().any(|et| token_match(qt, et))).count();
189    matches as f64 / query_tokens.len() as f64
190}
191
192/// Returns true when `query_token` exactly equals `entry_token`, or when
193/// `query_token` is a prefix of a compound `entry_token` whose next character
194/// is a hyphen or underscore.
195fn token_match(query_token: &str, entry_token: &str) -> bool {
196    if query_token == entry_token {
197        return true;
198    }
199    if query_token.len() < entry_token.len() {
200        let next_byte = entry_token.as_bytes()[query_token.len()];
201        (next_byte == b'-' || next_byte == b'_') && entry_token.starts_with(query_token)
202    } else {
203        false
204    }
205}
206
207/// Score a single memory entry against the tokenized user query.
208///
209/// Weights are tuned for coding-agent use: name and description are
210/// the strongest signals (they're written for retrieval); tags add
211/// categorical relevance; body is a weaker signal because it's noisy.
212fn compute_relevance(query_tokens: &[String], entry: &MemoryEntry) -> f64 {
213    let name_tokens = tokenize(&entry.name);
214    let desc_tokens = tokenize(&entry.description);
215    let tag_tokens: Vec<String> = entry.tags.iter().flat_map(|t| tokenize(t)).collect();
216    let body_tokens = tokenize(&entry.body);
217
218    let name_score = token_overlap(query_tokens, &name_tokens) * 0.35;
219    let desc_score = token_overlap(query_tokens, &desc_tokens) * 0.30;
220    let tag_score = token_overlap(query_tokens, &tag_tokens) * 0.25;
221    let body_score = token_overlap(query_tokens, &body_tokens) * 0.05;
222
223    let status_mult = match entry.status {
224        MemoryStatus::Working => 1.0,
225        MemoryStatus::NeedsFix => 0.5,
226        MemoryStatus::Deprecated => 0.0,
227    };
228
229    let usage_boost = ((entry.times_used as f64 + 1.0).ln() / 10.0).min(0.10);
230
231    (name_score + desc_score + tag_score + body_score) * status_mult + usage_boost
232}
telos_agent/knowledge/memory/index/query.rs

telos_agent/knowledge/memory/index/
query.rs
