riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(embedding,vision): candle embeddings with offline fallback, on-demand clipboard vision capture, and concurrency audit

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-07 06:36:09 +01:00
1 parent 5bd8b1587a
commit e4a0fe72df
47 files changed
+6292 -3503
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/search.rs
+77 -25
View File
@@ -27,6 +27,8 @@ pub struct MemoryIndex {
pub body_field: Field,
pub type_field: Field,
pub namespace_field: Field,
pub query_parser: tantivy::query::QueryParser,
}
impl MemoryIndex {
@@ -45,14 +47,14 @@ impl MemoryIndex {
let index = Index::open_in_dir(&index_dir)
.or_else(|_| Index::create_in_dir(&index_dir, schema.clone()))?;
let mut writer = index.writer(50_000_000)?;
writer.delete_all_documents()?;
writer.commit()?;
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,
@@ -63,6 +65,7 @@ impl MemoryIndex {
body_field,
type_field,
namespace_field,
query_parser,
})
}
@@ -76,14 +79,14 @@ impl MemoryIndex {
let schema = schema_builder.build();
let index = Index::create_in_ram(schema);
let mut writer = index.writer(50_000_000)?;
writer.delete_all_documents()?;
writer.commit()?;
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,
@@ -94,9 +97,20 @@ impl MemoryIndex {
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))
}
@@ -138,10 +152,18 @@ impl MemoryIndex {
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, t.acceptance_criteria.iter().map(|c| c.description.as_str()).collect::<Vec<_>>().join("\n")),
self.body_field => format!("{}\n{}", t.description, criteria_body),
self.type_field => "task",
self.namespace_field => "global"
);
@@ -172,6 +194,7 @@ impl MemoryIndex {
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());
@@ -181,6 +204,7 @@ impl MemoryIndex {
tantivy::TantivyError::SystemError(format!("First commit error: {}; Second commit error: {}", first_err, second_err))
})?;
}
let _ = reader.reload();
}
Ok(())
})
@@ -208,11 +232,9 @@ impl MemoryIndex {
namespace: Option<&str>,
) -> tantivy::Result<Vec<SearchResultTuple>> {
let searcher = self.reader.searcher();
let query_parser = tantivy::query::QueryParser::for_index(
&self.index,
vec![self.title_field, self.body_field],
);
let q = query_parser.parse_query(query).or_else(|_| {
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 {
@@ -223,17 +245,26 @@ impl MemoryIndex {
.collect();
let safe_query = sanitized.trim();
if safe_query.is_empty() {
query_parser.parse_query("*")
is_wildcard_fallback = true;
self.query_parser.parse_query("*")
} else {
query_parser.parse_query(safe_query)
self.query_parser.parse_query(safe_query)
}
})?;
let target_limit = if is_wildcard_fallback { 10 } else { 50 };
let fetch_limit = if namespace.is_some() {
// Fetch deeper candidate pool so other namespaces do not starve the target namespace
1000
} else {
target_limit
};
let top_docs = searcher.search(
&q,
&tantivy::collector::TopDocs::with_limit(50).order_by_score(),
&tantivy::collector::TopDocs::with_limit(fetch_limit).order_by_score(),
)?;
let mut results = Vec::with_capacity(top_docs.len());
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
@@ -267,6 +298,9 @@ impl MemoryIndex {
continue;
}
results.push((id, doc_type, title, body, score));
if results.len() >= target_limit {
break;
}
}
Ok(results)
}
@@ -319,59 +353,73 @@ impl MemoryIndex {
pub fn add_entity_sync(&self, e: &Entity) {
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.add_document(doc!(
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());
let _ = writer.delete_all_documents();
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());
let _res = writer.add_document(doc!(
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());
let _ = writer.add_document(doc!(
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());
let _ = writer.add_document(doc!(
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);
}
@@ -393,6 +441,7 @@ mod tests {
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();
@@ -408,6 +457,7 @@ mod tests {
dependencies: vec![],
parent_id: None,
expires_at: None,
..Default::default()
};
let _ = index.index_task(&task).await.unwrap();
@@ -419,6 +469,7 @@ mod tests {
updated_at: 0,
tags: vec![],
embedding: None,
..Default::default()
};
let _ = index.index_snippet(&snippet).await.unwrap();
@@ -431,6 +482,7 @@ mod tests {
status: "accepted".to_string(),
supersedes: None,
timestamp: 0,
..Default::default()
};
let _ = index.index_adr(&adr).await.unwrap();