writerslogic_warm/core/engine/
query.rs

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

writerslogic_warm/core/engine/
query.rs
