1use std::sync::Arc;
2
3use anyhow::Result;
4use serde_json::{Map, Value};
5use locus_core_rs::domain::contracts::NodeStore;
6use locus_core_rs::domain::models::AvecState;
7
8use crate::application::memory_aggregate::MemoryAggregateService;
9use crate::application::memory_explain::MemoryExplainService;
10use crate::application::manual_compression::ManualCompressionService;
11use crate::application::memory_recall::MemoryRecallService;
12use crate::application::memory_schema::MemorySchemaService;
13use crate::application::memory_transform::MemoryTransformService;
14use crate::domain::ai::AiProviderRegistry;
15use crate::domain::compression::ManualCompressionRequest;
16use crate::domain::memory::{
17 MemoryAggregateRequest, MemoryAggregateResult, MemoryExplainRequest, MemoryExplainResult,
18 MemoryFilter, MemoryGroupBy, MemoryRecallRequest, MemoryRecallResult, MemorySchemaResult,
19 MemoryScope, MemoryTransformRequest, MemoryTransformResult, clamp_nodes,
20};
21
22#[derive(Debug, Clone)]
23pub struct MemoryRecallWithExplainResult {
24 pub recall: MemoryRecallResult,
25 pub explain: MemoryExplainResult,
26}
27
28#[derive(Debug, Clone, Default)]
29pub struct MemoryDailyRollupRequest {
30 pub scope: MemoryScope,
31 pub filter: MemoryFilter,
32 pub max_days: usize,
33 pub max_nodes: usize,
34}
35
36#[derive(Debug, Clone)]
37pub struct MemoryTransformThenRecallRequest {
38 pub transform: MemoryTransformRequest,
39 pub recall: MemoryRecallRequest,
40}
41
42#[derive(Debug, Clone)]
43pub struct MemoryTransformThenRecallResult {
44 pub transform: MemoryTransformResult,
45 pub recall: MemoryRecallResult,
46}
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub enum CompositeRole {
50 User,
51 Model,
52 Document,
53 Conversation,
54}
55
56impl CompositeRole {
57 fn as_str(self) -> &'static str {
58 match self {
59 Self::User => "user",
60 Self::Model => "model",
61 Self::Document => "document",
62 Self::Conversation => "conversation",
63 }
64 }
65}
66
67#[derive(Debug, Clone)]
68pub struct CompositeInputItem {
69 pub role: CompositeRole,
70 pub text: String,
71 pub avec_override: Option<AvecState>,
72 pub context: Vec<CompositeInputItem>,
73}
74
75#[derive(Debug, Clone, Default)]
76pub struct CompositeRoleAvecOverrides {
77 pub user: Option<AvecState>,
78 pub model: Option<AvecState>,
79 pub document: Option<AvecState>,
80 pub conversation: Option<AvecState>,
81}
82
83impl CompositeRoleAvecOverrides {
84 fn resolve(&self, role: CompositeRole) -> Option<AvecState> {
85 match role {
86 CompositeRole::User => self.user,
87 CompositeRole::Model => self.model,
88 CompositeRole::Document => self.document,
89 CompositeRole::Conversation => self.conversation,
90 }
91 }
92}
93
94#[derive(Debug, Clone)]
95pub struct CompositeNodeFromTextOptions {
96 pub role_avec: CompositeRoleAvecOverrides,
97 pub global_avec: Option<AvecState>,
98 pub allow_llm_avec_fallback: bool,
99 pub max_recursion_depth: usize,
100}
101
102impl Default for CompositeNodeFromTextOptions {
103 fn default() -> Self {
104 Self {
105 role_avec: CompositeRoleAvecOverrides::default(),
106 global_avec: None,
107 allow_llm_avec_fallback: false,
108 max_recursion_depth: 5,
109 }
110 }
111}
112
113#[derive(Debug, Clone, Default)]
114pub struct CompositeNodeFromTextRequest {
115 pub items: Vec<CompositeInputItem>,
116 pub options: CompositeNodeFromTextOptions,
117}
118
119#[derive(Debug, Clone, Default)]
120pub struct CompositeNodeFromTextResult {
121 pub content: Value,
122 pub resolved_avec_count: usize,
123 pub unresolved_avec_count: usize,
124 pub requires_llm_avec: bool,
125}
126
127#[derive(Debug, Clone, Default)]
128struct CompositeBuildStats {
129 resolved_avec_count: usize,
130 unresolved_avec_count: usize,
131}
132
133pub struct MemoryCompositionService {
134 store: Arc<dyn NodeStore>,
135 recall: MemoryRecallService,
136 explain: MemoryExplainService,
137 aggregate: MemoryAggregateService,
138 schema: MemorySchemaService,
139}
140
141impl MemoryCompositionService {
142 pub fn new(store: Arc<dyn NodeStore>) -> Self {
143 Self {
144 store: store.clone(),
145 recall: MemoryRecallService::new(store.clone()),
146 explain: MemoryExplainService::new(store.clone()),
147 aggregate: MemoryAggregateService::new(store),
148 schema: MemorySchemaService::new(),
149 }
150 }
151
152 pub async fn recall_with_explain(
153 &self,
154 request: &MemoryRecallRequest,
155 ) -> Result<MemoryRecallWithExplainResult> {
156 let recall = self.recall.execute(request).await?;
157 let explain = self
158 .explain
159 .execute(&MemoryExplainRequest {
160 recall: request.clone(),
161 })
162 .await?;
163
164 Ok(MemoryRecallWithExplainResult { recall, explain })
165 }
166
167 pub async fn daily_rollup(
168 &self,
169 request: &MemoryDailyRollupRequest,
170 ) -> Result<MemoryAggregateResult> {
171 let max_days = if request.max_days == 0 {
172 30
173 } else {
174 request.max_days
175 };
176 let max_nodes = clamp_nodes(if request.max_nodes == 0 {
177 5000
178 } else {
179 request.max_nodes
180 });
181
182 self.aggregate
183 .execute(&MemoryAggregateRequest {
184 scope: request.scope.clone(),
185 filter: request.filter.clone(),
186 group_by: MemoryGroupBy::DateDay,
187 max_groups: max_days,
188 max_nodes,
189 })
190 .await
191 }
192
193 pub fn capability_bundle(&self) -> MemorySchemaResult {
194 self.schema.execute()
195 }
196
197 pub async fn transform_then_recall_verify(
198 &self,
199 providers: Arc<dyn AiProviderRegistry>,
200 request: &MemoryTransformThenRecallRequest,
201 ) -> Result<MemoryTransformThenRecallResult> {
202 let transform_service = MemoryTransformService::new(self.store.clone(), providers);
203 let transform = transform_service.execute(&request.transform).await?;
204 let recall = self.recall.execute(&request.recall).await?;
205
206 Ok(MemoryTransformThenRecallResult { transform, recall })
207 }
208
209 pub fn build_content_from_text(
210 &self,
211 request: &CompositeNodeFromTextRequest,
212 ) -> Result<CompositeNodeFromTextResult> {
213 let max_depth = request.options.max_recursion_depth.clamp(1, 5);
214 let compressor = ManualCompressionService::new();
215 let mut stats = CompositeBuildStats::default();
216
217 let mut root = Map::new();
218 for (idx, item) in request.items.iter().enumerate() {
219 let key = format!("entry_{idx}(.95)");
220 let value = build_composite_entry(
221 item,
222 1,
223 max_depth,
224 &request.options,
225 &compressor,
226 &mut stats,
227 )?;
228 root.insert(key, value);
229 }
230
231 let requires_llm = stats.unresolved_avec_count > 0;
232 if requires_llm && !request.options.allow_llm_avec_fallback {
233 anyhow::bail!(
234 "unable to resolve AVEC for {} item(s); provide overrides or enable llm fallback",
235 stats.unresolved_avec_count
236 );
237 }
238
239 Ok(CompositeNodeFromTextResult {
240 content: Value::Object(root),
241 resolved_avec_count: stats.resolved_avec_count,
242 unresolved_avec_count: stats.unresolved_avec_count,
243 requires_llm_avec: requires_llm,
244 })
245 }
246}
247
248fn build_composite_entry(
249 item: &CompositeInputItem,
250 depth: usize,
251 max_depth: usize,
252 options: &CompositeNodeFromTextOptions,
253 compressor: &ManualCompressionService,
254 stats: &mut CompositeBuildStats,
255) -> Result<Value> {
256 if depth > max_depth {
257 anyhow::bail!("composite context depth exceeded max depth of {max_depth}");
258 }
259
260 let resolved_avec = item
261 .avec_override
262 .or_else(|| options.role_avec.resolve(item.role))
263 .or(options.global_avec);
264
265 if resolved_avec.is_some() {
266 stats.resolved_avec_count += 1;
267 } else {
268 stats.unresolved_avec_count += 1;
269 }
270
271 let compressed = compressor.execute(&ManualCompressionRequest {
272 text: item.text.clone(),
273 ..Default::default()
274 });
275
276 let mut entry = Map::new();
277 entry.insert(
278 "role(.99)".to_string(),
279 Value::String(item.role.as_str().to_string()),
280 );
281 entry.insert("text(.70)".to_string(), Value::String(item.text.clone()));
282 entry.insert(
283 "anchor_topic(.86)".to_string(),
284 Value::String(compressed.anchor_topic),
285 );
286 entry.insert(
287 "key_points(.82)".to_string(),
288 Value::Array(
289 compressed
290 .key_points
291 .into_iter()
292 .map(Value::String)
293 .collect(),
294 ),
295 );
296
297 if let Some(avec) = resolved_avec {
298 let mut avec_obj = Map::new();
299 avec_obj.insert("stability(.99)".to_string(), Value::from(avec.stability as f64));
300 avec_obj.insert("friction(.99)".to_string(), Value::from(avec.friction as f64));
301 avec_obj.insert("logic(.99)".to_string(), Value::from(avec.logic as f64));
302 avec_obj.insert("autonomy(.99)".to_string(), Value::from(avec.autonomy as f64));
303 avec_obj.insert("psi(.99)".to_string(), Value::from(avec.psi() as f64));
304 entry.insert("resolved_avec(.95)".to_string(), Value::Object(avec_obj));
305 }
306
307 if !item.context.is_empty() {
308 let mut children = Map::new();
309 for (idx, child) in item.context.iter().enumerate() {
310 let child_key = format!("context_{idx}(.90)");
311 children.insert(
312 child_key,
313 build_composite_entry(child, depth + 1, max_depth, options, compressor, stats)?,
314 );
315 }
316 entry.insert("context(.88)".to_string(), Value::Object(children));
317 }
318
319 Ok(Value::Object(entry))
320}
321
322#[cfg(test)]
323mod tests {
324 use std::sync::Arc;
325
326 use anyhow::Result;
327 use async_trait::async_trait;
328 use chrono::{Duration, Utc};
329 use serde_json::Value;
330 use locus_core_rs::application::validation::TreeSitterValidator;
331 use locus_core_rs::domain::contracts::NodeValidator;
332 use locus_core_rs::domain::models::{AvecState, SttpNode};
333 use locus_core_rs::parsing::{
334 SttpContentSlice, SttpDocumentBuilder, SttpDocumentMetadata, SttpNodeParser,
335 };
336 use locus_core_rs::{InMemoryNodeStore, NodeStore};
337
338 use super::{
339 CompositeInputItem, CompositeNodeFromTextOptions, CompositeNodeFromTextRequest,
340 CompositeRole, MemoryCompositionService, MemoryDailyRollupRequest,
341 MemoryTransformThenRecallRequest,
342 };
343 use crate::domain::ai::{AiCapability, AiProvider, EmbedRequest, ScoreAvecRequest};
344 use crate::domain::memory::{
345 FallbackPolicy, MemoryFilter, MemoryRecallRequest, MemoryScoring, MemoryTransformOperation,
346 MemoryTransformRequest, RetrievalPath,
347 };
348 use crate::infrastructure::registry::InMemoryAiProviderRegistry;
349
350 struct MockEmbeddingProvider;
351
352 #[async_trait]
353 impl AiProvider for MockEmbeddingProvider {
354 fn provider_id(&self) -> &str {
355 "mock"
356 }
357
358 fn capabilities(&self) -> &'static [AiCapability] {
359 &[AiCapability::SemanticEmbedding]
360 }
361
362 async fn embed_semantic(&self, _request: &EmbedRequest) -> Result<Vec<f32>> {
363 Ok(vec![0.2, 0.3, 0.4])
364 }
365
366 async fn embed_avec(&self, _request: &EmbedRequest) -> Result<Vec<f32>> {
367 Ok(vec![0.2, 0.3, 0.4])
368 }
369
370 async fn score_avec(&self, _request: &ScoreAvecRequest) -> Result<AvecState> {
371 Ok(AvecState::zero())
372 }
373 }
374
375 #[tokio::test]
376 async fn recall_with_explain_returns_both_results() {
377 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
378 store
379 .upsert_node_async(test_node("s-recipe", Utc::now(), "keyword in payload"))
380 .await
381 .expect("upsert should succeed");
382
383 let service = MemoryCompositionService::new(store);
384 let result = service
385 .recall_with_explain(&MemoryRecallRequest {
386 query_text: Some("keyword".to_string()),
387 ..Default::default()
388 })
389 .await
390 .expect("composed recall should succeed");
391
392 assert!(!result.explain.stages.is_empty());
393 assert!(result.recall.retrieved <= result.recall.nodes.len());
394 }
395
396 #[tokio::test]
397 async fn recall_with_explain_marks_lexical_fallback_on_empty_policy() {
398 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
399 store
400 .upsert_node_async(test_node("s-fallback", Utc::now(), "payload without match"))
401 .await
402 .expect("upsert should succeed");
403
404 let service = MemoryCompositionService::new(store);
405 let result = service
406 .recall_with_explain(&MemoryRecallRequest {
407 query_text: Some("needle".to_string()),
408 filter: MemoryFilter {
409 has_embedding: Some(true),
410 ..Default::default()
411 },
412 scoring: MemoryScoring {
413 fallback_policy: FallbackPolicy::OnEmpty,
414 ..Default::default()
415 },
416 ..Default::default()
417 })
418 .await
419 .expect("composed recall should succeed");
420
421 assert_eq!(result.recall.retrieval_path, RetrievalPath::LexicalFallback);
422 assert!(result.explain.fallback_triggered);
423 }
424
425 #[tokio::test]
426 async fn daily_rollup_groups_by_day() {
427 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
428 let now = Utc::now();
429 store
430 .upsert_node_async(test_node("s-rollup", now - Duration::days(1), "a"))
431 .await
432 .expect("upsert should succeed");
433 store
434 .upsert_node_async(test_node("s-rollup", now, "b"))
435 .await
436 .expect("upsert should succeed");
437
438 let service = MemoryCompositionService::new(store);
439 let result = service
440 .daily_rollup(&MemoryDailyRollupRequest {
441 max_days: 10,
442 max_nodes: 100,
443 ..Default::default()
444 })
445 .await
446 .expect("daily rollup should succeed");
447
448 assert!(result.total_groups >= 2);
449 }
450
451 #[test]
452 fn capability_bundle_exposes_schema() {
453 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
454 let service = MemoryCompositionService::new(store);
455 let schema = service.capability_bundle();
456
457 assert_eq!(schema.schema_version, "locus-sdk.memory.v3");
458 assert!(schema
459 .transform_operations
460 .contains(&"embed_backfill".to_string()));
461 }
462
463 #[tokio::test]
464 async fn transform_then_recall_verify_returns_both_sides() {
465 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
466 store
467 .upsert_node_async(test_node("s-verify", Utc::now(), "verification payload"))
468 .await
469 .expect("upsert should succeed");
470
471 let service = MemoryCompositionService::new(store.clone());
472
473 let mut providers = InMemoryAiProviderRegistry::new();
474 providers.register(MockEmbeddingProvider);
475
476 let result = service
477 .transform_then_recall_verify(
478 Arc::new(providers),
479 &MemoryTransformThenRecallRequest {
480 transform: MemoryTransformRequest {
481 operation: MemoryTransformOperation::EmbedBackfill,
482 dry_run: false,
483 max_nodes: 100,
484 batch_size: 10,
485 ..Default::default()
486 },
487 recall: MemoryRecallRequest {
488 query_text: Some("verification".to_string()),
489 ..Default::default()
490 },
491 },
492 )
493 .await
494 .expect("transform then recall should succeed");
495
496 assert_eq!(result.transform.failed, 0);
497 assert_eq!(result.transform.updated, 1);
498 assert!(result.recall.retrieved <= result.recall.nodes.len());
499 }
500
501 #[test]
502 fn build_content_from_text_resolves_avec_from_role_then_global() {
503 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
504 let service = MemoryCompositionService::new(store);
505
506 let role_state = AvecState {
507 stability: 0.6,
508 friction: 0.2,
509 logic: 0.8,
510 autonomy: 0.7,
511 };
512 let global_state = AvecState {
513 stability: 0.4,
514 friction: 0.3,
515 logic: 0.6,
516 autonomy: 0.5,
517 };
518
519 let result = service
520 .build_content_from_text(&CompositeNodeFromTextRequest {
521 items: vec![
522 CompositeInputItem {
523 role: CompositeRole::User,
524 text: "policy retrieval stability".to_string(),
525 avec_override: None,
526 context: Vec::new(),
527 },
528 CompositeInputItem {
529 role: CompositeRole::Document,
530 text: "technical writeup migration".to_string(),
531 avec_override: None,
532 context: Vec::new(),
533 },
534 ],
535 options: CompositeNodeFromTextOptions {
536 role_avec: super::CompositeRoleAvecOverrides {
537 user: Some(role_state),
538 ..Default::default()
539 },
540 global_avec: Some(global_state),
541 ..Default::default()
542 },
543 })
544 .expect("composite build should succeed");
545
546 assert_eq!(result.resolved_avec_count, 2);
547 assert_eq!(result.unresolved_avec_count, 0);
548 assert!(!result.requires_llm_avec);
549 }
550
551 #[test]
552 fn build_content_from_text_fails_without_resolved_avec_when_llm_disabled() {
553 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
554 let service = MemoryCompositionService::new(store);
555
556 let err = service
557 .build_content_from_text(&CompositeNodeFromTextRequest {
558 items: vec![CompositeInputItem {
559 role: CompositeRole::Conversation,
560 text: "user asked then model replied".to_string(),
561 avec_override: None,
562 context: Vec::new(),
563 }],
564 options: CompositeNodeFromTextOptions {
565 allow_llm_avec_fallback: false,
566 ..Default::default()
567 },
568 })
569 .expect_err("missing avec should fail when llm fallback is disabled");
570
571 assert!(err.to_string().contains("unable to resolve AVEC"));
572 }
573
574 #[test]
575 fn build_content_from_text_enforces_depth_limit() {
576 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
577 let service = MemoryCompositionService::new(store);
578
579 let leaf = CompositeInputItem {
580 role: CompositeRole::User,
581 text: "leaf".to_string(),
582 avec_override: Some(AvecState::zero()),
583 context: Vec::new(),
584 };
585
586 let depth2 = CompositeInputItem {
587 role: CompositeRole::User,
588 text: "depth2".to_string(),
589 avec_override: Some(AvecState::zero()),
590 context: vec![leaf],
591 };
592
593 let depth1 = CompositeInputItem {
594 role: CompositeRole::User,
595 text: "depth1".to_string(),
596 avec_override: Some(AvecState::zero()),
597 context: vec![depth2],
598 };
599
600 let err = service
601 .build_content_from_text(&CompositeNodeFromTextRequest {
602 items: vec![depth1],
603 options: CompositeNodeFromTextOptions {
604 max_recursion_depth: 2,
605 ..Default::default()
606 },
607 })
608 .expect_err("depth overflow should fail");
609
610 assert!(err.to_string().contains("depth exceeded"));
611 }
612
613 #[test]
614 fn composite_content_parses_and_validates_under_strict_profile() {
615 let store: Arc<dyn NodeStore> = Arc::new(InMemoryNodeStore::new());
616 let service = MemoryCompositionService::new(store);
617
618 let request = CompositeNodeFromTextRequest {
619 items: vec![CompositeInputItem {
620 role: CompositeRole::Conversation,
621 text: "user asked for deterministic recall and model proposed fallback policy".to_string(),
622 avec_override: Some(AvecState {
623 stability: 0.8,
624 friction: 0.2,
625 logic: 0.85,
626 autonomy: 0.75,
627 }),
628 context: vec![CompositeInputItem {
629 role: CompositeRole::Document,
630 text: "system notes include lexical fallback and ranked retrieval".to_string(),
631 avec_override: Some(AvecState {
632 stability: 0.7,
633 friction: 0.3,
634 logic: 0.8,
635 autonomy: 0.7,
636 }),
637 context: Vec::new(),
638 }],
639 }],
640 options: CompositeNodeFromTextOptions {
641 allow_llm_avec_fallback: false,
642 ..Default::default()
643 },
644 };
645
646 let result = service
647 .build_content_from_text(&request)
648 .expect("composite content build should succeed");
649
650 let content = match result.content {
651 Value::Object(map) => map,
652 other => panic!("expected object content, got {other}"),
653 };
654 let avec = AvecState {
655 stability: 0.80,
656 friction: 0.20,
657 logic: 0.85,
658 autonomy: 0.75,
659 };
660 let metadata = SttpDocumentMetadata::new("sdk-composite-session")
661 .with_timestamp(
662 chrono::TimeZone::with_ymd_and_hms(&Utc, 2026, 5, 3, 0, 0, 0)
663 .single()
664 .expect("valid timestamp"),
665 )
666 .with_context_summary("sdk composite conformance")
667 .with_avec(avec, avec);
668 let raw_node = SttpDocumentBuilder::new(metadata)
669 .merge(SttpContentSlice::from_confidence_map(content).expect("content slice"))
670 .expect("merge content")
671 .build()
672 .expect("build document")
673 .render_canonical();
674
675 let validator = TreeSitterValidator::new();
676 let validation = validator.validate(&raw_node);
677 assert!(validation.is_valid, "validation failed: {:?}", validation.error);
678
679 let parser = SttpNodeParser::new();
680 let parse = parser.try_parse_strict_typed_ir(&raw_node, "sdk-composite-session");
681 assert!(parse.success, "strict typed-ir parse failed: {:?}", parse.error);
682 }
683
684 fn test_node(session_id: &str, timestamp: chrono::DateTime<Utc>, raw: &str) -> SttpNode {
685 let state = AvecState {
686 stability: 0.6,
687 friction: 0.4,
688 logic: 0.8,
689 autonomy: 0.7,
690 };
691
692 SttpNode {
693 raw: raw.to_string(),
694 session_id: session_id.to_string(),
695 tier: "raw".to_string(),
696 timestamp,
697 compression_depth: 1,
698 parent_node_id: None,
699 sync_key: format!("{}:{}", session_id, timestamp.timestamp_nanos_opt().unwrap_or_default()),
700 updated_at: timestamp,
701 source_metadata: None,
702 context_summary: Some(raw.to_string()),
703 semantic_tags: None,
704 semantic_links: None,
705 embedding_dimensions: None,
706 embedding_model: None,
707 embedding: None,
708 embedded_at: None,
709 user_avec: state,
710 model_avec: state,
711 compression_avec: Some(state),
712 rho: 0.9,
713 kappa: 0.8,
714 psi: 2.5,
715 }
716 }
717}