use crate::db::CodeChunk;
use crate::search::tokenizer::tokenize_code;
use anyhow::{Context, Result};
use sha2::{Digest, Sha256};
use std::path::Path;
use tantivy::collector::TopDocs;
use tantivy::query::QueryParser;
use tantivy::schema::*;
use tantivy::{doc, Index, IndexWriter, ReloadPolicy, TantivyDocument, Term};
const HEAP_SIZE: usize = 50_000_000;
#[derive(Debug, Clone)]
pub struct BM25SearchResult {
pub chunk_id: String,
pub file_path: String,
pub content: String,
pub start_line: usize,
pub end_line: usize,
pub chunk_type: String,
pub language: String,
pub symbol_name: Option<String>,
pub score: f32,
pub parent_symbol: Option<String>,
pub signature: Option<String>,
pub doc_comment: Option<String>,
pub module_path: Option<String>,
}
pub struct BM25Index {
index: Index,
reader: tantivy::IndexReader,
fields: IndexFields,
}
struct IndexFields {
chunk_id: Field,
file_path: Field,
content: Field,
content_tokens: Field, start_line: Field,
end_line: Field,
chunk_type: Field,
language: Field,
symbol_name: Field,
}
impl BM25Index {
pub fn new(index_path: &Path) -> Result<Self> {
let schema = build_schema();
let fields = IndexFields {
chunk_id: schema.get_field("chunk_id").unwrap(),
file_path: schema.get_field("file_path").unwrap(),
content: schema.get_field("content").unwrap(),
content_tokens: schema.get_field("content_tokens").unwrap(),
start_line: schema.get_field("start_line").unwrap(),
end_line: schema.get_field("end_line").unwrap(),
chunk_type: schema.get_field("chunk_type").unwrap(),
language: schema.get_field("language").unwrap(),
symbol_name: schema.get_field("symbol_name").unwrap(),
};
let index = if index_path.exists() {
Index::open_in_dir(index_path)?
} else {
std::fs::create_dir_all(index_path)?;
Index::create_in_dir(index_path, schema.clone())?
};
let reader = index
.reader_builder()
.reload_policy(ReloadPolicy::OnCommitWithDelay)
.try_into()?;
Ok(Self {
index,
reader,
fields,
})
}
pub fn index_chunks(&self, chunks: &[CodeChunk]) -> Result<()> {
let mut writer: IndexWriter<TantivyDocument> = self.index.writer(HEAP_SIZE)?;
for chunk in chunks {
let tokens = tokenize_code(&chunk.content);
let token_text = tokens.join(" ");
let doc = doc!(
self.fields.chunk_id => chunk.id.clone(),
self.fields.file_path => chunk.file_path.clone(),
self.fields.content => chunk.content.clone(),
self.fields.content_tokens => token_text,
self.fields.start_line => chunk.start_line as u64,
self.fields.end_line => chunk.end_line as u64,
self.fields.chunk_type => chunk.chunk_type.clone(),
self.fields.language => chunk.language.clone(),
self.fields.symbol_name => chunk.symbol_name.clone().unwrap_or_default(),
);
writer.add_document(doc)?;
}
writer.commit()?;
Ok(())
}
pub fn clear(&self) -> Result<()> {
let mut writer: IndexWriter<TantivyDocument> = self.index.writer(HEAP_SIZE)?;
writer.delete_all_documents()?;
writer.commit()?;
Ok(())
}
pub fn delete_by_chunk_ids(&self, chunk_ids: &[String]) -> Result<()> {
if chunk_ids.is_empty() {
return Ok(());
}
let mut writer: IndexWriter<TantivyDocument> = self.index.writer(HEAP_SIZE)?;
for id in chunk_ids {
let term = Term::from_field_text(self.fields.chunk_id, id);
writer.delete_term(term);
}
writer.commit()?;
Ok(())
}
pub fn search(&self, query: &str, limit: usize) -> Result<Vec<BM25SearchResult>> {
let searcher = self.reader.searcher();
let query_tokens = tokenize_code(query);
let query_text = query_tokens.join(" ");
let query_parser = QueryParser::for_index(
&self.index,
vec![
self.fields.content_tokens,
self.fields.content,
self.fields.symbol_name,
],
);
let query = query_parser
.parse_query(&query_text)
.context("Failed to parse query")?;
let top_docs = searcher
.search(&query, &TopDocs::with_limit(limit))
.context("Search failed")?;
let mut results = Vec::new();
for (score, doc_address) in top_docs {
let doc: TantivyDocument = searcher.doc(doc_address)?;
let chunk_id = doc
.get_first(self.fields.chunk_id)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let file_path = doc
.get_first(self.fields.file_path)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let content = doc
.get_first(self.fields.content)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let start_line = doc
.get_first(self.fields.start_line)
.and_then(|v| v.as_u64())
.unwrap_or(0) as usize;
let end_line = doc
.get_first(self.fields.end_line)
.and_then(|v| v.as_u64())
.unwrap_or(0) as usize;
let chunk_type = doc
.get_first(self.fields.chunk_type)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let language = doc
.get_first(self.fields.language)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let symbol_name = doc
.get_first(self.fields.symbol_name)
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.filter(|s| !s.is_empty());
results.push(BM25SearchResult {
chunk_id,
file_path,
content,
start_line,
end_line,
chunk_type,
language,
symbol_name,
score,
parent_symbol: None,
signature: None,
doc_comment: None,
module_path: None,
});
}
Ok(results)
}
pub fn schema_signature() -> Result<String> {
let schema = build_schema();
let entries: Vec<FieldEntry> = schema.fields().map(|(_, e)| e.clone()).collect();
let json = serde_json::to_string(&entries)?;
let mut hasher = Sha256::new();
hasher.update(json.as_bytes());
Ok(hex::encode(hasher.finalize()))
}
}
fn build_schema() -> Schema {
let mut schema_builder = Schema::builder();
schema_builder.add_text_field("chunk_id", STRING | STORED);
schema_builder.add_text_field("file_path", STRING | STORED);
schema_builder.add_text_field("content", TEXT | STORED);
schema_builder.add_text_field("content_tokens", TEXT); schema_builder.add_u64_field("start_line", STORED);
schema_builder.add_u64_field("end_line", STORED);
schema_builder.add_text_field("chunk_type", STRING | STORED);
schema_builder.add_text_field("language", STRING | STORED);
schema_builder.add_text_field("symbol_name", STRING | STORED);
schema_builder.build()
}