do-memory-core 0.1.34

//! Cascading retrieval pipeline (WG-131).
//!
//! Implements a 4-tier retrieval cascade:
//! 1. BM25 keyword index (CPU-local, no API calls)
//! 2. HDC hyperdimensional encoding (CPU-local, no API calls)
//! 3. ConceptGraph ontology expansion (CPU-local, no API calls)
//! 4. API embedding fallback (external API call)
//!
//! The cascade eliminates 50-70% of embedding API calls by satisfying
//! queries from CPU-local tiers before falling back to the API.

mod concept_graph;
pub use concept_graph::ConceptGraph;

mod types;
pub use types::{CascadeConfig, CascadeResult, TierResult};

/// Cascading retrieval orchestrator.
///
/// Coordinates the 4-tier retrieval pipeline, falling back to API
/// only when CPU-local tiers cannot satisfy the query.
pub struct CascadeRetriever {
    config: CascadeConfig,
    /// Episode data indexed for retrieval (id -> text).
    episode_data: Vec<(String, String)>,
    /// Concept graph for ontology-based term expansion (Tier 3).
    #[cfg(feature = "csm")]
    concept_graph: ConceptGraph,
    #[cfg(feature = "csm")]
    bm25_index: super::Bm25Index,
    #[cfg(feature = "csm")]
    hdc_encoder: super::HdcEncoder,
    #[cfg(feature = "csm")]
    hdc_vectors: Vec<(String, super::HVec10240)>,
}

impl CascadeRetriever {
    /// Create a new cascade retriever with given configuration.
    pub fn new(config: CascadeConfig) -> Self {
        Self {
            config,
            episode_data: Vec::new(),
            #[cfg(feature = "csm")]
            concept_graph: ConceptGraph::from_embedded(),
            #[cfg(feature = "csm")]
            bm25_index: super::Bm25Index::new(),
            #[cfg(feature = "csm")]
            hdc_encoder: super::HdcEncoder::new(),
            #[cfg(feature = "csm")]
            hdc_vectors: Vec::new(),
        }
    }

    /// Create a new cascade retriever with default configuration.
    #[must_use]
    pub fn default_config() -> Self {
        Self::new(CascadeConfig::default())
    }

    /// Tokenize text for BM25 indexing/search.
    #[cfg(feature = "csm")]
    fn tokenize(text: &str) -> Vec<String> {
        // Use default tokenization: not code-aware, lowercase enabled
        super::HdcEncoder::tokenize(text, false, true)
    }

    /// Add an episode to the retrieval index.
    ///
    /// This indexes the episode in BM25 and encodes it for HDC similarity search.
    /// When the `csm` feature is not enabled, this just stores the episode data
    /// for later retrieval.
    pub fn add_episode(&mut self, id: &str, text: &str) {
        self.episode_data.push((id.to_string(), text.to_string()));

        #[cfg(feature = "csm")]
        {
            // Tokenize and add to BM25 index
            let tokens = Self::tokenize(text);
            self.bm25_index.add_document(id, &tokens);

            // Encode and store HDC vector
            let hdc_vector = self.hdc_encoder.encode(text);
            self.hdc_vectors.push((id.to_string(), hdc_vector));
        }
    }

    /// Clear all indexed episodes.
    pub fn clear(&mut self) {
        self.episode_data.clear();

        #[cfg(feature = "csm")]
        {
            self.bm25_index.clear();
            self.hdc_vectors.clear();
        }
    }

    /// Get the number of indexed episodes.
    #[must_use]
    pub fn len(&self) -> usize {
        self.episode_data.len()
    }

