poimen-memory/crates/mem-cli/src/opensearch_client.rs

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<String>,
    client: reqwest::Client,
    cache: Arc<RwLock<SearchCache>>,
}

#[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<String>,
    pub method: String,  // "semantic", "lexical", or "hybrid"
}

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct HybridSearchResult {
    pub results: Vec<SearchResult>,
    pub total: usize,
    pub query: String,
    pub search_method: String,
}

struct SearchCache {
    queries: std::collections::HashMap<String, (HybridSearchResult, std::time::Instant)>,
    ttl_secs: u64,
}

impl OpenSearchClient {
    /// Create new OpenSearch client
    pub fn new(hosts: Vec<String>) -> 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<String>,
        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<Vec<(String, f32, String, String, Vec<String>)>> {
        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<String> = 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<Vec<(String, f32, String, String, Vec<String>)>> {
        // 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<String>)>,
        jwt_token: &str,
        limit: usize,
        weights: &HybridWeights,
    ) -> Result<HybridSearchResult> {
        // 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<String>)>,
        lexical: Vec<(String, f32, String, String, Vec<String>)>,
        limit: usize,
        weights: &HybridWeights,
    ) -> Vec<SearchResult> {
        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::<Vec<_>>();

        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::<Vec<_>>();

        // Combine with weighted average
        let mut combined: HashMap<String, (f32, String, String, Vec<String>)> = 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<bool> {
        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);
    }
}