dkp_core/search/
index.rs

1use serde::{Deserialize, Serialize};
2use tantivy::{
3    collector::TopDocs,
4    doc,
5    query::QueryParser,
6    schema::{Field, SchemaBuilder, Value, STORED, STRING, TEXT},
7    Index, TantivyDocument,
8};
9
10use crate::{error::DkpResult, pack::loader::Pack, DkpError};
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}
dkp_core/search/index.rs

dkp_core/search/
index.rs
