stasis/infrastructure/memory/
locus_memory_operations.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use locus_sdk::application::memory_evict::MemoryEvictService;
5use locus_sdk::domain::evict::{
6    InboundReferencesPreview as LocusInboundReferencesPreview,
7    MemoryEvictMode as LocusEvictMode, MemoryEvictRecord as LocusEvictRecord,
8    MemoryEvictRequest as LocusEvictRequest,
9};
10use locus_sdk::prelude::{
11    AiProviderRegistry, MemoryAggregateRequest as LocusAggregateRequest, MemoryAggregateService,
12    MemoryCompositionService, MemoryDailyRollupRequest, MemoryGroupBy, MemorySchemaService,
13    MemoryTransformOperation as LocusTransformOperation,
14    MemoryTransformRequest as LocusTransformRequest, MemoryTransformService,
15};
16
17use crate::domain::errors::{Result, StasisError};
18use crate::infrastructure::memory::locus_memory_mapping::{map_filter, map_scope};
19use crate::infrastructure::memory::locus_node_store_factory::LocusMemoryStore;
20use crate::ports::outbound::memory::memory_models::{
21    MemoryAggregateRequest, MemoryAggregateResponse, MemoryEvictMode, MemoryEvictRecord,
22    MemoryEvictRequest, MemoryEvictResponse, MemoryInboundReferencesPreview, MemoryRollupRequest,
23    MemoryRollupResponse, MemorySchemaResponse, MemoryTransformOperation, MemoryTransformRequest,
24    MemoryTransformResponse,
25};
26use crate::ports::outbound::memory::memory_operations::MemoryOperations;
27
28pub struct LocusMemoryOperations {
29    memory: Arc<LocusMemoryStore>,
30    aggregate: MemoryAggregateService,
31    composition: MemoryCompositionService,
32    schema: MemorySchemaService,
33    providers: Option<Arc<dyn AiProviderRegistry>>,
34}
35
36impl LocusMemoryOperations {
37    pub fn new(memory: Arc<LocusMemoryStore>, providers: Option<Arc<dyn AiProviderRegistry>>) -> Self {
38        Self {
39            aggregate: MemoryAggregateService::new(memory.node_store.clone()),
40            composition: MemoryCompositionService::new(memory.node_store.clone()),
41            schema: MemorySchemaService::new(),
42            memory,
43            providers,
44        }
45    }
46}
47
48#[async_trait]
49impl MemoryOperations for LocusMemoryOperations {
50    async fn aggregate(&self, request: &MemoryAggregateRequest) -> Result<MemoryAggregateResponse> {
51        let result = self
52            .aggregate
53            .execute(&LocusAggregateRequest {
54                scope: map_scope(&request.scope),
55                group_by: MemoryGroupBy::DateDay,
56                max_groups: request.max_groups,
57                max_nodes: request.max_nodes,
58                ..Default::default()
59            })
60            .await
61            .map_err(|e| StasisError::PortFailure(format!("locus aggregate failed: {e}")))?;
62
63        Ok(MemoryAggregateResponse {
64            total_groups: result.total_groups,
65            scanned_nodes: result.scanned_nodes,
66        })
67    }
68
69    async fn transform(&self, request: &MemoryTransformRequest) -> Result<MemoryTransformResponse> {
70        let providers = self.providers.clone().ok_or_else(|| {
71            StasisError::PortFailure("locus transform requires ai provider registry".to_string())
72        })?;
73
74        let service = MemoryTransformService::new(self.memory.node_store.clone(), providers)
75            .with_semantic_index(self.memory.semantic_index.clone());
76        let result = service
77            .execute(&LocusTransformRequest {
78                scope: map_scope(&request.scope),
79                filter: map_filter(&request.filter),
80                operation: map_transform_operation(request.operation),
81                dry_run: request.dry_run,
82                batch_size: request.batch_size,
83                max_nodes: request.max_nodes,
84                provider_id: request.provider_id.clone(),
85                model: request.model.clone(),
86            })
87            .await
88            .map_err(|e| StasisError::PortFailure(format!("locus transform failed: {e}")))?;
89
90        Ok(MemoryTransformResponse {
91            scanned: result.scanned,
92            selected: result.selected,
93            updated: result.updated,
94            skipped: result.skipped,
95            failed: result.failed,
96            duplicate: result.duplicate,
97            failures: result.failures,
98        })
99    }
100
101    async fn rollup(&self, request: &MemoryRollupRequest) -> Result<MemoryRollupResponse> {
102        let result = self
103            .composition
104            .daily_rollup(&MemoryDailyRollupRequest {
105                scope: map_scope(&request.scope),
106                max_days: request.max_days,
107                max_nodes: request.max_nodes,
108                ..Default::default()
109            })
110            .await
111            .map_err(|e| StasisError::PortFailure(format!("locus daily rollup failed: {e}")))?;
112
113        Ok(MemoryRollupResponse {
114            total_groups: result.total_groups,
115            scanned_nodes: result.scanned_nodes,
116        })
117    }
118
119    async fn schema(&self) -> Result<MemorySchemaResponse> {
120        let schema = self.schema.execute();
121        Ok(MemorySchemaResponse {
122            schema_version: schema.schema_version,
123            sort_fields: schema.sort_fields,
124            filter_fields: schema.filter_fields,
125            group_by_fields: schema.group_by_fields,
126            fallback_policies: schema.fallback_policies,
127            strictness_modes: schema.strictness_modes,
128            transform_operations: schema.transform_operations,
129            evict_operations: schema.evict_operations,
130        })
131    }
132
133    async fn evict(&self, request: &MemoryEvictRequest) -> Result<MemoryEvictResponse> {
134        let service = MemoryEvictService::new(self.memory.node_store.clone())
135            .with_semantic_index(self.memory.semantic_index.clone());
136        let result = service
137            .execute(&LocusEvictRequest {
138                mode: map_evict_mode(request.mode),
139                scope: map_scope(&request.scope),
140                filter: map_filter(&request.filter),
141                sync_keys: request.sync_keys.clone(),
142                node_ids: request.node_ids.clone(),
143                dry_run: request.dry_run,
144                force: request.force,
145                max_nodes: request.max_nodes,
146                include_calibration: request.include_calibration,
147                include_checkpoints: request.include_checkpoints,
148            })
149            .await
150            .map_err(|e| StasisError::PortFailure(format!("locus evict failed: {e}")))?;
151
152        Ok(MemoryEvictResponse {
153            dry_run: result.dry_run,
154            deleted: result.deleted,
155            blocked: result.blocked,
156            not_found: result.not_found,
157            skipped: result.skipped,
158            would_delete: result.would_delete,
159            calibrations_deleted: result.calibrations_deleted,
160            checkpoints_deleted: result.checkpoints_deleted,
161            records: result.records.iter().map(map_evict_record).collect(),
162        })
163    }
164}
165
166fn map_transform_operation(value: MemoryTransformOperation) -> LocusTransformOperation {
167    match value {
168        MemoryTransformOperation::EmbedBackfill => LocusTransformOperation::EmbedBackfill,
169        MemoryTransformOperation::ReindexEmbeddings => LocusTransformOperation::ReindexEmbeddings,
170        MemoryTransformOperation::EmbedTagBackfill => LocusTransformOperation::EmbedTagBackfill,
171        MemoryTransformOperation::ReindexTagEmbeddings => LocusTransformOperation::ReindexTagEmbeddings,
172    }
173}
174
175fn map_evict_mode(value: MemoryEvictMode) -> LocusEvictMode {
176    match value {
177        MemoryEvictMode::BySyncKeys => LocusEvictMode::BySyncKeys,
178        MemoryEvictMode::ByNodeIds => LocusEvictMode::ByNodeIds,
179        MemoryEvictMode::ByFilter => LocusEvictMode::ByFilter,
180        MemoryEvictMode::PurgeSession => LocusEvictMode::PurgeSession,
181    }
182}
183
184fn map_evict_record(record: &LocusEvictRecord) -> MemoryEvictRecord {
185    MemoryEvictRecord {
186        node_id: record.node_id.clone(),
187        sync_key: record.sync_key.clone(),
188        status: record.status.clone(),
189        reason: record.reason.clone(),
190        inbound_references: record
191            .inbound_references
192            .as_ref()
193            .map(map_inbound_references),
194    }
195}
196
197fn map_inbound_references(value: &LocusInboundReferencesPreview) -> MemoryInboundReferencesPreview {
198    MemoryInboundReferencesPreview {
199        child_parent_links: value.child_parent_links.clone(),
200        incoming_semantic_refs: value.incoming_semantic_refs.clone(),
201    }
202}
stasis/infrastructure/memory/locus_memory_operations.rs

stasis/infrastructure/memory/
locus_memory_operations.rs
