751 lines
27 KiB
Rust
751 lines
27 KiB
Rust
use crate::embedding::{cosine_similarity, generate_embedding_async, generate_embeddings_async};
|
|
use crate::models::{Adr, Entity, Snippet, Task};
|
|
use crate::state::MemoryState;
|
|
use std::sync::{Arc, Mutex};
|
|
use tantivy::schema::*;
|
|
use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc};
|
|
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct SearchResult {
|
|
pub id: String,
|
|
pub doc_type: String,
|
|
pub title: String,
|
|
pub body: String,
|
|
pub score: f32,
|
|
}
|
|
|
|
pub type SearchResultTuple = (String, String, String, String, f32);
|
|
|
|
fn calculate_time_decay(timestamp_sec: Option<u64>) -> f32 {
|
|
let Some(ts) = timestamp_sec else { return 1.0 };
|
|
if ts == 0 {
|
|
return 1.0;
|
|
} // Avoid decaying default unset timestamps
|
|
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_secs();
|
|
|
|
if ts > now {
|
|
return 1.0;
|
|
}
|
|
|
|
let age_days = (now - ts) as f32 / 86400.0;
|
|
// 1% decay per day
|
|
(-0.01 * age_days).exp()
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct MemoryIndex {
|
|
pub index: Index,
|
|
pub reader: IndexReader,
|
|
writer: Arc<Mutex<IndexWriter>>,
|
|
needs_commit: Arc<std::sync::atomic::AtomicBool>,
|
|
|
|
// Schema fields
|
|
pub id_field: Field,
|
|
pub title_field: Field,
|
|
pub body_field: Field,
|
|
pub type_field: Field,
|
|
pub namespace_field: Field,
|
|
|
|
pub query_parser: tantivy::query::QueryParser,
|
|
}
|
|
|
|
impl MemoryIndex {
|
|
pub fn new(store_dir: &std::path::Path) -> tantivy::Result<Self> {
|
|
let mut schema_builder = Schema::builder();
|
|
let id_field = schema_builder.add_text_field("id", STRING | STORED);
|
|
let title_field = schema_builder.add_text_field("title", TEXT | STORED);
|
|
let body_field = schema_builder.add_text_field("body", TEXT | STORED);
|
|
let type_field = schema_builder.add_text_field("type", STRING | STORED);
|
|
let namespace_field = schema_builder.add_text_field("namespace", STRING | STORED);
|
|
let schema = schema_builder.build();
|
|
|
|
let index_dir = store_dir.join("tantivy_index");
|
|
std::fs::create_dir_all(&index_dir)
|
|
.map_err(|e| tantivy::TantivyError::SystemError(e.to_string()))?;
|
|
let index = Index::open_in_dir(&index_dir)
|
|
.or_else(|_| Index::create_in_dir(&index_dir, schema.clone()))?;
|
|
|
|
let writer = index.writer(50_000_000)?;
|
|
let reader = index
|
|
.reader_builder()
|
|
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
|
.try_into()?;
|
|
|
|
let query_parser =
|
|
tantivy::query::QueryParser::for_index(&index, vec![title_field, body_field]);
|
|
|
|
Ok(Self {
|
|
index,
|
|
reader,
|
|
writer: Arc::new(Mutex::new(writer)),
|
|
needs_commit: Arc::new(std::sync::atomic::AtomicBool::new(false)),
|
|
id_field,
|
|
title_field,
|
|
body_field,
|
|
type_field,
|
|
namespace_field,
|
|
query_parser,
|
|
})
|
|
}
|
|
|
|
pub fn new_in_ram() -> tantivy::Result<Self> {
|
|
let mut schema_builder = Schema::builder();
|
|
let id_field = schema_builder.add_text_field("id", STRING | STORED);
|
|
let title_field = schema_builder.add_text_field("title", TEXT | STORED);
|
|
let body_field = schema_builder.add_text_field("body", TEXT | STORED);
|
|
let type_field = schema_builder.add_text_field("type", STRING | STORED);
|
|
let namespace_field = schema_builder.add_text_field("namespace", STRING | STORED);
|
|
let schema = schema_builder.build();
|
|
|
|
let index = Index::create_in_ram(schema);
|
|
let writer = index.writer(50_000_000)?;
|
|
let reader = index
|
|
.reader_builder()
|
|
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
|
.try_into()?;
|
|
|
|
let query_parser =
|
|
tantivy::query::QueryParser::for_index(&index, vec![title_field, body_field]);
|
|
|
|
Ok(Self {
|
|
index,
|
|
reader,
|
|
writer: Arc::new(Mutex::new(writer)),
|
|
needs_commit: Arc::new(std::sync::atomic::AtomicBool::new(false)),
|
|
id_field,
|
|
title_field,
|
|
body_field,
|
|
type_field,
|
|
namespace_field,
|
|
query_parser,
|
|
})
|
|
}
|
|
|
|
pub fn clear(&self) -> tantivy::Result<()> {
|
|
let mut writer = self
|
|
.writer
|
|
.lock()
|
|
.map_err(|e| tantivy::TantivyError::SystemError(e.to_string()))?;
|
|
writer.delete_all_documents()?;
|
|
writer.commit()?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
|
self.index_entities_batch(std::slice::from_ref(e))
|
|
}
|
|
|
|
pub fn index_entities_batch(
|
|
&self,
|
|
entities: &[Entity],
|
|
) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
|
let writer = Arc::clone(&self.writer);
|
|
let id_field = self.id_field;
|
|
let needs_commit = Arc::clone(&self.needs_commit);
|
|
|
|
let docs: Vec<(String, tantivy::TantivyDocument)> = entities
|
|
.iter()
|
|
.map(|e| {
|
|
let id_val = e.name.clone();
|
|
let doc = doc!(
|
|
self.id_field => e.name.as_str(),
|
|
self.title_field => e.name.as_str(),
|
|
self.body_field => e.observations.join(" "),
|
|
self.type_field => "entity",
|
|
self.namespace_field => e.namespace.as_str()
|
|
);
|
|
(id_val, doc)
|
|
})
|
|
.collect();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let writer_guard = writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
for (id_val, doc) in docs {
|
|
writer_guard.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
|
writer_guard.add_document(doc)?;
|
|
}
|
|
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
pub fn index_task(&self, t: &Task) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
|
let writer = Arc::clone(&self.writer);
|
|
let id_field = self.id_field;
|
|
let id_val = t.id.clone();
|
|
let needs_commit = Arc::clone(&self.needs_commit);
|
|
|
|
let criteria_body = t
|
|
.acceptance_criteria
|
|
.iter()
|
|
.fold(String::new(), |mut acc, c| {
|
|
if !acc.is_empty() {
|
|
acc.push('\n');
|
|
}
|
|
acc.push_str(&c.description);
|
|
acc
|
|
});
|
|
|
|
let doc = doc!(
|
|
self.id_field => t.id.as_str(),
|
|
self.title_field => t.title.as_str(),
|
|
self.body_field => format!("{}\n{}", t.description, criteria_body),
|
|
self.type_field => "task",
|
|
self.namespace_field => "global"
|
|
);
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
|
writer.add_document(doc)?;
|
|
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
pub fn index_clipboard_image(
|
|
&self,
|
|
id: &str,
|
|
file_path: &str,
|
|
ocr_text: &str,
|
|
) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
|
let writer = Arc::clone(&self.writer);
|
|
let id_field = self.id_field;
|
|
let id_val = id.to_string();
|
|
let needs_commit = Arc::clone(&self.needs_commit);
|
|
|
|
let doc = doc!(
|
|
self.id_field => id,
|
|
self.title_field => format!("Clipboard Screenshot: {}", file_path),
|
|
self.body_field => ocr_text,
|
|
self.type_field => "clipboard_image",
|
|
self.namespace_field => "global"
|
|
);
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
|
writer.add_document(doc)?;
|
|
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
pub fn delete_document(&self, id: &str) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
|
let writer = Arc::clone(&self.writer);
|
|
let id_field = self.id_field;
|
|
let id_val = id.to_string();
|
|
let needs_commit = Arc::clone(&self.needs_commit);
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
|
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
pub async fn commit(&self) -> tantivy::Result<()> {
|
|
let writer = Arc::clone(&self.writer);
|
|
let needs_commit = Arc::clone(&self.needs_commit);
|
|
let reader = self.reader.clone();
|
|
tokio::task::spawn_blocking(move || {
|
|
if needs_commit.swap(false, std::sync::atomic::Ordering::SeqCst) {
|
|
let mut writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
if let Err(first_err) = writer.commit() {
|
|
std::thread::sleep(std::time::Duration::from_millis(100));
|
|
writer.commit().map_err(|second_err| {
|
|
tantivy::TantivyError::SystemError(format!(
|
|
"First commit error: {}; Second commit error: {}",
|
|
first_err, second_err
|
|
))
|
|
})?;
|
|
}
|
|
let _ = reader.reload();
|
|
}
|
|
Ok(())
|
|
})
|
|
.await
|
|
.unwrap_or_else(|_| {
|
|
Err(tantivy::TantivyError::SystemError(
|
|
"Commit task panicked".to_string(),
|
|
))
|
|
})
|
|
}
|
|
|
|
pub fn commit_if_needed(&self) -> tantivy::Result<()> {
|
|
if self
|
|
.needs_commit
|
|
.swap(false, std::sync::atomic::Ordering::SeqCst)
|
|
{
|
|
let mut writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.commit()?;
|
|
let _ = self.reader.reload();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub fn search(
|
|
&self,
|
|
query: &str,
|
|
namespace: Option<&str>,
|
|
) -> tantivy::Result<Vec<SearchResultTuple>> {
|
|
let searcher = self.reader.searcher();
|
|
let mut is_wildcard_fallback = false;
|
|
|
|
let q = self.query_parser.parse_query(query).or_else(|_| {
|
|
let sanitized: String = query
|
|
.chars()
|
|
.map(|c| match c {
|
|
'+' | '-' | '&' | '|' | '!' | '(' | ')' | '{' | '}' | '[' | ']' | '^' | '"'
|
|
| '~' | '*' | '?' | ':' | '\\' | '/' => ' ',
|
|
_ => c,
|
|
})
|
|
.collect();
|
|
let safe_query = sanitized.trim();
|
|
if safe_query.is_empty() {
|
|
is_wildcard_fallback = true;
|
|
self.query_parser.parse_query("*")
|
|
} else {
|
|
self.query_parser.parse_query(safe_query)
|
|
}
|
|
})?;
|
|
|
|
let target_limit = if is_wildcard_fallback { 10 } else { 50 };
|
|
let final_query: Box<dyn tantivy::query::Query> = if let Some(ns) = namespace {
|
|
let ns_term_query: Box<dyn tantivy::query::Query> =
|
|
Box::new(tantivy::query::TermQuery::new(
|
|
tantivy::Term::from_field_text(self.namespace_field, ns),
|
|
tantivy::schema::IndexRecordOption::Basic,
|
|
));
|
|
let ns_filter: Box<dyn tantivy::query::Query> = if ns != "global" {
|
|
let global_term_query = Box::new(tantivy::query::TermQuery::new(
|
|
tantivy::Term::from_field_text(self.namespace_field, "global"),
|
|
tantivy::schema::IndexRecordOption::Basic,
|
|
));
|
|
Box::new(tantivy::query::BooleanQuery::new(vec![
|
|
(tantivy::query::Occur::Should, ns_term_query),
|
|
(tantivy::query::Occur::Should, global_term_query),
|
|
]))
|
|
} else {
|
|
ns_term_query
|
|
};
|
|
|
|
Box::new(tantivy::query::BooleanQuery::new(vec![
|
|
(tantivy::query::Occur::Must, q),
|
|
(tantivy::query::Occur::Must, ns_filter),
|
|
]))
|
|
} else {
|
|
q
|
|
};
|
|
|
|
let top_docs = searcher.search(
|
|
&final_query,
|
|
&tantivy::collector::TopDocs::with_limit(target_limit).order_by_score(),
|
|
)?;
|
|
|
|
let mut results = Vec::with_capacity(target_limit);
|
|
for (score, doc_address) in top_docs {
|
|
let retrieved_doc = searcher.doc::<tantivy::TantivyDocument>(doc_address)?;
|
|
let id = retrieved_doc
|
|
.get_first(self.id_field)
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
let doc_type = retrieved_doc
|
|
.get_first(self.type_field)
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
let title = retrieved_doc
|
|
.get_first(self.title_field)
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
let body = retrieved_doc
|
|
.get_first(self.body_field)
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
results.push((id, doc_type, title, body, score));
|
|
}
|
|
Ok(results)
|
|
}
|
|
|
|
pub fn index_snippet(&self, s: &Snippet) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
|
let writer = Arc::clone(&self.writer);
|
|
let id_field = self.id_field;
|
|
let id_val = s.name.clone();
|
|
let needs_commit = Arc::clone(&self.needs_commit);
|
|
|
|
let doc = doc!(
|
|
self.id_field => s.name.as_str(),
|
|
self.title_field => s.name.as_str(),
|
|
self.body_field => format!("{} {}\n{}", s.language, s.description, s.code),
|
|
self.type_field => "snippet",
|
|
self.namespace_field => "global"
|
|
);
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
|
writer.add_document(doc)?;
|
|
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
pub fn index_adr(&self, a: &Adr) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
|
let writer = Arc::clone(&self.writer);
|
|
let id_field = self.id_field;
|
|
let id_val = a.id.clone();
|
|
let needs_commit = Arc::clone(&self.needs_commit);
|
|
|
|
let doc = doc!(
|
|
self.id_field => a.id.as_str(),
|
|
self.title_field => a.title.as_str(),
|
|
self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
|
|
self.type_field => "adr",
|
|
self.namespace_field => "global"
|
|
);
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
|
writer.add_document(doc)?;
|
|
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
pub fn add_entity_sync(&self, e: &Entity) {
|
|
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &e.name));
|
|
if let Err(err) = writer.add_document(doc!(
|
|
self.id_field => e.name.as_str(),
|
|
self.title_field => e.name.as_str(),
|
|
self.body_field => e.observations.join(" "),
|
|
self.type_field => "entity",
|
|
self.namespace_field => e.namespace.as_str()
|
|
)) {
|
|
tracing::warn!("Sync entity indexing failed for '{}': {}", e.name, err);
|
|
}
|
|
self.needs_commit
|
|
.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
}
|
|
|
|
pub fn delete_all(&self) {
|
|
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
if let Err(err) = writer.delete_all_documents() {
|
|
tracing::warn!("delete_all_documents failed: {}", err);
|
|
}
|
|
self.needs_commit
|
|
.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
}
|
|
|
|
pub fn add_task_sync(&self, t: &Task) {
|
|
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &t.id));
|
|
if let Err(err) = writer.add_document(doc!(
|
|
self.id_field => t.id.as_str(),
|
|
self.title_field => t.title.as_str(),
|
|
self.body_field => t.description.as_str(),
|
|
self.type_field => "task",
|
|
self.namespace_field => "global"
|
|
)) {
|
|
tracing::warn!("Sync task indexing failed for '{}': {}", t.id, err);
|
|
}
|
|
self.needs_commit
|
|
.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
}
|
|
|
|
pub fn add_snippet_sync(&self, s: &Snippet) {
|
|
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &s.name));
|
|
if let Err(err) = writer.add_document(doc!(
|
|
self.id_field => s.name.as_str(),
|
|
self.title_field => s.name.as_str(),
|
|
self.body_field => format!("{} {}", s.language, s.description),
|
|
self.type_field => "snippet",
|
|
self.namespace_field => "global"
|
|
)) {
|
|
tracing::warn!("Sync snippet indexing failed for '{}': {}", s.name, err);
|
|
}
|
|
self.needs_commit
|
|
.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
}
|
|
|
|
pub fn add_adr_sync(&self, a: &Adr) {
|
|
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
|
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &a.id));
|
|
if let Err(err) = writer.add_document(doc!(
|
|
self.id_field => a.id.as_str(),
|
|
self.title_field => a.title.as_str(),
|
|
self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
|
|
self.type_field => "adr",
|
|
self.namespace_field => "global"
|
|
)) {
|
|
tracing::warn!("Sync ADR indexing failed for '{}': {}", a.id, err);
|
|
}
|
|
self.needs_commit
|
|
.store(true, std::sync::atomic::Ordering::SeqCst);
|
|
}
|
|
}
|
|
|
|
pub struct SearchService {
|
|
state: Arc<MemoryState>,
|
|
}
|
|
|
|
impl SearchService {
|
|
pub fn new(state: Arc<MemoryState>) -> Self {
|
|
Self { state }
|
|
}
|
|
|
|
pub async fn semantic_search(
|
|
&self,
|
|
query: &str,
|
|
filter_namespace: Option<&str>,
|
|
limit: usize,
|
|
) -> crate::error::Result<Vec<SearchResult>> {
|
|
let query_emb = generate_embedding_async(query.to_string())
|
|
.await
|
|
.unwrap_or_default();
|
|
if query_emb.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let mut results = Vec::new();
|
|
let mut uncached_texts = Vec::new();
|
|
let mut uncached_meta = Vec::new();
|
|
|
|
self.state.code.snippets.read_with(|snips| {
|
|
for snippet in snips.iter() {
|
|
if let Some(ref emb) = snippet.embedding {
|
|
let base_sim = cosine_similarity(&query_emb, emb);
|
|
let sim = base_sim * calculate_time_decay(Some(snippet.updated_at));
|
|
results.push(SearchResult {
|
|
id: snippet.name.clone(),
|
|
doc_type: "snippet".to_string(),
|
|
title: snippet.name.clone(),
|
|
body: snippet.description.clone(),
|
|
score: sim,
|
|
});
|
|
} else if uncached_texts.len() < 50 {
|
|
uncached_texts.push(format!(
|
|
"{} {} {}",
|
|
snippet.name, snippet.description, snippet.code
|
|
));
|
|
uncached_meta.push((
|
|
snippet.name.clone(),
|
|
"snippet".to_string(),
|
|
snippet.description.clone(),
|
|
Some(snippet.updated_at),
|
|
));
|
|
}
|
|
}
|
|
});
|
|
|
|
self.state.read_graph(|graph| {
|
|
for entity in graph.entities.values() {
|
|
if let Some(ns) = filter_namespace
|
|
&& entity.namespace != ns {
|
|
continue;
|
|
}
|
|
let obs = entity.observations.join("; ");
|
|
let desc = format!("{}: {}", entity.entity_type, obs);
|
|
if let Some(ref emb) = entity.embedding {
|
|
let base_sim = cosine_similarity(&query_emb, emb);
|
|
let sim = base_sim * calculate_time_decay(entity.updated_at);
|
|
results.push(SearchResult {
|
|
id: entity.name.clone(),
|
|
doc_type: "entity".to_string(),
|
|
title: entity.name.clone(),
|
|
body: desc,
|
|
score: sim,
|
|
});
|
|
} else if uncached_texts.len() < 50 {
|
|
uncached_texts.push(format!("{} {} {}", entity.name, entity.entity_type, obs));
|
|
uncached_meta.push((
|
|
entity.name.clone(),
|
|
"entity".to_string(),
|
|
desc,
|
|
entity.updated_at,
|
|
));
|
|
}
|
|
}
|
|
});
|
|
|
|
self.state.code.error_fixes.read_with(|fixes| {
|
|
for fix in fixes.iter() {
|
|
if let Some(ref emb) = fix.embedding {
|
|
let base_sim = cosine_similarity(&query_emb, emb);
|
|
let sim = base_sim * calculate_time_decay(Some(fix.timestamp));
|
|
results.push(SearchResult {
|
|
id: fix.signature.clone(),
|
|
doc_type: "error_fix".to_string(),
|
|
title: fix.signature.clone(),
|
|
body: fix.solution.clone(),
|
|
score: sim,
|
|
});
|
|
} else if uncached_texts.len() < 50 {
|
|
uncached_texts.push(format!("{} {}", fix.signature, fix.solution));
|
|
uncached_meta.push((
|
|
fix.signature.clone(),
|
|
"error_fix".to_string(),
|
|
fix.solution.clone(),
|
|
Some(fix.timestamp),
|
|
));
|
|
}
|
|
}
|
|
});
|
|
|
|
if !uncached_texts.is_empty()
|
|
&& let Ok(embeddings) = generate_embeddings_async(uncached_texts).await
|
|
{
|
|
for (emb, (title, doc_type, body, ts)) in embeddings.into_iter().zip(uncached_meta) {
|
|
let base_sim = cosine_similarity(&query_emb, &emb);
|
|
let sim = base_sim * calculate_time_decay(ts);
|
|
results.push(SearchResult {
|
|
id: title.clone(),
|
|
doc_type,
|
|
title,
|
|
body,
|
|
score: sim,
|
|
});
|
|
}
|
|
}
|
|
|
|
results.sort_by(|a, b| {
|
|
b.score
|
|
.partial_cmp(&a.score)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
results.truncate(limit);
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
pub async fn keyword_search(
|
|
&self,
|
|
query: &str,
|
|
filter_namespace: Option<&str>,
|
|
limit: usize,
|
|
) -> crate::error::Result<Vec<SearchResult>> {
|
|
let idx = self.state.get_search_index().await;
|
|
let matches = idx
|
|
.search(query, filter_namespace)
|
|
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
|
|
|
let mut results = Vec::new();
|
|
for (id, doc_type, title, body, score) in matches.into_iter().take(limit) {
|
|
results.push(SearchResult {
|
|
id,
|
|
doc_type,
|
|
title,
|
|
body,
|
|
score,
|
|
});
|
|
}
|
|
Ok(results)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tempfile::TempDir;
|
|
|
|
#[tokio::test]
|
|
async fn test_search_index_and_retrieve() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let index = MemoryIndex::new(temp_dir.path()).unwrap();
|
|
|
|
let entity = Entity {
|
|
name: "TestEntity".to_string(),
|
|
entity_type: "Component".to_string(),
|
|
observations: vec!["This is a test observation".to_string()],
|
|
namespace: "global".to_string(),
|
|
git_branch: None,
|
|
..Default::default()
|
|
};
|
|
let _ = index.index_entity(&entity).await.unwrap();
|
|
|
|
let task = Task {
|
|
id: "task-1".to_string(),
|
|
title: "Test Task".to_string(),
|
|
description: "Test task description".to_string(),
|
|
status: "open".to_string(),
|
|
created_at: 0,
|
|
updated_at: 0,
|
|
git_branch: None,
|
|
acceptance_criteria: vec![],
|
|
dependencies: vec![],
|
|
parent_id: None,
|
|
expires_at: None,
|
|
..Default::default()
|
|
};
|
|
let _ = index.index_task(&task).await.unwrap();
|
|
|
|
let snippet = Snippet {
|
|
name: "test_snippet".to_string(),
|
|
code: "fn main() {}".to_string(),
|
|
language: "rust".to_string(),
|
|
description: "A test snippet".to_string(),
|
|
updated_at: 0,
|
|
tags: vec![],
|
|
embedding: None,
|
|
..Default::default()
|
|
};
|
|
let _ = index.index_snippet(&snippet).await.unwrap();
|
|
|
|
let adr = Adr {
|
|
id: "adr-1".to_string(),
|
|
title: "Test ADR".to_string(),
|
|
context: "Test context".to_string(),
|
|
decision: "Test decision".to_string(),
|
|
consequence: "Test consequence".to_string(),
|
|
status: "accepted".to_string(),
|
|
supersedes: None,
|
|
timestamp: 0,
|
|
..Default::default()
|
|
};
|
|
let _ = index.index_adr(&adr).await.unwrap();
|
|
|
|
index.commit().await.unwrap();
|
|
index.reader.reload().unwrap();
|
|
|
|
// Test search
|
|
let results = index.search("observation", None).unwrap();
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, "TestEntity");
|
|
assert_eq!(results[0].1, "entity");
|
|
|
|
let results = index.search("task", None).unwrap();
|
|
assert!(results.iter().any(|r| r.0 == "task-1"));
|
|
|
|
let results = index.search("snippet", None).unwrap();
|
|
assert!(results.iter().any(|r| r.0 == "test_snippet"));
|
|
|
|
let results = index.search("decision", None).unwrap();
|
|
assert!(results.iter().any(|r| r.0 == "adr-1"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_search_malformed_query() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let index = MemoryIndex::new(temp_dir.path()).unwrap();
|
|
|
|
// Malformed lucene query (unclosed parenthesis) - sanitized gracefully
|
|
let result = index.search("title: (unclosed", None);
|
|
assert!(result.is_ok());
|
|
|
|
// Another malformed query (unclosed quote) - sanitized gracefully
|
|
let result2 = index.search("title: \"unclosed", None);
|
|
assert!(result2.is_ok());
|
|
}
|
|
}
|