lean_ctx/core/knowledge/
query.rs1use chrono::{DateTime, Utc};
2
3use super::ranking::{build_token_index, sort_fact_for_output};
4use super::types::{KnowledgeFact, ProjectKnowledge};
5
6impl ProjectKnowledge {
7 pub fn recall(&self, query: &str) -> Vec<&KnowledgeFact> {
8 let q = query.to_lowercase();
9 let terms: Vec<&str> = q.split_whitespace().collect();
10 if terms.is_empty() {
11 return Vec::new();
12 }
13
14 let index = build_token_index(&self.facts, true);
15 let mut match_counts: std::collections::HashMap<usize, usize> =
16 std::collections::HashMap::new();
17 for term in &terms {
18 if let Some(indices) = index.get(*term) {
19 for &idx in indices {
20 if self.facts[idx].is_current() {
21 *match_counts.entry(idx).or_insert(0) += 1;
22 }
23 }
24 }
25 }
26
27 let mut results: Vec<(&KnowledgeFact, f32)> = match_counts
28 .into_iter()
29 .map(|(idx, count)| {
30 let f = &self.facts[idx];
31 let relevance = (count as f32 / terms.len() as f32) * f.quality_score();
32 (f, relevance)
33 })
34 .collect();
35
36 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
37 results.into_iter().map(|(f, _)| f).collect()
38 }
39
40 pub fn recall_by_category(&self, category: &str) -> Vec<&KnowledgeFact> {
41 self.facts
42 .iter()
43 .filter(|f| f.category == category && f.is_current())
44 .collect()
45 }
46
47 pub fn recall_at_time(&self, query: &str, at: DateTime<Utc>) -> Vec<&KnowledgeFact> {
48 let q = query.to_lowercase();
49 let terms: Vec<&str> = q.split_whitespace().collect();
50 if terms.is_empty() {
51 return Vec::new();
52 }
53
54 let index = build_token_index(&self.facts, false);
55 let mut match_counts: std::collections::HashMap<usize, usize> =
56 std::collections::HashMap::new();
57 for term in &terms {
58 if let Some(indices) = index.get(*term) {
59 for &idx in indices {
60 if self.facts[idx].was_valid_at(at) {
61 *match_counts.entry(idx).or_insert(0) += 1;
62 }
63 }
64 }
65 }
66
67 let mut results: Vec<(&KnowledgeFact, f32)> = match_counts
68 .into_iter()
69 .map(|(idx, count)| {
70 let f = &self.facts[idx];
71 (f, count as f32 / terms.len() as f32)
72 })
73 .collect();
74
75 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
76 results.into_iter().map(|(f, _)| f).collect()
77 }
78
79 pub fn timeline(&self, category: &str) -> Vec<&KnowledgeFact> {
80 let mut facts: Vec<&KnowledgeFact> = self
81 .facts
82 .iter()
83 .filter(|f| f.category == category)
84 .collect();
85 facts.sort_by_key(|x| x.created_at);
86 facts
87 }
88
89 pub fn list_rooms(&self) -> Vec<(String, usize)> {
90 let mut categories: std::collections::BTreeMap<String, usize> =
91 std::collections::BTreeMap::new();
92 for f in &self.facts {
93 if f.is_current() {
94 *categories.entry(f.category.clone()).or_insert(0) += 1;
95 }
96 }
97 categories.into_iter().collect()
98 }
99
100 pub fn recall_for_output(&mut self, query: &str, limit: usize) -> (Vec<KnowledgeFact>, usize) {
101 let q = query.to_lowercase();
102 let terms: Vec<&str> = q.split_whitespace().filter(|t| !t.is_empty()).collect();
103 if terms.is_empty() {
104 return (Vec::new(), 0);
105 }
106
107 let index = build_token_index(&self.facts, true);
108 let mut match_counts: std::collections::HashMap<usize, usize> =
109 std::collections::HashMap::new();
110 for term in &terms {
111 if let Some(indices) = index.get(*term) {
112 for &idx in indices {
113 if self.facts[idx].is_current() {
114 *match_counts.entry(idx).or_insert(0) += 1;
115 }
116 }
117 }
118 }
119
120 struct Scored {
121 idx: usize,
122 relevance: f32,
123 }
124
125 let mut scored: Vec<Scored> = match_counts
126 .into_iter()
127 .map(|(idx, count)| {
128 let f = &self.facts[idx];
129 let relevance = (count as f32 / terms.len() as f32) * f.confidence;
130 Scored { idx, relevance }
131 })
132 .collect();
133
134 scored.sort_by(|a, b| {
135 b.relevance
136 .partial_cmp(&a.relevance)
137 .unwrap_or(std::cmp::Ordering::Equal)
138 .then_with(|| sort_fact_for_output(&self.facts[a.idx], &self.facts[b.idx]))
139 });
140
141 let total = scored.len();
142 scored.truncate(limit);
143
144 let now = Utc::now();
145 let mut out: Vec<KnowledgeFact> = Vec::new();
146 for s in scored {
147 if let Some(f) = self.facts.get_mut(s.idx) {
148 f.retrieval_count = f.retrieval_count.saturating_add(1);
149 f.last_retrieved = Some(now);
150 out.push(f.clone());
151 }
152 }
153
154 (out, total)
155 }
156
157 pub fn recall_by_category_for_output(
158 &mut self,
159 category: &str,
160 limit: usize,
161 ) -> (Vec<KnowledgeFact>, usize) {
162 let mut idxs: Vec<usize> = self
163 .facts
164 .iter()
165 .enumerate()
166 .filter(|(_, f)| f.is_current() && f.category == category)
167 .map(|(i, _)| i)
168 .collect();
169
170 idxs.sort_by(|a, b| sort_fact_for_output(&self.facts[*a], &self.facts[*b]));
171
172 let total = idxs.len();
173 idxs.truncate(limit);
174
175 let now = Utc::now();
176 let mut out = Vec::new();
177 for idx in idxs {
178 if let Some(f) = self.facts.get_mut(idx) {
179 f.retrieval_count = f.retrieval_count.saturating_add(1);
180 f.last_retrieved = Some(now);
181 out.push(f.clone());
182 }
183 }
184
185 (out, total)
186 }
187}