holographic_memory/core/engine/
query.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4use super::HmsCore;
5use crate::core::entangled::EntangledHVec;
6use crate::core::hopfield;
7use crate::core::types::{IndexStatus, QueryExplanation, RetrievalResult, StorageHealth};
8use rayon::prelude::*;
9
10impl HmsCore {
11    /// Query the memory system for the k most similar vectors.
12    pub fn query(&self, query_vec: &EntangledHVec, k: u32) -> Vec<RetrievalResult> {
13        if query_vec.dim != self.dimensions || k == 0 {
14            return Vec::new();
15        }
16        // Documents memorized with deferred encoding join the vector index on first use; a
17        // failure leaves them pending (and absent from this query) and the next call retries.
18        // `try_read`: a query must keep progressing while compaction or a mutation holds the
19        // store for writing, so a contended lock means "no catch-up this time".
20        if self
21            .documents
22            .try_read()
23            .is_some_and(|store| store.pending_vector_count() > 0)
24        {
25            if let Err(error) = self.encode_pending_documents() {
26                tracing::warn!(%error, "deferred document encoding failed; pending chunks stay outside this query");
27            }
28        }
29
30        let _transaction = self.mutation_gate.read();
31        self.shards
32            .read()
33            .query(query_vec, k.min(10000), self.dimensions)
34    }
35
36    /// The registry of agents admitted by verified access credentials.
37    #[cfg(feature = "provenance")]
38    pub fn agent_registry(
39        &self,
40    ) -> &parking_lot::RwLock<crate::core::provenance::access::AgentRegistry> {
41        &self.agents
42    }
43
44    /// [`Self::query`], refused unless `agent_did` holds a current read grant.
45    /// The caller must already have authenticated control of `agent_did`.
46    #[cfg(feature = "provenance")]
47    pub fn query_as(
48        &self,
49        agent_did: &str,
50        query_vec: &EntangledHVec,
51        k: u32,
52    ) -> anyhow::Result<Vec<RetrievalResult>> {
53        self.agents
54            .read()
55            .authorize(agent_did, crate::core::provenance::access::Permission::Read)?;
56        Ok(self.query(query_vec, k))
57    }
58
59    pub fn explain_query(&self, query_vec: &EntangledHVec, k: u32) -> QueryExplanation {
60        let mut explanation = self
61            .shards
62            .read()
63            .explain_query(query_vec, k, self.dimensions);
64        let (pending, lexical) = self
65            .documents
66            .try_read()
67            .map(|store| (store.pending_vector_count(), store.lexical_vector_count()))
68            .unwrap_or((0, 0));
69        if pending > 0 || lexical > 0 {
70            explanation.rationale.push_str(&format!(
71                " {pending} document chunk(s) await deferred encoding (a query, \
72                 encode_pending_documents or maintain_indices installs them); {lexical} \
73                 lexical-only chunk(s) are outside vector search."
74            ));
75        }
76        explanation
77    }
78
79    pub fn index_status(&self) -> IndexStatus {
80        let count = self.vector_count().min(u32::MAX as u64) as u32;
81        let nsg = self.nsg_trained();
82        let ivf = self.ivf_trained();
83        let nsg_recommended = !nsg
84            && count > 0
85            && self.config.nsg.auto_threshold > 0
86            && count as usize >= self.config.nsg.auto_threshold;
87        let ivf_recommended = self.config.ivf.enabled
88            && !ivf
89            && count > 0
90            && self.config.ivf.auto_threshold > 0
91            && count as usize >= self.config.ivf.auto_threshold;
92        let recommendation = if nsg_recommended && ivf_recommended {
93            "train NSG and IVF indices"
94        } else if nsg_recommended {
95            "train the NSG index"
96        } else if ivf_recommended {
97            "train the IVF index"
98        } else {
99            "indices are current for the configured thresholds"
100        };
101        IndexStatus {
102            vector_count: count,
103            nsg_trained: nsg,
104            ivf_trained: ivf,
105            nsg_training_recommended: nsg_recommended,
106            ivf_training_recommended: ivf_recommended,
107            recommendation: recommendation.to_string(),
108        }
109    }
110
111    pub fn storage_health(&self) -> StorageHealth {
112        let stats = self.arena.stats();
113        StorageHealth {
114            format_version: stats.format_version,
115            segment_count: stats.segment_count.min(u32::MAX as usize) as u32,
116            used_bytes: stats.used_bytes as f64,
117            capacity_bytes: stats.capacity_bytes as f64,
118        }
119    }
120
121    /// Train any index whose configured threshold has been reached, then return
122    /// the resulting lifecycle state. Calls are idempotent when indices are current.
123    pub fn maintain_indices(&self) -> anyhow::Result<IndexStatus> {
124        self.maybe_auto_shard(self.vector_count())?;
125        self.encode_pending_documents()?;
126        let before = self.index_status();
127        if before.ivf_training_recommended {
128            self.train_ivf()?;
129        }
130        if before.nsg_training_recommended {
131            self.train_nsg()?;
132        }
133        let document_training_recommended = self.documents.read().dense_index_recommended();
134        if document_training_recommended {
135            self.train_document_index()?;
136        }
137        self.maintain_vector_index(crate::core::documents::DEFAULT_GRAPH_MIN_ROWS)?;
138        Ok(self.index_status())
139    }
140
141    /// Run maintenance on a dedicated thread; training publishes only a current snapshot.
142    pub fn maintain_indices_background(
143        self: &std::sync::Arc<Self>,
144    ) -> std::thread::JoinHandle<anyhow::Result<IndexStatus>> {
145        let core = std::sync::Arc::clone(self);
146        std::thread::spawn(move || core.maintain_indices())
147    }
148
149    /// Flush persistent vector data to durable storage.
150    pub fn flush(&self) -> anyhow::Result<()> {
151        let _transaction = self.mutation_gate.write();
152        self.arena.flush()?;
153        self.persist_indices(&self.shards.read())
154    }
155
156    /// Energy-based associative retrieval using Hopfield-Fenchel-Young dynamics.
157    ///
158    /// Unlike `query` (which returns a fixed top-k by similarity), this uses
159    /// sparse entmax attention to naturally determine how many results are
160    /// relevant. Returns at most `max_results` patterns with non-zero
161    /// Hopfield attention weight.
162    pub fn query_hopfield(
163        &self,
164        query_vec: &EntangledHVec,
165        max_results: u32,
166    ) -> Vec<RetrievalResult> {
167        let shards = self.shards.read();
168        let patterns = shards.collect_all_patterns();
169        let config = &self.config.hopfield;
170
171        if config.max_iter > 1 {
172            hopfield::hopfield_query_iterative(query_vec, &patterns, config, max_results as usize)
173        } else {
174            hopfield::hopfield_query(query_vec, &patterns, config, max_results as usize)
175        }
176    }
177
178    /// Process multiple queries in parallel using rayon.
179    pub fn query_batch(&self, queries: &[EntangledHVec], k: u32) -> Vec<Vec<RetrievalResult>> {
180        queries.par_iter().map(|q| self.query(q, k)).collect()
181    }
182
183    /// Analyze components of a vector by finding its nearest neighbors.
184    /// Filters by similarity threshold from QueryConfig (default 0.05).
185    pub fn analyze_components(&self, vector: &EntangledHVec) -> Vec<RetrievalResult> {
186        let cfg = &self.config.query;
187        let neighbors = self.query(vector, cfg.component_max_neighbors);
188        neighbors
189            .into_iter()
190            .filter(|r| r.similarity > cfg.component_similarity_threshold)
191            .collect()
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    #[test]
200    fn diagnostics_report_route_lifecycle_and_storage() -> anyhow::Result<()> {
201        let directory = tempfile::tempdir()?;
202        let hms = HmsCore::new(
203            1024,
204            Some(directory.path().to_string_lossy().into_owned()),
205            None,
206        )?;
207        let query = EntangledHVec::new_deterministic(1024, 1);
208        let explanation = hms.explain_query(&query, 5);
209        assert_eq!(explanation.route, "brute_force");
210        assert_eq!(hms.index_status().vector_count, 0);
211        assert_eq!(
212            hms.storage_health().format_version,
213            crate::core::storage::STORAGE_FORMAT_VERSION
214        );
215        hms.flush()?;
216        Ok(())
217    }
218}
holographic_memory/core/engine/query.rs

holographic_memory/core/engine/
query.rs
