perf: optimize store batch serialization, zero-clone semantic search, and lock contention

server/src/search.rs

@@ -1,4 +1,6 @@
+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};
@@ -432,6 +434,154 @@ impl MemoryIndex {
     }
 }
 
+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 sim = cosine_similarity(&query_emb, emb);
+                    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(),
+                    ));
+                }
+            }
+        });
+
+        self.state.read_graph(|graph| {
+            for entity in graph.entities.values() {
+                if let Some(ns) = filter_namespace {
+                    if entity.namespace != ns {
+                        continue;
+                    }
+                }
+                let obs = entity.observations.join("; ");
+                let desc = format!("{}: {}", entity.entity_type, obs);
+                if let Some(ref emb) = entity.embedding {
+                    let sim = cosine_similarity(&query_emb, emb);
+                    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));
+                }
+            }
+        });
+
+        self.state.code.error_fixes.read_with(|fixes| {
+            for fix in fixes.iter() {
+                if let Some(ref emb) = fix.embedding {
+                    let sim = cosine_similarity(&query_emb, emb);
+                    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(),
+                    ));
+                }
+            }
+        });
+
+        if !uncached_texts.is_empty()
+            && let Ok(embeddings) = generate_embeddings_async(uncached_texts).await
+        {
+            for (emb, (title, doc_type, body)) in embeddings.into_iter().zip(uncached_meta) {
+                let sim = cosine_similarity(&query_emb, &emb);
+                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::*;