clawhdf5/crates/clawhdf5-agent/src/search.rs

//! Search and agents_md methods for HDF5Memory.

use std::path::Path;

use std::collections::HashSet;

use crate::bm25;
use crate::confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence};
use crate::hybrid;
use crate::reranker::{ReRankConfig, RerankInput, rerank};
use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult};

/// Options for [`HDF5Memory::search`].
///
/// [`SearchOptions::new`] is plain hybrid search with the tuned default
/// fusion — the same as `hybrid_search_with(.., hybrid::DEFAULT_FUSION, k)`.
/// Every stage beyond that is opt-in.
#[derive(Debug, Clone)]
pub struct SearchOptions {
    /// Number of results to return.
    pub k: usize,
    /// How the vector and keyword stages are combined.
    pub fusion: hybrid::Fusion,
    /// Only consider records whose `source_channel` is one of these. The
    /// filter applies *before* ranking, so a filtered search still returns up
    /// to `k` results and scores are normalised over the records it can
    /// return. `None` searches everything; an empty list matches nothing.
    pub source_channels: Option<Vec<String>>,
    /// Re-rank a candidate pool by retrieval relevance, recency, source
    /// authority and activation — the pipeline the OpenClaw backend runs.
    pub rerank: Option<ReRankConfig>,
    /// Candidates retrieved for re-ranking; 0 means `max(3k, 10)`.
    pub rerank_pool: usize,
    /// Drop low-confidence results (after re-ranking, when that is on).
    pub confidence: Option<ConfidenceConfig>,
    /// The time recency is measured from, in seconds since the epoch.
    /// `None` uses the system clock.
    pub now: Option<f64>,
}

impl SearchOptions {
    pub fn new(k: usize) -> Self {
        Self {
            k,
            fusion: hybrid::DEFAULT_FUSION,
            source_channels: None,
            rerank: None,
            rerank_pool: 0,
            confidence: None,
            now: None,
        }
    }

    pub fn with_fusion(mut self, fusion: hybrid::Fusion) -> Self {
        self.fusion = fusion;
        self
    }

    /// Search only records from these source channels.
    pub fn with_sources<S: Into<String>>(mut self, channels: impl IntoIterator<Item = S>) -> Self {
        self.source_channels = Some(channels.into_iter().map(Into::into).collect());
        self
    }

    pub fn with_rerank(mut self, config: ReRankConfig) -> Self {
        self.rerank = Some(config);
        self
    }

    pub fn with_confidence(mut self, config: ConfidenceConfig) -> Self {
        self.confidence = Some(config);
        self
    }

    /// Measure recency from `now` (seconds since the epoch) instead of the
    /// system clock — for reproducible results and tests.
    pub fn at_time(mut self, now: f64) -> Self {
        self.now = Some(now);
        self
    }
}

impl Default for SearchOptions {
    fn default() -> Self {
        Self::new(10)
    }
}

