1use serde::{Deserialize, Serialize};
2use tantivy::{
3 Index, TantivyDocument,
4 collector::TopDocs,
5 doc,
6 query::QueryParser,
7 schema::{Field, STORED, STRING, SchemaBuilder, TEXT, Value},
8};
9
10use crate::{DkpError, error::DkpResult, pack::loader::Pack};
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct SearchResult {
14 pub id: String,
15 pub asset_type: String,
16 pub title: String,
17 pub excerpt: String,
18 pub score: f32,
19 pub source_file: String,
20}
21
22struct Fields {
23 id: Field,
24 asset_type: Field,
25 title: Field,
26 body: Field,
27 source_file: Field,
28}
29
30pub struct SearchIndex {
31 index: Index,
32 fields: Fields,
33}
34
35impl SearchIndex {
36 pub fn build(pack: &Pack) -> DkpResult<Self> {
37 let mut builder = SchemaBuilder::new();
38 let id = builder.add_text_field("id", STRING | STORED);
39 let asset_type = builder.add_text_field("asset_type", STRING | STORED);
40 let title = builder.add_text_field("title", TEXT | STORED);
41 let body = builder.add_text_field("body", TEXT | STORED);
42 let source_file = builder.add_text_field("source_file", STRING | STORED);
43 let schema = builder.build();
44 let fields = Fields {
45 id,
46 asset_type,
47 title,
48 body,
49 source_file,
50 };
51
52 let index = Index::create_in_ram(schema);
53 let mut writer = index
54 .writer(15_000_000)
55 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
56
57 if let Some(gf) = pack.load_glossary()? {
58 for t in &gf.terms {
59 let body = format!(
60 "{} {} {}",
61 t.definition,
62 t.aliases.join(" "),
63 t.tags.join(" ")
64 );
65 writer
66 .add_document(doc!(
67 fields.id => t.id.as_str(),
68 fields.asset_type => "term",
69 fields.title => t.term.as_str(),
70 fields.body => body.as_str(),
71 fields.source_file => "machine/glossary.json",
72 ))
73 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
74 }
75 }
76
77 for c in &pack.load_chunks()? {
78 let body = format!("{} {}", c.chunk_text, c.tags.join(" "));
79 writer
80 .add_document(doc!(
81 fields.id => c.id.as_str(),
82 fields.asset_type => "chunk",
83 fields.title => c.title.as_str(),
84 fields.body => body.as_str(),
85 fields.source_file => "machine/retrieval_chunks.jsonl",
86 ))
87 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
88 }
89
90 if let Some(rf) = pack.load_rules()? {
91 for r in &rf.rules {
92 let body = format!("{} {}", r.description, r.tags.join(" "));
93 writer
94 .add_document(doc!(
95 fields.id => r.id.as_str(),
96 fields.asset_type => "rule",
97 fields.title => r.title.as_str(),
98 fields.body => body.as_str(),
99 fields.source_file => "machine/rules.json",
100 ))
101 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
102 }
103 }
104
105 if let Some(cf) = pack.load_constraints()? {
106 for c in cf.all_constraints() {
107 let body = format!("{} {}", c.description, c.tags.join(" "));
108 writer
109 .add_document(doc!(
110 fields.id => c.id.as_str(),
111 fields.asset_type => "constraint",
112 fields.title => c.title.as_str(),
113 fields.body => body.as_str(),
114 fields.source_file => "machine/constraints.json",
115 ))
116 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
117 }
118 }
119
120 writer
121 .commit()
122 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
123
124 Ok(Self { index, fields })
125 }
126
127 pub fn search(&self, query: &str, limit: usize) -> DkpResult<Vec<SearchResult>> {
128 let reader = self
129 .index
130 .reader()
131 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
132 let searcher = reader.searcher();
133
134 let parser = QueryParser::for_index(&self.index, vec![self.fields.title, self.fields.body]);
135
136 let parsed = parser
137 .parse_query(query)
138 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
139
140 let top_docs: Vec<(f32, _)> = searcher
141 .search(&parsed, &TopDocs::with_limit(limit).order_by_score())
142 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
143
144 let mut results = Vec::with_capacity(top_docs.len());
145 for (score, addr) in top_docs {
146 let doc: TantivyDocument = searcher
147 .doc(addr)
148 .map_err(|e| DkpError::SearchIndex(e.to_string()))?;
149
150 let get = |f: Field| -> String {
151 doc.get_first(f)
152 .and_then(|v| v.as_str())
153 .unwrap_or("")
154 .to_string()
155 };
156
157 let body = get(self.fields.body);
158 let excerpt = if body.len() > 120 {
159 format!("{}…", &body[..120])
160 } else {
161 body
162 };
163
164 results.push(SearchResult {
165 id: get(self.fields.id),
166 asset_type: get(self.fields.asset_type),
167 title: get(self.fields.title),
168 excerpt,
169 score,
170 source_file: get(self.fields.source_file),
171 });
172 }
173
174 Ok(results)
175 }
176}
177
178#[cfg(test)]
179mod tests {
180 use super::*;
181 use tempfile::TempDir;
182
183 fn pack_with_glossary_and_chunks(tmp: &TempDir) -> Pack {
184 std::fs::write(
185 tmp.path().join("manifest.json"),
186 r#"{
187 "spec": "1.0.0",
188 "name": "test-pack",
189 "version": "1.0.0",
190 "domain": "testing",
191 "audience": "internal",
192 "intended_use": "unit tests",
193 "known_limitations": "none",
194 "update_date": "2026-01-01"
195 }"#,
196 )
197 .unwrap();
198 let pack = Pack::open(tmp.path()).unwrap();
199 std::fs::create_dir_all(pack.machine_dir()).unwrap();
200 std::fs::write(
201 pack.machine_file("glossary.json"),
202 r#"{"terms": [{"id": "t1", "term": "Widget", "definition": "A mechanical device used in manufacturing."}]}"#,
203 )
204 .unwrap();
205 let jsonl = concat!(
206 r#"{"id":"c1","title":"Onboarding","chunk_text":"How to onboard a new employee","source_ref":"generated"}"#,
207 "\n",
208 r#"{"id":"c2","title":"Offboarding","chunk_text":"How to offboard a departing employee","source_ref":"generated"}"#,
209 "\n",
210 );
211 std::fs::write(pack.machine_file("retrieval_chunks.jsonl"), jsonl).unwrap();
212 pack
213 }
214
215 #[test]
216 fn query_matches_only_relevant_chunk() {
217 let tmp = TempDir::new().unwrap();
218 let pack = pack_with_glossary_and_chunks(&tmp);
219 let index = SearchIndex::build(&pack).unwrap();
220
221 let results = index.search("onboard", 10).unwrap();
222 assert!(!results.is_empty());
223 assert!(results.iter().any(|r| r.id == "c1"));
224 assert!(!results.iter().any(|r| r.id == "c2"));
225 }
226
227 #[test]
228 fn query_matches_glossary_term() {
229 let tmp = TempDir::new().unwrap();
230 let pack = pack_with_glossary_and_chunks(&tmp);
231 let index = SearchIndex::build(&pack).unwrap();
232
233 let results = index.search("widget", 10).unwrap();
234 assert!(
235 results
236 .iter()
237 .any(|r| r.id == "t1" && r.asset_type == "term")
238 );
239 }
240
241 #[test]
242 fn no_match_query_returns_empty_without_panicking() {
243 let tmp = TempDir::new().unwrap();
244 let pack = pack_with_glossary_and_chunks(&tmp);
245 let index = SearchIndex::build(&pack).unwrap();
246
247 let results = index.search("zzznonexistentzzz", 10).unwrap();
248 assert!(results.is_empty());
249 }
250
251 #[test]
252 fn empty_pack_index_builds_and_returns_no_results() {
253 let tmp = TempDir::new().unwrap();
254 std::fs::write(
255 tmp.path().join("manifest.json"),
256 r#"{
257 "spec": "1.0.0",
258 "name": "empty-pack",
259 "version": "1.0.0",
260 "domain": "testing",
261 "audience": "internal",
262 "intended_use": "unit tests",
263 "known_limitations": "none",
264 "update_date": "2026-01-01"
265 }"#,
266 )
267 .unwrap();
268 let pack = Pack::open(tmp.path()).unwrap();
269 let index = SearchIndex::build(&pack).unwrap();
270
271 let results = index.search("anything", 10).unwrap();
272 assert!(results.is_empty());
273 }
274
275 #[test]
276 fn search_limit_is_respected() {
277 let tmp = TempDir::new().unwrap();
278 let pack = pack_with_glossary_and_chunks(&tmp);
279 let index = SearchIndex::build(&pack).unwrap();
280
281 let results = index.search("employee", 1).unwrap();
282 assert!(results.len() <= 1);
283 }
284}