use serde::{Deserialize, Serialize};
use tantivy::{
collector::TopDocs,
doc,
query::QueryParser,
schema::{Field, SchemaBuilder, Value, STORED, STRING, TEXT},
Index, TantivyDocument,
};
use crate::{error::DkpResult, pack::loader::Pack, DkpError};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResult {
pub id: String,
pub asset_type: String,
pub title: String,
pub excerpt: String,
pub score: f32,
pub source_file: String,
}
struct Fields {
id: Field,
asset_type: Field,
title: Field,
body: Field,
source_file: Field,
}
pub struct SearchIndex {
index: Index,
fields: Fields,
}
impl SearchIndex {
pub fn build(pack: &Pack) -> DkpResult<Self> {
let mut builder = SchemaBuilder::new();
let id = builder.add_text_field("id", STRING | STORED);
let asset_type = builder.add_text_field("asset_type", STRING | STORED);
let title = builder.add_text_field("title", TEXT | STORED);
let body = builder.add_text_field("body", TEXT | STORED);
let source_file = builder.add_text_field("source_file", STRING | STORED);
let schema = builder.build();
let fields = Fields {
id,
asset_type,
title,
body,
source_file,
};
let index = Index::create_in_ram(schema);
let mut writer = index
.writer(15_000_000)
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
if let Some(gf) = pack.load_glossary()? {
for t in &gf.terms {
let body = format!(
"{} {} {}",
t.definition,
t.aliases.join(" "),
t.tags.join(" ")
);
writer
.add_document(doc!(
fields.id => t.id.as_str(),
fields.asset_type => "term",
fields.title => t.term.as_str(),
fields.body => body.as_str(),
fields.source_file => "machine/glossary.json",
))
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
}
}
for c in &pack.load_chunks()? {
let body = format!("{} {}", c.chunk_text, c.tags.join(" "));
writer
.add_document(doc!(
fields.id => c.id.as_str(),
fields.asset_type => "chunk",
fields.title => c.title.as_str(),
fields.body => body.as_str(),
fields.source_file => "machine/retrieval_chunks.jsonl",
))
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
}
if let Some(rf) = pack.load_rules()? {
for r in &rf.rules {
let body = format!("{} {}", r.description, r.tags.join(" "));
writer
.add_document(doc!(
fields.id => r.id.as_str(),
fields.asset_type => "rule",
fields.title => r.title.as_str(),
fields.body => body.as_str(),
fields.source_file => "machine/rules.json",
))
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
}
}
if let Some(cf) = pack.load_constraints()? {
for c in cf.all_constraints() {
let body = format!("{} {}", c.description, c.tags.join(" "));
writer
.add_document(doc!(
fields.id => c.id.as_str(),
fields.asset_type => "constraint",
fields.title => c.title.as_str(),
fields.body => body.as_str(),
fields.source_file => "machine/constraints.json",
))
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
}
}
writer
.commit()
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
Ok(Self { index, fields })
}
pub fn search(&self, query: &str, limit: usize) -> DkpResult<Vec<SearchResult>> {
let reader = self
.index
.reader()
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
let searcher = reader.searcher();
let parser = QueryParser::for_index(&self.index, vec![self.fields.title, self.fields.body]);
let parsed = parser
.parse_query(query)
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
let top_docs: Vec<(f32, _)> = searcher
.search(&parsed, &TopDocs::with_limit(limit).order_by_score())
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
let mut results = Vec::with_capacity(top_docs.len());
for (score, addr) in top_docs {
let doc: TantivyDocument = searcher
.doc(addr)
.map_err(|e| DkpError::SearchIndex(e.to_string()))?;
let get = |f: Field| -> String {
doc.get_first(f)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string()
};
let body = get(self.fields.body);
let excerpt = if body.len() > 120 {
format!("{}…", &body[..120])
} else {
body
};
results.push(SearchResult {
id: get(self.fields.id),
asset_type: get(self.fields.asset_type),
title: get(self.fields.title),
excerpt,
score,
source_file: get(self.fields.source_file),
});
}
Ok(results)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn pack_with_glossary_and_chunks(tmp: &TempDir) -> Pack {
std::fs::write(
tmp.path().join("manifest.json"),
r#"{
"spec": "1.0.0",
"name": "test-pack",
"version": "1.0.0",
"domain": "testing",
"audience": "internal",
"intended_use": "unit tests",
"known_limitations": "none",
"update_date": "2026-01-01"
}"#,
)
.unwrap();
let pack = Pack::open(tmp.path()).unwrap();
std::fs::create_dir_all(pack.machine_dir()).unwrap();
std::fs::write(
pack.machine_file("glossary.json"),
r#"{"terms": [{"id": "t1", "term": "Widget", "definition": "A mechanical device used in manufacturing."}]}"#,
)
.unwrap();
let jsonl = concat!(
r#"{"id":"c1","title":"Onboarding","chunk_text":"How to onboard a new employee","source_ref":"generated"}"#,
"\n",
r#"{"id":"c2","title":"Offboarding","chunk_text":"How to offboard a departing employee","source_ref":"generated"}"#,
"\n",
);
std::fs::write(pack.machine_file("retrieval_chunks.jsonl"), jsonl).unwrap();
pack
}
#[test]
fn query_matches_only_relevant_chunk() {
let tmp = TempDir::new().unwrap();
let pack = pack_with_glossary_and_chunks(&tmp);
let index = SearchIndex::build(&pack).unwrap();
let results = index.search("onboard", 10).unwrap();
assert!(!results.is_empty());
assert!(results.iter().any(|r| r.id == "c1"));
assert!(!results.iter().any(|r| r.id == "c2"));
}
#[test]
fn query_matches_glossary_term() {
let tmp = TempDir::new().unwrap();
let pack = pack_with_glossary_and_chunks(&tmp);
let index = SearchIndex::build(&pack).unwrap();
let results = index.search("widget", 10).unwrap();
assert!(results
.iter()
.any(|r| r.id == "t1" && r.asset_type == "term"));
}
#[test]
fn no_match_query_returns_empty_without_panicking() {
let tmp = TempDir::new().unwrap();
let pack = pack_with_glossary_and_chunks(&tmp);
let index = SearchIndex::build(&pack).unwrap();
let results = index.search("zzznonexistentzzz", 10).unwrap();
assert!(results.is_empty());
}
#[test]
fn empty_pack_index_builds_and_returns_no_results() {
let tmp = TempDir::new().unwrap();
std::fs::write(
tmp.path().join("manifest.json"),
r#"{
"spec": "1.0.0",
"name": "empty-pack",
"version": "1.0.0",
"domain": "testing",
"audience": "internal",
"intended_use": "unit tests",
"known_limitations": "none",
"update_date": "2026-01-01"
}"#,
)
.unwrap();
let pack = Pack::open(tmp.path()).unwrap();
let index = SearchIndex::build(&pack).unwrap();
let results = index.search("anything", 10).unwrap();
assert!(results.is_empty());
}
#[test]
fn search_limit_is_respected() {
let tmp = TempDir::new().unwrap();
let pack = pack_with_glossary_and_chunks(&tmp);
let index = SearchIndex::build(&pack).unwrap();
let results = index.search("employee", 1).unwrap();
assert!(results.len() <= 1);
}
}