    /// Check if the index is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.episode_data.is_empty()
    }

    /// Execute the cascading retrieval pipeline.
    ///
    /// When the `csm` feature is enabled, this implements a 4-tier cascade:
    /// 1. BM25 keyword search (CPU-local, 0 API calls)
    /// 2. HDC similarity search (CPU-local, 0 API calls)
    /// 3. ConceptGraph expansion (CPU-local, 0 API calls)
    /// 4. API fallback (requires external embedding call)
    ///
    /// Without `csm`, returns empty results (placeholder behavior).
    pub fn retrieve(&self, query: &str) -> CascadeResult {
        #[cfg(feature = "csm")]
        {
            self.retrieve_with_csm(query)
        }

        #[cfg(not(feature = "csm"))]
        {
            tracing::warn!(
                "CSM feature not enabled; cascade retrieval returns empty results. \
                 Enable the `csm` feature for BM25/HDC/ConceptGraph retrieval."
            );
            let _ = query;
            CascadeResult {
                episode_ids: Vec::new(),
                scores: Vec::new(),
                contributing_tiers: Vec::new(),
                api_calls: 0,
            }
        }
    }

    /// Full cascade implementation using CSM components.
    #[cfg(feature = "csm")]
    fn retrieve_with_csm(&self, query: &str) -> CascadeResult {
        use super::{compute_weights, merge_results};

        // Tier 1: BM25 keyword search
        let bm25_results = self.retrieve_bm25(query);

        // Check if BM25 produced sufficient results
        if bm25_results.sufficient {
            return CascadeResult {
                episode_ids: bm25_results.ids(),
                scores: bm25_results.scores(),
                contributing_tiers: vec!["bm25".to_string()],
                api_calls: 0,
            };
        }

        // Tier 2: HDC similarity search
        let hdc_results = self.retrieve_hdc(query);

        // Check if HDC produced sufficient results (or merge with BM25)
        if self.config.merge_results && !bm25_results.is_empty() {
            // Merge BM25 and HDC results with query-length-dependent weights
            let weights = compute_weights(query.len());
            let merged = merge_results(&bm25_results.results, &hdc_results.results, weights);

            // Check if merged results are sufficient
            if merged.len() >= self.config.min_results {
                return CascadeResult {
                    episode_ids: merged.iter().map(|(id, _)| id.clone()).collect(),
                    scores: merged.iter().map(|(_, s)| *s).collect(),
                    contributing_tiers: vec!["bm25".to_string(), "hdc".to_string()],
                    api_calls: 0,
                };
            }
        } else if hdc_results.sufficient {
            return CascadeResult {
                episode_ids: hdc_results.ids(),
                scores: hdc_results.scores(),
                contributing_tiers: vec!["hdc".to_string()],
                api_calls: 0,
            };
        }

        // Tier 3: ConceptGraph expansion (optional)
        if self.config.enable_concept_expansion {
            let concept_results = self.retrieve_concept_graph(query);

            if concept_results.sufficient {
                return CascadeResult {
                    episode_ids: concept_results.ids(),
                    scores: concept_results.scores(),
                    contributing_tiers: vec!["concept_graph".to_string()],
                    api_calls: 0,
                };
            }
        }

        // Tier 4: API fallback - mark that we need an API call
        // Return best available results with api_calls = 1 indicator
        let best_results: Vec<(String, f32)> = if self.config.merge_results {
            let weights = compute_weights(query.len());
            merge_results(&bm25_results.results, &hdc_results.results, weights)
        } else if !hdc_results.is_empty() {
            hdc_results.results.clone()
        } else {
            bm25_results.results.clone()
        };

        CascadeResult {
            episode_ids: best_results.iter().map(|(id, _)| id.clone()).collect(),
            scores: best_results.iter().map(|(_, s)| *s).collect(),
            contributing_tiers: if best_results.is_empty() {
                vec!["none".to_string()]
            } else {
                // Preserve original tier attributions + note that API fallback was needed
                let mut tiers = Vec::new();
                if !bm25_results.is_empty() {
                    tiers.push("bm25".to_string());
                }
                if !hdc_results.is_empty() {
                    tiers.push("hdc".to_string());
                }
                tiers.push("api_fallback_needed".to_string());
                tiers
            },
            api_calls: 1, // Indicates API call would be needed
        }
    }

    /// BM25 keyword search (Tier 1).
    #[cfg(feature = "csm")]
    fn retrieve_bm25(&self, query: &str) -> TierResult {
        // Tokenize query for BM25 search
        let query_tokens = Self::tokenize(query);
        let raw_results = self.bm25_index.search(&query_tokens, self.config.top_k);

        // Normalize BM25 scores to 0.0-1.0 range
        let results = super::normalize_scores(&raw_results);

        // Determine if results are sufficient
        let sufficient = results.len() >= self.config.min_results
            && results
                .iter()
                .any(|(_, s)| *s >= self.config.bm25_threshold);

        TierResult {
            tier: "bm25".to_string(),
            results,
            sufficient,
        }
    }

    /// HDC hyperdimensional similarity search (Tier 2).
    #[cfg(feature = "csm")]
    fn retrieve_hdc(&self, query: &str) -> TierResult {
        // Encode query to HDC vector
        let query_vector = self.hdc_encoder.encode(query);

        // Compute similarities with all indexed vectors
        let mut similarities: Vec<(String, f32)> = self
            .hdc_vectors
            .iter()
            .map(|(id, vec)| {
                // Use cosine similarity (normalized hamming distance)
                let sim = query_vector.cosine_similarity(vec);
                (id.clone(), sim)
            })
            .collect();

        // Select top-k by similarity (highest first)
        // Optimization: O(N + k log k) instead of O(N log N)
        let top_k = self.config.top_k;
        let similarities = crate::search::select_top_k(&mut similarities, top_k, |a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.0.cmp(&b.0))
        });

        // Determine if results are sufficient
        let sufficient = similarities.len() >= self.config.min_results
            && similarities
                .iter()
                .any(|(_, s)| *s >= self.config.hdc_threshold);

        TierResult {
            tier: "hdc".to_string(),
            results: similarities,
            sufficient,
        }
    }

    /// ConceptGraph expansion search (Tier 3).
    ///
    /// Uses the embedded coding-agent domain ontology to expand query terms
    /// and match against indexed episode data using expanded terminology.
    ///
    /// # How it works
    ///
    /// 1. Expand query terms using the ontology (e.g., "auth" → "authentication")
    /// 2. Match expanded terms against episode text
    /// 3. Score results based on term overlap density
    #[cfg(feature = "csm")]
    fn retrieve_concept_graph(&self, query: &str) -> TierResult {
        // Expand query terms using the ontology
        let expanded_terms = self.concept_graph.expand_terms(query);

        if expanded_terms.is_empty() {
            return TierResult {
                tier: "concept_graph".to_string(),
                results: Vec::new(),
                sufficient: false,
            };
        }

        // Match expanded terms against episode data
        let mut scored: Vec<(String, f32)> = self
            .episode_data
            .iter()
            .map(|(id, text)| {
                let text_lower = text.to_lowercase();
                let match_count = expanded_terms
                    .iter()
                    .filter(|term| text_lower.split_whitespace().any(|w| w == term.as_str()))
                    .count();

                // Score based on term overlap density
                let score = if expanded_terms.is_empty() {
                    0.0
                } else {
                    match_count as f32 / expanded_terms.len() as f32
                };

                (id.clone(), score)
            })
            .filter(|(_, s)| *s >= self.config.concept_graph_threshold)
            .collect();

        // Select top-k by score using the shared select_top_k utility
        // (consistent with the HDC tier's approach)
        let top_k = self.config.top_k;
        let scored = crate::search::select_top_k(&mut scored, top_k, |a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.0.cmp(&b.0))
        });

        let sufficient = scored.len() >= self.config.min_results;

        TierResult {
            tier: "concept_graph".to_string(),
            results: scored,
            sufficient,
        }
    }

    /// Get the configuration for this retriever.
    pub fn config(&self) -> &CascadeConfig {
        &self.config
    }

    /// Estimate the probability that a query would require an API call.
    ///
    /// Returns a value in [0.0, 1.0] where:
    /// - 0.0 means CPU-local tiers (BM25/HDC/ConceptGraph) are very likely to suffice
    /// - 1.0 means an API embedding call is almost certainly needed
    ///
    /// Heuristic: short keyword-rich queries resolve via BM25 (low probability);
    /// long abstract queries with few known terms need semantic embedding (high probability).
    pub fn estimate_api_call_probability(&self, query: &str) -> f32 {
        let len = query.len() as f32;
        let word_count = query.split_whitespace().count() as f32;

        // Base probability from query length — short queries favor BM25
        let length_factor: f32 = if len < 20.0 {
            0.1
        } else if len < 50.0 {
            0.25
        } else if len < 100.0 {
            0.5
        } else {
            0.7
        };

        // Keyword density — queries with many short words are more BM25-friendly
        let avg_word_len = if word_count > 0.0 {
            len / word_count
        } else {
            10.0
        };
        let keyword_factor: f32 = if avg_word_len < 5.0 {
            0.0 // Short words = good keyword match candidates
        } else if avg_word_len < 8.0 {
            0.15
        } else {
            0.3 // Long words = more semantic, harder for BM25
        };

        // Concept-level boost — code-like tokens (identifiers, paths) are BM25-friendly
        let code_token_count = query
            .split_whitespace()
            .filter(|w| w.contains('_') || w.contains("::") || w.contains('/'))
            .count() as f32;
        let code_factor: f32 = if word_count > 0.0 && code_token_count / word_count > 0.3 {
            -0.15 // Many code tokens boost BM25 relevance
        } else {
            0.0
        };

        (length_factor + keyword_factor + code_factor).clamp(0.0, 1.0)
    }
}

#[cfg(test)]
mod tests;
#[cfg(feature = "csm")]
pub mod weights;
#[cfg(feature = "csm")]
pub use weights::compute_tier_weights;