a3s_code_core/durable_memory/
semantic.rs

1mod indexing;
2
3pub(crate) use indexing::SemanticRefreshEmbeddingCache;
4
5use super::semantic_binding::{
6    invalid, DurableMemorySemanticBindingV1, DurableMemorySemanticError,
7    DurableMemorySemanticRecallPolicy,
8};
9use crate::embedding::{
10    EmbeddingError, EmbeddingExecutor, EmbeddingExecutorConfig, EmbeddingInput, EmbeddingProvider,
11};
12use a3s_memory::repository::{MemoryNamespace, MemoryNode, MemoryRepository, MemoryStatus};
13use a3s_memory::vector::{
14    VectorIndex, VectorIndexObservation, VectorIndexStatus, VectorMutationConsistency,
15    VectorSearchRequest,
16};
17use sha2::{Digest, Sha256};
18use std::collections::HashSet;
19use std::sync::Arc;
20use tokio::sync::Mutex;
21use tokio_util::sync::CancellationToken;
22
23const PARTITION_ID_DOMAIN: &str = "a3s.code.memory.semantic-partition.v1";
24const RECORD_ID_DOMAIN: &str = "a3s.code.memory.semantic-record.v1";
25const RECORD_SCHEMA_V1: &str = "a3s.code.memory.semantic-record.v1";
26const QUERY_ID: &str = "durable-memory-semantic-query";
27const LABEL_SCHEMA: &str = "a3s.memory.semantic.schema";
28const LABEL_NODE_ID: &str = "a3s.memory.semantic.node_id";
29const LABEL_NODE_REVISION: &str = "a3s.memory.semantic.node_revision";
30const LABEL_CONTENT_DIGEST: &str = "a3s.memory.semantic.content_digest";
31
32/// Explicit host-owned embedding generation plus caller-owned vector index.
33///
34/// Construction and namespace replacement are inert with respect to background
35/// work. The host decides whether to invoke refresh directly or install an
36/// owned schedule; Code owns bounded query execution and candidate verification.
37#[derive(Clone)]
38pub struct DurableMemorySemanticRecall {
39    binding: DurableMemorySemanticBindingV1,
40    serving_generation_digest: String,
41    executor: EmbeddingExecutor,
42    index: Arc<dyn VectorIndex>,
43    refresh_lock: Arc<Mutex<()>>,
44}
45
46impl std::fmt::Debug for DurableMemorySemanticRecall {
47    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        formatter
49            .debug_struct("DurableMemorySemanticRecall")
50            .field("binding", &self.binding)
51            .finish_non_exhaustive()
52    }
53}
54
55impl DurableMemorySemanticRecall {
56    pub fn new(
57        authority_digest: impl Into<String>,
58        provider: Arc<dyn EmbeddingProvider>,
59        embedding_config: EmbeddingExecutorConfig,
60        index: Arc<dyn VectorIndex>,
61        policy: DurableMemorySemanticRecallPolicy,
62    ) -> Result<Self, DurableMemorySemanticError> {
63        let executor = EmbeddingExecutor::new(provider, embedding_config)?;
64        let binding = DurableMemorySemanticBindingV1::new(
65            authority_digest.into(),
66            executor.descriptor().clone(),
67            embedding_config,
68            index.descriptor().clone(),
69            policy,
70        )?;
71        let serving_generation_digest = binding.serving_generation_digest()?;
72        Ok(Self {
73            binding,
74            serving_generation_digest,
75            executor,
76            index,
77            refresh_lock: Arc::new(Mutex::new(())),
78        })
79    }
80
81    pub fn binding(&self) -> &DurableMemorySemanticBindingV1 {
82        &self.binding
83    }
84
85    /// Return the backend's latest locally cached compatibility status.
86    ///
87    /// Durable correctness paths use [`Self::observe_index`] instead.
88    pub fn index_status(&self) -> VectorIndexStatus {
89        self.index.status()
90    }
91
92    /// Read one exact, fallible index status/history observation.
93    pub async fn observe_index(
94        &self,
95    ) -> Result<VectorIndexObservation, DurableMemorySemanticError> {
96        self.index.observe().await.map_err(Into::into)
97    }
98
99    /// Return the strongest partition-mutation ordering advertised by the
100    /// injected vector backend.
101    pub fn mutation_consistency(&self) -> VectorMutationConsistency {
102        self.index.mutation_consistency()
103    }
104
105    pub(super) fn serving_generation_digest(&self) -> &str {
106        &self.serving_generation_digest
107    }
108
109    pub(super) async fn query_verified(
110        &self,
111        repository: &dyn MemoryRepository,
112        namespace: &MemoryNamespace,
113        text: &str,
114        cancellation: CancellationToken,
115    ) -> Result<Vec<SemanticRecallCandidate>, DurableMemorySemanticError> {
116        check_cancellation(&cancellation)?;
117        let execution = self
118            .executor
119            .embed(
120                vec![EmbeddingInput::new(QUERY_ID, text.to_string())],
121                cancellation.clone(),
122            )
123            .await?;
124        let vector = execution
125            .vectors
126            .into_iter()
127            .next()
128            .ok_or_else(|| invalid("query.embedding", "provider returned no vector"))?;
129        check_cancellation(&cancellation)?;
130
131        let partition = semantic_partition_id(namespace, &self.serving_generation_digest);
132        let observed = tokio::select! {
133            result = self.observe_index() => result?,
134            _ = cancellation.cancelled() => return Err(EmbeddingError::Cancelled.into()),
135        };
136        let request =
137            VectorSearchRequest::new(vector.values, self.binding.policy().candidate_limit())
138                .with_partition(&partition)
139                .with_label(LABEL_SCHEMA, RECORD_SCHEMA_V1);
140        let result = tokio::select! {
141            result = self.index.search(request) => result?,
142            _ = cancellation.cancelled() => return Err(EmbeddingError::Cancelled.into()),
143        };
144        let after_search = tokio::select! {
145            result = self.observe_index() => result?,
146            _ = cancellation.cancelled() => return Err(EmbeddingError::Cancelled.into()),
147        };
148        if result.status != observed.status || after_search != observed {
149            return Err(DurableMemorySemanticError::IndexRevisionChanged);
150        }
151
152        let mut known_nodes = HashSet::new();
153        let mut candidates = Vec::new();
154        for hit in result.hits {
155            check_cancellation(&cancellation)?;
156            if hit.partition != partition
157                || !hit.score.is_finite()
158                || !(-1.0..=1.0).contains(&hit.score)
159                || hit.score < self.binding.policy().min_score()
160            {
161                continue;
162            }
163            let Some(node_id) = hit.labels.get(LABEL_NODE_ID) else {
164                continue;
165            };
166            let Some(node_revision) = hit
167                .labels
168                .get(LABEL_NODE_REVISION)
169                .and_then(|value| value.parse::<u64>().ok())
170            else {
171                continue;
172            };
173            let Some(content_digest) = hit.labels.get(LABEL_CONTENT_DIGEST) else {
174                continue;
175            };
176            if hit.labels.get(LABEL_SCHEMA).map(String::as_str) != Some(RECORD_SCHEMA_V1)
177                || hit.id != semantic_record_id(&partition, node_id, node_revision, content_digest)
178                || !known_nodes.insert(node_id.clone())
179            {
180                continue;
181            }
182            let node = tokio::select! {
183                result = repository.get(namespace, node_id) => result?,
184                _ = cancellation.cancelled() => return Err(EmbeddingError::Cancelled.into()),
185            };
186            let Some(node) = node else {
187                continue;
188            };
189            if node.status != MemoryStatus::Active
190                || node.revision != node_revision
191                || digest(&node.content) != *content_digest
192            {
193                continue;
194            }
195            candidates.push(SemanticRecallCandidate {
196                node,
197                score: hit.score,
198            });
199        }
200        check_cancellation(&cancellation)?;
201        let after_verification = tokio::select! {
202            result = self.observe_index() => result?,
203            _ = cancellation.cancelled() => return Err(EmbeddingError::Cancelled.into()),
204        };
205        if after_verification != observed {
206            return Err(DurableMemorySemanticError::IndexRevisionChanged);
207        }
208        candidates.sort_by(|left, right| {
209            right
210                .score
211                .total_cmp(&left.score)
212                .then_with(|| right.node.updated_at.cmp(&left.node.updated_at))
213                .then_with(|| left.node.id.cmp(&right.node.id))
214        });
215        candidates.truncate(self.binding.policy().candidate_limit());
216        Ok(candidates)
217    }
218}
219
220#[derive(Clone)]
221pub(super) struct SemanticRecallCandidate {
222    pub(super) node: MemoryNode,
223    pub(super) score: f32,
224}
225
226fn semantic_partition_id(namespace: &MemoryNamespace, serving_generation_digest: &str) -> String {
227    let mut hasher = Sha256::new();
228    hasher.update(PARTITION_ID_DOMAIN.as_bytes());
229    for (label, value) in [
230        (b"tenant".as_slice(), namespace.tenant_id()),
231        (b"principal".as_slice(), namespace.principal_id()),
232        (b"scope".as_slice(), namespace.scope_id()),
233    ] {
234        hasher.update([0]);
235        hasher.update(label);
236        hasher.update([0]);
237        hasher.update(Sha256::digest(value.as_bytes()));
238    }
239    hasher.update([0]);
240    hasher.update(b"serving-generation");
241    hasher.update([0]);
242    hasher.update(serving_generation_digest.as_bytes());
243    format!("a3s-memory-semantic-{:x}", hasher.finalize())
244}
245
246fn semantic_record_id(
247    partition: &str,
248    node_id: &str,
249    node_revision: u64,
250    content_digest: &str,
251) -> String {
252    let mut hasher = Sha256::new();
253    hasher.update(RECORD_ID_DOMAIN.as_bytes());
254    hasher.update([0]);
255    hasher.update(partition.as_bytes());
256    hasher.update([0]);
257    hasher.update(Sha256::digest(node_id.as_bytes()));
258    hasher.update([0]);
259    hasher.update(node_revision.to_le_bytes());
260    hasher.update([0]);
261    hasher.update(content_digest.as_bytes());
262    format!("a3s-memory-semantic-record-{:x}", hasher.finalize())
263}
264
265fn digest(content: &str) -> String {
266    format!("sha256:{:x}", Sha256::digest(content.as_bytes()))
267}
268
269fn check_cancellation(cancellation: &CancellationToken) -> Result<(), DurableMemorySemanticError> {
270    if cancellation.is_cancelled() {
271        Err(EmbeddingError::Cancelled.into())
272    } else {
273        Ok(())
274    }
275}
a3s_code_core/durable_memory/semantic.rs

a3s_code_core/durable_memory/
semantic.rs
