edgehdf5_memory/
search.rs

1//! Search and agents_md methods for HDF5Memory.
2
3use std::path::Path;
4
5use crate::bm25;
6use crate::hybrid;
7use crate::{HDF5Memory, MemoryError, Result, SearchResult};
8
9impl HDF5Memory {
10    /// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
11    pub fn hybrid_search(
12        &mut self,
13        query_embedding: &[f32],
14        query_text: &str,
15        vector_weight: f32,
16        keyword_weight: f32,
17        k: usize,
18    ) -> Vec<SearchResult> {
19        let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones);
20        let scored = hybrid::hybrid_search(
21            query_embedding,
22            query_text,
23            &self.cache.embeddings,
24            &self.cache.chunks,
25            &self.cache.tombstones,
26            &bm25,
27            vector_weight,
28            keyword_weight,
29            k,
30        );
31        let mut results: Vec<SearchResult> = scored
32            .into_iter()
33            .map(|(idx, score)| {
34                let w = self.cache.activation_weights[idx];
35                SearchResult {
36                    score: score * w.sqrt(),
37                    chunk: self.cache.chunks[idx].clone(),
38                    index: idx,
39                    timestamp: self.cache.timestamps[idx],
40                    source_channel: self.cache.source_channels[idx].clone(),
41                    activation: w,
42                }
43            })
44            .collect();
45        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
46
47        let hit_indices: Vec<usize> = results.iter().map(|r| r.index).collect();
48        self.apply_hebbian_boost(&hit_indices);
49        self.flush().ok();
50
51        results
52    }
53
54    fn apply_hebbian_boost(&mut self, hit_indices: &[usize]) {
55        for &idx in hit_indices {
56            self.cache.activation_weights[idx] += self.config.hebbian_boost;
57        }
58    }
59
60    /// Get the chunk text for a memory entry by index.
61    pub fn get_chunk(&self, index: usize) -> Option<&str> {
62        if index < self.cache.chunks.len() && self.cache.tombstones[index] == 0 {
63            Some(&self.cache.chunks[index])
64        } else {
65            None
66        }
67    }
68
69    /// Generate an AGENTS.md string from current memory state.
70    pub fn generate_agents_md(&self) -> String {
71        crate::agents_md::generate(&self.config, &self.cache, &self.sessions, &self.knowledge)
72    }
73
74    /// Write AGENTS.md to disk alongside the .h5 file.
75    pub fn write_agents_md(&self) -> Result<()> {
76        let md = self.generate_agents_md();
77        let md_path = self.config.path.with_extension("agents.md");
78        std::fs::write(&md_path, md).map_err(MemoryError::Io)
79    }
80
81    /// Read AGENTS.md from disk (if it exists).
82    pub fn read_agents_md(path: &Path) -> Result<String> {
83        let md_path = path.with_extension("agents.md");
84        std::fs::read_to_string(&md_path).map_err(MemoryError::Io)
85    }
86}
edgehdf5_memory/search.rs

edgehdf5_memory/
search.rs