impl HDF5Memory {
    /// Vector + keyword scoring stage of [`HDF5Memory::search`].
    ///
    /// Without the `hnsw` feature this is a full linear cosine scan (the exact
    /// previous behaviour, also used as the correctness oracle in tests). With
    /// `hnsw` enabled and an index available, the vector candidates come from an
    /// approximate-nearest-neighbour search over an over-fetched pool, then merge
    /// with BM25 via the shared [`hybrid::merge_vector_keyword`].
    ///
    /// `exclude`, when given, marks records that must not be returned (1 =
    /// excluded; it covers tombstones too). The index is over-fetched in
    /// proportion to how much the mask removes. Surfacing `pool` candidates
    /// costs the index roughly `pool × M` distance evaluations, while an exact
    /// scan of the allowed records costs one each — so whenever that scan is
    /// the cheaper of the two it is used instead, and it is also the fallback
    /// if the pool comes back with too few allowed hits (the allowed records
    /// sit away from the query). A filtered search never comes back short.
    #[cfg(feature = "hnsw")]
    fn vector_keyword_search(
        &mut self,
        query_embedding: &[f32],
        query_text: &str,
        bm25: &bm25::BM25Index,
        fusion: hybrid::Fusion,
        k: usize,
        exclude: Option<&[u8]>,
    ) -> Vec<(usize, f32)> {
        self.ensure_hnsw_fresh();
        let n = self.cache.len();
        // Over-fetch so the merge sees a useful vector pool. `ef` is
        // configurable, but the pool the fusion stage sees is not tied to it:
        // a caller lowering `ef` for speed should not silently narrow what
        // fusion has to work with.
        let mut pool = (k * 8).max(64);
        let mut allowed = n;
        if let Some(ex) = exclude {
            allowed = ex.iter().filter(|&&e| e == 0).count();
            if allowed == 0 {
                return Vec::new();
            }
            // Expect `pool` allowed hits if the filter is independent of the
            // query's neighbourhood.
            pool = pool.saturating_mul(n).div_ceil(allowed);
            if allowed <= pool.saturating_mul(self.hnsw_m()) {
                return self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex);
            }
        }
        match self.hnsw.as_ref() {
            Some(index) if !index.is_empty() && index.dimension() == query_embedding.len() => {
                let ef = self.hnsw_ef_search(k).max(pool);
                let candidates = index.search(query_embedding, pool, ef);
                // A quantised index returns approximate distances, and no
                // amount of `ef` fixes that — the loss is in the distances,
                // not the graph. Re-score the pool against the cache's exact
                // embeddings, which cost nothing extra to keep: recall then
                // matches an f32 index. See `BENCHMARKS.md`.
                let exact = index.storage() == clawhdf5_ann::Storage::Int8;
                let vec_scores: Vec<(usize, f32)> = candidates
                    .into_iter()
                    .filter(|(id, _)| exclude.is_none_or(|ex| ex[*id] == 0))
                    .map(|(id, dist)| {
                        let score = if exact {
                            crate::vector_search::cosine_similarity(
                                query_embedding,
                                &self.cache.embeddings[id],
                            )
                        } else {
                            1.0 - dist
                        };
                        (id, score)
                    })
                    .collect();
                // Fusion normalises over every keyword match, so it needs all
                // the scores — but not ranked.
                let mut kw_scores = bm25.scores(query_text);
                if let Some(ex) = exclude {
                    if vec_scores.len() < k.min(allowed) {
                        // The allowed records are not where the index looked.
                        return self.exact_masked_search(
                            query_embedding,
                            query_text,
                            bm25,
                            fusion,
                            k,
                            ex,
                        );
                    }
                    kw_scores.retain(|(id, _)| ex[*id] == 0);
                }
                hybrid::fuse(vec_scores, kw_scores, fusion, k)
            }
            _ => match exclude {
                Some(ex) => {
                    self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex)
                }
                None => hybrid::hybrid_search_fused(
                    query_embedding,
                    query_text,
                    &self.cache.embeddings,
                    &self.cache.chunks,
                    &self.cache.tombstones,
                    bm25,
                    fusion,
                    k,
                ),
            },
        }
    }

    #[cfg(not(feature = "hnsw"))]
    fn vector_keyword_search(
        &mut self,
        query_embedding: &[f32],
        query_text: &str,
        bm25: &bm25::BM25Index,
        fusion: hybrid::Fusion,
        k: usize,
        exclude: Option<&[u8]>,
    ) -> Vec<(usize, f32)> {
        match exclude {
            Some(ex) => self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex),
            None => hybrid::hybrid_search_fused(
                query_embedding,
                query_text,
                &self.cache.embeddings,
                &self.cache.chunks,
                &self.cache.tombstones,
                bm25,
                fusion,
                k,
            ),
        }
    }

    /// Exact hybrid search over the records `exclude` leaves (0 = allowed).
    fn exact_masked_search(
        &self,
        query_embedding: &[f32],
        query_text: &str,
        bm25: &bm25::BM25Index,
        fusion: hybrid::Fusion,
        k: usize,
        exclude: &[u8],
    ) -> Vec<(usize, f32)> {
        let vec_scores =
            hybrid::exact_vector_scores(query_embedding, &self.cache.embeddings, exclude);
        let mut kw_scores = bm25.scores(query_text);
        kw_scores.retain(|(id, _)| exclude.get(*id) == Some(&0));
        hybrid::fuse(vec_scores, kw_scores, fusion, k)
    }

    /// The exclusion mask for a source-channel filter: 1 for a tombstoned
    /// record or one from a channel not in `channels`.
    fn source_mask(&self, channels: &[String]) -> Vec<u8> {
        let allowed: HashSet<&str> = channels.iter().map(String::as_str).collect();
        self.cache
            .source_channels
            .iter()
            .zip(&self.cache.tombstones)
            .map(|(ch, &t)| u8::from(t != 0 || !allowed.contains(ch.as_str())))
            .collect()
    }

    /// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
    pub fn hybrid_search(
        &mut self,
        query_embedding: &[f32],
        query_text: &str,
        vector_weight: f32,
        keyword_weight: f32,
        k: usize,
    ) -> Vec<SearchResult> {
        self.hybrid_search_with(
            query_embedding,
            query_text,
            hybrid::Fusion::Weighted {
                vector: vector_weight,
                keyword: keyword_weight,
            },
            k,
        )
    }

    /// [`HDF5Memory::hybrid_search`] with the fusion method chosen explicitly.
    ///
    /// [`hybrid::DEFAULT_FUSION`] is what the weighted form defaults to;
    /// [`hybrid::Fusion::Rrf`] combines the two stages by rank instead of by
    /// score.
    pub fn hybrid_search_with(
        &mut self,
        query_embedding: &[f32],
        query_text: &str,
        fusion: hybrid::Fusion,
        k: usize,
    ) -> Vec<SearchResult> {
        self.search(
            query_embedding,
            query_text,
            &SearchOptions::new(k).with_fusion(fusion),
        )
    }

    /// Hybrid search with optional source filtering, re-ranking and
    /// confidence rejection — see [`SearchOptions`].
    ///
    /// Stages, in order: vector + keyword retrieval over the records the
    /// source filter allows; fusion; scaling by Hebbian activation; re-ranking
    /// (if on) of a `rerank_pool` of candidates; confidence rejection (if on);
    /// the top `k`. The records returned with a positive score get their
    /// Hebbian boost.
    pub fn search(
        &mut self,
        query_embedding: &[f32],
        query_text: &str,
        options: &SearchOptions,
    ) -> Vec<SearchResult> {
        let k = options.k;
        let fetch = match options.rerank {
            Some(_) if options.rerank_pool > 0 => options.rerank_pool.max(k),
            Some(_) => k.saturating_mul(3).max(10),
            None => k,
        };
        let exclude = options
            .source_channels
            .as_deref()
            .map(|channels| self.source_mask(channels));

        // The keyword index lives for the life of the store and is updated
        // incrementally. Take it out for the duration of the call so the
        // vector stage can borrow `self` mutably, then put it back.
        self.ensure_bm25_fresh();
        let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
        let scored = self.vector_keyword_search(
            query_embedding,
            query_text,
            &bm25,
            options.fusion,
            fetch,
            exclude.as_deref(),
        );
        self.bm25 = Some(bm25);

        let mut results: Vec<SearchResult> = scored
            .into_iter()
            .map(|(idx, score)| {
                let w = self.cache.activation_weights[idx];
                SearchResult {
                    score: score * w.sqrt(),
                    chunk: self.cache.chunks[idx].clone(),
                    index: idx,
                    timestamp: self.cache.timestamps[idx],
                    source_channel: self.cache.source_channels[idx].clone(),
                    activation: w,
                }
            })
            .collect();
        // Ties broken by index so results (and therefore which records get
        // boosted) don't depend on HashMap iteration order upstream.
        results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then(a.index.cmp(&b.index))
        });

        if let Some(config) = &options.rerank {
            results = Self::rerank_results(results, config, options.now);
        }
        if let Some(config) = &options.confidence {
            let scored: Vec<ScoredResult> = results
                .iter()
                .map(|r| ScoredResult {
                    index: r.index,
                    score: r.score,
                })
                .collect();
            let keep: HashSet<usize> = reject_low_confidence(&scored, config)
                .into_iter()
                .map(|r| r.index)
                .collect();
            results.retain(|r| keep.contains(&r.index));
        }
        results.truncate(k);

        // Only reinforce records that actually matched. When fewer than `k`
        // records are relevant, the rest of the list is zero-score filler;
        // boosting it would teach the store that arbitrary records are
        // important just because they were nearby in iteration order.
        let hit_indices: Vec<usize> = results
            .iter()
            .filter(|r| r.score > 0.0)
            .map(|r| r.index)
            .collect();
        self.apply_hebbian_boost(&hit_indices);

        results
    }

    /// Reorder by the re-ranker's combined score, which also becomes each
    /// result's `score`.
    fn rerank_results(
        results: Vec<SearchResult>,
        config: &ReRankConfig,
        now: Option<f64>,
    ) -> Vec<SearchResult> {
        let now = now.unwrap_or_else(|| {
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_secs_f64())
                .unwrap_or(0.0)
        });
        let inputs: Vec<RerankInput> = results
            .iter()
            .map(|r| RerankInput {
                index: r.index,
                timestamp: r.timestamp,
                source_channel: r.source_channel.clone(),
                raw_activation: r.activation,
                relevance: r.score,
            })
            .collect();
        let mut by_index: std::collections::HashMap<usize, SearchResult> =
            results.into_iter().map(|r| (r.index, r)).collect();
        rerank(&inputs, config, now)
            .into_iter()
            .filter_map(|rr| {
                let mut r = by_index.remove(&rr.index)?;
                r.score = rr.combined_score;
                Some(r)
            })
            .collect()
    }

    /// Reinforce the records a query returned. The new weights are persisted by
    /// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not
    /// by rewriting the whole store inside the query, which is what made
    /// `hybrid_search` cost O(store size) in disk I/O. They are a ranking hint,
    /// not user data: a crash before the next checkpoint only forgets the
    /// boosts since the last one.
    fn apply_hebbian_boost(&mut self, hit_indices: &[usize]) {
        if hit_indices.is_empty() || self.config.hebbian_boost == 0.0 {
            return;
        }
        for &idx in hit_indices {
            let w = &mut self.cache.activation_weights[idx];
            *w = (*w + self.config.hebbian_boost).min(MAX_ACTIVATION_WEIGHT);
        }
        self.activations_dirty = true;
    }

    /// Get the chunk text for a memory entry by index.
    pub fn get_chunk(&self, index: usize) -> Option<&str> {
        if index < self.cache.chunks.len() && self.cache.tombstones[index] == 0 {
            Some(&self.cache.chunks[index])
        } else {
            None
        }
    }

    /// Generate an AGENTS.md string from current memory state.
    pub fn generate_agents_md(&self) -> String {
        crate::agents_md::generate(&self.config, &self.cache, &self.sessions, &self.knowledge)
    }

    /// Write AGENTS.md to disk alongside the .h5 file.
    pub fn write_agents_md(&self) -> Result<()> {
        let md = self.generate_agents_md();
        let md_path = self.config.path.with_extension("agents.md");
        std::fs::write(&md_path, md).map_err(MemoryError::Io)
    }

    /// Read AGENTS.md from disk (if it exists).
    pub fn read_agents_md(path: &Path) -> Result<String> {
        let md_path = path.with_extension("agents.md");
        std::fs::read_to_string(&md_path).map_err(MemoryError::Io)
    }
}