use anyhow::{anyhow, Result}; use serde_json::{json, Value}; use std::sync::Arc; use tokio::sync::RwLock; /// OpenSearch client for hybrid search (semantic + lexical) pub struct OpenSearchClient { hosts: Vec, client: reqwest::Client, cache: Arc>, } #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub struct SearchResult { pub id: String, pub chunk: String, pub score: f32, pub source: String, pub level: String, pub breadcrumb: Vec, pub method: String, // "semantic", "lexical", or "hybrid" } #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub struct HybridSearchResult { pub results: Vec, pub total: usize, pub query: String, pub search_method: String, } struct SearchCache { queries: std::collections::HashMap, ttl_secs: u64, } impl OpenSearchClient { /// Create new OpenSearch client pub fn new(hosts: Vec) -> Self { let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(30)) .build() .expect("Failed to create HTTP client"); Self { hosts, client, cache: Arc::new(RwLock::new(SearchCache { queries: std::collections::HashMap::new(), ttl_secs: 300, // 5 minute cache })), } } /// Get the primary host fn primary_host(&self) -> &str { &self.hosts[0] } /// Index a document (called on vault changes) pub async fn index_document( &self, doc_id: &str, content: &str, source: &str, level: &str, breadcrumb: Vec, jwt_token: &str, ) -> Result<()> { let url = format!( "https://{}/vault-*/_doc/{}", self.primary_host(), doc_id ); let body = json!({ "content": content, "source": source, "level": level, "breadcrumb": breadcrumb, "indexed_at": chrono::Utc::now().to_rfc3339(), }); let response = self .client .put(&url) .header("Authorization", format!("Bearer {}", jwt_token)) .json(&body) .send() .await?; if !response.status().is_success() { return Err(anyhow!( "OpenSearch index failed: {} {}", response.status(), response.text().await.unwrap_or_default() )); } // Invalidate cache after indexing self.cache.write().await.queries.clear(); Ok(()) } /// BM25 lexical search via OpenSearch async fn lexical_search( &self, query: &str, limit: usize, jwt_token: &str, ) -> Result)>> { let url = format!("https://{}/vault-*/_search", self.primary_host()); let search_body = json!({ "size": limit * 2, "query": { "multi_match": { "query": query, "fields": ["content^2", "source", "breadcrumb"], "fuzziness": "AUTO", "operator": "or" } }, "_source": ["content", "source", "level", "breadcrumb"] }); let response = self .client .get(&url) .header("Authorization", format!("Bearer {}", jwt_token)) .header("Content-Type", "application/json") .json(&search_body) .send() .await?; if !response.status().is_success() { return Err(anyhow!( "OpenSearch search failed: {} {}", response.status(), response.text().await.unwrap_or_default() )); } let result: Value = response.json().await?; let mut results = Vec::new(); if let Some(hits) = result["hits"]["hits"].as_array() { for hit in hits { let score = hit["_score"].as_f64().unwrap_or(0.0) as f32; let source = &hit["_source"]; let id = hit["_id"].as_str().unwrap_or("").to_string(); let chunk = source["content"].as_str().unwrap_or("").to_string(); let src = source["source"].as_str().unwrap_or("").to_string(); let level = source["level"].as_str().unwrap_or("L0").to_string(); let breadcrumb: Vec = source["breadcrumb"] .as_array() .map(|arr| { arr.iter() .filter_map(|v| v.as_str().map(|s| s.to_string())) .collect() }) .unwrap_or_default(); results.push((id, score, chunk, src, breadcrumb)); } } Ok(results) } /// Semantic search via pgvector (called from memory service) /// This is separate - pgvector search happens in PostgreSQL pub async fn semantic_search( &self, embedding: &[f32], limit: usize, jwt_token: &str, ) -> Result)>> { // NOTE: This is actually handled by pgvector in PostgreSQL // This method is a placeholder for consistency // The actual semantic search happens in crates/mem-cli/src/http_server.rs Err(anyhow!( "Semantic search must be done via pgvector in PostgreSQL, not OpenSearch" )) } /// Hybrid search: combine lexical (OpenSearch) + semantic (pgvector) pub async fn hybrid_search( &self, query: &str, semantic_results: Vec<(String, f32, String, String, Vec)>, jwt_token: &str, limit: usize, weights: &HybridWeights, ) -> Result { // Check cache { let cache = self.cache.read().await; if let Some((cached, timestamp)) = cache.queries.get(query) { if timestamp.elapsed().as_secs() < cache.ttl_secs { return Ok(cached.clone()); } } } // Perform lexical search let lexical_results = self .lexical_search(query, limit, jwt_token) .await .unwrap_or_default(); // Combine results let combined = self.combine_results( semantic_results, lexical_results, limit, weights, ); let result = HybridSearchResult { results: combined, total: limit, query: query.to_string(), search_method: "hybrid".to_string(), }; // Cache result { let mut cache = self.cache.write().await; cache.queries.insert(query.to_string(), (result.clone(), std::time::Instant::now())); } Ok(result) } /// Combine semantic and lexical results with reranking fn combine_results( &self, semantic: Vec<(String, f32, String, String, Vec)>, lexical: Vec<(String, f32, String, String, Vec)>, limit: usize, weights: &HybridWeights, ) -> Vec { use std::collections::HashMap; // Normalize scores to 0-1 let sem_max = semantic.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max); let lex_max = lexical.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max); let sem_norm = semantic.into_iter().map(|(id, s, chunk, src, bc)| { let normalized = if sem_max > 0.0 { s / sem_max } else { 0.0 }; (id, normalized, chunk, src, bc) }).collect::>(); let lex_norm = lexical.into_iter().map(|(id, s, chunk, src, bc)| { let normalized = if lex_max > 0.0 { s / lex_max } else { 0.0 }; (id, normalized, chunk, src, bc) }).collect::>(); // Combine with weighted average let mut combined: HashMap)> = HashMap::new(); for (id, sem_score, chunk, src, bc) in sem_norm { let lex_score = lex_norm .iter() .find(|(lid, _, _, _, _)| lid == &id) .map(|(_, s, _, _, _)| *s) .unwrap_or(0.0); let final_score = weights.semantic * sem_score + weights.lexical * lex_score; combined.insert(id, (final_score, chunk, src, bc)); } // Add lexical-only results for (id, lex_score, chunk, src, bc) in lex_norm { if !combined.contains_key(&id) { let final_score = weights.lexical * lex_score; combined.insert(id, (final_score, chunk, src, bc)); } } // Sort and take top-k let mut results: Vec<_> = combined .into_iter() .map(|(id, (score, chunk, src, bc))| SearchResult { id, chunk, score, source: src, level: "L1".to_string(), breadcrumb: bc, method: "hybrid".to_string(), }) .collect(); results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap()); results.truncate(limit); results } /// Health check pub async fn health(&self, jwt_token: &str) -> Result { let url = format!("https://{}/_cluster/health", self.primary_host()); let response = self .client .get(&url) .header("Authorization", format!("Bearer {}", jwt_token)) .send() .await?; Ok(response.status().is_success()) } } #[derive(Clone, Debug)] pub struct HybridWeights { pub semantic: f32, // 0.6 = 60% pub lexical: f32, // 0.4 = 40% } impl Default for HybridWeights { fn default() -> Self { Self { semantic: 0.6, lexical: 0.4, } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_hybrid_weights_sum() { let weights = HybridWeights::default(); assert!((weights.semantic + weights.lexical - 1.0).abs() < 0.01); } #[test] fn test_combine_results_ranking() { let client = OpenSearchClient::new(vec!["localhost:9200".to_string()]); let semantic = vec![ ( "doc1".to_string(), 0.9, "deployment content".to_string(), "deploy.md".to_string(), vec!["runbooks".to_string()], ), ( "doc2".to_string(), 0.7, "networking content".to_string(), "network.md".to_string(), vec!["docs".to_string()], ), ]; let lexical = vec![ ( "doc1".to_string(), 0.95, "deployment content".to_string(), "deploy.md".to_string(), vec!["runbooks".to_string()], ), ]; let weights = HybridWeights::default(); let results = client.combine_results(semantic, lexical, 10, &weights); assert_eq!(results.len(), 2); assert_eq!(results[0].id, "doc1"); // doc1 has both semantic and lexical scores assert!(results[0].score > results[1].score); } }